Rework unloadability fix (#68550)

* Rework the unloadability fix

The recent fix had a problem when GC collected the managed Assembly
loaded via the Load override or the Resolving event before we added the
reference between the related native LoaderAllocators. The managed
LoaderAllocator of the ALC we've loaded the Assembly into was collected
and we couldn't create the managed Type objects for the interfaces from
that assembly in the GetInterfaces call on a type that implemented
those.

This change fixes it by moving the creation of reference between the
native LoaderAllocators to the RuntimeInvokeHostAssemblyResolver where
we still have a live managed reference to the loaded Assembly.

* Move clearing of the m_pAssemblyLoaderAllocator to a later stage

* Throw exception on collectible assembly in non-collectible ALC

* Update the ResolvedFromDifferentContext test

Make it test that we throw an exception when the parent ALC is not
collectible.

Author: janvorli

src/coreclr/binder/assemblybindercommon.cpp

@@ -30,6 +30,7 @@
 extern HRESULT RuntimeInvokeHostAssemblyResolver(INT_PTR pManagedAssemblyLoadContextToBindWithin,
                                                  BINDER_SPACE::AssemblyName *pAssemblyName,
                                                  DefaultAssemblyBinder *pDefaultBinder,
+                                                 AssemblyBinder *pBinder,
                                                  BINDER_SPACE::Assembly **ppLoadedAssembly);
 
 #endif // !defined(DACCESS_COMPILE)
@@ -1153,9 +1154,10 @@ namespace BINDER_SPACE
 
 #if !defined(DACCESS_COMPILE)
 HRESULT AssemblyBinderCommon::BindUsingHostAssemblyResolver(/* in */ INT_PTR pManagedAssemblyLoadContextToBindWithin,
-                                                            /* in */ AssemblyName       *pAssemblyName,
+                                                            /* in */ AssemblyName *pAssemblyName,
                                                             /* in */ DefaultAssemblyBinder *pDefaultBinder,
-                                                            /* out */ Assembly           **ppAssembly)
+                                                            /* in */ AssemblyBinder *pBinder,
+                                                            /* out */ Assembly **ppAssembly)
 {
     HRESULT hr = E_FAIL;
 
@@ -1164,7 +1166,7 @@ HRESULT AssemblyBinderCommon::BindUsingHostAssemblyResolver(/* in */ INT_PTR pMa
     // RuntimeInvokeHostAssemblyResolver will perform steps 2-4 of CustomAssemblyBinder::BindAssemblyByName.
     BINDER_SPACE::Assembly *pLoadedAssembly = NULL;
     hr = RuntimeInvokeHostAssemblyResolver(pManagedAssemblyLoadContextToBindWithin,
-                                           pAssemblyName, pDefaultBinder, &pLoadedAssembly);
+                                           pAssemblyName, pDefaultBinder, pBinder, &pLoadedAssembly);
     if (SUCCEEDED(hr))
     {
         _ASSERTE(pLoadedAssembly != NULL);

src/coreclr/binder/customassemblybinder.cpp

@@ -72,7 +72,8 @@ HRESULT CustomAssemblyBinder::BindUsingAssemblyName(BINDER_SPACE::AssemblyName*
             // of what to do next. The host-overridden binder can either fail the bind or return reference to an existing assembly
             // that has been loaded.
             //
-            hr = AssemblyBinderCommon::BindUsingHostAssemblyResolver(GetManagedAssemblyLoadContext(), pAssemblyName, m_pDefaultBinder, &pCoreCLRFoundAssembly);
+            hr = AssemblyBinderCommon::BindUsingHostAssemblyResolver(GetManagedAssemblyLoadContext(), pAssemblyName,
+                                                                     m_pDefaultBinder, this, &pCoreCLRFoundAssembly);
             if (SUCCEEDED(hr))
             {
                 // We maybe returned an assembly that was bound to a different AssemblyBinder instance.
@@ -224,9 +225,13 @@ void CustomAssemblyBinder::PrepareForLoadContextRelease(INT_PTR ptrManagedStrong
 
     // We cannot delete the binder here as it is used indirectly when comparing assemblies with the same binder
     // It will be deleted when the LoaderAllocator will be deleted
-    // But we can release the LoaderAllocator as we are no longer using it here
+    // We need to keep the LoaderAllocator pointer set as it still may be needed for creating references between the
+    // native LoaderAllocators of two collectible contexts in case the AssemblyLoadContext.Unload was called on the current
+    // context before returning from its AssemblyLoadContext.Load override or the context's Resolving event.
+    // But we need to release the LoaderAllocator so that it doesn't prevent completion of the final phase of unloading in
+    // some cases. It is safe to do as the AssemblyLoaderAllocator is guaranteed to be alive at least until the 
+    // CustomAssemblyBinder::ReleaseLoadContext is called, where we NULL this pointer.
     m_pAssemblyLoaderAllocator->Release();
-    m_pAssemblyLoaderAllocator = NULL;
 
     // Destroy the strong handle to the LoaderAllocator in order to let it reach its finalizer
     DestroyHandle(reinterpret_cast<OBJECTHANDLE>(m_loaderAllocatorHandle));
@@ -251,6 +256,10 @@ void CustomAssemblyBinder::ReleaseLoadContext()
     handle = reinterpret_cast<OBJECTHANDLE>(m_ptrManagedStrongAssemblyLoadContext);
     DestroyHandle(handle);
     SetManagedAssemblyLoadContext(NULL);
+
+    // The AssemblyLoaderAllocator is in a process of shutdown and should not be used 
+    // after this point.
+    m_pAssemblyLoaderAllocator = NULL;
 }
 
 #endif // !defined(DACCESS_COMPILE)

src/coreclr/binder/defaultassemblybinder.cpp

@@ -85,7 +85,7 @@ HRESULT DefaultAssemblyBinder::BindUsingAssemblyName(BINDER_SPACE::AssemblyName
         if (pManagedAssemblyLoadContext != NULL)
         {
             hr = AssemblyBinderCommon::BindUsingHostAssemblyResolver(pManagedAssemblyLoadContext, pAssemblyName,
-                                                                     NULL, &pCoreCLRFoundAssembly);
+                                                                     NULL, this, &pCoreCLRFoundAssembly);
             if (SUCCEEDED(hr))
             {
                 // We maybe returned an assembly that was bound to a different AssemblyLoadContext instance.

src/coreclr/binder/inc/assemblybindercommon.hpp

@@ -52,6 +52,7 @@ namespace BINDER_SPACE
         static HRESULT BindUsingHostAssemblyResolver (/* in */ INT_PTR pManagedAssemblyLoadContextToBindWithin,
                                                       /* in */ AssemblyName       *pAssemblyName,
                                                       /* in */ DefaultAssemblyBinder *pDefaultBinder,
+                                                      /* in */ AssemblyBinder *pBinder,
                                                       /* out */ Assembly           **ppAssembly);
 
         static HRESULT BindUsingPEImage(/* in */  AssemblyBinder     *pBinder,

src/coreclr/vm/appdomain.cpp

@@ -3370,37 +3370,6 @@ BOOL AppDomain::AddFileToCache(AssemblySpec* pSpec, PEAssembly * pPEAssembly, BO
         }
     }
 
-    // If the ALC for the result (pPEAssembly) is collectible and not the same as the ALC for the request (pSpec),
-    // which can happen when extension points (AssemblyLoadContext.Load, AssemblyLoadContext.Resolving) resolve the
-    // assembly, we need to ensure the result ALC is not collected before the request ALC. We do this by adding an
-    // explicit reference between the LoaderAllocators corresponding to the ALCs.
-    //
-    // To get the LoaderAllocators, we rely on the DomainAssembly corresponding to:
-    // - Parent assembly for the request
-    //   - For loads via explicit load or reflection, this will be NULL. In these cases, the request ALC should not
-    //    implicitly (as in not detectable via managed references) depend on the result ALC, so it is fine to not
-    //    add the reference.
-    //   - For loads via assembly references, this will never be NULL.
-    // - Result assembly for the result
-    //   - For dynamic assemblies, there is no host assembly, so we don't have the result assembly / LoaderAllocator.
-    //     We currently block resolving to dynamic assemblies, so we simply assert that we have a host assembly.
-    //   - For non-dynamic assemblies, we should be able to get the host assembly.
-    DomainAssembly *pParentAssembly = pSpec->GetParentAssembly();
-    BINDER_SPACE::Assembly *pBinderSpaceAssembly = pPEAssembly->GetHostAssembly();
-    _ASSERTE(pBinderSpaceAssembly != NULL);
-    DomainAssembly *pResultAssembly = pBinderSpaceAssembly->GetDomainAssembly();
-    if ((pParentAssembly != NULL) && (pResultAssembly != NULL))
-    {
-        LoaderAllocator *pParentAssemblyLoaderAllocator = pParentAssembly->GetLoaderAllocator();
-        LoaderAllocator *pResultAssemblyLoaderAllocator = pResultAssembly->GetLoaderAllocator();
-        _ASSERTE(pParentAssemblyLoaderAllocator);
-        _ASSERTE(pResultAssemblyLoaderAllocator);
-        if (pResultAssemblyLoaderAllocator->IsCollectible())
-        {
-            pParentAssemblyLoaderAllocator->EnsureReference(pResultAssemblyLoaderAllocator);
-        }
-    }
-
     return TRUE;
 }
 
@@ -4937,7 +4906,7 @@ AppDomain::AssemblyIterator::Next_Unlocked(
 #if !defined(DACCESS_COMPILE)
 
 // Returns S_OK if the assembly was successfully loaded
-HRESULT RuntimeInvokeHostAssemblyResolver(INT_PTR pManagedAssemblyLoadContextToBindWithin, BINDER_SPACE::AssemblyName *pAssemblyName, DefaultAssemblyBinder *pDefaultBinder, BINDER_SPACE::Assembly **ppLoadedAssembly)
+HRESULT RuntimeInvokeHostAssemblyResolver(INT_PTR pManagedAssemblyLoadContextToBindWithin, BINDER_SPACE::AssemblyName *pAssemblyName, DefaultAssemblyBinder *pDefaultBinder, AssemblyBinder *pBinder, BINDER_SPACE::Assembly **ppLoadedAssembly)
 {
     CONTRACTL
     {
@@ -5115,6 +5084,22 @@ HRESULT RuntimeInvokeHostAssemblyResolver(INT_PTR pManagedAssemblyLoadContextToB
                 COMPlusThrowHR(COR_E_INVALIDOPERATION, IDS_HOST_ASSEMBLY_RESOLVER_DYNAMICALLY_EMITTED_ASSEMBLIES_UNSUPPORTED, name);
             }
 
+            // For collectible assemblies, ensure that the parent loader allocator keeps the assembly's loader allocator
+            // alive for all its lifetime.
+            if (pDomainAssembly->IsCollectible())
+            {
+                LoaderAllocator *pResultAssemblyLoaderAllocator = pDomainAssembly->GetLoaderAllocator();
+                LoaderAllocator *pParentLoaderAllocator = pBinder->GetLoaderAllocator();
+                if (pParentLoaderAllocator == NULL)
+                {
+                    // The AssemblyLoadContext for which we are resolving the the Assembly is not collectible.
+                    COMPlusThrow(kNotSupportedException, W("NotSupported_CollectibleBoundNonCollectible")); 
+                }
+
+                _ASSERTE(pResultAssemblyLoaderAllocator);
+                pParentLoaderAllocator->EnsureReference(pResultAssemblyLoaderAllocator);
+            }
+
             pResolvedAssembly = pLoadedPEAssembly->GetHostAssembly();
         }
 

src/tests/Loader/CollectibleAssemblies/ResolvedFromDifferentContext/ResolvedFromDifferentContext.cs

@@ -30,8 +30,8 @@ protected override Assembly Load(AssemblyName assemblyName)
         if (assemblyName.Name == "TestInterface")
         {
             AssemblyLoadContext alc1 = new AssemblyLoadContext("Dependencies", true);
-            Console.WriteLine($"Loading TestInterface by alc {alc1} for alc {this}");
-            Assembly a = alc1.LoadFromAssemblyPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, @"..\TestInterface\TestInterface.dll"));
+            Console.WriteLine($"Loading TestInterface by alc {alc1} for {(IsCollectible ? "collectible" : "non-collectible")} alc {this}");
+            Assembly a = alc1.LoadFromAssemblyPath(Test.GetTestAssemblyPath(@"..\TestInterface\TestInterface.dll"));
             interfaceAssemblyRef = new WeakReference(a);
             return a;
         }
@@ -45,13 +45,18 @@ class Test
     static AssemblyLoadContext alc1 = null;
     static WeakReference interfaceAssemblyRef = null;
 
+    public static string GetTestAssemblyPath(string subPath)
+    {
+        return Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), subPath);
+    }
+
     [MethodImpl(MethodImplOptions.NoInlining)]
-    public static Assembly LoadUsingResolvingEvent()
+    private static Assembly LoadUsingResolvingEvent(bool collectibleParent)
     {
         alc1 = new AssemblyLoadContext("Dependencies", true);
-        AssemblyLoadContext alc2 = new AssemblyLoadContext("Test1", true);
+        AssemblyLoadContext alc2 = new AssemblyLoadContext("Test1", collectibleParent);
         alc2.Resolving += Alc2_Resolving;
-        Assembly assembly = alc2.LoadFromAssemblyPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, @"..\TestClass\TestClass.dll"));
+        Assembly assembly = alc2.LoadFromAssemblyPath(Test.GetTestAssemblyPath(@"..\TestClass\TestClass.dll"));
 
         Type t = assembly.GetType("TestClass.Class");
         Console.WriteLine($"Type {t} obtained");
@@ -70,8 +75,8 @@ private static Assembly Alc2_Resolving(AssemblyLoadContext arg1, AssemblyName ar
         Console.WriteLine($"Resolving event by alc {alc1} for alc {arg1}");
         if (alc1 != null && arg2.Name == "TestInterface")
         {
-            Console.WriteLine($"Loading TestInterface by alc {alc1} for alc {arg1}");
-            Assembly a = alc1.LoadFromAssemblyPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, @"..\TestInterface\TestInterface.dll"));
+            Console.WriteLine($"Loading TestInterface by alc {alc1} for {(arg1.IsCollectible ? "collectible" : "non-collectible")} alc {arg1}");
+            Assembly a = alc1.LoadFromAssemblyPath(Test.GetTestAssemblyPath(@"..\TestInterface\TestInterface.dll"));
             interfaceAssemblyRef = new WeakReference(a);
             return a;
         }
@@ -80,10 +85,10 @@ private static Assembly Alc2_Resolving(AssemblyLoadContext arg1, AssemblyName ar
     }
 
     [MethodImpl(MethodImplOptions.NoInlining)]
-    public static Assembly LoadUsingLoadOverride()
+    private static Assembly LoadUsingLoadOverride(bool collectibleParent)
     {
-        TestAssemblyLoadContext alc2 = new TestAssemblyLoadContext("Test2", true);
-        Assembly assembly = alc2.LoadFromAssemblyPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, @"..\TestClass\TestClass.dll"));
+        TestAssemblyLoadContext alc2 = new TestAssemblyLoadContext("Test2", collectibleParent);
+        Assembly assembly = alc2.LoadFromAssemblyPath(Test.GetTestAssemblyPath(@"..\TestClass\TestClass.dll"));
 
         Type t = assembly.GetType("TestClass.Class");
 
@@ -94,18 +99,33 @@ public static Assembly LoadUsingLoadOverride()
         return assembly;
     }
 
+    private enum TestCase
+    {
+        ResolvingEvent,
+        LoadOverride,
+        ResolvingEventInNonCollectible,
+        LoadOverrideInNonCollectible
+    }
+
     [MethodImpl(MethodImplOptions.NoInlining)]
-    private static WeakReference TestDependencies(string testCase)
+    private static WeakReference TestDependencies(TestCase testCase)
     {
-        Assembly assembly;
+        Assembly assembly = null;
 
-        if (testCase == "ResolvingEvent")
-        {
-            assembly = LoadUsingResolvingEvent();
-        }
-        else
+        switch (testCase)
         {
-            assembly = LoadUsingLoadOverride();
+            case TestCase.ResolvingEvent:
+                assembly = LoadUsingResolvingEvent(collectibleParent: true);
+                break;
+            case TestCase.LoadOverride:
+                assembly = LoadUsingLoadOverride(collectibleParent: true);
+                break;
+            case TestCase.ResolvingEventInNonCollectible:
+                assembly = LoadUsingResolvingEvent(collectibleParent: false);
+                break;
+            case TestCase.LoadOverrideInNonCollectible:
+                assembly = LoadUsingLoadOverride(collectibleParent: false);
+                break;
         }
 
         for (int i = 0; interfaceAssemblyRef.IsAlive && (i < 10); i++)
@@ -124,45 +144,78 @@ private static WeakReference TestDependencies(string testCase)
         return new WeakReference(assembly);
     }
 
-    public static int TestFullUnload(string testCase)
+    private static bool ShouldThrow(TestCase testCase)
     {
-        Console.WriteLine($"Running test case {testCase}");
+        return (testCase == TestCase.LoadOverrideInNonCollectible) || (testCase == TestCase.ResolvingEventInNonCollectible);
+    }
 
-        WeakReference assemblyRef = TestDependencies(testCase);
-        if (assemblyRef == null)
-        {
-            return 101;
-        }
+    private static int TestFullUnload(TestCase testCase)
+    {
+        Console.WriteLine($"Running test case {testCase}");
 
-        for (int i = 0; assemblyRef.IsAlive && (i < 10); i++)
+        try
         {
-            GC.Collect();
-            GC.WaitForPendingFinalizers();
+            WeakReference assemblyRef = TestDependencies(testCase);
+            if (assemblyRef == null)
+            {
+                return 101;
+            }
+
+            for (int i = 0; assemblyRef.IsAlive && (i < 10); i++)
+            {
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
+            }
+
+            if (assemblyRef.IsAlive)
+            {
+                Console.WriteLine("Failed to unload alc2");
+                return 102;
+            }
+
+            if (interfaceAssemblyRef.IsAlive)
+            {
+                Console.WriteLine("Failed to unload alc1");
+                return 103;
+            }
+
+            Console.WriteLine();
         }
-
-        if (assemblyRef.IsAlive)
+        catch (System.IO.FileLoadException e)
         {
-            Console.WriteLine("Failed to unload alc2");
-            return 102;
+            if (!ShouldThrow(testCase))
+            {
+                Console.WriteLine("Failure - unexpected exception");
+                return 104;
+            }
+            if ((e.InnerException == null) || e.InnerException.GetType() != typeof(System.NotSupportedException))
+            {
+                Console.WriteLine($"Failure - unexpected exception type {e.InnerException}");
+                return 105;
+            }
+
+            return 100;
         }
 
-        if (interfaceAssemblyRef.IsAlive)
+        if (ShouldThrow(testCase))
         {
-            Console.WriteLine("Failed to unload alc1");
-            return 103;
+            Console.WriteLine("Failure - resolved collectible assembly into non-collectible context without throwing exception");
+            return 106;
         }
 
-        Console.WriteLine();
         return 100;
-         
     }
 
     public static int Main(string[] args)
     {
         int status = 100;
-        foreach (string testCase in new string[] {"LoadOverride", "ResolvingEvent"})
+        foreach (TestCase testCase in Enum.GetValues(typeof(TestCase)))
         {
             status = TestFullUnload(testCase);
+            if (status != 100)
+            {
+                break;
+            }
         }
 
         return status;