[RuntimeAsync] Respect IsFlowSuppressed when capturing execution context for suspension.

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -32,6 +32,9 @@ internal struct ExecutionAndSyncBlockStore
         public void Push()
         {
             _thread = Thread.CurrentThread;
+            // Here we get the execution context for synchronous restoring,
+            // not for flowing across suspension to potentially another thread.
+            // Therefore we do not need to worry about IsFlowSuppressed
             _previousExecutionCtx = _thread._executionContext;
             _previousSyncCtx = _thread._synchronizationContext;
         }
@@ -191,6 +194,9 @@ private struct RuntimeAsyncAwaitState
             public void CaptureContexts()
             {
                 Thread curThread = Thread.CurrentThreadAssumedInitialized;
+                // Here we get the execution context for presenting to the notifier,
+                // not for flowing across suspension to potentially another thread.
+                // Therefore we do not need to worry about IsFlowSuppressed
                 ExecutionContext = curThread._executionContext;
                 SynchronizationContext = curThread._synchronizationContext;
             }
@@ -788,15 +794,21 @@ private static ValueTask ValueTaskFromException(Exception ex)
             return new ValueTask<T?>(TaskFromException<T>(ex));
         }
 
+        // Called when capturing execution context for suspension.
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private static ExecutionContext? CaptureExecutionContext()
         {
-            return Thread.CurrentThreadAssumedInitialized._executionContext;
+            return ExecutionContext.CaptureForSuspension(Thread.CurrentThreadAssumedInitialized);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private static void RestoreExecutionContext(ExecutionContext? previousExecCtx)
         {
+            if (previousExecCtx == ExecutionContext.DefaultFlowSuppressed)
+            {
+                return;
+            }
+
             Thread thread = Thread.CurrentThreadAssumedInitialized;
             ExecutionContext? currentExecCtx = thread._executionContext;
             if (previousExecCtx != currentExecCtx)
@@ -809,6 +821,9 @@ private static void RestoreExecutionContext(ExecutionContext? previousExecCtx)
         private static void CaptureContexts(out ExecutionContext? execCtx, out SynchronizationContext? syncCtx)
         {
             Thread thread = Thread.CurrentThreadAssumedInitialized;
+            // Here we get the execution context for synchronous restoring,
+            // not for flowing across suspension to potentially another thread.
+            // Therefore we do not need to worry about IsFlowSuppressed
             execCtx = thread._executionContext;
             syncCtx = thread._synchronizationContext;
         }

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs

@@ -155,7 +155,7 @@ private static IAsyncStateMachineBox GetStateMachineBox<TStateMachine>(
             [NotNull] ref Task<TResult>? taskField)
             where TStateMachine : IAsyncStateMachine
         {
-            ExecutionContext? currentContext = ExecutionContext.Capture();
+            ExecutionContext? currentContext = ExecutionContext.CaptureForSuspension(Thread.CurrentThread);
 
             IAsyncStateMachineBox result;
 
@@ -360,7 +360,7 @@ private void MoveNext(Thread? threadPoolThread)
                 }
 
                 ExecutionContext? context = Context;
-                if (context == null)
+                if (context == ExecutionContext.DefaultFlowSuppressed)
                 {
                     Debug.Assert(StateMachine != null);
                     StateMachine.MoveNext();

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs

@@ -154,7 +154,7 @@ private static IAsyncStateMachineBox GetStateMachineBox<TStateMachine>(
             [NotNull] ref StateMachineBox? boxFieldRef)
             where TStateMachine : IAsyncStateMachine
         {
-            ExecutionContext? currentContext = ExecutionContext.Capture();
+            ExecutionContext? currentContext = ExecutionContext.CaptureForSuspension(Thread.CurrentThread);
 
             // Check first for the most common case: not the first yield in an async method.
             // In this case, the first yield will have already "boxed" the state machine in
@@ -403,7 +403,7 @@ public void MoveNext()
             {
                 ExecutionContext? context = Context;
 
-                if (context is null)
+                if (context == ExecutionContext.DefaultFlowSuppressed)
                 {
                     Debug.Assert(StateMachine is not null, $"Null {nameof(StateMachine)}");
                     StateMachine.MoveNext();

src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs

@@ -20,7 +20,9 @@ namespace System.Threading
     public sealed class ExecutionContext : IDisposable, ISerializable
     {
         internal static readonly ExecutionContext Default = new ExecutionContext();
-        private static ExecutionContext? s_defaultFlowSuppressed;
+#pragma warning disable CA1825, IDE0300 // Avoid unnecessary zero-length array allocations
+        internal static readonly ExecutionContext DefaultFlowSuppressed = new ExecutionContext(AsyncLocalValueMap.Empty, new IAsyncLocal[0], isFlowSuppressed: true);
+#pragma warning restore CA1825, IDE0300
 
         private readonly IAsyncLocalValueMap? m_localValues;
         private readonly IAsyncLocal[]? m_localChangeNotifications;
@@ -79,17 +81,30 @@ public void GetObjectData(SerializationInfo info, StreamingContext context)
             return Thread.CurrentThread._executionContext;
         }
 
+        // Capture for flowing/restoring across suspension points.
+        // Respects m_isFlowSuppressed, but avoids 'null' -> 'Default' -> 'null' reinterpretation.
+        internal static ExecutionContext? CaptureForSuspension(Thread currentThread)
+        {
+            Debug.Assert(Thread.CurrentThread == currentThread);
+
+            ExecutionContext? executionContext = currentThread._executionContext;
+            if (executionContext?.m_isFlowSuppressed == true)
+            {
+                executionContext = DefaultFlowSuppressed;
+            }
+
+            return executionContext;
+        }
+
         private ExecutionContext? ShallowClone(bool isFlowSuppressed)
         {
             Debug.Assert(isFlowSuppressed != m_isFlowSuppressed);
 
             if (m_localValues == null || AsyncLocalValueMap.IsEmpty(m_localValues))
             {
-#pragma warning disable CA1825, IDE0300 // Avoid unnecessary zero-length array allocations
                 return isFlowSuppressed ?
-                    (s_defaultFlowSuppressed ??= new ExecutionContext(AsyncLocalValueMap.Empty, new IAsyncLocal[0], isFlowSuppressed: true)) :
+                    DefaultFlowSuppressed :
                     null; // implies the default context
-#pragma warning restore CA1825, IDE0300
             }
 
             return new ExecutionContext(m_localValues, m_localChangeNotifications, isFlowSuppressed);
@@ -245,7 +260,7 @@ internal static void RestoreInternal(ExecutionContext? executionContext)
             }
         }
 
-        internal static void RunFromThreadPoolDispatchLoop(Thread threadPoolThread, ExecutionContext executionContext, ContextCallback callback, object state)
+        internal static void RunFromThreadPoolDispatchLoop(Thread threadPoolThread, ExecutionContext? executionContext, ContextCallback callback, object state)
         {
             Debug.Assert(threadPoolThread == Thread.CurrentThread);
             CheckThreadPoolAndContextsAreDefault();

src/tests/async/execution-context/execution-context.cs

@@ -10,11 +10,17 @@
 public class Async2ExecutionContext
 {
     [Fact]
-    public static void TestEntryPoint()
+    public static void TestDefaultFlow()
     {
         Test().GetAwaiter().GetResult();
     }
 
+    [Fact]
+    public static void TestSuppressedFlow()
+    {
+        TestNoFlowOuter().GetAwaiter().GetResult();
+    }
+
     public static AsyncLocal<long?> s_local = new AsyncLocal<long?>();
     private static async Task Test()
     {
@@ -51,6 +57,33 @@ private static async Task Test()
         Assert.Equal(46, s_local.Value);
     }
 
+    private static async Task TestNoFlowOuter()
+    {
+        s_local.Value = 7;
+        await TestNoFlowInner();
+        // by default exec context should flow, even if inner frames suppress the flow
+        Assert.Equal(7, s_local.Value);
+    }
+
+    private static async Task TestNoFlowInner()
+    {
+        ExecutionContext.SuppressFlow();
+
+        s_local.Value = 42;
+        // returns synchronously, context stays the same.
+        await ChangeThenReturn();
+        Assert.Equal(42, s_local.Value);
+
+        // returns asynchronously, context should not flow.
+        // the value is technically nondeterministic,
+        // but in our current implementation it will be 12345
+        await ChangeYieldThenReturn();
+        Assert.Equal(12345, s_local.Value);
+
+        // NB: no need to restore flow here as we will
+        //     be popping to the parent context anyways.
+    }
+
     [MethodImpl(MethodImplOptions.NoInlining)]
     private static async Task ChangeThenThrow()
     {
@@ -64,6 +97,16 @@ private static async Task ChangeThenReturn()
         s_local.Value = 123;
     }
 
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static async Task ChangeYieldThenReturn()
+    {
+        s_local.Value = 12345;
+        // restore flow so that state is not cleared by Yield
+        ExecutionContext.RestoreFlow();
+        await Task.Yield();
+        Assert.Equal(12345, s_local.Value);
+    }
+
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private static async Task ChangeThenThrowInlined()
     {