Skip save/restore of contexts in task-returning thunks; mark thunk infrastructure NonVersionable

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

@@ -5,6 +5,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Reflection;
 using System.Runtime.InteropServices;
+using System.Runtime.Versioning;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Sources;
@@ -207,6 +208,7 @@ public static partial class AsyncHelpers
         // task-returning thunk, or DispatchContinuations. A pointer to this
         // state is kept in the runtime async TLS. This storage method avoids
         // costly write barriers on the hot path of suspension/resumption.
+        [NonVersionable]
         private ref struct RuntimeAsyncStackState
         {
             // The following are the possible introducers of asynchrony into a chain of awaits.
@@ -221,39 +223,12 @@ private ref struct RuntimeAsyncStackState
             public ExecutionContext? LeafExecutionContext;
             public SynchronizationContext? LeafSynchronizationContext;
 
-            // When we enter the root of the async chain (either an async thunk
-            // or DispatchContinuations), the contexts are captured into these
-            // fields.
-            public ExecutionContext? RootExecutionContext;
-            public SynchronizationContext? RootSynchronizationContext;
-
             public unsafe RuntimeAsyncStackState* Next;
-
-            public void Push(Thread thread)
-            {
-                RootExecutionContext = thread._executionContext;
-                RootSynchronizationContext = thread._synchronizationContext;
-            }
-
-            public void Pop(Thread thread)
-            {
-                // The common case is that these have not changed, so avoid the cost of a write barrier if not needed.
-                if (RootSynchronizationContext != thread._synchronizationContext)
-                {
-                    // Restore changed SynchronizationContext back to previous
-                    thread._synchronizationContext = RootSynchronizationContext;
-                }
-
-                ExecutionContext? currentExecutionCtx = thread._executionContext;
-                if (RootExecutionContext != currentExecutionCtx)
-                {
-                    ExecutionContext.RestoreChangedContextToThread(thread, RootExecutionContext, currentExecutionCtx);
-                }
-            }
         }
 
         // Used during suspensions to hold the continuation chain and on what we are waiting.
-        // Methods like FinalizeTaskReturningThunk will unlink the state and wrap into a Task.
+        // Methods like CreateRuntimeAsyncTask will unlink the state and wrap into a Task.
+        [NonVersionable]
         private unsafe struct RuntimeAsyncAwaitState
         {
             public Continuation? SentinelContinuation;
@@ -282,22 +257,51 @@ public void CaptureContexts()
 
             // At the start of an async chain (task-returning thunk or DispatchContinuations) this function
             // is called
+            [NonVersionable]
             public void Push(RuntimeAsyncStackState* stackState)
             {
                 stackState->Next = StackState;
                 StackState = stackState;
-                stackState->Push(CurrentThread ??= Thread.CurrentThread);
+                CurrentThread ??= Thread.CurrentThread;
             }
 
             // This function is called at the end of an async chain
+            [NonVersionable]
             public void Pop()
             {
                 Debug.Assert(CurrentThread != null);
-                StackState->Pop(CurrentThread);
                 StackState = StackState->Next;
             }
         }
 
+        private struct AsyncContexts
+        {
+            private SynchronizationContext? _synchronizationContext;
+            private ExecutionContext? _executionContext;
+
+            public void Push(Thread thread)
+            {
+                _synchronizationContext = thread._synchronizationContext;
+                _executionContext = thread._executionContext;
+            }
+
+            public void Pop(Thread thread)
+            {
+                // The common case is that these have not changed, so avoid the cost of a write barrier if not needed.
+                if (_synchronizationContext != thread._synchronizationContext)
+                {
+                    // Restore changed SynchronizationContext back to previous
+                    thread._synchronizationContext = _synchronizationContext;
+                }
+
+                ExecutionContext? currentExecutionCtx = thread._executionContext;
+                if (_executionContext != currentExecutionCtx)
+                {
+                    ExecutionContext.RestoreChangedContextToThread(thread, _executionContext, currentExecutionCtx);
+                }
+            }
+        }
+
         [ThreadStatic]
         private static RuntimeAsyncAwaitState t_runtimeAsyncAwaitState;
 
@@ -917,6 +921,9 @@ private unsafe void DispatchContinuations()
                 ref RuntimeAsyncAwaitState awaitState = ref t_runtimeAsyncAwaitState;
                 awaitState.Push(&stackState);
 
+                AsyncContexts contexts = default;
+                contexts.Push(awaitState.CurrentThread!);
+
                 ref AsyncDispatcherInfo* refDispatcherInfo = ref AsyncDispatcherInfo.t_current;
 
                 AsyncDispatcherInfo asyncDispatcherInfo;
@@ -948,6 +955,7 @@ private unsafe void DispatchContinuations()
                             newContinuation.Next = nextContinuation;
                             HandleSuspended(ref awaitState);
 
+                            contexts.Pop(awaitState.CurrentThread!);
                             awaitState.Pop();
                             refDispatcherInfo = asyncDispatcherInfo.Next;
                             return;
@@ -964,6 +972,7 @@ private unsafe void DispatchContinuations()
                                 TrySetCanceled(oce.CancellationToken, oce) :
                                 TrySetException(ex);
 
+                            contexts.Pop(awaitState.CurrentThread!);
                             awaitState.Pop();
                             refDispatcherInfo = asyncDispatcherInfo.Next;
 
@@ -983,6 +992,7 @@ private unsafe void DispatchContinuations()
                     {
                         bool successfullySet = TrySetResult(m_result);
 
+                        contexts.Pop(awaitState.CurrentThread!);
                         awaitState.Pop();
                         refDispatcherInfo = asyncDispatcherInfo.Next;
 
@@ -996,6 +1006,7 @@ private unsafe void DispatchContinuations()
 
                     if (QueueContinuationFollowUpActionIfNecessary(asyncDispatcherInfo.NextContinuation))
                     {
+                        contexts.Pop(awaitState.CurrentThread!);
                         awaitState.Pop();
                         refDispatcherInfo = asyncDispatcherInfo.Next;
                         return;
@@ -1005,6 +1016,7 @@ private unsafe void DispatchContinuations()
                     {
                         SetContinuationState(asyncDispatcherInfo.NextContinuation);
 
+                        contexts.Pop(awaitState.CurrentThread!);
                         awaitState.Pop();
                         refDispatcherInfo = asyncDispatcherInfo.Next;
 
@@ -1025,6 +1037,9 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
                 ref RuntimeAsyncAwaitState awaitState = ref t_runtimeAsyncAwaitState;
                 awaitState.Push(&stackState);
 
+                AsyncContexts contexts = default;
+                contexts.Push(awaitState.CurrentThread!);
+
                 ref AsyncDispatcherInfo* refDispatcherInfo = ref AsyncDispatcherInfo.t_current;
 
                 AsyncDispatcherInfo asyncDispatcherInfo;
@@ -1063,6 +1078,7 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
                             RuntimeAsyncInstrumentationHelpers.AwaitSuspendedRuntimeAsyncContext(ref asyncDispatcherInfo, flags, curContinuation, newContinuation, awaitState.SentinelContinuation!.Next);
                             InstrumentedHandleSuspended(flags, ref awaitState);
 
+                            contexts.Pop(awaitState.CurrentThread!);
                             awaitState.Pop();
                             refDispatcherInfo = asyncDispatcherInfo.Next;
                             return;
@@ -1083,6 +1099,7 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
                                 TrySetCanceled(oce.CancellationToken, oce) :
                                 TrySetException(ex);
 
+                            contexts.Pop(awaitState.CurrentThread!);
                             awaitState.Pop();
                             refDispatcherInfo = asyncDispatcherInfo.Next;
 
@@ -1106,6 +1123,7 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
 
                         bool successfullySet = TrySetResult(m_result);
 
+                        contexts.Pop(awaitState.CurrentThread!);
                         awaitState.Pop();
                         refDispatcherInfo = asyncDispatcherInfo.Next;
 
@@ -1121,6 +1139,7 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
                     {
                         RuntimeAsyncInstrumentationHelpers.QueueSuspendedRuntimeAsyncContext(ref asyncDispatcherInfo, flags, asyncDispatcherInfo.NextContinuation);
 
+                        contexts.Pop(awaitState.CurrentThread!);
                         awaitState.Pop();
                         refDispatcherInfo = asyncDispatcherInfo.Next;
                         return;
@@ -1265,30 +1284,30 @@ private static void FinalizeRuntimeAsyncTask<T>(ref RuntimeAsyncAwaitState state
 #pragma warning disable CA1859
         // When a Task-returning thunk gets a continuation result
         // it calls here to make a Task that awaits on the current async state.
-        private static Task<T?> FinalizeTaskReturningThunk<T>(ref RuntimeAsyncAwaitState state)
+        private static Task<T?> CreateRuntimeAsyncTask<T>(ref RuntimeAsyncAwaitState state)
         {
             RuntimeAsyncTask<T?> result = new();
             FinalizeRuntimeAsyncTask(ref state, result!);
             return result;
         }
 
-        private static Task FinalizeTaskReturningThunk(ref RuntimeAsyncAwaitState state)
+        private static Task CreateRuntimeAsyncTask(ref RuntimeAsyncAwaitState state)
         {
             RuntimeAsyncTask<VoidTaskResult> result = new();
             FinalizeRuntimeAsyncTask(ref state, result!);
             return result;
         }
 
-        private static ValueTask<T?> FinalizeValueTaskReturningThunk<T>(ref RuntimeAsyncAwaitState state)
+        private static ValueTask<T?> CreateRuntimeAsyncValueTask<T>(ref RuntimeAsyncAwaitState state)
         {
             // We only come to these methods in the expensive case (already
             // suspended), so ValueTask optimization here is not relevant.
-            return new ValueTask<T?>(FinalizeTaskReturningThunk<T>(ref state));
+            return new ValueTask<T?>(CreateRuntimeAsyncTask<T>(ref state));
         }
 
-        private static ValueTask FinalizeValueTaskReturningThunk(ref RuntimeAsyncAwaitState state)
+        private static ValueTask CreateRuntimeAsyncValueTask(ref RuntimeAsyncAwaitState state)
         {
-            return new ValueTask(FinalizeTaskReturningThunk(ref state));
+            return new ValueTask(CreateRuntimeAsyncTask(ref state));
         }
 
         private static Task<T?> TaskFromException<T>(Exception ex)
@@ -1544,19 +1563,6 @@ private static void FinishSuspensionWithContinuationContext(ref object continuat
             }
         }
 
-        [StackTraceHidden]
-        internal static T CompletedTaskResult<T>(Task<T> task)
-        {
-            TaskAwaiter.ValidateEnd(task);
-            return task.ResultOnSuccess;
-        }
-
-        [StackTraceHidden]
-        internal static void CompletedTask(Task task)
-        {
-            TaskAwaiter.ValidateEnd(task);
-        }
-
         // Instrumentation helpers called from InstrumentedDispatchContinuations.
         // These methods should not throw - exceptions would break the dispatch loop.
         internal static class RuntimeAsyncInstrumentationHelpers

src/coreclr/inc/readytorun.h

@@ -19,10 +19,10 @@
 //  src/coreclr/nativeaot/Runtime/inc/ModuleHeaders.h
 // If you update this, ensure you run `git grep MINIMUM_READYTORUN_MAJOR_VERSION`
 // and handle pending work.
-#define READYTORUN_MAJOR_VERSION 24
+#define READYTORUN_MAJOR_VERSION 25
 #define READYTORUN_MINOR_VERSION 0x0000
 
-#define MINIMUM_READYTORUN_MAJOR_VERSION 24
+#define MINIMUM_READYTORUN_MAJOR_VERSION 25
 
 // R2R Version 2.1 adds the InliningInfo section
 // R2R Version 2.2 adds the ProfileDataInfo section
@@ -64,6 +64,7 @@
 // R2R Version 22 changes NativeVarInfo encoding to include CALL_RETURN_VALUE
 // R2R Version 23 changes delegate layout to have target before methodPtr
 // R2R Version 24 changes ARM32 virtual stub dispatch hidden parameter register to R12
+// R2R Version 25 renames runtime async infrastructure members and makes thunk-used members NonVersionable
 
 struct READYTORUN_CORE_HEADER
 {

src/coreclr/nativeaot/Runtime/inc/ModuleHeaders.h

@@ -11,7 +11,7 @@ struct ReadyToRunHeaderConstants
 {
     static const uint32_t Signature = 0x00525452; // 'RTR'
 
-    static const uint32_t CurrentMajorVersion = 24;
+    static const uint32_t CurrentMajorVersion = 25;
     static const uint32_t CurrentMinorVersion = 0;
 };
 

src/coreclr/tools/Common/Internal/Runtime/ModuleHeaders.cs

@@ -15,7 +15,7 @@ internal struct ReadyToRunHeaderConstants
     {
         public const uint Signature = 0x00525452; // 'RTR'
 
-        public const ushort CurrentMajorVersion = 24;
+        public const ushort CurrentMajorVersion = 25;
         public const ushort CurrentMinorVersion = 0;
     }
 #if READYTORUN

src/coreclr/tools/Common/TypeSystem/IL/Stubs/AsyncThunks.cs

@@ -11,7 +11,7 @@ public static class AsyncThunkILEmitter
         // Emits a thunk that wraps an async method to return a Task or ValueTask.
         // The thunk calls the async method, and if it completes synchronously,
         // it returns a completed Task/ValueTask. If the async method suspends,
-        // it calls FinalizeTaskReturningThunk/FinalizeValueTaskReturningThunk method to get the Task/ValueTask.
+        // it calls CreateRuntimeAsyncTask/CreateRuntimeAsyncValueTask method to get the Task/ValueTask.
 
         // The emitted code matches method EmitTaskReturningThunk in CoreCLR VM.
         public static MethodIL EmitTaskReturningThunk(MethodDesc taskReturningMethod, MethodDesc asyncMethod)
@@ -191,35 +191,35 @@ public static MethodIL EmitTaskReturningThunk(MethodDesc taskReturningMethod, Me
 
                 codestream.EmitLabel(suspendedLabel);
 
-                MethodDesc finalizeTaskReturningThunkMd;
+                MethodDesc createRuntimeAsyncTaskMd;
                 if (logicalReturnType != null)
                 {
-                    MethodSignature finalizeReturningThunkSignature = new MethodSignature(
+                    MethodSignature createRuntimeAsyncTaskSignature = new MethodSignature(
                         MethodSignatureFlags.Static,
                         genericParameterCount: 1,
                         returnType: ((MetadataType)returnType.GetTypeDefinition()).MakeInstantiatedType(context.GetSignatureVariable(0, true)),
                         parameters: [awaitStateType.MakeByRefType()]
                     );
 
-                    finalizeTaskReturningThunkMd = asyncHelpersType
-                        .GetKnownMethod(isValueTask ? "FinalizeValueTaskReturningThunk"u8 : "FinalizeTaskReturningThunk"u8, finalizeReturningThunkSignature)
+                    createRuntimeAsyncTaskMd = asyncHelpersType
+                        .GetKnownMethod(isValueTask ? "CreateRuntimeAsyncValueTask"u8 : "CreateRuntimeAsyncTask"u8, createRuntimeAsyncTaskSignature)
                         .MakeInstantiatedMethod(new Instantiation(logicalReturnType));
                 }
                 else
                 {
-                    MethodSignature finalizeReturningThunkSignature = new MethodSignature(
+                    MethodSignature createRuntimeAsyncTaskSignature = new MethodSignature(
                         MethodSignatureFlags.Static,
                         genericParameterCount: 0,
                         returnType: returnType,
                         parameters: [awaitStateType.MakeByRefType()]
                     );
 
-                    finalizeTaskReturningThunkMd = asyncHelpersType
-                        .GetKnownMethod(isValueTask ? "FinalizeValueTaskReturningThunk"u8 : "FinalizeTaskReturningThunk"u8, finalizeReturningThunkSignature);
+                    createRuntimeAsyncTaskMd = asyncHelpersType
+                        .GetKnownMethod(isValueTask ? "CreateRuntimeAsyncValueTask"u8 : "CreateRuntimeAsyncTask"u8, createRuntimeAsyncTaskSignature);
                 }
 
                 codestream.EmitLdLoc(refAwaitStateLocal);
-                codestream.Emit(ILOpcode.call, emitter.NewToken(finalizeTaskReturningThunkMd));
+                codestream.Emit(ILOpcode.call, emitter.NewToken(createRuntimeAsyncTaskMd));
                 codestream.EmitStLoc(returnTaskLocal);
                 codestream.Emit(ILOpcode.leave, returnTaskLabel);
 

src/coreclr/vm/asyncthunks.cpp

@@ -93,7 +93,7 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncCallVariant, MetaSig&
     //     if (AsyncHelpers.AsyncCallContinuation() == null)
     //       return Task.FromResult(result);
     //
-    //     return FinalizeTaskReturningThunk(ref awaitState);
+    //     return CreateRuntimeAsyncTask(ref awaitState);
     //   }
     //   catch (Exception ex)
     //   {
@@ -229,29 +229,29 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncCallVariant, MetaSig&
 
         pCode->EmitLabel(suspendedLabel);
 
-        int finalizeTaskReturningThunkToken;
+        int createRuntimeAsyncTaskToken;
         if (logicalResultLocal != UINT_MAX)
         {
             MethodDesc* md;
             if (isValueTask)
-                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINALIZE_VALUETASK_RETURNING_THUNK_1);
+                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CREATE_RUNTIME_ASYNC_VALUE_TASK_1);
             else
-                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINALIZE_TASK_RETURNING_THUNK_1);
+                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CREATE_RUNTIME_ASYNC_TASK_1);
 
             md = FindOrCreateAssociatedMethodDesc(md, md->GetMethodTable(), FALSE, Instantiation(&thLogicalRetType, 1), FALSE);
-            finalizeTaskReturningThunkToken = GetTokenForGenericMethodCallWithAsyncReturnType(pCode, md);
+            createRuntimeAsyncTaskToken = GetTokenForGenericMethodCallWithAsyncReturnType(pCode, md);
         }
         else
         {
             MethodDesc* md;
             if (isValueTask)
-                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINALIZE_VALUETASK_RETURNING_THUNK);
+                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CREATE_RUNTIME_ASYNC_VALUE_TASK);
             else
-                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINALIZE_TASK_RETURNING_THUNK);
-            finalizeTaskReturningThunkToken = pCode->GetToken(md);
+                md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__CREATE_RUNTIME_ASYNC_TASK);
+            createRuntimeAsyncTaskToken = pCode->GetToken(md);
         }
         pCode->EmitLDLOC(refAwaitStateLocal);
-        pCode->EmitCALL(finalizeTaskReturningThunkToken, 1, 1);
+        pCode->EmitCALL(createRuntimeAsyncTaskToken, 1, 1);
         pCode->EmitSTLOC(returnTaskLocal);
         pCode->EmitLEAVE(returnTaskLabel);