Add test case to make sure ValueTasks work properly with a race condi…

…tion between `IsCompleted` and `OnCompleted`. Changed AwaitHelper to use `ManualResetEventSlim` instead of `Monitor.Wait`.
dotnet · Feb 17, 2023 · 43138da · 43138da
1 parent b3d8e81
commit 43138da
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 24 deletions.
diff --git a/src/BenchmarkDotNet/Helpers/AwaitHelper.cs b/src/BenchmarkDotNet/Helpers/AwaitHelper.cs
@@ -2,6 +2,7 @@
 using System.Linq;
 using System.Reflection;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace BenchmarkDotNet.Helpers
@@ -10,54 +11,55 @@ public static class AwaitHelper
     {
         private class ValueTaskWaiter
         {
-            // We use thread static field so that multiple threads can use individual lock object and callback.
+            // We use thread static field so that each thread uses its own individual callback and reset event.
             [ThreadStatic]
             private static ValueTaskWaiter ts_current;
             internal static ValueTaskWaiter Current => ts_current ??= new ValueTaskWaiter();
 
+            // We cache the callback to prevent allocations for memory diagnoser.
             private readonly Action awaiterCallback;
-            private bool awaiterCompleted;
+            private readonly ManualResetEventSlim resetEvent;
 
             private ValueTaskWaiter()
             {
-                awaiterCallback = AwaiterCallback;
-            }
-
-            private void AwaiterCallback()
-            {
-                lock (this)
-                {
-                    awaiterCompleted = true;
-                    System.Threading.Monitor.Pulse(this);
-                }
+                resetEvent = new ();
+                awaiterCallback = resetEvent.Set;
             }
 
             // Hook up a callback instead of converting to Task to prevent extra allocations on each benchmark run.
             internal void Wait(ConfiguredValueTaskAwaitable.ConfiguredValueTaskAwaiter awaiter)
             {
-                lock (this)
+                resetEvent.Reset();
+                awaiter.UnsafeOnCompleted(awaiterCallback);
+
+                // The fastest way to wait for completion is to spin a bit before waiting on the event. This is the same logic that Task.GetAwaiter().GetResult() uses.
+                var spinner = new SpinWait();
+                while (!resetEvent.IsSet)
                 {
-                    awaiterCompleted = false;
-                    awaiter.UnsafeOnCompleted(awaiterCallback);
-                    // Check if the callback executed synchronously before blocking.
-                    if (!awaiterCompleted)
+                    if (spinner.NextSpinWillYield)
                     {
-                        System.Threading.Monitor.Wait(this);
+                        resetEvent.Wait();
+                        return;
                     }
+                    spinner.SpinOnce();
                 }
             }
 
             internal void Wait<T>(ConfiguredValueTaskAwaitable<T>.ConfiguredValueTaskAwaiter awaiter)
             {
-                lock (this)
+                resetEvent.Reset();
+                awaiter.UnsafeOnCompleted(awaiterCallback);
+
+                // The fastest way to wait for completion is to spin a bit before waiting on the event. This is the same logic that Task.GetAwaiter().GetResult() uses.
+                var spinner = new SpinWait();
+                while (!resetEvent.IsSet)
                 {
-                    awaiterCompleted = false;
-                    awaiter.UnsafeOnCompleted(awaiterCallback);
-                    // Check if the callback executed synchronously before blocking.
-                    if (!awaiterCompleted)
+                    if (spinner.NextSpinWillYield)
                     {
-                        System.Threading.Monitor.Wait(this);
+                        resetEvent.Wait();
+                        return;
                     }
+                    spinner.SpinOnce();
                 }
             }
         }

diff --git a/tests/BenchmarkDotNet.IntegrationTests/AsyncBenchmarksTests.cs b/tests/BenchmarkDotNet.IntegrationTests/AsyncBenchmarksTests.cs
@@ -22,6 +22,23 @@ internal class ValueTaskSource<T> : IValueTaskSource<T>, IValueTaskSource
         public void SetResult(T result) => _core.SetResult(result);
     }
 
+    // This is used to test the case of ValueTaskAwaiter.IsCompleted returns false, then OnCompleted invokes the callback immediately because it happened to complete between the 2 calls.
+    internal class ValueTaskSourceCallbackOnly<T> : IValueTaskSource<T>, IValueTaskSource
+    {
+        private ManualResetValueTaskSourceCore<T> _core;
+
+        T IValueTaskSource<T>.GetResult(short token) => _core.GetResult(token);
+        void IValueTaskSource.GetResult(short token) => _core.GetResult(token);
+        // Always return pending state so OnCompleted will be called.
+        ValueTaskSourceStatus IValueTaskSource<T>.GetStatus(short token) => ValueTaskSourceStatus.Pending;
+        ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => ValueTaskSourceStatus.Pending;
+        void IValueTaskSource<T>.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags) => _core.OnCompleted(continuation, state, token, flags);
+        void IValueTaskSource.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags) => _core.OnCompleted(continuation, state, token, flags);
+        public void Reset() => _core.Reset();
+        public short Token => _core.Version;
+        public void SetResult(T result) => _core.SetResult(result);
+    }
+
     public class AsyncBenchmarksTests : BenchmarkTestExecutor
     {
         public AsyncBenchmarksTests(ITestOutputHelper output) : base(output) { }
@@ -41,6 +58,9 @@ public void TaskReturningMethodsAreAwaited()
             }
         }
 
+        [Fact]
+        public void TaskReturningMethodsAreAwaited_AlreadyComplete() => CanExecute<TaskImmediateMethods>();
+
         public class TaskDelayMethods
         {
             private readonly ValueTaskSource<int> valueTaskSource = new ();
@@ -89,5 +109,58 @@ public ValueTask<int> ReturningGenericValueTaskBackByIValueTaskSource()
                 return new ValueTask<int>(valueTaskSource, valueTaskSource.Token);
             }
         }
+
+        public class TaskImmediateMethods
+        {
+            private readonly ValueTaskSource<int> valueTaskSource = new ();
+            private readonly ValueTaskSourceCallbackOnly<int> valueTaskSourceCallbackOnly = new ();
+
+            [Benchmark]
+            public Task ReturningTask() => Task.CompletedTask;
+
+            [Benchmark]
+            public ValueTask ReturningValueTask() => new ValueTask();
+
+            [Benchmark]
+            public ValueTask ReturningValueTaskBackByIValueTaskSource()
+            {
+                valueTaskSource.Reset();
+                valueTaskSource.SetResult(default);
+                return new ValueTask(valueTaskSource, valueTaskSource.Token);
+            }
+
+            [Benchmark]
+            public ValueTask ReturningValueTaskBackByIValueTaskSource_ImmediateCallback()
+            {
+                valueTaskSourceCallbackOnly.Reset();
+                valueTaskSourceCallbackOnly.SetResult(default);
+                return new ValueTask(valueTaskSourceCallbackOnly, valueTaskSourceCallbackOnly.Token);
+            }
+
+            [Benchmark]
+            public async Task Awaiting() => await Task.CompletedTask;
+
+            [Benchmark]
+            public Task<int> ReturningGenericTask() => ReturningTask().ContinueWith(_ => default(int));
+
+            [Benchmark]
+            public ValueTask<int> ReturningGenericValueTask() => new ValueTask<int>(ReturningGenericTask());
+
+            [Benchmark]
+            public ValueTask<int> ReturningGenericValueTaskBackByIValueTaskSource()
+            {
+                valueTaskSource.Reset();
+                valueTaskSource.SetResult(default);
+                return new ValueTask<int>(valueTaskSource, valueTaskSource.Token);
+            }
+
+            [Benchmark]
+            public ValueTask<int> ReturningGenericValueTaskBackByIValueTaskSource_ImmediateCallback()
+            {
+                valueTaskSourceCallbackOnly.Reset();
+                valueTaskSourceCallbackOnly.SetResult(default);
+                return new ValueTask<int>(valueTaskSourceCallbackOnly, valueTaskSourceCallbackOnly.Token);
+            }
+        }
     }
 }