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

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:
```csharp
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

Author: jakobbotsch

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

@@ -556,7 +556,7 @@ public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskComplet
                 sentinelContinuation.Next = null;
 
                 // Head continuation should be the result of async call to AwaitAwaiter or UnsafeAwaitAwaiter.
-                // These never have special continuation handling.
+                // These never have special continuation context handling.
                 const ContinuationFlags continueFlags =
                     ContinuationFlags.ContinueOnCapturedSynchronizationContext |
                     ContinuationFlags.ContinueOnThreadPool |
@@ -774,36 +774,41 @@ private static void CaptureContexts(out ExecutionContext? execCtx, out Synchroni
             syncCtx = thread._synchronizationContext;
         }
 
+        // Restore contexts onto current Thread. If "resumed" then this is not the first starting call for the async method.
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void RestoreContexts(bool suspended, ExecutionContext? previousExecCtx, SynchronizationContext? previousSyncCtx)
+        private static void RestoreContexts(bool resumed, ExecutionContext? previousExecCtx, SynchronizationContext? previousSyncCtx)
         {
-            Thread thread = Thread.CurrentThreadAssumedInitialized;
-            if (!suspended && previousSyncCtx != thread._synchronizationContext)
+            if (!resumed)
             {
-                thread._synchronizationContext = previousSyncCtx;
-            }
+                Thread thread = Thread.CurrentThreadAssumedInitialized;
+                if (previousSyncCtx != thread._synchronizationContext)
+                {
+                    thread._synchronizationContext = previousSyncCtx;
+                }
 
-            ExecutionContext? currentExecCtx = thread._executionContext;
-            if (previousExecCtx != currentExecCtx)
-            {
-                ExecutionContext.RestoreChangedContextToThread(thread, previousExecCtx, currentExecCtx);
+                ExecutionContext? currentExecCtx = thread._executionContext;
+                if (previousExecCtx != currentExecCtx)
+                {
+                    ExecutionContext.RestoreChangedContextToThread(thread, previousExecCtx, currentExecCtx);
+                }
             }
         }
 
-        private static void CaptureContinuationContext(SynchronizationContext syncCtx, ref object context, ref ContinuationFlags flags)
+        private static void CaptureContinuationContext(ref object continuationContext, ref ContinuationFlags flags)
         {
+            SynchronizationContext? syncCtx = Thread.CurrentThreadAssumedInitialized._synchronizationContext;
             if (syncCtx != null && syncCtx.GetType() != typeof(SynchronizationContext))
             {
                 flags |= ContinuationFlags.ContinueOnCapturedSynchronizationContext;
-                context = syncCtx;
+                continuationContext = syncCtx;
                 return;
             }
 
             TaskScheduler? sched = TaskScheduler.InternalCurrent;
             if (sched != null && sched != TaskScheduler.Default)
             {
                 flags |= ContinuationFlags.ContinueOnCapturedTaskScheduler;
-                context = sched;
+                continuationContext = sched;
                 return;
             }
 

src/coreclr/inc/corinfo.h

@@ -722,6 +722,7 @@ enum CorInfoOptions
                                                CORINFO_GENERICS_CTXT_FROM_METHODDESC |
                                                CORINFO_GENERICS_CTXT_FROM_METHODTABLE),
     CORINFO_GENERICS_CTXT_KEEP_ALIVE        = 0x00000100, // Keep the generics context alive throughout the method even if there is no explicit use, and report its location to the CLR
+    CORINFO_ASYNC_SAVE_CONTEXTS             = 0x00000200, // Runtime async method must save and restore contexts
 };
 
 //

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* a802fbbf-3e14-4b34-a348-5fba9fd756d4 */
-    0xa802fbbf,
-    0x3e14,
-    0x4b34,
-    {0xa3, 0x48, 0x5f, 0xba, 0x9f, 0xd7, 0x56, 0xd4}
+constexpr GUID JITEEVersionIdentifier = { /* f0752445-e116-444d-98ea-aaa7cbc30baa */
+    0xf0752445,
+    0xe116,
+    0x444d,
+    {0x98, 0xea, 0xaa, 0xa7, 0xcb, 0xc3, 0x0b, 0xaa}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H