Fix Assembly reference in unloadability

src/coreclr/vm/appdomain.cpp

@@ -3352,20 +3352,53 @@ BOOL AppDomain::AddFileToCache(AssemblySpec* pSpec, PEAssembly * pPEAssembly, BO
     }
     CONTRACTL_END;
 
-    GCX_PREEMP();
-    DomainCacheCrstHolderForGCCoop holder(this);
-
-    // !!! suppress exceptions
-    if(!m_AssemblyCache.StorePEAssembly(pSpec, pPEAssembly) && !fAllowFailure)
     {
-        // TODO: Disabling the below assertion as currently we experience
-        // inconsistency on resolving the Microsoft.Office.Interop.MSProject.dll
-        // This causes below assertion to fire and crashes the VS. This issue
-        // is being tracked with Dev10 Bug 658555. Brought back it when this bug
-        // is fixed.
-        // _ASSERTE(FALSE);
+        GCX_PREEMP();
+        DomainCacheCrstHolderForGCCoop holder(this);
+
+        // !!! suppress exceptions
+        if(!m_AssemblyCache.StorePEAssembly(pSpec, pPEAssembly) && !fAllowFailure)
+        {
+            // TODO: Disabling the below assertion as currently we experience
+            // inconsistency on resolving the Microsoft.Office.Interop.MSProject.dll
+            // This causes below assertion to fire and crashes the VS. This issue
+            // is being tracked with Dev10 Bug 658555. Brought back it when this bug
+            // is fixed.
+            // _ASSERTE(FALSE);
+
+            EEFileLoadException::Throw(pSpec, FUSION_E_CACHEFILE_FAILED, NULL);
+        }
+    }
 
-        EEFileLoadException::Throw(pSpec, FUSION_E_CACHEFILE_FAILED, NULL);
+    // 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;

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

@@ -0,0 +1,15 @@
+using TestInterface;
+
+namespace TestClass
+{
+    public class Class : TestInterface.Interface
+    {
+        public void MainTest()
+        {
+            TestInterface.Class.Test();
+        }
+        public void Test()
+        {
+        }
+    }
+}

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

@@ -0,0 +1,17 @@
+using System;
+
+namespace TestInterface
+{
+    public interface Interface
+    {
+        void Test();
+    }
+
+    public class Class
+    {
+        public static void Test()
+        {
+            Console.WriteLine("TestInterface.Class.Test");
+        }
+    }
+}

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

@@ -0,0 +1,170 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+// When an assembly A1 from an unloadable AssemblyLoadContext 1 is referenced
+// by another assembly A2 from an unloadable AssemblyLoadContext 2, e.g.
+// when a class from A2 implements an interface from A1 and there are no
+// instances of types from A1 and the managed assembly type and the related
+// AssemblyLoadContext is gone too, the AssemblyLoadContext 1 cannot
+// get collected, it would result in later crashes, null references
+// etc.
+
+using System;
+using System.Runtime.Loader;
+using System.Reflection;
+using System.Runtime.CompilerServices;
+using System.IO;
+
+class TestAssemblyLoadContext : AssemblyLoadContext
+{
+    WeakReference interfaceAssemblyRef = null;
+
+    public TestAssemblyLoadContext(string name, bool isCollectible) : base(name,  isCollectible)
+    {
+    }
+
+    public WeakReference InterfaceAssemblyRef { get { return interfaceAssemblyRef; } }
+
+    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"));
+            interfaceAssemblyRef = new WeakReference(a);
+            return a;
+        }
+
+        return null;
+    }
+}
+
+class Test
+{
+    static AssemblyLoadContext alc1 = null;
+    static WeakReference interfaceAssemblyRef = null;
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static Assembly LoadUsingResolvingEvent()
+    {
+        alc1 = new AssemblyLoadContext("Dependencies", true);
+        AssemblyLoadContext alc2 = new AssemblyLoadContext("Test1", true);
+        alc2.Resolving += Alc2_Resolving;
+        Assembly assembly = alc2.LoadFromAssemblyPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, @"..\TestClass\TestClass.dll"));
+
+        Type t = assembly.GetType("TestClass.Class");
+        Console.WriteLine($"Type {t} obtained");
+
+        MethodInfo mi = t.GetMethod("MainTest");
+        Console.WriteLine($"Method {mi} obtained");
+
+        alc1 = null;
+        Console.WriteLine("Load done");
+
+        return assembly;
+    }
+
+    private static Assembly Alc2_Resolving(AssemblyLoadContext arg1, AssemblyName arg2)
+    {
+        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"));
+            interfaceAssemblyRef = new WeakReference(a);
+            return a;
+        }
+
+        return null;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static Assembly LoadUsingLoadOverride()
+    {
+        TestAssemblyLoadContext alc2 = new TestAssemblyLoadContext("Test2", true);
+        Assembly assembly = alc2.LoadFromAssemblyPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, @"..\TestClass\TestClass.dll"));
+
+        Type t = assembly.GetType("TestClass.Class");
+
+        Console.WriteLine($"Load done, type {t} obtained");
+
+        interfaceAssemblyRef = alc2.InterfaceAssemblyRef;
+
+        return assembly;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static WeakReference TestDependencies(string testCase)
+    {
+        Assembly assembly;
+
+        if (testCase == "ResolvingEvent")
+        {
+            assembly = LoadUsingResolvingEvent();
+        }
+        else
+        {
+            assembly = LoadUsingLoadOverride();
+        }
+
+        for (int i = 0; interfaceAssemblyRef.IsAlive && (i < 10); i++)
+        {
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+        }
+
+        Type t = assembly.GetType("TestClass.Class");
+        Type[] ifaces = t.GetInterfaces();
+        ConstructorInfo ctor = t.GetConstructor(new Type[] { });
+        object instance = ctor.Invoke(new object[0]);
+        MethodInfo mi = t.GetMethod("MainTest");
+        mi.Invoke(instance, null);
+
+        return new WeakReference(assembly);
+    }
+
+    public static int TestFullUnload(string testCase)
+    {
+        Console.WriteLine($"Running test case {testCase}");
+
+        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();
+        return 100;
+
+    }
+
+    public static int Main(string[] args)
+    {
+        int status = 100;
+        foreach (string testCase in new string[] {"LoadOverride", "ResolvingEvent"})
+        {
+            status = TestFullUnload(testCase);
+        }
+
+        return status;
+    }
+}

src/tests/Loader/CollectibleAssemblies/ResolvedFromDifferentContext/ResolvedFromDifferentContext.csproj

@@ -0,0 +1,9 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <CLRTestPriority>1</CLRTestPriority>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="ResolvedFromDifferentContext.cs" />
+  </ItemGroup>
+</Project>

src/tests/Loader/CollectibleAssemblies/ResolvedFromDifferentContext/TestClass.csproj

@@ -0,0 +1,16 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <OutputType>Library</OutputType>
+    <CLRTestKind>SharedLibrary</CLRTestKind>
+  </PropertyGroup>
+
+  <ItemGroup>
+   <Compile Include="Class.cs" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="TestInterface.csproj" />
+  </ItemGroup>
+
+</Project>

src/tests/Loader/CollectibleAssemblies/ResolvedFromDifferentContext/TestInterface.csproj

@@ -0,0 +1,12 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <OutputType>Library</OutputType>
+    <CLRTestKind>SharedLibrary</CLRTestKind>
+  </PropertyGroup>
+
+  <ItemGroup>
+   <Compile Include="Interface.cs" />
+  </ItemGroup>
+
+</Project>