Merge branch 'palliate-assembly-loading-errors' into 'main'

Improvements to extension loading

See merge request Sharpmake/sharpmake!545

Author: baudronp

Sharpmake.Application/Program.cs

@@ -10,6 +10,7 @@
 using System.Linq;
 using System.Reflection;
 using System.Runtime.InteropServices;
+using System.Runtime.Loader;
 using System.Threading;
 using System.Threading.Tasks;
 using Sharpmake.Generators;
@@ -438,7 +439,7 @@ private static void LogSharpmakeExtensionLoaded(Assembly extensionAssembly)
                 return;
 
             GetAssemblyInfo(extensionAssembly, out var extensionName, out var _, out var extensionVersion, out var extensionLocation);
-            LogWriteLine("    {0} {1} loaded from '{2}'", extensionName, extensionVersion, extensionLocation);
+            LogWriteLine("    {0} {1} loaded from '{2}' in assembly load context '{3}'", extensionName, extensionVersion, extensionLocation, AssemblyLoadContext.GetLoadContext(extensionAssembly).Name);
         }
 
         private static void CreateBuilderAndGenerate(BuildContext.BaseBuildContext buildContext, Argument parameters, bool generateDebugSolution)

Sharpmake/BuilderExtension.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.Loader;
 
 namespace Sharpmake
 {
@@ -40,7 +41,7 @@ private void AppDomain_AssemblyLoad(object sender, AssemblyLoadEventArgs args)
         /// <param name="extensionAssembly">The <see cref="Assembly"/> to scan.</param>
         private void RegisterExtensionAssembly(Assembly extensionAssembly)
         {
-            if (extensionAssembly.ReflectionOnly)
+            if (ExtensionLoader.IsTempAssembly(extensionAssembly))
                 return;
 
             // Ignores if the assembly does not declare itself as a Sharpmake extension.

Sharpmake/ExtensionLoader.cs

@@ -2,10 +2,12 @@
 // Licensed under the Apache 2.0 License. See LICENSE.md in the project root for license information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.IO;
+using System.Linq;
 using System.Reflection;
-using Microsoft.CodeAnalysis.CSharp.Syntax;
+using System.Runtime.Loader;
 
 namespace Sharpmake
 {
@@ -19,44 +21,14 @@ namespace Sharpmake
     /// </remarks>
     public class ExtensionLoader : IDisposable
     {
+        public const string TempContextNamePrefix = "Temp extension check AssemblyLoadContext";
         public class ExtensionChecker : MarshalByRefObject
         {
-            private readonly Dictionary<string, bool> _loadedAssemblies = new Dictionary<string, bool>(StringComparer.OrdinalIgnoreCase);
-
             public bool IsSharpmakeExtension(string assemblyPath)
             {
-                if (assemblyPath == null)
-                    throw new ArgumentNullException(nameof(assemblyPath));
-                if (!File.Exists(assemblyPath))
-                    throw new FileNotFoundException("Cannot find the assembly DLL.", assemblyPath);
-
-                lock (_loadedAssemblies)
-                {
-                    // If the assembly has already been loaded, check if it is.
-                    bool result;
-                    if (_loadedAssemblies.TryGetValue(assemblyPath, out result))
-                        return result;
-
-                    try
-                    {
-                        Assembly assembly = Assembly.LoadFrom(assemblyPath);
-                        result = IsSharpmakeExtension(assembly);
-                    }
-                    catch (BadImageFormatException)
-                    {
-                        // This is either a native C/C++ assembly, or there is a x86/x64 mismatch
-                        // that prevents it to load. Sharpmake platforms have no reason to not be
-                        // AnyCPU so just assume that it's not a Sharpmake extension.
-                        result = false;
-                    }
-                    catch (Exception ex)
-                    {
-                        throw new Error("An unexpected error has occurred while loading a potential Sharpmake extension assembly: {0}", ex.Message);
-                    }
-
-                    _loadedAssemblies.Add(assemblyPath, result);
-                    return result;
-                }
+                ExtensionLoader extensionLoader = new();
+                bool result =  extensionLoader.LoadExtension(assemblyPath) != null;
+                return result;
             }
 
             public static bool IsSharpmakeExtension(Assembly assembly)
@@ -68,6 +40,12 @@ public static bool IsSharpmakeExtension(Assembly assembly)
             }
         }
 
+        public static bool IsTempAssembly(Assembly assembly)
+        {
+            return AssemblyLoadContext.GetLoadContext(assembly)?.Name?.StartsWith(ExtensionLoader.TempContextNamePrefix) ?? false;
+        }
+
+        private static ConcurrentDictionary<string, byte> _nonExtensionAssemblyPaths = new();
         /// <summary>
         /// Releases the remote <see cref="AppDomain"/> if one was created.
         /// </summary>
@@ -92,23 +70,75 @@ public Assembly LoadExtension(string assemblyPath, bool fastLoad)
         /// Loads a Sharpmake extension assembly.
         /// </summary>
         /// <param name="assemblyPath">The path of the assembly that contains the Sharpmake extension.</param>
-        /// <returns>The loaded extension's <see cref="Assembly"/>.</returns>
+        /// <returns>The loaded extension's <see cref="Assembly"/> or null if it's not a Sharpmake extension.</returns>
         public Assembly LoadExtension(string assemblyPath)
         {
             if (assemblyPath == null)
                 throw new ArgumentNullException(nameof(assemblyPath));
 
-            try
+            if (!File.Exists(assemblyPath))
+                throw new FileNotFoundException("Cannot find the assembly DLL.", assemblyPath);
+            
+            if (_nonExtensionAssemblyPaths.Keys.Contains(assemblyPath))
+                return null;
+
+            // If we've already loaded the assembly in some loading context, check the attributes without loading it again
+            foreach (AssemblyLoadContext alc in AssemblyLoadContext.All)
             {
-                Assembly assembly = Assembly.LoadFrom(assemblyPath);
-                if (!ExtensionChecker.IsSharpmakeExtension(assembly))
+                Assembly assembly = alc.Assemblies.FirstOrDefault(a => AssemblyHasSamePath(a, assemblyPath));
+                if (assembly != null)
+                {
+                    if (ExtensionChecker.IsSharpmakeExtension(assembly))
+                        return assembly;
+                    
+                    _nonExtensionAssemblyPaths.TryAdd(assemblyPath, 0);
+
                     return null;
-                return assembly;
+                }
+            }
+
+            // Check the attributes via an unloadable context (reflexion only is not a thing anymore in .net 6)
+            var contextName = $"{TempContextNamePrefix} - {Path.GetFileNameWithoutExtension(assemblyPath)}";
+            AssemblyLoadContext tempLoadContext = new(contextName, true);
+
+            bool isSharpmakeExtension;
+
+            try
+            {
+                var tempAssembly = tempLoadContext.LoadFromAssemblyPath(assemblyPath);
+                isSharpmakeExtension = ExtensionChecker.IsSharpmakeExtension(tempAssembly);
+                // log unloading of contexts, only log info for Sharpmake extensions
+                Action<string, object[]> logger = isSharpmakeExtension ? Builder.Instance.LogWriteLine : Builder.Instance.DebugWriteLine;
+                tempLoadContext.Unloading += (context) => logger($"    [ExtensionLoader] Attempting to unload temporary context {context.Name}", new object[]{});
             }
             catch (BadImageFormatException)
             {
+                // This is either a native C/C++ assembly, or there is a x86/x64 mismatch
+                // that prevents it to load. Sharpmake platforms have no reason to not be
+                // AnyCPU so just assume that it's not a Sharpmake extension.
                 return null;
             }
+            finally
+            {
+                // Unload the AssemblyLoadContext
+                // NOTE: the unloading event is fired on calling unload for the context, the assembly might not be unloaded
+                //  if there are references to it (make sure it is not used on temp assembly loading)
+                tempLoadContext.Unload();
+            }
+
+            if (isSharpmakeExtension)
+                return AssemblyLoadContext.Default.LoadFromAssemblyPath(assemblyPath);
+
+            _nonExtensionAssemblyPaths.TryAdd(assemblyPath, 0);
+
+            return null;
+
+            static bool AssemblyHasSamePath(Assembly assembly, string assemblyPath)
+            {
+                return assembly != null
+                       && !string.IsNullOrEmpty(assemblyPath) && !assembly.IsDynamic &&
+                       string.Equals(assembly.Location, assemblyPath, StringComparison.Ordinal);
+            }
         }
 
         [Obsolete("This can't work on .net 6. Dead code", error: true)]

Sharpmake/PlatformRegistry.cs

@@ -104,7 +104,7 @@ public static void RegisterExtensionAssembly(Assembly extensionAssembly)
             if (extensionAssembly == null)
                 throw new ArgumentNullException(nameof(extensionAssembly));
 
-            if (extensionAssembly.ReflectionOnly)
+            if (ExtensionLoader.IsTempAssembly(extensionAssembly))
                 return;
 
             // Don't support loading dynamically compiled assemblies