Try a new approach

Author: kg

src/coreclr/vm/dllimportcallback.cpp

@@ -221,6 +221,22 @@ UMEntryThunkData * GetMostRecentUMEntryThunkData()
 }
 #endif
 
+static void PatchUMEntryPrecodeIfNecessary(UMEntryThunkData * pUMEntryThunkData, PCODE callTarget)
+{
+#ifdef FEATURE_INTERPRETER
+    // HACK: We currently rely on the thread-local stash in TheUMEntryPrestubWorker for proper operation
+    //  of reverse p/invoke IL stubs in the interpreter, so we can't patch the precode in this scenario.
+    // Instead, every reverse p/invoke will go through TheUMEntryPrestubWorker.
+    MethodDesc *pMD = NonVirtualEntry2MethodDesc(callTarget);
+    if (pMD && pMD->GetInterpreterCode())
+        return;
+#endif
+
+    // NOTE: This may be run concurrently by multiple threads, but it should be safe since we will just be
+    //  setting the target of the precode to the same thing multiple times in a row.
+    pUMEntryThunkData->PatchPrecode();
+}
+
 PCODE TheUMEntryPrestubWorker(UMEntryThunkData * pUMEntryThunkData)
 {
     STATIC_CONTRACT_THROWS;
@@ -240,23 +256,36 @@ PCODE TheUMEntryPrestubWorker(UMEntryThunkData * pUMEntryThunkData)
     if (pUMEntryThunkData->IsCollectedDelegate())
         CallbackOnCollectedDelegate(pUMEntryThunkData);
 
-    INSTALL_MANAGED_EXCEPTION_DISPATCHER;
-    // this method is called by stubs which are called by managed code,
-    // so we need an unwind and continue handler so that our internal
-    // exceptions don't leak out into managed code.
-    INSTALL_UNWIND_AND_CONTINUE_HANDLER;
-
+    PCODE callTarget = (PCODE)0;
 #ifdef FEATURE_INTERPRETER
     // HACK: Stash the entry thunk data address so that the interpreter can recover it.
     // The InterpreterStub overwrites the register that normally would contain this address.
     t_MostRecentUMEntryThunkData = pUMEntryThunkData;
+
+    // See whether another thread has already prepared the IL stub for the reverse p/invoke.
+    callTarget = pUMEntryThunkData->GetCachedCallTarget();
 #endif
-    pUMEntryThunkData->RunTimeInit();
 
-    UNINSTALL_UNWIND_AND_CONTINUE_HANDLER;
-    UNINSTALL_MANAGED_EXCEPTION_DISPATCHER;
+    // We don't have a pre-prepared reverse p/invoke, and might be racing with another thread.
+    if (!callTarget)
+    {
+        INSTALL_MANAGED_EXCEPTION_DISPATCHER;
+        // this method is called by stubs which are called by managed code,
+        // so we need an unwind and continue handler so that our internal
+        // exceptions don't leak out into managed code.
+        INSTALL_UNWIND_AND_CONTINUE_HANDLER;
+
+        pUMEntryThunkData->RunTimeInit();
+        callTarget = pUMEntryThunkData->GetCachedCallTarget();
+        // We're likely to be the first thread to have gotten here for this UM entry thunk, so make
+        //  sure that the precode is patched to point to the IL stub instead of us, if necessary.
+        PatchUMEntryPrecodeIfNecessary(pUMEntryThunkData, callTarget);
+
+        UNINSTALL_UNWIND_AND_CONTINUE_HANDLER;
+        UNINSTALL_MANAGED_EXCEPTION_DISPATCHER;
+    }
 
-    return (PCODE)pUMEntryThunkData->GetCode();
+    return callTarget;
 }
 
 UMEntryThunkData* UMEntryThunkData::CreateUMEntryThunk()

src/coreclr/vm/dllimportcallback.h

@@ -134,7 +134,7 @@ class UMEntryThunk : private StubPrecode
     PTR_UMEntryThunkData GetData() const
     {
         LIMITED_METHOD_CONTRACT;
-        
+
         return dac_cast<PTR_UMEntryThunkData>(GetSecretParam());
     }
 };
@@ -148,6 +148,9 @@ class UMEntryThunkData
     // if m_pObjectHandle is non-NULL, this field is still set to help with diagnostic of call on collected delegate crashes
     // but it may not have the correct value.
     PCODE                   m_pManagedTarget;
+    // The native code to tailcall. May or may not match what the UM entry precode has been patched to call directly.
+    // This is populated by RunTimeInit at the end of execution.
+    PCODE                   m_pCachedCallTarget;
 
     // This is used for debugging and profiling.
     PTR_MethodDesc          m_pMD;
@@ -200,6 +203,7 @@ class UMEntryThunkData
         m_pManagedTarget = pManagedTarget;
         m_pObjectHandle     = pObjectHandle;
         m_pUMThunkMarshInfo = pUMThunkMarshInfo;
+        m_pCachedCallTarget = (PCODE)0;
 
         m_pMD = pMD;
 
@@ -214,7 +218,20 @@ class UMEntryThunkData
 
     void Terminate();
 
-    VOID RunTimeInit()
+    PCODE GetCachedCallTarget()
+    {
+        STANDARD_VM_CONTRACT;
+        return m_pCachedCallTarget;
+    }
+
+    void PatchPrecode()
+    {
+        STANDARD_VM_CONTRACT;
+
+        m_pUMEntryThunk->SetTargetUnconditional(m_pCachedCallTarget);
+    }
+
+    void RunTimeInit()
     {
         STANDARD_VM_CONTRACT;
 
@@ -227,7 +244,8 @@ class UMEntryThunkData
         if (m_pObjectHandle == NULL && m_pManagedTarget == (TADDR)0)
             m_pManagedTarget = m_pMD->GetMultiCallableAddrOfCode();
 
-        m_pUMEntryThunk->SetTargetUnconditional(m_pUMThunkMarshInfo->GetExecStubEntryPoint());
+        // FIXME: Do we need to use atomics here?
+        m_pCachedCallTarget = m_pUMThunkMarshInfo->GetExecStubEntryPoint();
 
 #ifdef _DEBUG
         m_state = kRunTimeInited;