Fix profiler ReJIT of Runtime Async method bodies

src/coreclr/vm/codeversion.cpp

@@ -1621,6 +1621,21 @@ HRESULT CodeVersionManager::SetActiveILCodeVersions(ILCodeVersion* pActiveVersio
         *pMethodDescs = CDynArray<MethodDesc*>();
 
         MethodDesc* pLoadedMethodDesc = pActiveVersions[i].GetModule()->LookupMethodDef(pActiveVersions[i].GetMethodDef());
+
+        // A Task- or ValueTask-returning method may also have an async variant with native code
+        // compiled from the same IL. Include the async variant in addition to the primary method.
+        if (pLoadedMethodDesc != NULL && pLoadedMethodDesc->ReturnsTaskOrValueTask())
+        {
+            MethodDesc* pAsyncVariant =
+                pLoadedMethodDesc->GetMethodTable()->GetParallelMethodDesc(pLoadedMethodDesc, AsyncVariantLookup::Async);
+            if (pAsyncVariant != NULL &&
+                FAILED(hr = CodeVersionManager::EnumerateClosedMethodDescs(pAsyncVariant, pMethodDescs, &errorRecords)))
+            {
+                _ASSERTE(hr == E_OUTOFMEMORY);
+                return hr;
+            }
+        }
+
         if (FAILED(hr = CodeVersionManager::EnumerateClosedMethodDescs(pLoadedMethodDesc, pMethodDescs, &errorRecords)))
         {
             _ASSERTE(hr == E_OUTOFMEMORY);
@@ -2046,14 +2061,10 @@ HRESULT CodeVersionManager::EnumerateClosedMethodDescs(
     // Ok, now the case of a generic function (or function on generic class), which
     // is loaded, and may thus have compiled instantiations.
     // It's impossible to get to any other kind of domain from the profiling API
-    Module* pModule = pMD->GetModule();
-    mdMethodDef methodDef = pMD->GetMemberDef();
-
     // Module is unshared, so just use the module's domain to find instantiations.
     hr = EnumerateDomainClosedMethodDescs(
         AppDomain::GetCurrentDomain(),
-        pModule,
-        methodDef,
+        pMD,
         pClosedMethodDescs,
         pUnsupportedMethodErrors);
 
@@ -2069,8 +2080,7 @@ HRESULT CodeVersionManager::EnumerateClosedMethodDescs(
 // static
 HRESULT CodeVersionManager::EnumerateDomainClosedMethodDescs(
     AppDomain * pAppDomainToSearch,
-    Module* pModuleContainingMethodDef,
-    mdMethodDef methodDef,
+    MethodDesc* pMethodDesc,
     CDynArray<MethodDesc*> * pClosedMethodDescs,
     CDynArray<CodePublishError> * pUnsupportedMethodErrors)
 {
@@ -2081,12 +2091,14 @@ HRESULT CodeVersionManager::EnumerateDomainClosedMethodDescs(
         MODE_PREEMPTIVE;
         CAN_TAKE_LOCK;
         PRECONDITION(CheckPointer(pAppDomainToSearch, NULL_OK));
-        PRECONDITION(CheckPointer(pModuleContainingMethodDef));
+        PRECONDITION(CheckPointer(pMethodDesc));
         PRECONDITION(CheckPointer(pClosedMethodDescs));
         PRECONDITION(CheckPointer(pUnsupportedMethodErrors));
     }
     CONTRACTL_END;
 
+    Module* pModuleContainingMethodDef = pMethodDesc->GetModule();
+    mdMethodDef methodDef = pMethodDesc->GetMemberDef();
     _ASSERTE(methodDef != mdTokenNil);
 
     HRESULT hr;
@@ -2100,11 +2112,20 @@ HRESULT CodeVersionManager::EnumerateDomainClosedMethodDescs(
     {
         assemFlags = (AssemblyIterationFlags)(kIncludeAvailableToProfilers | kIncludeExecution);
     }
-    LoadedMethodDescIterator it(
-        pAppDomainToSearch,
-        pModuleContainingMethodDef,
-        methodDef,
-        assemFlags);
+    LoadedMethodDescIterator it;
+    if (pMethodDesc->ReturnsTaskOrValueTask() || pMethodDesc->IsAsyncVariantMethod())
+    {
+        it.StartForAsyncVariant(
+            pAppDomainToSearch,
+            pModuleContainingMethodDef,
+            methodDef,
+            pMethodDesc,
+            assemFlags);
+    }
+    else
+    {
+        it.Start(pAppDomainToSearch, pModuleContainingMethodDef, methodDef, assemFlags);
+    }
     CollectibleAssemblyHolder<Assembly *> pAssembly;
     while (it.Next(pAssembly.This()))
     {

src/coreclr/vm/codeversion.h

@@ -642,8 +642,7 @@ class CodeVersionManager
     static HRESULT EnumerateClosedMethodDescs(MethodDesc* pMD, CDynArray<MethodDesc*> * pClosedMethodDescs, CDynArray<CodePublishError> * pUnsupportedMethodErrors);
     static HRESULT EnumerateDomainClosedMethodDescs(
         AppDomain * pAppDomainToSearch,
-        Module* pModuleContainingMethodDef,
-        mdMethodDef methodDef,
+        MethodDesc* pMethodDesc,
         CDynArray<MethodDesc*> * pClosedMethodDescs,
         CDynArray<CodePublishError> * pUnsupportedMethodErrors);
     static HRESULT GetNonVersionableError(MethodDesc* pMD);

src/coreclr/vm/methoditer.cpp

@@ -136,6 +136,9 @@ BOOL LoadedMethodDescIterator::Next(
             goto ADVANCE_METHOD;
         if (m_methodIteratorEntry->GetMethod()->GetMemberDef() != m_md)
             goto ADVANCE_METHOD;
+        if (m_filterAsyncVariant &&
+            !m_methodIteratorEntry->GetMethod()->MatchesAsyncVariantLookup(m_asyncVariantLookup))
+            goto ADVANCE_METHOD;
     }
     else if (m_startedNonGenericMethod)
     {
@@ -195,7 +198,9 @@ MethodDesc *LoadedMethodDescIterator::Current()
 
     MethodTable *pMT = m_typeIteratorEntry->GetTypeHandle().GetMethodTable()->GetCanonicalMethodTable();
     _ASSERTE(pMT != NULL);
-    return pMT->GetParallelMethodDesc(m_mainMD);
+    return m_filterAsyncVariant
+        ? pMT->GetParallelMethodDesc(m_mainMD, m_asyncVariantLookup)
+        : pMT->GetParallelMethodDesc(m_mainMD);
 }
 
 void
@@ -220,6 +225,8 @@ LoadedMethodDescIterator::Start(
     m_md = md;
     m_pAppDomain = pAppDomain;
     m_fFirstTime = TRUE;
+    m_asyncVariantLookup = AsyncVariantLookup::Ordinary;
+    m_filterAsyncVariant = false;
 
     _ASSERTE(pAppDomain != NULL);
     _ASSERTE(TypeFromToken(m_md) == mdtMethodDef);
@@ -238,11 +245,27 @@ LoadedMethodDescIterator::Start(
     m_mainMD = pMethodDesc;
 }
 
+void
+LoadedMethodDescIterator::StartForAsyncVariant(
+    AppDomain     *pAppDomain,
+    Module          *pModule,
+    mdMethodDef     md,
+    MethodDesc      *pMethodDesc,
+    AssemblyIterationFlags assemblyIterationFlags)
+{
+    Start(pAppDomain, pModule, md, assemblyIterationFlags);
+    m_mainMD = pMethodDesc;
+    m_asyncVariantLookup = pMethodDesc->GetMatchingAsyncVariantLookup();
+    m_filterAsyncVariant = true;
+}
+
 LoadedMethodDescIterator::LoadedMethodDescIterator(void)
 {
     LIMITED_METHOD_CONTRACT;
     m_mainMD = NULL;
     m_module = NULL;
     m_md = mdTokenNil;
     m_pAppDomain = NULL;
+    m_asyncVariantLookup = AsyncVariantLookup::Ordinary;
+    m_filterAsyncVariant = false;
 }

src/coreclr/vm/methoditer.h

@@ -33,6 +33,8 @@ class LoadedMethodDescIterator
     mdMethodDef  m_md;
     MethodDesc * m_mainMD;
     AppDomain *  m_pAppDomain;
+    AsyncVariantLookup m_asyncVariantLookup;
+    bool m_filterAsyncVariant;
 
     // The following hold the state of the iteration....
     // Yes we iterate everything for the moment - we need
@@ -69,6 +71,12 @@ class LoadedMethodDescIterator
                mdMethodDef md,
                AssemblyIterationFlags assemIterationFlags = (AssemblyIterationFlags)(kIncludeLoaded | kIncludeExecution));
     void Start(AppDomain * pAppDomain, Module *pModule, mdMethodDef md, MethodDesc *pDesc);
+    void StartForAsyncVariant(
+        AppDomain * pAppDomain,
+        Module *pModule,
+        mdMethodDef md,
+        MethodDesc *pDesc,
+        AssemblyIterationFlags assemIterationFlags);
 
     LoadedMethodDescIterator(
         AppDomain * pAppDomain,

src/tests/profiler/native/rejitprofiler/rejitprofiler.cpp

@@ -26,6 +26,24 @@ using std::vector;
 #define INFO(MSG) _MESSAGE("INFO", MSG)
 #define FAIL(MSG) _MESSAGE("FAIL", MSG)
 
+static bool StartsWith(const String& value, const String& prefix)
+{
+    if (value.Length() < prefix.Length())
+    {
+        return false;
+    }
+
+    for (size_t i = 0; i < prefix.Length(); i++)
+    {
+        if (value[i] != prefix[i])
+        {
+            return false;
+        }
+    }
+
+    return true;
+}
+
 #ifdef __clang__
 #pragma clang diagnostic ignored "-Wnull-arithmetic"
 #endif // __clang__
@@ -35,11 +53,14 @@ ReJITProfiler::ReJITProfiler() : Profiler(),
     _failures(0),
     _rejits(0),
     _reverts(0),
+    _runtimeAsyncRejits(0),
     _inlinings(),
     _triggerFuncId(0),
     _targetFuncId(0),
     _targetModuleId(0),
-    _targetMethodDef(mdTokenNil)
+    _targetMethodDef(mdTokenNil),
+    _runtimeAsyncTargetModuleId(0),
+    _runtimeAsyncTargetMethodDef(mdTokenNil)
 {
 
 }
@@ -110,8 +131,9 @@ HRESULT ReJITProfiler::Shutdown()
     }
 
     INFO(L" rejit count=" << _rejits << L" expected rejit count=" << expectedRejitCount);
+    INFO(L" runtime async rejit count=" << _runtimeAsyncRejits << L" expected runtime async rejit count=2");
 
-    if(_failures == 0 && _rejits == expectedRejitCount)
+    if(_failures == 0 && _rejits == expectedRejitCount && _runtimeAsyncRejits == 2)
     {
         printf("PROFILER TEST PASSES\n");
     }
@@ -209,7 +231,16 @@ bool ReJITProfiler::FunctionSeen(FunctionID functionId)
         }
     }
 
-    if (functionName == TargetMethodName && EndsWith(moduleName, TargetModuleName))
+    if (StartsWith(functionName, RuntimeAsyncTargetMethodName) && EndsWith(moduleName, TargetModuleName))
+    {
+        _runtimeAsyncTargetModuleId = moduleId;
+        _runtimeAsyncTargetMethodDef = GetMethodDefForFunction(functionId);
+        INFO(L"Runtime async target native version compiled. FunctionID=" << std::hex << functionId
+            << L", ModuleID=" << _runtimeAsyncTargetModuleId
+            << L", MethodDef=" << _runtimeAsyncTargetMethodDef);
+        return true;
+    }
+    else if (functionName == TargetMethodName && EndsWith(moduleName, TargetModuleName))
     {
         INFO(L"Found function id for target method");
         _targetFuncId = functionId;
@@ -237,6 +268,23 @@ bool ReJITProfiler::FunctionSeen(FunctionID functionId)
         INFO(L"ModuleID=" << std::hex << _targetModuleId << L" and MethodDef=" << std::hex << _targetMethodDef);
         _profInfo10->RequestRevert(1, &_targetModuleId, &_targetMethodDef, nullptr);
     }
+    else if (functionName == RuntimeAsyncReJITTriggerMethodName && EndsWith(moduleName, TargetModuleName))
+    {
+        INFO(L"Runtime async ReJIT trigger method jitting finished: " << functionName);
+        INFO(L"Requesting ReJIT with inliners for runtime async method. ModuleID=" << std::hex
+            << _runtimeAsyncTargetModuleId << L", MethodDef=" << _runtimeAsyncTargetMethodDef);
+
+        HRESULT hr = _profInfo10->RequestReJITWithInliners(
+            COR_PRF_REJIT_BLOCK_INLINING | COR_PRF_REJIT_INLINING_CALLBACKS,
+            1,
+            &_runtimeAsyncTargetModuleId,
+            &_runtimeAsyncTargetMethodDef);
+        if (FAILED(hr))
+        {
+            _failures++;
+            FAIL(L"RequestReJITWithInliners failed for runtime async target with hr=" << std::hex << hr);
+        }
+    }
 
     return false;
 }
@@ -317,16 +365,27 @@ HRESULT STDMETHODCALLTYPE ReJITProfiler::ReJITCompilationStarted(FunctionID func
 {
     SHUTDOWNGUARD();
 
-    INFO(L"Saw a ReJIT for function " << GetFunctionIDName(functionId));
-    _rejits++;
+    String functionName = GetFunctionIDName(functionId);
+    if (StartsWith(functionName, RuntimeAsyncTargetMethodName))
+    {
+        INFO(L"Saw a runtime async target ReJIT. FunctionID=" << std::hex << functionId << L", ReJITID=" << rejitId);
+        _runtimeAsyncRejits++;
+    }
+    else
+    {
+        INFO(L"Saw a ReJIT for function " << functionName);
+        _rejits++;
+    }
     return S_OK;
 }
 
 HRESULT STDMETHODCALLTYPE ReJITProfiler::GetReJITParameters(ModuleID moduleId, mdMethodDef methodId, ICorProfilerFunctionControl* pFunctionControl)
 {
     SHUTDOWNGUARD();
 
-    String functionName = GetFunctionIDName(GetFunctionIDFromToken(moduleId, methodId, false));
+    String functionName = methodId == _runtimeAsyncTargetMethodDef
+        ? RuntimeAsyncTargetMethodName
+        : GetFunctionIDName(GetFunctionIDFromToken(moduleId, methodId, false));
     INFO(L"Starting to build IL for method " << functionName);
     COMPtrHolder<IUnknown> pUnk;
     HRESULT hr = _profInfo10->GetModuleMetaData(moduleId, ofWrite, IID_IMetaDataEmit2, &pUnk);

src/tests/profiler/native/rejitprofiler/rejitprofiler.h

@@ -56,14 +56,19 @@ class ReJITProfiler : public Profiler
     std::atomic<int> _failures;
     std::atomic<int> _rejits;
     std::atomic<int> _reverts;
+    std::atomic<int> _runtimeAsyncRejits;
     std::map<FunctionID, std::shared_ptr<std::unordered_set<FunctionID>>> _inlinings;
     FunctionID _triggerFuncId;
     FunctionID _targetFuncId;
     ModuleID _targetModuleId;
     mdMethodDef _targetMethodDef;
+    ModuleID _runtimeAsyncTargetModuleId;
+    mdMethodDef _runtimeAsyncTargetMethodDef;
 
     const String ReJITTriggerMethodName = WCHAR("TriggerReJIT");
     const String RevertTriggerMethodName = WCHAR("TriggerRevert");
+    const String RuntimeAsyncReJITTriggerMethodName = WCHAR("TriggerRuntimeAsyncReJIT");
+    const String RuntimeAsyncTargetMethodName = WCHAR("RuntimeAsyncTarget");
     const String TargetMethodName = WCHAR("InlineeTarget");
     const String TargetModuleName = WCHAR("rejit.dll");
 };