From d344fff2f063aa242c41da6b4a5ee2ddb95e1d6a Mon Sep 17 00:00:00 2001
From: David Mason <davmason@microsoft.com>
Date: Tue, 25 Jun 2019 20:15:02 -0700
Subject: [PATCH] Fix ReJIT failures (#25377)

---
 src/inc/profilepriv.h            |  1 +
 src/vm/codeversion.cpp           |  1 +
 src/vm/proftoeeinterfaceimpl.cpp | 39 +++++++++++++++++++++++++++++++++++++++
 src/vm/proftoeeinterfaceimpl.h   |  2 ++
 src/vm/proftoeeinterfaceimpl.inl |  5 +++--
 5 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/inc/profilepriv.h b/src/inc/profilepriv.h
index a5e2398..91791f5 100644
--- a/src/inc/profilepriv.h
+++ b/src/inc/profilepriv.h
@@ -146,6 +146,7 @@ GVAL_DECL(ProfControlBlock, g_profControlBlock);
 #define COR_PRF_CALLBACKSTATE_INCALLBACK                 0x1
 #define COR_PRF_CALLBACKSTATE_IN_TRIGGERS_SCOPE          0x2
 #define COR_PRF_CALLBACKSTATE_FORCEGC_WAS_CALLED         0x4
+#define COR_PRF_CALLBACKSTATE_REJIT_WAS_CALLED           0x8
 //
 //---------------------------------------------------------------
 
diff --git a/src/vm/codeversion.cpp b/src/vm/codeversion.cpp
index 6d7cbe8..c90189b 100644
--- a/src/vm/codeversion.cpp
+++ b/src/vm/codeversion.cpp
@@ -2348,6 +2348,7 @@ HRESULT CodeVersionManager::PublishNativeCodeVersion(MethodDesc* pMethod, Native
         {
             MODE_ANY;
         }
+        NOTHROW;
     }
     CONTRACTL_END;
 
diff --git a/src/vm/proftoeeinterfaceimpl.cpp b/src/vm/proftoeeinterfaceimpl.cpp
index f238723..cff1721 100644
--- a/src/vm/proftoeeinterfaceimpl.cpp
+++ b/src/vm/proftoeeinterfaceimpl.cpp
@@ -6805,6 +6805,12 @@ HRESULT ProfToEEInterfaceImpl::RequestReJITWithInliners(
 
     // Remember the profiler is doing this, as that means we must never detach it!
     g_profControlBlock.pProfInterface->SetUnrevertiblyModifiedILFlag();
+
+    HRESULT hr = SetupThreadForReJIT();
+    if (FAILED(hr))
+    {
+        return hr;
+    }
     
     GCX_PREEMP();
     return ReJitManager::RequestReJIT(cFunctions, moduleIds, methodIds, static_cast<COR_PRF_REJIT_FLAGS>(dwRejitFlags));
@@ -8655,6 +8661,27 @@ HRESULT ProfToEEInterfaceImpl::GetReJITIDs(
     return ReJitManager::GetReJITIDs(pMD, cReJitIds, pcReJitIds, reJitIds);
 }
 
+
+HRESULT ProfToEEInterfaceImpl::SetupThreadForReJIT()
+{
+    LIMITED_METHOD_CONTRACT;
+
+    HRESULT hr = S_OK;
+    EX_TRY
+    {
+        if (GetThread() == NULL)
+        {
+            SetupThread();
+        }
+
+        Thread *pThread = GetThread();
+        pThread->SetProfilerCallbackStateFlags(COR_PRF_CALLBACKSTATE_REJIT_WAS_CALLED);
+    }
+    EX_CATCH_HRESULT(hr);
+
+    return hr;
+}
+
 HRESULT ProfToEEInterfaceImpl::RequestReJIT(ULONG       cFunctions,   // in
                                             ModuleID    moduleIds[],  // in
                                             mdMethodDef methodIds[])  // in
@@ -8704,6 +8731,12 @@ HRESULT ProfToEEInterfaceImpl::RequestReJIT(ULONG       cFunctions,   // in
     // Remember the profiler is doing this, as that means we must never detach it!
     g_profControlBlock.pProfInterface->SetUnrevertiblyModifiedILFlag();
     
+    HRESULT hr = SetupThreadForReJIT();
+    if (FAILED(hr))
+    {
+        return hr;
+    }
+    
     GCX_PREEMP();
     return ReJitManager::RequestReJIT(cFunctions, moduleIds, methodIds, static_cast<COR_PRF_REJIT_FLAGS>(0));
 }
@@ -8761,6 +8794,12 @@ HRESULT ProfToEEInterfaceImpl::RequestRevert(ULONG       cFunctions,  // in
         memset(rgHrStatuses, 0, sizeof(HRESULT) * cFunctions);
         _ASSERTE(S_OK == rgHrStatuses[0]);
     }
+    
+    HRESULT hr = SetupThreadForReJIT();
+    if (FAILED(hr))
+    {
+        return hr;
+    }
 
     GCX_PREEMP();
     return ReJitManager::RequestRevert(cFunctions, moduleIds, methodIds, rgHrStatuses);
diff --git a/src/vm/proftoeeinterfaceimpl.h b/src/vm/proftoeeinterfaceimpl.h
index 335b4f7..bc4b13d 100644
--- a/src/vm/proftoeeinterfaceimpl.h
+++ b/src/vm/proftoeeinterfaceimpl.h
@@ -659,6 +659,8 @@ protected:
 
     HRESULT EnumJITedFunctionsHelper(ProfilerFunctionEnum ** ppEnum, IJitManager ** ppJitMgr);
 
+    HRESULT SetupThreadForReJIT();
+
 #ifdef _TARGET_X86_
     HRESULT ProfilerEbpWalker(Thread * pThreadToSnapshot, LPCONTEXT pctxSeed, StackSnapshotCallback * callback, void * clientData);
 #endif //_TARGET_X86_
diff --git a/src/vm/proftoeeinterfaceimpl.inl b/src/vm/proftoeeinterfaceimpl.inl
index 58a3770..6652585 100644
--- a/src/vm/proftoeeinterfaceimpl.inl
+++ b/src/vm/proftoeeinterfaceimpl.inl
@@ -63,9 +63,10 @@ inline BOOL AreCallbackStateFlagsSet(DWORD dwFlags)
     BOOL fRet;
     BEGIN_GETTHREAD_ALLOWED_IN_NO_THROW_REGION;
     DWORD dwProfilerCallbackFullStateFlags = pThread->GetProfilerCallbackFullState();
-    if ((dwProfilerCallbackFullStateFlags & COR_PRF_CALLBACKSTATE_FORCEGC_WAS_CALLED) != 0)
+    if (((dwProfilerCallbackFullStateFlags & COR_PRF_CALLBACKSTATE_FORCEGC_WAS_CALLED) != 0)
+        || ((dwProfilerCallbackFullStateFlags & COR_PRF_CALLBACKSTATE_REJIT_WAS_CALLED) != 0))
     {
-        // Threads on which ForceGC() was successfully called should be treated just
+        // Threads on which ForceGC() or RequestReJIT() was successfully called should be treated just
         // like native threads.  Profiler can do whatever it wants
         return TRUE;
     }
-- 
2.7.4