[RuntimeAsync] Keep LoaderAllocator objects alive from continuations

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

@@ -583,6 +583,44 @@ private static Continuation AllocContinuation(Continuation prevContinuation, nui
             return newContinuation;
         }
 
+        private static unsafe Continuation AllocContinuationMethod(Continuation prevContinuation, nuint numGCRefs, nuint dataSize, MethodDesc* method)
+        {
+            LoaderAllocator loaderAllocator = RuntimeMethodHandle.GetLoaderAllocator(new RuntimeMethodHandleInternal((IntPtr)method));
+            object?[] gcData;
+            if (loaderAllocator != null)
+            {
+                gcData = new object[numGCRefs + 1];
+                gcData[numGCRefs] = loaderAllocator;
+            }
+            else
+            {
+                gcData = new object[numGCRefs];
+            }
+
+            Continuation newContinuation = new Continuation { Data = new byte[dataSize], GCData = gcData };
+            prevContinuation.Next = newContinuation;
+            return newContinuation;
+        }
+
+        private static unsafe Continuation AllocContinuationClass(Continuation prevContinuation, nuint numGCRefs, nuint dataSize, MethodTable* methodTable)
+        {
+            IntPtr loaderAllocatorHandle = methodTable->GetLoaderAllocatorHandle();
+            object?[] gcData;
+            if (loaderAllocatorHandle != IntPtr.Zero)
+            {
+                gcData = new object[numGCRefs + 1];
+                gcData[numGCRefs] = GCHandle.FromIntPtr(loaderAllocatorHandle).Target;
+            }
+            else
+            {
+                gcData = new object[numGCRefs];
+            }
+
+            Continuation newContinuation = new Continuation { Data = new byte[dataSize], GCData = gcData };
+            prevContinuation.Next = newContinuation;
+            return newContinuation;
+        }
+
         private struct RuntimeAsyncAwaitState
         {
             public Continuation? SentinelContinuation;
@@ -640,7 +678,7 @@ private static Continuation UnlinkHeadContinuation(out AwaitableProxy awaitableP
             return head;
         }
 
-        private static async Task<T> FinalizeTaskReturningThunk<T>(Continuation continuation)
+        private static async Task<T?> FinalizeTaskReturningThunk<T>(Continuation continuation)
         {
             Continuation finalContinuation = new Continuation();
 
@@ -670,7 +708,7 @@ private static async Task<T> FinalizeTaskReturningThunk<T>(Continuation continua
                     Debug.Assert(finalResult == finalContinuation);
                     if (IsReferenceOrContainsReferences<T>())
                     {
-                        return (T)finalResult.GCData![0];
+                        return (T?)finalResult.GCData![0];
                     }
                     else
                     {
@@ -701,7 +739,7 @@ private static async Task FinalizeTaskReturningThunk(Continuation continuation)
             }
         }
 
-        private static async ValueTask<T> FinalizeValueTaskReturningThunk<T>(Continuation continuation)
+        private static async ValueTask<T?> FinalizeValueTaskReturningThunk<T>(Continuation continuation)
         {
             Continuation finalContinuation = new Continuation();
 
@@ -731,7 +769,7 @@ private static async ValueTask<T> FinalizeValueTaskReturningThunk<T>(Continuatio
                     Debug.Assert(finalResult == finalContinuation);
                     if (IsReferenceOrContainsReferences<T>())
                     {
-                        return (T)finalResult.GCData![0];
+                        return (T?)finalResult.GCData![0];
                     }
                     else
                     {
@@ -1148,6 +1186,9 @@ public TypeHandle GetArrayElementTypeHandle()
         /// </summary>
         [MethodImpl(MethodImplOptions.InternalCall)]
         public extern MethodTable* GetMethodTableMatchingParentClass(MethodTable* parent);
+
+        [MethodImpl(MethodImplOptions.InternalCall)]
+        public extern IntPtr GetLoaderAllocatorHandle();
     }
 
     // Subset of src\vm\methodtable.h
@@ -1353,6 +1394,6 @@ internal sealed unsafe class Continuation
         // set, and otherwise at GCData[1].
         //
         public byte[]? Data;
-        public object[]? GCData;
+        public object?[]? GCData;
     }
 }

src/coreclr/inc/corinfo.h

@@ -596,6 +596,8 @@ enum CorInfoHelpFunc
     CORINFO_HELP_DISPATCH_INDIRECT_CALL,    // CFG: Validate and dispatch to pointer
 
     CORINFO_HELP_ALLOC_CONTINUATION,
+    CORINFO_HELP_ALLOC_CONTINUATION_METHOD,
+    CORINFO_HELP_ALLOC_CONTINUATION_CLASS,
     CORINFO_HELP_RESUME_OSR,
 
     CORINFO_HELP_COUNT,
@@ -1730,6 +1732,9 @@ struct CORINFO_ASYNC2_INFO
     CORINFO_FIELD_HANDLE continuationDataFldHnd;
     // 'GCData' field
     CORINFO_FIELD_HANDLE continuationGCDataFldHnd;
+    // Whether or not the continuation needs to be alloated through the
+    // helper that also takes a method handle
+    bool continuationsNeedMethodHandle;
 };
 
 // Flags passed from JIT to runtime.

src/coreclr/inc/jithelpers.h

@@ -343,6 +343,8 @@
 #endif
 
     DYNAMICJITHELPER(CORINFO_HELP_ALLOC_CONTINUATION, NULL, METHOD__RUNTIME_HELPERS__ALLOC_CONTINUATION)
+    DYNAMICJITHELPER(CORINFO_HELP_ALLOC_CONTINUATION_METHOD, NULL, METHOD__RUNTIME_HELPERS__ALLOC_CONTINUATION_METHOD)
+    DYNAMICJITHELPER(CORINFO_HELP_ALLOC_CONTINUATION_CLASS, NULL, METHOD__RUNTIME_HELPERS__ALLOC_CONTINUATION_CLASS)
     JITHELPER(CORINFO_HELP_RESUME_OSR, JIT_ResumeOSR, CORINFO_HELP_SIG_REG_ONLY)
 
 #undef JITHELPER

src/coreclr/jit/async.cpp

@@ -164,7 +164,7 @@ PhaseStatus Async2Transformation::Run()
     {
         for (GenTree* tree : LIR::AsRange(block))
         {
-            if (tree->IsCall() && tree->AsCall()->IsAsync2())
+            if (tree->IsCall() && tree->AsCall()->IsAsync2() && !tree->AsCall()->IsTailCall())
             {
                 JITDUMP(FMT_BB " contains await(s)\n", block->bbNum);
                 worklist.Push(block);
@@ -588,11 +588,9 @@ void Async2Transformation::Transform(
     block->GetFalseEdge()->setLikelihood(1 - RESUME_SUSPEND_LIKELIHOOD);
 
     // Allocate continuation
-    returnedContinuation           = m_comp->gtNewLclvNode(m_returnedContinuationVar, TYP_REF);
-    GenTree*     gcRefsCountNode   = m_comp->gtNewIconNode((ssize_t)gcRefsCount, TYP_I_IMPL);
-    GenTree*     dataSizeNode      = m_comp->gtNewIconNode((ssize_t)dataSize, TYP_I_IMPL);
-    GenTreeCall* allocContinuation = m_comp->gtNewHelperCallNode(CORINFO_HELP_ALLOC_CONTINUATION, TYP_REF,
-                                                                 returnedContinuation, gcRefsCountNode, dataSizeNode);
+    returnedContinuation = m_comp->gtNewLclvNode(m_returnedContinuationVar, TYP_REF);
+
+    GenTreeCall* allocContinuation = CreateAllocContinuationCall(life, returnedContinuation, gcRefsCount, dataSize);
 
     m_comp->compCurBB = suspendBB;
     m_comp->fgMorphTree(allocContinuation);
@@ -1100,6 +1098,50 @@ void Async2Transformation::Transform(
     m_resumptionBBs.push_back(resumeBB);
 }
 
+GenTreeCall* Async2Transformation::CreateAllocContinuationCall(AsyncLiveness& life, GenTree* prevContinuation, unsigned gcRefsCount, unsigned dataSize)
+{
+    GenTree* gcRefsCountNode = m_comp->gtNewIconNode((ssize_t)gcRefsCount, TYP_I_IMPL);
+    GenTree* dataSizeNode = m_comp->gtNewIconNode((ssize_t)dataSize, TYP_I_IMPL);
+    // If VM requests that we report the method handle, or if we have a shared generic context method handle
+    // that is live here, then we need to call a different helper.
+    GenTree* methodHandleArg = nullptr;
+    GenTree* classHandleArg = nullptr;
+    if (((m_comp->info.compMethodInfo->options & CORINFO_GENERICS_CTXT_FROM_METHODDESC) != 0) && life.IsLive(m_comp->info.compTypeCtxtArg))
+    {
+        methodHandleArg = m_comp->gtNewLclvNode(m_comp->info.compTypeCtxtArg, TYP_I_IMPL);
+    }
+    else if (((m_comp->info.compMethodInfo->options & CORINFO_GENERICS_CTXT_FROM_METHODTABLE) != 0) && life.IsLive(m_comp->info.compTypeCtxtArg))
+    {
+        classHandleArg = m_comp->gtNewLclvNode(m_comp->info.compTypeCtxtArg, TYP_I_IMPL);
+    }
+    else if (m_async2Info.continuationsNeedMethodHandle)
+    {
+        methodHandleArg = m_comp->gtNewIconEmbMethHndNode(m_comp->info.compMethodHnd);
+    }
+
+    if (methodHandleArg != nullptr)
+    {
+        return m_comp->gtNewHelperCallNode(
+            CORINFO_HELP_ALLOC_CONTINUATION_METHOD,
+            TYP_REF,
+            prevContinuation, gcRefsCountNode, dataSizeNode, methodHandleArg);
+    }
+
+    if (classHandleArg != nullptr)
+    {
+        return m_comp->gtNewHelperCallNode(
+            CORINFO_HELP_ALLOC_CONTINUATION_CLASS,
+            TYP_REF,
+            prevContinuation, gcRefsCountNode, dataSizeNode, classHandleArg);
+    }
+
+    return m_comp->gtNewHelperCallNode(
+        CORINFO_HELP_ALLOC_CONTINUATION,
+        TYP_REF,
+        prevContinuation, gcRefsCountNode, dataSizeNode);
+}
+
+
 GenTreeIndir* Async2Transformation::LoadFromOffset(GenTree*     base,
                                                    unsigned     offset,
                                                    var_types    type,

src/coreclr/jit/async.h

@@ -45,6 +45,12 @@ class Async2Transformation
                    class AsyncLiveness&      life,
                    BasicBlock**              remainder);
 
+    GenTreeCall* CreateAllocContinuationCall(
+        AsyncLiveness& life,
+        GenTree* prevContinuation,
+        unsigned gcRefsCount,
+        unsigned int dataSize);
+
     GenTreeIndir* LoadFromOffset(GenTree*     base,
                                  unsigned     offset,
                                  var_types    type,

src/coreclr/jit/compiler.h

@@ -3106,7 +3106,7 @@ class Compiler
     GenTreeCall* gtNewIndCallNode(GenTree* addr, var_types type, const DebugInfo& di = DebugInfo());
 
     GenTreeCall* gtNewHelperCallNode(
-        unsigned helper, var_types type, GenTree* arg1 = nullptr, GenTree* arg2 = nullptr, GenTree* arg3 = nullptr);
+        unsigned helper, var_types type, GenTree* arg1 = nullptr, GenTree* arg2 = nullptr, GenTree* arg3 = nullptr, GenTree* arg4 = nullptr);
 
     GenTreeCall* gtNewRuntimeLookupHelperCallNode(CORINFO_RUNTIME_LOOKUP* pRuntimeLookup,
                                                   GenTree*                ctxTree,

src/coreclr/jit/compiler.hpp

@@ -1538,7 +1538,7 @@ inline GenTree* Compiler::gtNewIconEmbFldHndNode(CORINFO_FIELD_HANDLE fldHnd)
 //    New CT_HELPER node
 //
 inline GenTreeCall* Compiler::gtNewHelperCallNode(
-    unsigned helper, var_types type, GenTree* arg1, GenTree* arg2, GenTree* arg3)
+    unsigned helper, var_types type, GenTree* arg1, GenTree* arg2, GenTree* arg3, GenTree* arg4)
 {
     GenTreeCall* const result = gtNewCallNode(CT_HELPER, eeFindHelper(helper), type);
 
@@ -1557,6 +1557,12 @@ inline GenTreeCall* Compiler::gtNewHelperCallNode(
     result->gtInlineObservation = InlineObservation::CALLSITE_IS_CALL_TO_HELPER;
 #endif
 
+    if (arg4 != nullptr)
+    {
+        result->gtArgs.PushFront(this, NewCallArg::Primitive(arg4));
+        result->gtFlags |= arg4->gtFlags & GTF_ALL_EFFECT;
+    }
+
     if (arg3 != nullptr)
     {
         result->gtArgs.PushFront(this, NewCallArg::Primitive(arg3));

src/coreclr/tools/superpmi/superpmi-shared/agnostic.h

@@ -204,6 +204,7 @@ struct Agnostic_CORINFO_ASYNC2_INFO
     DWORDLONG continuationFlagsFldHnd;
     DWORDLONG continuationDataFldHnd;
     DWORDLONG continuationGCDataFldHnd;
+    DWORD continuationsNeedMethodHandle;
 };
 
 struct Agnostic_GetOSRInfo

src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp

@@ -4432,24 +4432,27 @@ void MethodContext::recGetAsync2Info(const CORINFO_ASYNC2_INFO* pAsync2Info)
         GetAsync2Info = new LightWeightMap<DWORD, Agnostic_CORINFO_ASYNC2_INFO>();
 
     Agnostic_CORINFO_ASYNC2_INFO value;
+    ZeroMemory(&value, sizeof(value));
+
     value.continuationClsHnd = CastHandle(pAsync2Info->continuationClsHnd);
     value.continuationNextFldHnd = CastHandle(pAsync2Info->continuationNextFldHnd);
     value.continuationResumeFldHnd = CastHandle(pAsync2Info->continuationResumeFldHnd);
     value.continuationStateFldHnd = CastHandle(pAsync2Info->continuationStateFldHnd);
     value.continuationFlagsFldHnd = CastHandle(pAsync2Info->continuationFlagsFldHnd);
     value.continuationDataFldHnd = CastHandle(pAsync2Info->continuationDataFldHnd);
     value.continuationGCDataFldHnd = CastHandle(pAsync2Info->continuationGCDataFldHnd);
+    value.continuationsNeedMethodHandle = pAsync2Info->continuationsNeedMethodHandle ? 1 : 0;
 
     GetAsync2Info->Add(0, value);
     DEBUG_REC(dmpGetAsync2Info(0, value));
 }
 void MethodContext::dmpGetAsync2Info(DWORD key, const Agnostic_CORINFO_ASYNC2_INFO& value)
 {
     printf("GetAsync2Info key %u value contClsHnd-%016" PRIX64 " contNextFldHnd-%016" PRIX64 " contResumeFldHnd-%016" PRIX64
-           " contStateFldHnd-%016" PRIX64 " contFlagsFldHnd-%016" PRIX64 " contDataFldHnd-%016" PRIX64 " contGCDataFldHnd-%016" PRIX64,
+           " contStateFldHnd-%016" PRIX64 " contFlagsFldHnd-%016" PRIX64 " contDataFldHnd-%016" PRIX64 " contGCDataFldHnd-%016" PRIX64 " contsNeedMethodHandle-%d",
         key, value.continuationClsHnd, value.continuationNextFldHnd, value.continuationResumeFldHnd,
         value.continuationStateFldHnd, value.continuationFlagsFldHnd, value.continuationDataFldHnd,
-        value.continuationGCDataFldHnd);
+        value.continuationGCDataFldHnd, value.continuationsNeedMethodHandle);
 }
 void MethodContext::repGetAsync2Info(CORINFO_ASYNC2_INFO* pAsync2InfoOut)
 {
@@ -4461,6 +4464,7 @@ void MethodContext::repGetAsync2Info(CORINFO_ASYNC2_INFO* pAsync2InfoOut)
     pAsync2InfoOut->continuationFlagsFldHnd = (CORINFO_FIELD_HANDLE)value.continuationFlagsFldHnd;
     pAsync2InfoOut->continuationDataFldHnd = (CORINFO_FIELD_HANDLE)value.continuationDataFldHnd;
     pAsync2InfoOut->continuationGCDataFldHnd = (CORINFO_FIELD_HANDLE)value.continuationGCDataFldHnd;
+    pAsync2InfoOut->continuationsNeedMethodHandle = value.continuationsNeedMethodHandle != 0;
     DEBUG_REP(dmpGetAsync2Info(0, value));
 }
 

src/coreclr/vm/appdomain.cpp

@@ -1088,7 +1088,6 @@ void SystemDomain::LoadBaseSystemClasses()
         // further loading of nonprimitive types may need casting support.
         // initialize cast cache here.
         CastCache::Initialize();
-        ECall::PopulateAsyncHelpers();
 
         // used by IsImplicitInterfaceOfSZArray
         CoreLibBinder::GetClass(CLASS__IENUMERABLEGENERIC);

src/coreclr/vm/comutilnative.cpp

@@ -1818,6 +1818,14 @@ FCIMPL2(MethodTable*, MethodTableNative::GetMethodTableMatchingParentClass, Meth
 }
 FCIMPLEND
 
+FCIMPL1(OBJECTHANDLE, MethodTableNative::GetLoaderAllocatorHandle, MethodTable *mt)
+{
+    FCALL_CONTRACT;
+
+    return mt->GetLoaderAllocatorObjectHandle();
+}
+FCIMPLEND
+
 extern "C" BOOL QCALLTYPE MethodTable_AreTypesEquivalent(MethodTable* mta, MethodTable* mtb)
 {
     QCALL_CONTRACT;

src/coreclr/vm/comutilnative.h

@@ -259,6 +259,7 @@ class MethodTableNative {
     static FCDECL1(UINT32, GetNumInstanceFieldBytes, MethodTable* mt);
     static FCDECL1(CorElementType, GetPrimitiveCorElementType, MethodTable* mt);
     static FCDECL2(MethodTable*, GetMethodTableMatchingParentClass, MethodTable* mt, MethodTable* parent);
+    static FCDECL1(OBJECTHANDLE, GetLoaderAllocatorHandle, MethodTable* mt);
 };
 
 extern "C" BOOL QCALLTYPE MethodTable_AreTypesEquivalent(MethodTable* mta, MethodTable* mtb);

src/coreclr/vm/corelib.h

@@ -661,7 +661,9 @@ DEFINE_METHOD(RUNTIME_HELPERS,      ALLOC_TAILCALL_ARG_BUFFER, AllocTailCallArgB
 DEFINE_METHOD(RUNTIME_HELPERS,      GET_TAILCALL_INFO,      GetTailCallInfo, NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS,      DISPATCH_TAILCALLS,     DispatchTailCalls,          NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS,      COPY_CONSTRUCT,         CopyConstruct, NoSig)
-DEFINE_METHOD(RUNTIME_HELPERS,      ALLOC_CONTINUATION,     AllocContinuation, NoSig)
+DEFINE_METHOD(RUNTIME_HELPERS,      ALLOC_CONTINUATION,        AllocContinuation, NoSig)
+DEFINE_METHOD(RUNTIME_HELPERS,      ALLOC_CONTINUATION_METHOD, AllocContinuationMethod, NoSig)
+DEFINE_METHOD(RUNTIME_HELPERS,      ALLOC_CONTINUATION_CLASS,  AllocContinuationClass, NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS,      FINALIZE_TASK_RETURNING_THUNK, FinalizeTaskReturningThunk, SM_Continuation_RetTask)
 DEFINE_METHOD(RUNTIME_HELPERS,      FINALIZE_TASK_RETURNING_THUNK_1, FinalizeTaskReturningThunk, GM_Continuation_RetTaskOfT)
 DEFINE_METHOD(RUNTIME_HELPERS,      FINALIZE_VALUETASK_RETURNING_THUNK, FinalizeValueTaskReturningThunk, SM_Continuation_RetValueTask)

src/coreclr/vm/ecall.cpp

@@ -96,14 +96,6 @@ void ECall::PopulateManagedStringConstructors()
     INDEBUG(fInitialized = true);
 }
 
-void ECall::PopulateAsyncHelpers()
-{
-    STANDARD_VM_CONTRACT;
-
-    MethodDesc* pMD = CoreLibBinder::GetMethod((BinderMethodID)(METHOD__RUNTIME_HELPERS__ALLOC_CONTINUATION));
-    PCODE pDest = pMD->GetMultiCallableAddrOfCode();
-    SetJitHelperFunction(CORINFO_HELP_ALLOC_CONTINUATION, pDest);
-}
 static CrstStatic gFCallLock;
 
 // This variable is used to force the compiler not to tailcall a function.

src/coreclr/vm/ecall.h

@@ -94,8 +94,6 @@ class ECall
 
         static void PopulateManagedStringConstructors();
 
-        static void PopulateAsyncHelpers();
-
 #ifdef DACCESS_COMPILE
         // Enumerates all gFCallMethods for minidumps.
         static void EnumFCallMethods();

src/coreclr/vm/ecalllist.h

@@ -363,6 +363,7 @@ FCFuncStart(gMethodTableFuncs)
     FCFuncElement("GetNumInstanceFieldBytes", MethodTableNative::GetNumInstanceFieldBytes)
     FCFuncElement("GetPrimitiveCorElementType", MethodTableNative::GetPrimitiveCorElementType)
     FCFuncElement("GetMethodTableMatchingParentClass", MethodTableNative::GetMethodTableMatchingParentClass)
+    FCFuncElement("GetLoaderAllocatorHandle", MethodTableNative::GetLoaderAllocatorHandle)
 FCFuncEnd()
 
 FCFuncStart(gStubHelperFuncs)

src/coreclr/vm/jitinterface.cpp

@@ -10188,6 +10188,7 @@ void CEEInfo::getAsync2Info(CORINFO_ASYNC2_INFO* pAsync2InfoOut)
     pAsync2InfoOut->continuationFlagsFldHnd = CORINFO_FIELD_HANDLE(CoreLibBinder::GetField(FIELD__CONTINUATION__FLAGS));
     pAsync2InfoOut->continuationDataFldHnd = CORINFO_FIELD_HANDLE(CoreLibBinder::GetField(FIELD__CONTINUATION__DATA));
     pAsync2InfoOut->continuationGCDataFldHnd = CORINFO_FIELD_HANDLE(CoreLibBinder::GetField(FIELD__CONTINUATION__GCDATA));
+    pAsync2InfoOut->continuationsNeedMethodHandle = m_pMethodBeingCompiled->GetLoaderAllocator()->CanUnload();
 }
 
 const char16_t * CEEInfo::getJitTimeLogFilename()