Add debug information for runtime async methods (#120303)

- Add new JIT-EE API to report back debug information about the generated state
  machine and continuations
- Refactor debug info storage on VM side to be more easily extensible. The new
  format has either a thin or fat header. The fat header is used when we have
  either uninstrumented bounds, patchpoint info, rich debug info or async debug
  info, and stores the blob sizes of all of those components in addition to the
  bounds and vars. It is indicated by the first field (size of bounds) having
  value 0, which is an uncommon value for this field.
- Add new async debug information to the storage on the VM side
- Implement new format in R2R as well, bump R2R major version (might as well do
  this now as we expect to need to store async debug info in R2R during .NET 11
  anyway)
- Remove `Continuation.Resume` in favor of a `Continuation.ResumeInfo` pointer
  that gives both the resumption stub and a `DiagnosticIP`. The `DiagnosticIP`
  can be used for async stackwalking and for keying into the async debug info
- Teach the JIT to emit "async resumption info tables" and update the async
  transformation to point `Continuation.ResumeInfo` at this table's entries on
  suspension
- Add a new `ICorDebugInfo::ASYNC` source type; emit mappings with this source
  type for generated suspension and resumption code

Author: jakobbotsch

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

@@ -78,11 +78,27 @@ internal enum ContinuationFlags
         ContinueOnCapturedTaskScheduler = 64,
     }
 
+    // Keep in sync with dataAsyncResumeInfo in the JIT
+    internal unsafe struct ResumeInfo
+    {
+        public delegate*<Continuation, ref byte, Continuation?> Resume;
+        // IP to use for diagnostics. Points into the jitted suspension code.
+        // For debug codegen the IP resolves via an ASYNC native->IL mapping to
+        // the IL AsyncHelpers.Await (or other async function) call which
+        // caused the suspension.
+        // For optimized codegen the mapping into the root method may be more
+        // approximate (e.g. because of inlining).
+        // For all codegens the offset of DiagnosticsIP matches
+        // DiagnosticNativeOffset for the corresponding AsyncSuspensionPoint in
+        // the debug info.
+        public void* DiagnosticIP;
+    }
+
 #pragma warning disable CA1852 // "Type can be sealed" -- no it cannot because the runtime constructs subtypes dynamically
     internal unsafe class Continuation
     {
         public Continuation? Next;
-        public delegate*<Continuation, ref byte, Continuation?> Resume;
+        public ResumeInfo* ResumeInfo;
         public ContinuationFlags Flags;
         public int State;
 
@@ -195,9 +211,8 @@ private interface IRuntimeAsyncTaskOps<T>
             static abstract ref byte GetResultStorage(T task);
         }
 
-        /// <summary>
-        /// Represents a wrapped runtime async operation.
-        /// </summary>
+        // Represents execution of a chain of suspended and resuming runtime
+        // async functions.
         private sealed class RuntimeAsyncTask<T> : Task<T>, ITaskCompletionAction
         {
             public RuntimeAsyncTask()
@@ -261,9 +276,8 @@ public static void PostToSyncContext(RuntimeAsyncTask<T> task, SynchronizationCo
             }
         }
 
-        /// <summary>
-        /// Represents a wrapped runtime async operation.
-        /// </summary>
+        // Represents execution of a chain of suspended and resuming runtime
+        // async functions.
         private sealed class RuntimeAsyncTask : Task, ITaskCompletionAction
         {
             public RuntimeAsyncTask()
@@ -329,35 +343,63 @@ public static void PostToSyncContext(RuntimeAsyncTask task, SynchronizationConte
 
         private static class RuntimeAsyncTaskCore
         {
+            [StructLayout(LayoutKind.Explicit)]
+            private unsafe ref struct DispatcherInfo
+            {
+                // Dispatcher info for next dispatcher present on stack, or
+                // null if none.
+                [FieldOffset(0)]
+                public DispatcherInfo* Next;
+
+                // Next continuation the dispatcher will process.
+#if TARGET_64BIT
+                [FieldOffset(8)]
+#else
+                [FieldOffset(4)]
+#endif
+                public Continuation? NextContinuation;
+            }
+
+            // Information about current task dispatching, to be used for async
+            // stackwalking.
+            [ThreadStatic]
+            private static unsafe DispatcherInfo* t_dispatcherInfo;
+
             public static unsafe void DispatchContinuations<T, TOps>(T task) where T : Task, ITaskCompletionAction where TOps : IRuntimeAsyncTaskOps<T>
             {
                 ExecutionAndSyncBlockStore contexts = default;
                 contexts.Push();
-                Continuation? continuation = TOps.GetContinuationState(task);
+
+                DispatcherInfo dispatcherInfo;
+                dispatcherInfo.Next = t_dispatcherInfo;
+                dispatcherInfo.NextContinuation = TOps.GetContinuationState(task);
+                t_dispatcherInfo = &dispatcherInfo;
 
                 while (true)
                 {
-                    Debug.Assert(continuation != null);
+                    Debug.Assert(dispatcherInfo.NextContinuation != null);
                     try
                     {
-                        ref byte resultLoc = ref continuation.Next != null ? ref continuation.Next.GetResultStorageOrNull() : ref TOps.GetResultStorage(task);
-                        Continuation? newContinuation = continuation.Resume(continuation, ref resultLoc);
+                        Continuation curContinuation = dispatcherInfo.NextContinuation;
+                        Continuation? nextContinuation = curContinuation.Next;
+                        dispatcherInfo.NextContinuation = nextContinuation;
+
+                        ref byte resultLoc = ref nextContinuation != null ? ref nextContinuation.GetResultStorageOrNull() : ref TOps.GetResultStorage(task);
+                        Continuation? newContinuation = curContinuation.ResumeInfo->Resume(curContinuation, ref resultLoc);
 
                         if (newContinuation != null)
                         {
-                            newContinuation.Next = continuation.Next;
+                            newContinuation.Next = nextContinuation;
                             HandleSuspended<T, TOps>(task);
                             contexts.Pop();
+                            t_dispatcherInfo = dispatcherInfo.Next;
                             return;
                         }
-
-                        continuation = continuation.Next;
                     }
                     catch (Exception ex)
                     {
-                        Debug.Assert(continuation != null);
-                        Continuation? nextContinuation = UnwindToPossibleHandler(continuation);
-                        if (nextContinuation == null)
+                        Continuation? handlerContinuation = UnwindToPossibleHandler(dispatcherInfo.NextContinuation);
+                        if (handlerContinuation == null)
                         {
                             // Tail of AsyncTaskMethodBuilderT.SetException
                             bool successfullySet = ex is OperationCanceledException oce ?
@@ -366,6 +408,8 @@ public static unsafe void DispatchContinuations<T, TOps>(T task) where T : Task,
 
                             contexts.Pop();
 
+                            t_dispatcherInfo = dispatcherInfo.Next;
+
                             if (!successfullySet)
                             {
                                 ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
@@ -374,17 +418,18 @@ public static unsafe void DispatchContinuations<T, TOps>(T task) where T : Task,
                             return;
                         }
 
-                        nextContinuation.SetException(ex);
-
-                        continuation = nextContinuation;
+                        handlerContinuation.SetException(ex);
+                        dispatcherInfo.NextContinuation = handlerContinuation;
                     }
 
-                    if (continuation == null)
+                    if (dispatcherInfo.NextContinuation == null)
                     {
                         bool successfullySet = TOps.SetCompleted(task);
 
                         contexts.Pop();
 
+                        t_dispatcherInfo = dispatcherInfo.Next;
+
                         if (!successfullySet)
                         {
                             ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
@@ -393,26 +438,23 @@ public static unsafe void DispatchContinuations<T, TOps>(T task) where T : Task,
                         return;
                     }
 
-                    if (QueueContinuationFollowUpActionIfNecessary<T, TOps>(task, continuation))
+                    if (QueueContinuationFollowUpActionIfNecessary<T, TOps>(task, dispatcherInfo.NextContinuation))
                     {
                         contexts.Pop();
+                        t_dispatcherInfo = dispatcherInfo.Next;
                         return;
                     }
                 }
             }
 
-            private static Continuation? UnwindToPossibleHandler(Continuation continuation)
+            private static Continuation? UnwindToPossibleHandler(Continuation? continuation)
             {
                 while (true)
                 {
-                    Continuation? nextContinuation = continuation.Next;
-                    if (nextContinuation == null)
-                        return null;
-
-                    if ((nextContinuation.Flags & ContinuationFlags.HasException) != 0)
-                        return nextContinuation;
+                    if (continuation == null || (continuation.Flags & ContinuationFlags.HasException) != 0)
+                        return continuation;
 
-                    continuation = nextContinuation;
+                    continuation = continuation.Next;
                 }
             }
 

src/coreclr/inc/cordebuginfo.h

@@ -45,7 +45,8 @@ class ICorDebugInfo
         STACK_EMPTY                = 0x02, // The stack is empty here
         CALL_SITE                  = 0x04, // This is a call site.
         NATIVE_END_OFFSET_UNKNOWN  = 0x08, // Indicates a epilog endpoint
-        CALL_INSTRUCTION           = 0x10  // The actual instruction of a call.
+        CALL_INSTRUCTION           = 0x10, // The actual instruction of a call.
+        ASYNC                      = 0x20, // Indicates suspension/resumption for an async call
 
     };
 
@@ -431,4 +432,30 @@ class ICorDebugInfo
         // Source information about the IL instruction in the inlinee
         SourceTypes Source;
     };
+
+    struct AsyncContinuationVarInfo
+    {
+        // IL number of variable (or one of the special IL numbers, like TYPECTXT_ILNUM)
+        uint32_t VarNumber;
+        // Offset in continuation object where this variable is stored
+        uint32_t Offset;
+    };
+
+    struct AsyncSuspensionPoint
+    {
+        // Offset of IP stored in ResumeInfo.DiagnosticIP. This offset maps to
+        // the IL call that resulted in the suspension point through an ASYNC
+        // mapping. Also used as a unique key for debug information about the
+        // suspension point. See ResumeInfo.DiagnosticIP in SPC for more info.
+        uint32_t DiagnosticNativeOffset;
+        // Count of AsyncContinuationVarInfo in array of locals starting where
+        // the previous suspension point's locals end.
+        uint32_t NumContinuationVars;
+    };
+
+    struct AsyncInfo
+    {
+        // Number of suspension points in the method.
+        uint32_t NumSuspensionPoints;
+    };
 };

src/coreclr/inc/corinfo.h

@@ -1735,8 +1735,8 @@ struct CORINFO_ASYNC_INFO
     CORINFO_CLASS_HANDLE continuationClsHnd;
     // 'Next' field
     CORINFO_FIELD_HANDLE continuationNextFldHnd;
-    // 'Resume' field
-    CORINFO_FIELD_HANDLE continuationResumeFldHnd;
+    // 'ResumeInfo' field
+    CORINFO_FIELD_HANDLE continuationResumeInfoFldHnd;
     // 'State' field
     CORINFO_FIELD_HANDLE continuationStateFldHnd;
     // 'Flags' field
@@ -2913,6 +2913,16 @@ class ICorStaticInfo
             uint32_t                          numMappings         // [IN] Number of rich mappings
             ) = 0;
 
+    // Report async debug information to EE.
+    // The arrays are expected to be allocated with allocateArray
+    // and ownership is transferred to the EE with this call.
+    virtual void reportAsyncDebugInfo(
+            ICorDebugInfo::AsyncInfo*             asyncInfo,         // [IN] Async method information
+            ICorDebugInfo::AsyncSuspensionPoint*  suspensionPoints,  // [IN] Array of async suspension points, indexed by state number
+            ICorDebugInfo::AsyncContinuationVarInfo* vars,           // [IN] Array of async continuation variable info
+            uint32_t                              numVars            // [IN] Number of entries in the async vars array
+            ) = 0;
+
     // Report back some metadata about the compilation to the EE -- for
     // example, metrics about the compilation.
     virtual void reportMetadata(
@@ -3341,7 +3351,7 @@ class ICorDynamicInfo : public ICorStaticInfo
             CORINFO_TAILCALL_HELPERS* pResult
             ) = 0;
 
-    virtual CORINFO_METHOD_HANDLE getAsyncResumptionStub() = 0;
+    virtual CORINFO_METHOD_HANDLE getAsyncResumptionStub(void** entryPoint) = 0;
 
     virtual CORINFO_CLASS_HANDLE getContinuationType(
         size_t dataSize,

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -453,6 +453,12 @@ void reportRichMappings(
           ICorDebugInfo::RichOffsetMapping* mappings,
           uint32_t numMappings) override;
 
+void reportAsyncDebugInfo(
+          ICorDebugInfo::AsyncInfo* asyncInfo,
+          ICorDebugInfo::AsyncSuspensionPoint* suspensionPoints,
+          ICorDebugInfo::AsyncContinuationVarInfo* vars,
+          uint32_t numVars) override;
+
 void reportMetadata(
           const char* key,
           const void* value,
@@ -660,7 +666,8 @@ CORINFO_CLASS_HANDLE getContinuationType(
           bool* objRefs,
           size_t objRefsSize) override;
 
-CORINFO_METHOD_HANDLE getAsyncResumptionStub() override;
+CORINFO_METHOD_HANDLE getAsyncResumptionStub(
+          void** entryPoint) override;
 
 bool convertPInvokeCalliToCall(
           CORINFO_RESOLVED_TOKEN* pResolvedToken,

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* 68e93e9d-dd28-49ca-9ebc-a01a54532bb3 */
-    0x68e93e9d,
-    0xdd28,
-    0x49ca,
-    {0x9e, 0xbc, 0xa0, 0x1a, 0x54, 0x53, 0x2b, 0xb3}
+constexpr GUID JITEEVersionIdentifier = { /* a802fbbf-3e14-4b34-a348-5fba9fd756d4 */
+    0xa802fbbf,
+    0x3e14,
+    0x4b34,
+    {0xa3, 0x48, 0x5f, 0xba, 0x9f, 0xd7, 0x56, 0xd4}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

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 16
+#define READYTORUN_MAJOR_VERSION 17
 #define READYTORUN_MINOR_VERSION 0x0000
 
-#define MINIMUM_READYTORUN_MAJOR_VERSION 16
+#define MINIMUM_READYTORUN_MAJOR_VERSION 17
 
 // R2R Version 2.1 adds the InliningInfo section
 // R2R Version 2.2 adds the ProfileDataInfo section
@@ -47,6 +47,7 @@
 // R2R Version 14 changed x86 code generation to use funclets
 // R2R Version 15 removes double to int/uint helper calls
 // R2R Version 16 replaces the compression format for debug boundaries with a new format that is smaller and more efficient to parse
+// R2R Version 17 adds support for producing "fat" debug information (that e.g. can include async debug info)
 
 struct READYTORUN_CORE_HEADER
 {

src/coreclr/jit/ICorJitInfo_names_generated.h

@@ -112,6 +112,7 @@ DEF_CLR_API(setBoundaries)
 DEF_CLR_API(getVars)
 DEF_CLR_API(setVars)
 DEF_CLR_API(reportRichMappings)
+DEF_CLR_API(reportAsyncDebugInfo)
 DEF_CLR_API(reportMetadata)
 DEF_CLR_API(allocateArray)
 DEF_CLR_API(freeArray)

src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp

@@ -1067,6 +1067,17 @@ void WrapICorJitInfo::reportRichMappings(
     API_LEAVE(reportRichMappings);
 }
 
+void WrapICorJitInfo::reportAsyncDebugInfo(
+          ICorDebugInfo::AsyncInfo* asyncInfo,
+          ICorDebugInfo::AsyncSuspensionPoint* suspensionPoints,
+          ICorDebugInfo::AsyncContinuationVarInfo* vars,
+          uint32_t numVars)
+{
+    API_ENTER(reportAsyncDebugInfo);
+    wrapHnd->reportAsyncDebugInfo(asyncInfo, suspensionPoints, vars, numVars);
+    API_LEAVE(reportAsyncDebugInfo);
+}
+
 void WrapICorJitInfo::reportMetadata(
           const char* key,
           const void* value,
@@ -1547,10 +1558,11 @@ CORINFO_CLASS_HANDLE WrapICorJitInfo::getContinuationType(
     return temp;
 }
 
-CORINFO_METHOD_HANDLE WrapICorJitInfo::getAsyncResumptionStub()
+CORINFO_METHOD_HANDLE WrapICorJitInfo::getAsyncResumptionStub(
+          void** entryPoint)
 {
     API_ENTER(getAsyncResumptionStub);
-    CORINFO_METHOD_HANDLE temp = wrapHnd->getAsyncResumptionStub();
+    CORINFO_METHOD_HANDLE temp = wrapHnd->getAsyncResumptionStub(entryPoint);
     API_LEAVE(getAsyncResumptionStub);
     return temp;
 }