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.

[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