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

[VSadov]

Fixes: #122052

=== async1 counterpart:

Regular/synchronous Capture/Restore of execution context should work unconditionally.

runtime/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilderCore.cs

Lines 30 to 33 in aa500a4

	// Store current ExecutionContext and SynchronizationContext as "previousXxx".
	// This allows us to restore them and undo any Context changes made in stateMachine.MoveNext
	// so that they won't "leak" out of the first await.
	ExecutionContext? previousExecutionCtx = currentThread._executionContext;

But captures on the suspension code path use ExecutionContext.Capture, that does not capture if the current context suppresses the flow.

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

Lines 158 to 167 in aa500a4

	ExecutionContext? currentContext = ExecutionContext.Capture();
	IAsyncStateMachineBox result;
	// 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
	// a strongly-typed manner into an AsyncStateMachineBox. It will already contain
	// the state machine as well as a MoveNextDelegate and a context. The only thing
	// we might need to do is update the context if that's changed since it was stored.
	if (taskField is AsyncStateMachineBox<TStateMachine> stronglyTypedBox)

[Copilot]

Pull request overview

This PR fixes issue #122052 by ensuring that ExecutionContext.IsFlowSuppressed() is properly respected when capturing execution context during async method suspension in the RuntimeAsync feature. The fix aligns the RuntimeAsync suspension code path with the existing behavior in AsyncMethodBuilderCore, ensuring that when execution context flow is suppressed, the captured context is treated as null (default execution context) rather than capturing the current context.

Key changes:

• Modified CaptureExecutionContext() to check m_isFlowSuppressed and return null when flow is suppressed
• Changed ExecutionContext.m_isFlowSuppressed visibility from private to internal to enable the check
• Enabled RuntimeAsync feature by default and removed test build restrictions

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs	Added flow suppression check to CaptureExecutionContext() method and clarifying comments to distinguish between contexts captured for synchronous restoration vs. suspension
src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs	Changed m_isFlowSuppressed field visibility from private to internal to allow access from AsyncHelpers
src/coreclr/inc/clrconfigvalues.h	Enabled RuntimeAsync feature by default (changed from 0 to 1)
src/tests/async/Directory.Build.targets	Removed conditional that disabled runtime async testing for non-NativeAOT builds

[VSadov]

TODO: undo enabling of async/tests before merging

[jakobbotsch]

ExecutionContext.Capture does something else, returning ExecutionContext.Default when the current execution context is null. Do we need to respect that?

[VSadov]

ExecutionContext.Capture has many callers with different purposes. For our purpose we only need to match what the AsyncStateMachineBox does. It looks like it stores whatever comes from ExecutionContext.Capture as-is, but when it comes to restoring it to a thread, all the paths check for Default and replace with null.

            private void MoveNext(Thread? threadPoolThread)
            {
                Debug.Assert(!IsCompleted);

                bool loggingOn = TplEventSource.Log.IsEnabled();
                if (loggingOn)
                {
                    TplEventSource.Log.TraceSynchronousWorkBegin(this.Id, CausalitySynchronousWork.Execution);
                }

                ExecutionContext? context = Context;
                if (context == null)
                {
                    Debug.Assert(StateMachine != null);
                    StateMachine.MoveNext();
                }
                else
                {
                    if (threadPoolThread is null)
                    {
                        ExecutionContext.RunInternal(context, s_callback, this);
                    }
                    else
                    {
                        ExecutionContext.RunFromThreadPoolDispatchLoop(threadPoolThread, context, s_callback, this);
                    }
                }

Both ExecutionContext.RunInternal and RunFromThreadPoolDispatchLoop have code that is:

            if (executionContext != null && executionContext.m_isDefault)
            {
                // Default is a null ExecutionContext internally
                executionContext = null;
            }

or

            // ThreadPool starts on Default Context so we don't need to save the "previous" state as we know it is Default (null)

            // Default is a null ExecutionContext internally
            if (executionContext != null && !executionContext.m_isDefault)
            {
                // Non-Default context to restore
                RestoreChangedContextToThread(threadPoolThread, contextToRestore: executionContext, currentContext: null);
            }

I think AsyncStateMachineBox could also store null for Default.

In fact, if we stored the result of ExecutionContext.Capture, we would need compensating code when restoring. Otherwise we hit asserts/crashes in a few places that expect an invariant "if context is not null, there must be variables in it".

[VSadov]

I’ll add an internal helper for async capture and use in AsyncStateMachineBox as well.

[jakobbotsch]

This would not be necessary if we used ExecutionContext.Capture.

[jakobbotsch]

Do we have a targeted test for this?

[VSadov]

added some test scenarios. @jakobbotsch @davidwrighton - PTAL. Thanks!

[jakobbotsch]

Add a new [Fact] for this test?

[jakobbotsch]

What about the non-T version of this?

Also cc @stephentoub

[VSadov]

non-T versions internally delegate to Task<VoidTaskResult>, same for ValueTask and ValueTask<T>

[VSadov]

Hold on. There are testcase failures in Libraries now.
I think I missed the subtle part that when the flow is suppressed, it does not mean "if capturing, capture default", which would make sense to me, it actually means "if restoring, do not restore".

The result is different if running a callback with an already dirty context.

That is why the Capture does its weird thing.

When context is applied to a thread, null means default state with no variables, but in the captured state null becomes Default (and when applied to a thread will clear its context), but suppressed context is captured as null and cause the callback to run without restoring at all.

[VSadov]

Basically, suppressing context flow will result in callback executing in a nondeterministic context that may depend on implementation details. (i.e. did we switch threads or ran "inline",...)

Indeed, the following produces randomly different output:

    internal class Program
    {
        static void Main(string[] args)
        {
            TestNoFlowOuter().GetAwaiter().GetResult();
        }

        public static AsyncLocal<long?> s_local = new AsyncLocal<long?>();

        private static async Task TestNoFlowOuter()
        {
            s_local.Value = 7;
            await TestNoFlowInner();
            System.Console.WriteLine(s_local.Value);
        }

        private static async Task TestNoFlowInner()
        {
            ExecutionContext.SuppressFlow();

            s_local.Value = 42;
            // returns synchronously, context stays the same.
            await ChangeThenReturn();
            System.Console.WriteLine(s_local.Value);

            // returns asynchronously, context should not flow.
            await ChangeYieldThenReturn();

// The following may print either False or True;
// I see mostly True in Debug, but sometimes False.
// In Release I see False, but would not be surprised if sometimes it prints True

            System.Console.WriteLine(s_local.Value == null); 
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static async Task ChangeThenReturn()
        {
            s_local.Value = 123;
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static async Task ChangeYieldThenReturn()
        {
            s_local.Value = 123;
            // restore flow so that state is not cleared by Yield
            ExecutionContext.RestoreFlow();
            await Task.Yield();
            System.Console.WriteLine(s_local.Value);
        }
    }

[VSadov]

Using null as a sentinel for flow-suppressed contexts would force us to map null to Default and then map back the Default to null and that is a per-frame expense that can add up.

I'll try using some other sentinel, so that the most common case, when the context is null, could be captured/restored as-is.

[VSadov]

The failures seem to be: #122228

[VSadov]

failures in Cryptography.Tests.MLKemImplementationTests are #122228