1 files changed, 22 insertions, 6 deletions
diff --git a/src/vm/multicorejitplayer.cpp b/src/vm/multicorejitplayer.cpp
index d7c2cec8a1..247fa0a14a 100644
--- a/src/vm/multicorejitplayer.cpp
+++ b/src/vm/multicorejitplayer.cpp
@@ -12,12 +12,10 @@
 
 #include "common.h"
 #include "vars.hpp"
-#include "security.h"
 #include "eeconfig.h"
 #include "dllimport.h"
 #include "comdelegate.h"
 #include "dbginterface.h"
-#include "listlock.inl"
 #include "stubgen.h"
 #include "eventtrace.h"
 #include "array.h"
@@ -103,7 +101,7 @@ void MulticoreJitCodeStorage::StoreMethodCode(MethodDesc * pMD, PCODE pCode)
 
 
 // Query from MakeJitWorker: Lookup stored JITted methods
-PCODE MulticoreJitCodeStorage::QueryMethodCode(MethodDesc * pMethod)
+PCODE MulticoreJitCodeStorage::QueryMethodCode(MethodDesc * pMethod, BOOL shouldRemoveCode)
 {
     STANDARD_VM_CONTRACT;
 
@@ -113,7 +111,7 @@ PCODE MulticoreJitCodeStorage::QueryMethodCode(MethodDesc * pMethod)
     {
         CrstHolder holder(& m_crstCodeMap);
     
-        if (m_nativeCodeMap.Lookup(pMethod, & code))
+        if (m_nativeCodeMap.Lookup(pMethod, & code) && shouldRemoveCode)
         {
             m_nReturned ++;
 
@@ -507,6 +505,23 @@ HRESULT MulticoreJitProfilePlayer::HandleModuleRecord(const ModuleRecord * pMod)
 }
 
 
+#ifndef DACCESS_COMPILE
+class MulticoreJitPrepareCodeConfig : public PrepareCodeConfig
+{
+public:
+    MulticoreJitPrepareCodeConfig(MethodDesc* pMethod) :
+        PrepareCodeConfig(NativeCodeVersion(pMethod), FALSE, FALSE)
+    {}
+    
+    virtual BOOL SetNativeCode(PCODE pCode, PCODE * ppAlternateCodeToUse)
+    {
+        MulticoreJitManager & mcJitManager = GetAppDomain()->GetMulticoreJitManager();
+        mcJitManager.GetMulticoreJitCodeStorage().StoreMethodCode(GetMethodDesc(), pCode);
+        return TRUE;
+    }
+};
+#endif
+
 // Call JIT to compile a method
 
 bool MulticoreJitProfilePlayer::CompileMethodDesc(Module * pModule, MethodDesc * pMD)
@@ -529,8 +544,9 @@ bool MulticoreJitProfilePlayer::CompileMethodDesc(Module * pModule, MethodDesc *
         // Reset the flag to allow managed code to be called in multicore JIT background thread from this routine
         ThreadStateNCStackHolder holder(-1, Thread::TSNC_CallingManagedCodeDisabled);
 
-        // MakeJitWorker calls back to MulticoreJitCodeStorage::StoreMethodCode under MethodDesc lock
-        pMD->MakeJitWorker(& header, CORJIT_FLAGS(CORJIT_FLAGS::CORJIT_FLAG_MCJIT_BACKGROUND));
+        // PrepareCode calls back to MulticoreJitCodeStorage::StoreMethodCode under MethodDesc lock
+        MulticoreJitPrepareCodeConfig config(pMD);
+        pMD->PrepareCode(&config);
 
         return true;
     }