Support moving continuations to SynchronizationContext/TaskScheduler (#117314)

This adds support for moving continuations after task awaits to the right
context:
* For unconfigured task awaits or `ConfigureAwait(true)`, continuations are
* posted to the captured `SynchronizationContext` or `TaskScheduler` by the BCL
* For `ConfigureAwait(false)` continuations are moved to the thread pool

The JIT inserts a call to a new `AsyncHelpers.CaptureContinuationContext` when
suspending because of a task await. That function either captures the current
`SynchronizationContext` or `TaskScheduler` into the continuation in a known
place. Later the BCL will fetch the context from the continuation object based
on a flag to determine whether the continuation needs to be posted somewhere
else or whether it can be inlined. There is a new `Task`-derived `ThunkTask`
that replaces the existing C#-async-function-with-a-loop and which handles the
continuations explicitly.

Inlining continuations on the same thread currently has no check on stack depth.
This might be necessary to add. There are some differences compared to async 1
with how continuations are executed. In particular some common cases that
resulted in unbounded stack usage with async 1 does _not_ result in increasing
stack usage with runtime async. Hence I am not 100% sure whether we need a
mitigation or not, but more investigation is needed from my side. I'd like to
leave that as a follow-up.

To get to the right behavior this PR disables code inlining of task-awaited
functions. This is a large concession, but getting the behavior right in the
face of inlining proved to be quite tricky with the JIT's current internal
representation. Work to reenable this will be done separately.

There is still one missing piece for correctness: saving and restoring the
`Thread._synchronizationContext` field around async calls. It runs into similar
trickiness when trying to represent the expected behavior (which is that the
field is restored around synchronous calls, but _not_ restored on resumption).

Author: jakobbotsch

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

@@ -1712,6 +1712,15 @@ enum CorInfoContinuationFlags
     // OSR method saved in the beginning of 'Data', or -1 if the continuation
     // belongs to a tier 0 method.
     CORINFO_CONTINUATION_OSR_IL_OFFSET_IN_DATA = 4,
+    // If this bit is set the continuation should continue on the thread
+    // pool.
+    CORINFO_CONTINUATION_CONTINUE_ON_THREAD_POOL = 8,
+    // If this bit is set the continuation has a SynchronizationContext
+    // that we should continue on.
+    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_SYNCHRONIZATION_CONTEXT = 16,
+    // If this bit is set the continuation has a TaskScheduler
+    // that we should continue on.
+    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_TASK_SCHEDULER = 32,
 };
 
 struct CORINFO_ASYNC_INFO
@@ -1737,6 +1746,7 @@ struct CORINFO_ASYNC_INFO
     CORINFO_METHOD_HANDLE captureExecutionContextMethHnd;
     // Method handle for AsyncHelpers.RestoreExecutionContext
     CORINFO_METHOD_HANDLE restoreExecutionContextMethHnd;
+    CORINFO_METHOD_HANDLE captureContinuationContextMethHnd;
 };
 
 // Flags passed from JIT to runtime.

src/coreclr/inc/corinfo.h

@@ -1040,6 +1040,13 @@ ContinuationLayout AsyncTransformation::LayOutContinuation(BasicBlock*
                 block->getTryIndex(), layout.ExceptionGCDataIndex);
     }
 
+    if (call->GetAsyncInfo().ContinuationContextHandling == ContinuationContextHandling::ContinueOnCapturedContext)
+    {
+        layout.ContinuationContextGCDataIndex = layout.GCRefsCount++;
+        JITDUMP("  Continuation continues on captured context; context will be at GC@+%02u in GC data\n",
+                layout.ContinuationContextGCDataIndex);
+    }
+
     if (call->GetAsyncInfo().ExecutionContextHandling == ExecutionContextHandling::AsyncSaveAndRestore)
     {
         layout.ExecContextGCDataIndex = layout.GCRefsCount++;
@@ -1200,13 +1207,16 @@ BasicBlock* AsyncTransformation::CreateSuspension(
     LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_comp, storeState));
 
     // Fill in 'flags'
-    unsigned continuationFlags = 0;
+    const AsyncCallInfo& callInfo          = call->GetAsyncInfo();
+    unsigned             continuationFlags = 0;
     if (layout.ReturnInGCData)
         continuationFlags |= CORINFO_CONTINUATION_RESULT_IN_GCDATA;
     if (block->hasTryIndex())
         continuationFlags |= CORINFO_CONTINUATION_NEEDS_EXCEPTION;
     if (m_comp->doesMethodHavePatchpoints() || m_comp->opts.IsOSR())
         continuationFlags |= CORINFO_CONTINUATION_OSR_IL_OFFSET_IN_DATA;
+    if (callInfo.ContinuationContextHandling == ContinuationContextHandling::ContinueOnThreadPool)
+        continuationFlags |= CORINFO_CONTINUATION_CONTINUE_ON_THREAD_POOL;
 
     newContinuation      = m_comp->gtNewLclvNode(m_newContinuationVar, TYP_REF);
     unsigned flagsOffset = m_comp->info.compCompHnd->getFieldOffset(m_asyncInfo->continuationFlagsFldHnd);
@@ -1386,6 +1396,51 @@ void AsyncTransformation::FillInGCPointersOnSuspension(const ContinuationLayout&
         }
     }
 
+    if (layout.ContinuationContextGCDataIndex != UINT_MAX)
+    {
+        // Insert call AsyncHelpers.CaptureContinuationContext(ref
+        // newContinuation.GCData[ContinuationContextGCDataIndex], ref newContinuation.Flags).
+        GenTree*     contextElementPlaceholder = m_comp->gtNewZeroConNode(TYP_BYREF);
+        GenTree*     flagsPlaceholder          = m_comp->gtNewZeroConNode(TYP_BYREF);
+        GenTreeCall* captureCall =
+            m_comp->gtNewCallNode(CT_USER_FUNC, m_asyncInfo->captureContinuationContextMethHnd, TYP_VOID);
+
+        captureCall->gtArgs.PushFront(m_comp, NewCallArg::Primitive(flagsPlaceholder));
+        captureCall->gtArgs.PushFront(m_comp, NewCallArg::Primitive(contextElementPlaceholder));
+
+        m_comp->compCurBB = suspendBB;
+        m_comp->fgMorphTree(captureCall);
+
+        LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_comp, captureCall));
+
+        // Now replace contextElementPlaceholder with actual address of the context element
+        LIR::Use use;
+        bool     gotUse = LIR::AsRange(suspendBB).TryGetUse(contextElementPlaceholder, &use);
+        assert(gotUse);
+
+        GenTree* objectArr = m_comp->gtNewLclvNode(objectArrLclNum, TYP_REF);
+        unsigned offset = OFFSETOF__CORINFO_Array__data + (layout.ContinuationContextGCDataIndex * TARGET_POINTER_SIZE);
+        GenTree* contextElementOffset =
+            m_comp->gtNewOperNode(GT_ADD, TYP_BYREF, objectArr, m_comp->gtNewIconNode((ssize_t)offset, TYP_I_IMPL));
+
+        LIR::AsRange(suspendBB).InsertBefore(contextElementPlaceholder, LIR::SeqTree(m_comp, contextElementOffset));
+        use.ReplaceWith(contextElementOffset);
+        LIR::AsRange(suspendBB).Remove(contextElementPlaceholder);
+
+        // And now replace flagsPlaceholder with actual address of the flags
+        gotUse = LIR::AsRange(suspendBB).TryGetUse(flagsPlaceholder, &use);
+        assert(gotUse);
+
+        newContinuation          = m_comp->gtNewLclvNode(m_newContinuationVar, TYP_REF);
+        unsigned flagsOffset     = m_comp->info.compCompHnd->getFieldOffset(m_asyncInfo->continuationFlagsFldHnd);
+        GenTree* flagsOffsetNode = m_comp->gtNewOperNode(GT_ADD, TYP_BYREF, newContinuation,
+                                                         m_comp->gtNewIconNode((ssize_t)flagsOffset, TYP_I_IMPL));
+
+        LIR::AsRange(suspendBB).InsertBefore(flagsPlaceholder, LIR::SeqTree(m_comp, flagsOffsetNode));
+        use.ReplaceWith(flagsOffsetNode);
+        LIR::AsRange(suspendBB).Remove(flagsPlaceholder);
+    }
+
     if (layout.ExecContextGCDataIndex != UINT_MAX)
     {
         GenTreeCall* captureExecContext =

src/coreclr/jit/async.cpp

@@ -18,14 +18,15 @@ struct LiveLocalInfo
 
 struct ContinuationLayout
 {
-    unsigned                             DataSize               = 0;
-    unsigned                             GCRefsCount            = 0;
-    ClassLayout*                         ReturnStructLayout     = nullptr;
-    unsigned                             ReturnSize             = 0;
-    bool                                 ReturnInGCData         = false;
-    unsigned                             ReturnValDataOffset    = UINT_MAX;
-    unsigned                             ExceptionGCDataIndex   = UINT_MAX;
-    unsigned                             ExecContextGCDataIndex = UINT_MAX;
+    unsigned                             DataSize                       = 0;
+    unsigned                             GCRefsCount                    = 0;
+    ClassLayout*                         ReturnStructLayout             = nullptr;
+    unsigned                             ReturnSize                     = 0;
+    bool                                 ReturnInGCData                 = false;
+    unsigned                             ReturnValDataOffset            = UINT_MAX;
+    unsigned                             ExceptionGCDataIndex           = UINT_MAX;
+    unsigned                             ExecContextGCDataIndex         = UINT_MAX;
+    unsigned                             ContinuationContextGCDataIndex = UINT_MAX;
     const jitstd::vector<LiveLocalInfo>& Locals;
 
     explicit ContinuationLayout(const jitstd::vector<LiveLocalInfo>& locals)

src/coreclr/jit/async.h

@@ -4430,6 +4430,7 @@ class Compiler
 #endif
         // This call is a task await
         PREFIX_IS_TASK_AWAIT = 0x00000080,
+        PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT = 0x00000100,
     };
 
     static void impValidateMemoryAccessOpcode(const BYTE* codeAddr, const BYTE* codeEndp, bool volatilePrefix);

src/coreclr/jit/compiler.h

@@ -4328,10 +4328,21 @@ enum class ExecutionContextHandling
     AsyncSaveAndRestore,
 };
 
+enum class ContinuationContextHandling
+{
+    // No special handling of SynchronizationContext/TaskScheduler is required.
+    None,
+    // Continue on SynchronizationContext/TaskScheduler
+    ContinueOnCapturedContext,
+    // Continue on thread pool thread
+    ContinueOnThreadPool,
+};
+
 // Additional async call info.
 struct AsyncCallInfo
 {
-    ExecutionContextHandling ExecutionContextHandling = ExecutionContextHandling::None;
+    ExecutionContextHandling    ExecutionContextHandling    = ExecutionContextHandling::None;
+    ContinuationContextHandling ContinuationContextHandling = ContinuationContextHandling::None;
 };
 
 // Return type descriptor of a GT_CALL node.

src/coreclr/jit/gentree.h

@@ -9094,17 +9094,18 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                 // many other places.  We unfortunately embed that knowledge here.
                 if (opcode != CEE_CALLI)
                 {
-                    bool isAwait = false;
-                    // TODO: The configVal should be wired to the actual implementation
-                    //       that control the flow of sync context.
-                    //       We do not have that yet.
-                    int configVal = -1; // -1 not configured, 0/1 configured to false/true
+                    bool isAwait   = false;
+                    int  configVal = -1; // -1 not configured, 0/1 configured to false/true
                     if (compIsAsync() && JitConfig.JitOptimizeAwait())
                     {
                         if (impMatchTaskAwaitPattern(codeAddr, codeEndp, &configVal))
                         {
                             isAwait = true;
                             prefixFlags |= PREFIX_IS_TASK_AWAIT;
+                            if (configVal != 0)
+                            {
+                                prefixFlags |= PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT;
+                            }
                         }
                     }
 

src/coreclr/jit/importer.cpp

@@ -701,17 +701,26 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
     {
         AsyncCallInfo asyncInfo;
 
-        JITDUMP("Call is an async ");
-
         if ((prefixFlags & PREFIX_IS_TASK_AWAIT) != 0)
         {
-            JITDUMP("task await\n");
+            JITDUMP("Call is an async task await\n");
 
             asyncInfo.ExecutionContextHandling = ExecutionContextHandling::SaveAndRestore;
+
+            if ((prefixFlags & PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT) != 0)
+            {
+                asyncInfo.ContinuationContextHandling = ContinuationContextHandling::ContinueOnCapturedContext;
+                JITDUMP("  Continuation continues on captured context\n");
+            }
+            else
+            {
+                asyncInfo.ContinuationContextHandling = ContinuationContextHandling::ContinueOnThreadPool;
+                JITDUMP("  Continuation continues on thread pool\n");
+            }
         }
         else
         {
-            JITDUMP("non-task await\n");
+            JITDUMP("Call is an async non-task await\n");
             // Only expected non-task await to see in IL is one of the AsyncHelpers.AwaitAwaiter variants.
             // These are awaits of custom awaitables, and they come with the behavior that the execution context
             // is captured and restored on suspension/resumption.
@@ -7884,6 +7893,14 @@ void Compiler::impMarkInlineCandidateHelper(GenTreeCall*           call,
         return;
     }
 
+    if (call->IsAsync() && (call->GetAsyncInfo().ContinuationContextHandling != ContinuationContextHandling::None))
+    {
+        // Cannot currently handle moving to captured context/thread pool when logically returning from inlinee.
+        //
+        inlineResult->NoteFatal(InlineObservation::CALLSITE_CONTINUATION_HANDLING);
+        return;
+    }
+
     // Ignore indirect calls, unless they are indirect virtual stub calls with profile info.
     //
     if (call->gtCallType == CT_INDIRECT)

src/coreclr/jit/importercalls.cpp

@@ -163,6 +163,7 @@ INLINE_OBSERVATION(RETURN_TYPE_MISMATCH,      bool,   "return type mismatch",
 INLINE_OBSERVATION(STFLD_NEEDS_HELPER,        bool,   "stfld needs helper",                   FATAL,       CALLSITE)
 INLINE_OBSERVATION(TOO_MANY_LOCALS,           bool,   "too many locals",                      FATAL,       CALLSITE)
 INLINE_OBSERVATION(PINVOKE_EH,                bool,   "PInvoke call site with EH",            FATAL,       CALLSITE)
+INLINE_OBSERVATION(CONTINUATION_HANDLING,     bool,   "Callsite needs continuation handling", FATAL,       CALLSITE)
 
 // ------ Call Site Performance -------
 

src/coreclr/jit/inline.def

@@ -728,6 +728,7 @@ DEFINE_METHOD(ASYNC_HELPERS,      FINALIZE_VALUETASK_RETURNING_THUNK_1, Finalize
 DEFINE_METHOD(ASYNC_HELPERS,      UNSAFE_AWAIT_AWAITER_1,    UnsafeAwaitAwaiter, GM_T_RetVoid)
 DEFINE_METHOD(ASYNC_HELPERS,      CAPTURE_EXECUTION_CONTEXT, CaptureExecutionContext, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      RESTORE_EXECUTION_CONTEXT, RestoreExecutionContext, NoSig)
+DEFINE_METHOD(ASYNC_HELPERS,      CAPTURE_CONTINUATION_CONTEXT, CaptureContinuationContext, NoSig)
 
 DEFINE_CLASS(SPAN_HELPERS,          System,                 SpanHelpers)
 DEFINE_METHOD(SPAN_HELPERS,         MEMSET,                 Fill, SM_RefByte_Byte_UIntPtr_RetVoid)