Suppress ExecutionContext flow and set BenchmarkSynchronizationContext in InProcess background threads. (#3115)

Author: timcassell

src/BenchmarkDotNet/Helpers/ExecutionContextHelper.cs

@@ -0,0 +1,24 @@
+namespace BenchmarkDotNet.Helpers;
+
+internal static class ExecutionContextHelper
+{
+#if NETCOREAPP
+    internal static AsyncFlowControl SuppressFlow()
+        => ExecutionContext.SuppressFlow();
+#else
+    // .Net Framework throws if ExecutionContext.SuppressFlow() is called when it's already suppressed.
+    internal static WrappedAsyncFlowControl SuppressFlow()
+        => ExecutionContext.IsFlowSuppressed() ? default : new WrappedAsyncFlowControl(ExecutionContext.SuppressFlow());
+
+    internal readonly struct WrappedAsyncFlowControl(AsyncFlowControl asyncFlowControl) : IDisposable
+    {
+        public void Dispose()
+        {
+            if (asyncFlowControl != default)
+            {
+                asyncFlowControl.Dispose();
+            }
+        }
+    }
+#endif
+}

src/BenchmarkDotNet/Toolchains/InProcess/Emit/InProcessEmitExecutor.cs

@@ -21,28 +21,33 @@ public async ValueTask<ExecuteResult> ExecuteAsync(ExecuteParameters executePara
             if (executeOnSeparateThread)
             {
                 var taskCompletionSource = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
-                var runThread = new Thread(async () =>
+                Thread runThread;
+                // Ensure the new thread is started with a fresh execution context.
+                using (ExecutionContextHelper.SuppressFlow())
                 {
-                    try
+                    runThread = new Thread(() =>
                     {
-                        taskCompletionSource.SetResult(await ExecuteCore(host, executeParameters).ConfigureAwait(false));
-                    }
-                    catch (Exception ex)
+                        // The engine yields immediately, so we need to set the synchronization context to ensure the benchmarks are executed on this dedicated thread, and not on a ThreadPool thread.
+                        using var context = BenchmarkSynchronizationContext.CreateAndSetCurrent();
+                        try
+                        {
+                            var executeTask = ExecuteCore(host, executeParameters);
+                            taskCompletionSource.SetResult(context.ExecuteUntilComplete(executeTask));
+                        }
+                        catch (Exception ex)
+                        {
+                            taskCompletionSource.SetException(ex);
+                        }
+                    });
+                    if (executeParameters.BenchmarkCase.Descriptor.WorkloadMethod.GetCustomAttributes<STAThreadAttribute>(false).Any()
+                        && OsDetector.IsWindows())
                     {
-                        taskCompletionSource.SetException(ex);
+                        runThread.SetApartmentState(ApartmentState.STA);
                     }
-                });
-
-                if (executeParameters.BenchmarkCase.Descriptor.WorkloadMethod.GetCustomAttributes<STAThreadAttribute>(false).Any()
-                    && OsDetector.IsWindows())
-                {
-                    runThread.SetApartmentState(ApartmentState.STA);
+                    runThread.IsBackground = true;
+                    runThread.Start();
                 }
 
-                runThread.IsBackground = true;
-
-                runThread.Start();
-
                 exitCode = await taskCompletionSource.Task.ConfigureAwait(true);
                 runThread.Join();
             }

src/BenchmarkDotNet/Toolchains/InProcess/NoEmit/InProcessNoEmitExecutor.cs

@@ -21,28 +21,33 @@ public async ValueTask<ExecuteResult> ExecuteAsync(ExecuteParameters executePara
             if (executeOnSeparateThread)
             {
                 var taskCompletionSource = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
-                var runThread = new Thread(async () =>
+                Thread runThread;
+                // Ensure the new thread is started with a fresh execution context.
+                using (ExecutionContextHelper.SuppressFlow())
                 {
-                    try
+                    runThread = new Thread(() =>
                     {
-                        taskCompletionSource.SetResult(await ExecuteCore(host, executeParameters, benchmarkActionFactory).ConfigureAwait(false));
-                    }
-                    catch (Exception ex)
+                        // The engine yields immediately, so we need to set the synchronization context to ensure the benchmarks are executed on this dedicated thread, and not on a ThreadPool thread.
+                        using var context = BenchmarkSynchronizationContext.CreateAndSetCurrent();
+                        try
+                        {
+                            var executeTask = ExecuteCore(host, executeParameters, benchmarkActionFactory);
+                            taskCompletionSource.SetResult(context.ExecuteUntilComplete(executeTask));
+                        }
+                        catch (Exception ex)
+                        {
+                            taskCompletionSource.SetException(ex);
+                        }
+                    });
+                    if (executeParameters.BenchmarkCase.Descriptor.WorkloadMethod.GetCustomAttributes<STAThreadAttribute>(false).Any()
+                        && OsDetector.IsWindows())
                     {
-                        taskCompletionSource.SetException(ex);
+                        runThread.SetApartmentState(ApartmentState.STA);
                     }
-                });
-
-                if (executeParameters.BenchmarkCase.Descriptor.WorkloadMethod.GetCustomAttributes<STAThreadAttribute>(false).Any()
-                    && OsDetector.IsWindows())
-                {
-                    runThread.SetApartmentState(ApartmentState.STA);
+                    runThread.IsBackground = true;
+                    runThread.Start();
                 }
 
-                runThread.IsBackground = true;
-
-                runThread.Start();
-
                 exitCode = await taskCompletionSource.Task.ConfigureAwait(true);
                 runThread.Join();
             }