JIT: Move Thread sync and exec context save/restore to happen around runtime async bodies

[jakobbotsch]

This fixes an issue where there was a behavioral difference between async1 and async2 for async to sync runtime async calls. Before this PR the JIT would always save and restore contexts around awaits of task-returning methods, regardless of whether they were implemented as runtime async functions or not. That does not match async1: in async1, the context save and restore happens around the body in each async method. The former logic was an optimization, but the optimization is only correct for runtime async to runtime async calls.

We cannot in general know if a callee is implemented by runtime async or not, so we have to move the context save/restore to happen around the body of all runtime async methods.

An example test program that saw the behavioral difference before is the following:

using System;
using System.Threading;
using System.Threading.Tasks;

public class Program
{
    static void Main()
    {
        Foo().GetAwaiter().GetResult();
    }

    static async Task Foo()
    {
        SynchronizationContext.SetSynchronizationContext(new MySyncCtx(123));
        Console.WriteLine(((MySyncCtx)SynchronizationContext.Current).Arg);
        await Bar();
        Console.WriteLine(((MySyncCtx)SynchronizationContext.Current).Arg);
    }

    static Task Bar()
    {
        SynchronizationContext.SetSynchronizationContext(new MySyncCtx(124));
        return Task.CompletedTask;
    }

    class MySyncCtx(int arg) : SynchronizationContext
    {
        public int Arg => arg;
    }
}

Async1: 123 124
Runtime async before: 123 123
Runtime async after: 123 124

[dotnet-policy-service]

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch
See info in area-owners.md if you want to be subscribed.

[jakobbotsch]

FYI @VSadov @davidwrighton

Still not ready -- needs more testing, and I need to introduce the test case. But this is a heads up that some changes are incoming around the save/restore behavior.

[davidwrighton]

This doesn't feel right around Sync/Execution save and restore at suspension points. It looks correct for save/restore around the initial call, but the async callsite on suspension behavior looks wrong to me. Let me see if I can write up a few test cases. Do we have a convenient async testbed for these sorts of tests? In general, I think save/restore at the top level makes sense, but using that context saved at the top level during a suspend as the "captured" context I think is wrong.

[jakobbotsch]

Let me see if I can write up a few test cases. Do we have a convenient async testbed for these sorts of tests?

src/tests/async is where they live, there is a synchronization context test.

using that context saved at the top level during a suspend as the "captured" context I think is wrong.

We aren't doing that, we are getting new contexts from Thread.CurrentThread for saving inside the continuation (and for restoring on resumption). Then we restore the thread ones from the ones we saved at the beginning.
(This path is currently pretty inefficient, we could have a single helper to save the current contexts and restore the old ones, but I left that for a future PR)

[davidwrighton]

@jakobbotsch Ah, I was confused by some of the terminology around the JIT here. What happens if there is an exception while restoring the old contexts on suspension? Are we still inside the try/finally which will ALSO attempt to restore the old contexts? Notably, restoring a context involves running arbitrary code since you can have an AsyncLocal with value changed notifications, and what is the behavior if they throw? I don't want us to try to run the ExecutionContext restore twice. I have no idea what that would do. Now that I've looked at this, I think otherwise its all good.

[jakobbotsch]

What happens if there is an exception while restoring the old contexts on suspension? Are we still inside the try/finally which will ALSO attempt to restore the old contexts? Notably, restoring a context involves running arbitrary code since you can have an AsyncLocal with value changed notifications, and what is the behavior if they throw? I don't want us to try to run the ExecutionContext restore twice. I have no idea what that would do. Now that I've looked at this, I think otherwise its all good.

No, suspension/resumption code is not inside the try clause. So if OnValuesChanges throws we won't retry the restore.
Still, it is a good point about exceptions during suspension. I opened #121465 to investigate that further.

[Copilot]

Pull Request Overview

This PR refactors the async context (ExecutionContext and SynchronizationContext) handling for runtime async methods. The changes move from a per-await-call context save/restore model to a method-level context management approach, simplifying the implementation and enabling runtime async support by default.

Key Changes

• Method-level context management: ExecutionContext and SynchronizationContext are now saved once at method entry and restored at method exit/exception, rather than around each individual async call.
• Simplified async call handling: Removed per-call suspended indicator logic and execution context tracking fields from AsyncCallInfo.
• Enabled by default: Runtime async support is now enabled by default via the RuntimeAsync configuration flag.

Reviewed Changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/tests/async/synchronization-context/synchronization-context.cs	Adds test for async-to-sync context switching behavior
src/tests/async/Directory.Build.targets	Removes file that disabled async tests for certain configurations
src/coreclr/vm/method.hpp	Adds RequiresAsyncContextSaveAndRestore() method and fixes typo in comment
src/coreclr/vm/jitinterface.cpp	Passes CORINFO_ASYNC_SAVE_CONTEXTS flag to JIT for eligible methods
src/coreclr/jit/morph.cpp	Refactors GetCustomRegister to return bool and updates well-known arg names
src/coreclr/jit/lclmorph.cpp	Removes async suspended indicator local handling
src/coreclr/jit/importercalls.cpp	Simplifies async call setup by removing per-call context save/restore logic
src/coreclr/jit/importer.cpp	Removes tailcall restrictions for task awaits and adds new context args to inline handling
src/coreclr/jit/gentree.h	Refactors AsyncCallInfo structure and renames well-known args
src/coreclr/jit/gentree.cpp	Updates GetCustomRegister signature and removes obsolete methods
src/coreclr/jit/flowgraph.cpp	Minor comment punctuation fix
src/coreclr/jit/fginline.cpp	Corrects type and preserves async flag for inlined async calls
src/coreclr/jit/compiler.hpp	Removes async suspended indicator local handling from visitor
src/coreclr/jit/compiler.h	Adds method-level context locals and removes per-call tracking field
src/coreclr/jit/async.h	Renames field and adds RestoreContexts method
src/coreclr/jit/async.cpp	Major refactoring: implements method-level context save/restore with try/fault wrapper
src/coreclr/inc/corinfo.h	Adds CORINFO_ASYNC_SAVE_CONTEXTS flag
src/coreclr/inc/clrconfigvalues.h	Enables runtime async by default
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs	Updates parameter names and removes sync context argument from continuation context capture

[AndyAyersMS]

JIT changes LGTM, left some notes on (adding) comments.

[VSadov]

Do we still need context store/restore in task-returning thunks after this?

CC: @eduardo-vp

[jakobbotsch]

Do we still need context store/restore in task-returning thunks after this?

CC: @eduardo-vp

No, I do not think it is needed. We can switch it to just ensure Thread.CurrentThread is initialized. I think I will leave it as a follow up.

[jakobbotsch]

Anything else here @AndyAyersMS?

[VSadov]

LGTM. Thanks!

[jakobbotsch]

/ba-g Timeouts