diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 0045c9e53dbf3d..78ef25f0e4eb66 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/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(TaskFromException(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; } diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs index 5a2ccb1d635493..500a28a8d5952a 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs @@ -155,7 +155,7 @@ private static IAsyncStateMachineBox GetStateMachineBox( [NotNull] ref Task? 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(); diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs index 150f22c1970325..2c68e487d93c73 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/PoolingAsyncValueTaskMethodBuilderT.cs @@ -154,7 +154,7 @@ private static IAsyncStateMachineBox GetStateMachineBox( [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(); diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs index affc292913fcf5..584df1d63b99d9 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs +++ b/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(); diff --git a/src/tests/async/execution-context/execution-context.cs b/src/tests/async/execution-context/execution-context.cs index 6455792a0d0f50..0eea918e8ca7ca 100644 --- a/src/tests/async/execution-context/execution-context.cs +++ b/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 s_local = new AsyncLocal(); 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() {