Resolved the loaded assemblies not being concidered by the assembly resolver

src/Sandbox/ConsoleApp.SourceGenerator.Tests/AssemblyResolverTests.cs

@@ -77,6 +77,54 @@ public void Request_For_A_Newer_Version_Does_Not_Resolve_To_An_Older_One()
             Assert.Throws<FileNotFoundException>(() => Assembly.Load(new AssemblyName("OutdatedDependency, Version=2.0.0.0")));
         }
 
+        [Fact]
+        public void Unsatisfied_Request_Does_Not_Load_The_Embedded_Assembly_Again()
+        {
+            byte[] dependency = Emit("ReloadedDependency", "1.0.0.0");
+            LoadInOwnContext(Emit("ReloadingGenerator", "1.0.0.0", ("SGF.Assembly::ReloadedDependency.dll", dependency)));
+
+            for (int i = 0; i < 3; i++)
+            {
+                Assert.Throws<FileNotFoundException>(() => Assembly.Load(new AssemblyName("ReloadedDependency, Version=2.0.0.0")));
+            }
+
+            Assert.Single(AppDomain.CurrentDomain.GetAssemblies(), a => a.GetName().Name == "ReloadedDependency");
+        }
+
+        [Fact]
+        public void Resource_Fallback_Unpacks_An_Embedded_Assembly_Not_Unpacked_Yet()
+        {
+            byte[] dependency = Emit("LateDependency", "1.0.0.0");
+            Assembly generator = LoadInOwnContext(Emit("LateGenerator", "1.0.0.0", ("SGF.Assembly::LateDependency.dll", dependency)));
+            ForgetUnpacking(generator, "LateDependency, Version=1.0.0.0");
+
+            Assembly resolved = Assembly.Load(new AssemblyName("LateDependency, Version=1.0.0.0"));
+
+            Assert.Equal(new Version(1, 0, 0, 0), resolved.GetName().Version);
+        }
+
+        [Fact]
+        public void Resource_Fallback_Does_Not_Resolve_To_An_Older_Version()
+        {
+            byte[] dependency = Emit("LateOutdatedDependency", "1.0.0.0");
+            Assembly generator = LoadInOwnContext(Emit("LateOutdatedGenerator", "1.0.0.0", ("SGF.Assembly::LateOutdatedDependency.dll", dependency)));
+            ForgetUnpacking(generator, "LateOutdatedDependency, Version=1.0.0.0");
+
+            Assert.Throws<FileNotFoundException>(() => Assembly.Load(new AssemblyName("LateOutdatedDependency, Version=2.0.0.0")));
+        }
+
+        // Puts the resolver back in the window between registering an assembly with resources and unpacking them,
+        // which another generator initialising on another thread can hit. Only the resource fallback covers it.
+        private static void ForgetUnpacking(Assembly generator, string dependency)
+        {
+            Type hoist = typeof(ConsoleAppSourceGeneratorHoist).BaseType!;
+            T Field<T>(string name) => (T)hoist.GetField(name, BindingFlags.NonPublic | BindingFlags.Static)!.GetValue(null)!;
+
+            AssemblyName name = new(dependency);
+            Field<HashSet<string>>("s_extractedResources").Remove($"{generator.FullName}|SGF.Assembly::{name.Name}.dll");
+            Field<Dictionary<AssemblyName, Assembly>>("s_loadedAssemblies").Remove(name);
+        }
+
         private static byte[] Emit(string name, string version, params (string Name, byte[] Content)[] resources)
         {
             CSharpCompilation compilation = CSharpCompilation.Create(
@@ -94,7 +142,7 @@ private static byte[] Emit(string name, string version, params (string Name, byt
         }
 
         // Roslyn loads each analyzer directory in a context of its own; the resolver sees them all through the AppDomain.
-        private static void LoadInOwnContext(byte[] image)
+        private static Assembly LoadInOwnContext(byte[] image)
             => new AssemblyLoadContext(name: null).LoadFromStream(new MemoryStream(image));
     }
 }

src/SourceGenerator.Foundations/Templates/SourceGeneratorHoistBase.cs

@@ -32,6 +32,9 @@ internal abstract class SourceGeneratorHoist
         private static bool s_isInitialized;
         private static readonly List<Assembly> s_assembliesWithResources;
         private static readonly Dictionary<AssemblyName, Assembly> s_loadedAssemblies;
+        private static readonly object s_loadedAssembliesLock;
+        private static readonly object s_assembliesWithResourcesLock;
+        private static readonly HashSet<string> s_extractedResources;
 
         static SourceGeneratorHoist()
         {
@@ -46,6 +49,9 @@ static SourceGeneratorHoist()
 
             s_assembliesWithResources = new List<Assembly>();
             s_loadedAssemblies = new Dictionary<AssemblyName, Assembly>(new AssemblyNameComparer());
+            s_loadedAssembliesLock = new object();
+            s_assembliesWithResourcesLock = new object();
+            s_extractedResources = new HashSet<string>();
             Initialize();
         }
 
@@ -99,11 +105,15 @@ private static void AddAssembly(Assembly assembly)
         {
             AssemblyName assemblyName = assembly.GetName();
 
-            if (s_loadedAssemblies.ContainsKey(assemblyName))
+            lock (s_loadedAssembliesLock)
             {
-                return;
+                if (s_loadedAssemblies.ContainsKey(assemblyName))
+                {
+                    return;
+                }
+
+                s_loadedAssemblies.Add(assemblyName, assembly);
             }
-            s_loadedAssemblies.Add(assemblyName, assembly);
 
             if (assembly.IsDynamic) return;
 
@@ -113,8 +123,19 @@ private static void AddAssembly(Assembly assembly)
 
             if (resources.Length == 0) return;
 
+            // remember assemblies that contain embedded resources so the resolver can
+            // search them later as a fallback
+            lock (s_assembliesWithResourcesLock)
+            {
+                s_assembliesWithResources.Add(assembly);
+            }
+
             foreach (string resource in resources)
             {
+                if (!TryMarkExtracted(assembly, resource))
+                {
+                    continue;
+                }
 
 #pragma warning disable RS1035 // Do not use APIs banned for analyzers
                 System.Console.WriteLine($"Extracting {resource} assembly from {assemblyName.Name}'s resources.");
@@ -134,21 +155,25 @@ private static void AddAssembly(Assembly assembly)
         {
             AssemblyName assemblyName = new(args.Name);
 
-            if (s_loadedAssemblies.TryGetValue(assemblyName, out Assembly? assembly))
+            KeyValuePair<AssemblyName, Assembly>[] loadedSnapshot;
+            lock (s_loadedAssembliesLock)
             {
-                return assembly;
+                if (s_loadedAssemblies.TryGetValue(assemblyName, out Assembly? assembly))
+                {
+                    return assembly;
+                }
+
+                loadedSnapshot = s_loadedAssemblies.ToArray();
             }
 
             // No exact version: the newest one loaded that still satisfies the request, as the runtime binds a
             // reference to a newer version. Never an older one: it can lack what the caller was built against.
             Assembly? newest = null;
             Version newestVersion = new(0, 0);
-            foreach (KeyValuePair<AssemblyName, Assembly> loaded in s_loadedAssemblies)
+            foreach (KeyValuePair<AssemblyName, Assembly> loaded in loadedSnapshot)
             {
                 Version version = loaded.Key.Version ?? new Version(0, 0);
-                bool sameName = string.Equals(loaded.Key.Name, assemblyName.Name);
-                bool satisfies = assemblyName.Version == null || version >= assemblyName.Version;
-                if (sameName && satisfies && (newest == null || version > newestVersion))
+                if (Satisfies(loaded.Key, assemblyName) && (newest == null || version > newestVersion))
                 {
                     newest = loaded.Value;
                     newestVersion = version;
@@ -160,19 +185,58 @@ private static void AddAssembly(Assembly assembly)
                 return newest;
             }
 
-            foreach (Assembly loadedAssembly in s_assembliesWithResources)
+            Assembly[] resourceSnapshot;
+            lock (s_assembliesWithResourcesLock)
+            {
+                resourceSnapshot = s_assembliesWithResources.ToArray();
+            }
+
+            // Reaches a resource AddAssembly has not unpacked yet (generators initialise concurrently). Same version
+            // rule as above: a copy that does not satisfy the request is skipped, so an older one is never returned.
+            // A resource is unpacked once: Assembly.Load(byte[]) loads a new copy on every call, and none unloads.
+            string resourceName = $"SGF.Assembly::{assemblyName.Name}.dll";
+            foreach (Assembly loadedAssembly in resourceSnapshot)
             {
-                string resourceName = $"SGF.Assembly::{assemblyName.Name}.dll";
+                if (!TryMarkExtracted(loadedAssembly, resourceName))
+                {
+                    continue;
+                }
+
                 if (TryExtractingAssembly(loadedAssembly, resourceName, out Assembly? extractedAssembly))
                 {
                     AddAssembly(extractedAssembly!);
-                    return extractedAssembly!;
+                    if (Satisfies(extractedAssembly!.GetName(), assemblyName))
+                    {
+                        return extractedAssembly;
+                    }
                 };
             }
 
             return null;
         }
 
+        /// <summary>
+        /// Claims a resource of an assembly for extraction. False if it was claimed before.
+        /// </summary>
+        private static bool TryMarkExtracted(Assembly assembly, string resourceName)
+        {
+            lock (s_assembliesWithResourcesLock)
+            {
+                return s_extractedResources.Add($"{assembly.FullName}|{resourceName}");
+            }
+        }
+
+        /// <summary>
+        /// True if <paramref name="candidate"/> can stand in for <paramref name="requested"/>: same name, and at least
+        /// the requested version when there is one.
+        /// </summary>
+        private static bool Satisfies(AssemblyName candidate, AssemblyName requested)
+        {
+            Version version = candidate.Version ?? new Version(0, 0);
+            return string.Equals(candidate.Name, requested.Name)
+                && (requested.Version == null || version >= requested.Version);
+        }
+
 
         /// <summary>
         /// Attempts to load an assembly that is contained within another assembly as a resource