Fix collectible NativeCallable UMThunkEntry lifetime (#20438)

* Fix collectible NativeCallable UMThunkEntry lifetime

The UMEntryThunk cache entries created for NativeCallable target methods
for collectible classes were not properly cleaned up at the unload time.
This change fixes that by adding UMEntryThunkCache on LoaderAllocator
and using it for entries belonging to NativeCallable targets on
collectible classes. The cache is created lazily.

* Reflect PR feedback

Remove the UMEntryThunk cache from the AppDomain and leave it just on
the LoaderAllocator.

6 files changed, 42 insertions, 43 deletions

diff --git a/src/vm/appdomain.cpp b/src/vm/appdomain.cpp
index 61f709a9ca..6233094904 100644
--- a/src/vm/appdomain.cpp
+++ b/src/vm/appdomain.cpp
@@ -3885,8 +3885,6 @@ AppDomain::AppDomain()
     memset(m_rpCLRTypes, 0, sizeof(m_rpCLRTypes));
 #endif // FEATURE_COMINTEROP
 
-    m_pUMEntryThunkCache = NULL;
-
     m_pAsyncPool = NULL;
     m_handleStore = NULL;
 
@@ -4222,12 +4220,6 @@ void AppDomain::Terminate()
     delete m_pDefaultContext;
     m_pDefaultContext = NULL;
 
-    if (m_pUMEntryThunkCache)
-    {
-        delete m_pUMEntryThunkCache;
-        m_pUMEntryThunkCache = NULL;
-    }
-
 #ifdef FEATURE_COMINTEROP
     if (m_pRCWCache)
     {
@@ -7546,31 +7538,6 @@ BOOL AppDomain::WasSystemAssemblyLoadEventSent(void)
 }
 
 #ifndef CROSSGEN_COMPILE
-// U->M thunks created in this domain and not associated with a delegate.
-UMEntryThunkCache *AppDomain::GetUMEntryThunkCache()
-{
-    CONTRACTL
-    {
-        THROWS;
-        GC_TRIGGERS;
-        MODE_ANY;
-        INJECT_FAULT(COMPlusThrowOM(););
-    }
-    CONTRACTL_END;
-
-    if (!m_pUMEntryThunkCache)
-    {
-        UMEntryThunkCache *pUMEntryThunkCache = new UMEntryThunkCache(this);
-
-        if (FastInterlockCompareExchangePointer(&m_pUMEntryThunkCache, pUMEntryThunkCache, NULL) != NULL)
-        {
-            // some thread swooped in and set the field
-            delete pUMEntryThunkCache;
-        }
-    }
-    _ASSERTE(m_pUMEntryThunkCache);
-    return m_pUMEntryThunkCache;
-}
 
 #ifdef FEATURE_COMINTEROP
 
diff --git a/src/vm/appdomain.hpp b/src/vm/appdomain.hpp
index 2c838d1d51..5819c340fc 100644
--- a/src/vm/appdomain.hpp
+++ b/src/vm/appdomain.hpp
@@ -67,7 +67,6 @@ class DomainModule;
 class DomainAssembly;
 struct InteropMethodTableData;
 class LoadLevelLimiter;
-class UMEntryThunkCache;
 class TypeEquivalenceHashTable;
 class StringArrayList;
 
@@ -3394,10 +3393,6 @@ private:
     // IL stub cache with fabricated MethodTable parented by a random module in this AD.
     ILStubCache         m_ILStubCache;
 
-    // U->M thunks created in this domain and not associated with a delegate.
-    // The cache is keyed by MethodDesc pointers.
-    UMEntryThunkCache *m_pUMEntryThunkCache;
-
     // The number of  times we have entered this AD
     ULONG m_dwThreadEnterCount;
     // The number of threads that have entered this AD, for ADU only
@@ -3475,8 +3470,6 @@ public:
     BOOL IsBindingModelLocked();
     BOOL LockBindingModel();
 
-    UMEntryThunkCache *GetUMEntryThunkCache();
-
     ILStubCache* GetILStubCache()
     {
         LIMITED_METHOD_CONTRACT;
diff --git a/src/vm/comdelegate.cpp b/src/vm/comdelegate.cpp
index cef90b034e..6d7c256147 100644
--- a/src/vm/comdelegate.cpp
+++ b/src/vm/comdelegate.cpp
@@ -1076,8 +1076,8 @@ PCODE COMDelegate::ConvertToCallback(MethodDesc* pMD)
     if (NDirect::MarshalingRequired(pMD, pMD->GetSig(), pMD->GetModule()))
         COMPlusThrow(kNotSupportedException, W("NotSupported_NonBlittableTypes"));
 
-    // Get UMEntryThunk from appdomain thunkcache cache.
-    UMEntryThunk *pUMEntryThunk = GetAppDomain()->GetUMEntryThunkCache()->GetUMEntryThunk(pMD);
+    // Get UMEntryThunk from the thunk cache.
+    UMEntryThunk *pUMEntryThunk = pMD->GetLoaderAllocator()->GetUMEntryThunkCache()->GetUMEntryThunk(pMD);
 
 #if defined(_TARGET_X86_) && !defined(FEATURE_STUBS_AS_IL)
 
diff --git a/src/vm/corhost.cpp b/src/vm/corhost.cpp
index f9785ab19f..03fbdb9df8 100644
--- a/src/vm/corhost.cpp
+++ b/src/vm/corhost.cpp
@@ -798,7 +798,7 @@ HRESULT CorHost2::_CreateDelegate(
     if (pMD==NULL || !pMD->IsStatic() || pMD->ContainsGenericVariables()) 
         ThrowHR(COR_E_MISSINGMETHOD);
 
-    UMEntryThunk *pUMEntryThunk = GetAppDomain()->GetUMEntryThunkCache()->GetUMEntryThunk(pMD);
+    UMEntryThunk *pUMEntryThunk = pMD->GetLoaderAllocator()->GetUMEntryThunkCache()->GetUMEntryThunk(pMD);
     *fnPtr = (INT_PTR)pUMEntryThunk->GetCode();
 
     END_DOMAIN_TRANSITION;
diff --git a/src/vm/loaderallocator.cpp b/src/vm/loaderallocator.cpp
index 46358f79d0..7033abf8f2 100644
--- a/src/vm/loaderallocator.cpp
+++ b/src/vm/loaderallocator.cpp
@@ -72,6 +72,8 @@ LoaderAllocator::LoaderAllocator()
     m_pJumpStubCache = NULL;
     m_IsCollectible = false;
 
+    m_pUMEntryThunkCache = NULL;
+
     m_nLoaderAllocator = InterlockedIncrement64((LONGLONG *)&LoaderAllocator::cLoaderAllocatorsCreated);
 }
 
@@ -1284,6 +1286,9 @@ void LoaderAllocator::Terminate()
         m_fGCPressure = false;
     }
 
+    delete m_pUMEntryThunkCache;
+    m_pUMEntryThunkCache = NULL;
+
     m_crstLoaderAllocator.Destroy();
     m_LoaderAllocatorReferences.RemoveAll();
 
@@ -1872,6 +1877,32 @@ void AssemblyLoaderAllocator::ReleaseManagedAssemblyLoadContext()
     }
 }
 
+// U->M thunks created in this LoaderAllocator and not associated with a delegate.
+UMEntryThunkCache *LoaderAllocator::GetUMEntryThunkCache()
+{
+    CONTRACTL
+    {
+        THROWS;
+        GC_TRIGGERS;
+        MODE_ANY;
+        INJECT_FAULT(COMPlusThrowOM(););
+    }
+    CONTRACTL_END;
+
+    if (!m_pUMEntryThunkCache)
+    {
+        UMEntryThunkCache *pUMEntryThunkCache = new UMEntryThunkCache(GetAppDomain());
+
+        if (FastInterlockCompareExchangePointer(&m_pUMEntryThunkCache, pUMEntryThunkCache, NULL) != NULL)
+        {
+            // some thread swooped in and set the field
+            delete pUMEntryThunkCache;
+        }
+    }
+    _ASSERTE(m_pUMEntryThunkCache);
+    return m_pUMEntryThunkCache;
+}
+
 #endif // !CROSSGEN_COMPILE
 
 #endif // !DACCESS_COMPILE
diff --git a/src/vm/loaderallocator.hpp b/src/vm/loaderallocator.hpp
index f047076829..beaa3e818a 100644
--- a/src/vm/loaderallocator.hpp
+++ b/src/vm/loaderallocator.hpp
@@ -130,6 +130,7 @@ class VirtualCallStubManager;
 template <typename ELEMENT>
 class ListLockEntryBase;
 typedef ListLockEntryBase<void*> ListLockEntry;
+class UMEntryThunkCache;
 
 class LoaderAllocator
 {
@@ -173,6 +174,10 @@ protected:
     BYTE *              m_pVSDHeapInitialAlloc;
     BYTE *              m_pCodeHeapInitialAlloc;
 
+    // U->M thunks that are not associated with a delegate.
+    // The cache is keyed by MethodDesc pointers.
+    UMEntryThunkCache * m_pUMEntryThunkCache;
+
 public:
     BYTE *GetVSDHeapInitialBlock(DWORD *pSize);
     BYTE *GetCodeHeapInitialBlock(const BYTE * loAddr, const BYTE * hiAddr, DWORD minimumSize, DWORD *pSize);
@@ -510,6 +515,9 @@ public:
         LIMITED_METHOD_CONTRACT;
         return m_pVirtualCallStubManager;
     }
+
+    UMEntryThunkCache *GetUMEntryThunkCache();
+
 #endif
 };  // class LoaderAllocator