Add cleanup of the TypeIDMap at unload time (#21943)

The TypeIDMap is stored in the AppDomain and contains two hash maps -
id to MethodTable and MethodTable to id. We were missing removing
entries for MethodTables that belong to a LoaderAllocator that's being
destroyed. Thus we were leaking some memory, but also causing potential
issue. When at some point after the LoaderAllocator destruction a
MethodTable gets the same address as one of the MethodTables that was
destroyed in the past and was also recorded in the TypeIDMap, we would
incorrectly reuse the id. That is problematic in case the old
MethodTable didn't require fat id and the new does or vice versa.
I've hit assert due to that while running System.Numerics.Vectors.Tests
inside an unloadable AssemblyLoadContext.

The implementation of the fix is very primitive. It is expected that we
will be able to get rid of the TypeIDMap in a near future and so it is
not worth optimizing.

5 files changed, 60 insertions, 0 deletions

diff --git a/src/vm/appdomain.cpp b/src/vm/appdomain.cpp
index 19035604d6..f0933afda5 100644
--- a/src/vm/appdomain.cpp
+++ b/src/vm/appdomain.cpp
@@ -7864,6 +7864,20 @@ PTR_MethodTable BaseDomain::LookupType(UINT32 id) {
     return pMT;
 }
 
+#ifndef DACCESS_COMPILE
+//---------------------------------------------------------------------------------------
+void BaseDomain::RemoveTypesFromTypeIDMap(LoaderAllocator* pLoaderAllocator)
+{
+    CONTRACTL {
+        NOTHROW;
+        GC_NOTRIGGER;
+        SO_TOLERANT;
+    } CONTRACTL_END;
+
+    m_typeIDMap.RemoveTypes(pLoaderAllocator);
+}
+#endif // DACCESS_COMPILE
+
 //---------------------------------------------------------------------------------------
 // 
 BOOL 
diff --git a/src/vm/appdomain.hpp b/src/vm/appdomain.hpp
index 9db5ad5584..c67f8af00c 100644
--- a/src/vm/appdomain.hpp
+++ b/src/vm/appdomain.hpp
@@ -1427,6 +1427,9 @@ public:
     UINT32 GetTypeID(PTR_MethodTable pMT);
     UINT32 LookupTypeID(PTR_MethodTable pMT);
     PTR_MethodTable LookupType(UINT32 id);
+#ifndef DACCESS_COMPILE
+    void RemoveTypesFromTypeIDMap(LoaderAllocator* pLoaderAllocator);
+#endif // DACCESS_COMPILE
 
 private:
     // I have yet to figure out an efficent way to get the number of handles 
diff --git a/src/vm/contractimpl.cpp b/src/vm/contractimpl.cpp
index c82bd1b45b..1a83712b3e 100644
--- a/src/vm/contractimpl.cpp
+++ b/src/vm/contractimpl.cpp
@@ -155,6 +155,41 @@ UINT32 TypeIDMap::GetTypeID(PTR_MethodTable pMT)
 #ifndef DACCESS_COMPILE 
 
 //------------------------------------------------------------------------
+// Remove all types that belong to the passed in LoaderAllocator
+void TypeIDMap::RemoveTypes(LoaderAllocator *pLoaderAllocator)
+{
+    CONTRACTL {
+        NOTHROW;
+        GC_NOTRIGGER;
+        SO_TOLERANT;
+    } CONTRACTL_END;
+
+    // Take the lock
+    CrstHolder lh(&m_lock);
+
+    for (HashMap::Iterator it = m_mtMap.begin(); !it.end(); ++it)
+    {
+        if (((MethodTable*)it.GetKey())->GetLoaderAllocator() == pLoaderAllocator)
+        {
+            // Note: the entry is just marked for deletion, the removal happens in Compact below
+            m_mtMap.DeleteValue(it.GetKey(), it.GetValue());
+        }
+    }
+
+    m_mtMap.Compact();
+
+    for (HashMap::Iterator it = m_idMap.begin(); !it.end(); ++it)
+    {
+        if (((MethodTable*)(it.GetValue() << 1))->GetLoaderAllocator() == pLoaderAllocator)
+        {
+            // Note: the entry is just marked for deletion, the removal happens in Compact below
+            m_idMap.DeleteValue(it.GetKey(), it.GetValue());
+        }
+    }
+
+    m_idMap.Compact();
+}
+//------------------------------------------------------------------------
 // If TRUE, it points to a matching entry.
 // If FALSE, it is at the insertion point.
 BOOL 
diff --git a/src/vm/contractimpl.h b/src/vm/contractimpl.h
index d7aefaf12c..d0652c1e20 100644
--- a/src/vm/contractimpl.h
+++ b/src/vm/contractimpl.h
@@ -533,6 +533,12 @@ public:
     // returns the new ID.
     UINT32 GetTypeID(PTR_MethodTable pMT);
 
+#ifndef DACCESS_COMPILE
+    //------------------------------------------------------------------------
+    // Remove all types that belong to the passed in LoaderAllocator
+    void RemoveTypes(LoaderAllocator* pLoaderAllocator);
+#endif // DACCESS_COMPILE
+
     //------------------------------------------------------------------------
     inline UINT32 GetCount()
         { LIMITED_METHOD_CONTRACT; return m_entryCount; }
diff --git a/src/vm/loaderallocator.cpp b/src/vm/loaderallocator.cpp
index e3092ff9f3..994ae95f97 100644
--- a/src/vm/loaderallocator.cpp
+++ b/src/vm/loaderallocator.cpp
@@ -485,6 +485,8 @@ LoaderAllocator * LoaderAllocator::GCLoaderAllocators_RemoveAssemblies(AppDomain
     {
         _ASSERTE(!pDomainLoaderAllocatorDestroyIterator->IsAlive());
 
+        GetAppDomain()->RemoveTypesFromTypeIDMap(pDomainLoaderAllocatorDestroyIterator);
+
         DomainAssemblyIterator domainAssemblyIt(pDomainLoaderAllocatorDestroyIterator->m_pFirstDomainAssemblyFromSameALCToDelete);
 
         // Release all assemblies from the same ALC