Add thread-safety and wait-blocking to BenchmarkSynchronizationContext.

Author: timcassell

src/BenchmarkDotNet/Engines/BenchmarkSynchronizationContext.cs

@@ -10,12 +10,40 @@ namespace BenchmarkDotNet.Engines;
 // Used to ensure async continuations are posted back to the same thread that the benchmark process was started on.
 [UsedImplicitly]
 [EditorBrowsable(EditorBrowsableState.Never)]
-public sealed class BenchmarkSynchronizationContext : SynchronizationContext, IDisposable
+public readonly ref struct BenchmarkSynchronizationContext : IDisposable
 {
-    private readonly SynchronizationContext previousContext;
+    private readonly BenchmarkDotNetSynchronizationContext context;
+
+    private BenchmarkSynchronizationContext(BenchmarkDotNetSynchronizationContext context)
+    {
+        this.context = context;
+    }
+
+    public static BenchmarkSynchronizationContext CreateAndSetCurrent()
+    {
+        var context = new BenchmarkDotNetSynchronizationContext(SynchronizationContext.Current);
+        SynchronizationContext.SetSynchronizationContext(context);
+        return new(context);
+    }
+
+    public void Dispose()
+        => context.Dispose();
+
+    public void ExecuteUntilComplete(ValueTask valueTask)
+        => context.ExecuteUntilComplete(valueTask);
+
+    public T ExecuteUntilComplete<T>(ValueTask<T> valueTask)
+        => context.ExecuteUntilComplete(valueTask);
+}
+
+internal sealed class BenchmarkDotNetSynchronizationContext : SynchronizationContext
+{
+    private readonly SynchronizationContext? previousContext;
     private readonly Queue<(SendOrPostCallback d, object? state)> queue = new();
+    private bool isDisposed;
+    volatile private bool isCompleted;
 
-    private BenchmarkSynchronizationContext(SynchronizationContext previousContext)
+    internal BenchmarkDotNetSynchronizationContext(SynchronizationContext? previousContext)
     {
         this.previousContext = previousContext;
     }
@@ -24,53 +52,133 @@ public override SynchronizationContext CreateCopy()
         => this;
 
     public override void Post(SendOrPostCallback d, object? state)
-        => queue.Enqueue((d ?? throw new ArgumentNullException(nameof(d)), state));
+    {
+        if (d is null) throw new ArgumentNullException(nameof(d));
 
-    public static BenchmarkSynchronizationContext CreateAndSetCurrent()
+        lock (queue)
+        {
+            ThrowIfDisposed();
+
+            queue.Enqueue((d, state));
+            Monitor.Pulse(queue);
+        }
+    }
+
+    private void ThrowIfDisposed()
     {
-        var context = new BenchmarkSynchronizationContext(Current);
-        SetSynchronizationContext(context);
-        return context;
+        if (isDisposed) throw new ObjectDisposedException(nameof(BenchmarkDotNetSynchronizationContext));
     }
 
-    public void Dispose()
-        => SetSynchronizationContext(previousContext);
+    internal void Dispose()
+    {
+        lock (queue)
+        {
+            ThrowIfDisposed();
+            isDisposed = true;
 
-    public void ExecuteUntilComplete(ValueTask valueTask)
+            // Flush any remaining posted callbacks.
+            while (TryDequeue(out var callbackAndState))
+            {
+                callbackAndState.d(callbackAndState.state);
+            }
+        }
+        SetSynchronizationContext(previousContext);
+    }
+
+    internal void ExecuteUntilComplete(ValueTask valueTask)
     {
-        var spinner = new SpinWait();
-        while (!valueTask.IsCompleted)
+        ThrowIfDisposed();
+
+        var awaiter = valueTask.GetAwaiter();
+        if (valueTask.IsCompleted)
         {
-            DoSpin(ref spinner);
+            awaiter.GetResult();
+            return;
         }
-        valueTask.GetAwaiter().GetResult();
+
+        isCompleted = false;
+        awaiter.UnsafeOnCompleted(OnCompleted);
+        ExecuteUntilComplete();
+        awaiter.GetResult();
     }
 
-    public T ExecuteUntilComplete<T>(ValueTask<T> valueTask)
+    internal T ExecuteUntilComplete<T>(ValueTask<T> valueTask)
     {
-        var spinner = new SpinWait();
-        while (!valueTask.IsCompleted)
+        ThrowIfDisposed();
+
+        var awaiter = valueTask.GetAwaiter();
+        if (valueTask.IsCompleted)
         {
-            DoSpin(ref spinner);
+            return awaiter.GetResult();
         }
-        return valueTask.GetAwaiter().GetResult();
+
+        isCompleted = false;
+        awaiter.UnsafeOnCompleted(OnCompleted);
+        ExecuteUntilComplete();
+        return awaiter.GetResult();
     }
 
-    private void DoSpin(ref SpinWait spinner)
+    private void OnCompleted()
     {
-        if (queue.Count <= 0)
+        isCompleted = true;
+        lock (queue)
         {
+            Monitor.Pulse(queue);
+        }
+    }
+
+    private void ExecuteUntilComplete()
+    {
+        var spinner = new SpinWait();
+        while (true)
+        {
+            if (TryDequeue(out var callbackAndState))
+            {
+                do
+                {
+                    callbackAndState.d(callbackAndState.state);
+                }
+                while (TryDequeue(out callbackAndState));
+                // Reset spinner after any posted callback is executed.
+                spinner = new();
+            }
+
+            if (isCompleted)
+            {
+                return;
+            }
+
+            if (spinner.NextSpinWillYield)
+            {
+                // Yield the thread and wait for completion or for a posted callback.
+                lock (queue)
+                {
+                    Monitor.Wait(queue);
+                }
+                // Reset the spinner.
+                spinner = new();
+                continue;
+            }
+
             spinner.SpinOnce();
-            return;
         }
+    }
 
-        do
+    private bool TryDequeue(out (SendOrPostCallback d, object? state) callbackAndState)
+    {
+        lock (queue)
         {
-            var (d, state) = queue.Dequeue();
-            d(state);
+#if NETSTANDARD2_0
+            if (queue.Count > 0)
+            {
+                callbackAndState = queue.Dequeue();
+                return true;
+            }
+            callbackAndState = default;
+            return false;
+#else
+            return queue.TryDequeue(out callbackAndState);
+#endif
         }
-        while (queue.Count > 0);
-        // Reset spinner after any posted callback is executed.
-        spinner = new();
     }
 }

src/BenchmarkDotNet/Templates/BenchmarkProgram.txt

@@ -59,12 +59,17 @@ namespace BenchmarkDotNet.Autogenerated
                 {
                     host.WriteLine("You have not specified benchmark id (an integer) so the first benchmark will be executed.");
                 }
-                using (global::BenchmarkDotNet.Engines.BenchmarkSynchronizationContext benchmarkSynchronizationContext = global::BenchmarkDotNet.Engines.BenchmarkSynchronizationContext.CreateAndSetCurrent())
+                global::BenchmarkDotNet.Engines.BenchmarkSynchronizationContext benchmarkSynchronizationContext = global::BenchmarkDotNet.Engines.BenchmarkSynchronizationContext.CreateAndSetCurrent();
+                try
                 {
                     global::System.Threading.Tasks.ValueTask runTask;
                     $BenchmarkRunCall$
                     benchmarkSynchronizationContext.ExecuteUntilComplete(runTask);
                 }
+                finally
+                {
+                    benchmarkSynchronizationContext.Dispose();
+                }
                 return 0;
             }
             catch (global::System.Exception oom) when (oom is global::System.OutOfMemoryException || oom is global::System.Reflection.TargetInvocationException reflection && reflection.InnerException is global::System.OutOfMemoryException)