Add an escape hatch for ALC isolation (#5098)

Restore the original load-with-funky-rules-into-the-default-ALC behavior under an environment-variable escape hatch. Fixes #5086.

This looks like a lot of code, but it's almost all undeletions.

src/Shared/CoreCLRAssemblyLoader.cs

@@ -1,10 +1,13 @@
 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using Microsoft.Build.Shared.FileSystem;
+using Microsoft.Build.Utilities;
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
 using System.Reflection;
+using System.Runtime.Loader;
 
 namespace Microsoft.Build.Shared
 {
@@ -14,8 +17,26 @@ namespace Microsoft.Build.Shared
     internal sealed class CoreClrAssemblyLoader
     {
         private readonly Dictionary<string, Assembly> _pathsToAssemblies = new Dictionary<string, Assembly>(StringComparer.OrdinalIgnoreCase);
+        private readonly Dictionary<string, Assembly> _namesToAssemblies = new Dictionary<string, Assembly>();
+        private readonly HashSet<string> _dependencyPaths = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
         private readonly object _guard = new object();
 
+        private bool _resolvingHandlerHookedUp = false;
+
+        private static readonly Version _currentAssemblyVersion = new Version(Microsoft.Build.Shared.MSBuildConstants.CurrentAssemblyVersion);
+
+        public void AddDependencyLocation(string fullPath)
+        {
+            if (fullPath == null)
+            {
+                throw new ArgumentNullException(nameof(fullPath));
+            }
+
+            lock (_guard)
+            {
+                _dependencyPaths.Add(fullPath);
+            }
+        }
 
         public Assembly LoadFromPath(string fullPath)
         {
@@ -31,6 +52,38 @@ public Assembly LoadFromPath(string fullPath)
             // folders in a NuGet package).
             fullPath = FileUtilities.NormalizePath(fullPath);
 
+            if (Traits.Instance.EscapeHatches.UseSingleLoadContext)
+            {
+                return LoadUsingLegacyDefaultContext(fullPath);
+            }
+            else
+            {
+                return LoadUsingPluginContext(fullPath);
+            }
+        }
+
+        private Assembly LoadUsingLegacyDefaultContext(string fullPath)
+        {
+            lock (_guard)
+            {
+                if (!_resolvingHandlerHookedUp)
+                {
+                    AssemblyLoadContext.Default.Resolving += TryResolveAssembly;
+                    _resolvingHandlerHookedUp = true;
+                }
+
+                Assembly assembly;
+                if (_pathsToAssemblies.TryGetValue(fullPath, out assembly))
+                {
+                    return assembly;
+                }
+
+                return LoadAndCache(AssemblyLoadContext.Default, fullPath);
+            }
+        }
+
+        private Assembly LoadUsingPluginContext(string fullPath)
+        {
             lock (_guard)
             {
                 Assembly assembly;
@@ -51,5 +104,98 @@ public Assembly LoadFromPath(string fullPath)
                 return assembly;
             }
         }
+
+        private Assembly TryGetWellKnownAssembly(AssemblyLoadContext context, AssemblyName assemblyName)
+        {
+            if (!MSBuildLoadContext.WellKnownAssemblyNames.Contains(assemblyName.Name))
+            {
+                return null;
+            }
+
+            // Ensure we are attempting to load a matching version
+            // of the Microsoft.Build.* assembly.
+            assemblyName.Version = _currentAssemblyVersion;
+
+            var searchPaths = new[] { Assembly.GetExecutingAssembly().Location };
+            return TryResolveAssemblyFromPaths(context, assemblyName, searchPaths);
+        }
+
+        private Assembly TryResolveAssembly(AssemblyLoadContext context, AssemblyName assemblyName)
+        {
+            lock (_guard)
+            {
+                Assembly assembly = TryGetWellKnownAssembly(context, assemblyName);
+
+                if (assembly != null)
+                {
+                    return assembly;
+                }
+
+                if (_namesToAssemblies.TryGetValue(assemblyName.FullName, out assembly))
+                {
+                    return assembly;
+                }
+
+                return TryResolveAssemblyFromPaths(context, assemblyName, _dependencyPaths);
+            }
+        }
+
+        private Assembly TryResolveAssemblyFromPaths(AssemblyLoadContext context, AssemblyName assemblyName, IEnumerable<string> searchPaths)
+        {
+            foreach (var cultureSubfolder in string.IsNullOrEmpty(assemblyName.CultureName)
+                // If no culture is specified, attempt to load directly from
+                // the known dependency paths.
+                ? new[] { string.Empty }
+                // Search for satellite assemblies in culture subdirectories
+                // of the assembly search directories, but fall back to the
+                // bare search directory if that fails.
+                : new[] { assemblyName.CultureName, string.Empty })
+            {
+                foreach (var searchPath in searchPaths)
+                {
+                    foreach (var extension in MSBuildLoadContext.Extensions)
+                    {
+                        var candidatePath = Path.Combine(searchPath,
+                            cultureSubfolder,
+                            $"{assemblyName.Name}.{extension}");
+
+                        if (IsAssemblyAlreadyLoaded(candidatePath) ||
+                            !FileSystems.Default.FileExists(candidatePath))
+                        {
+                            continue;
+                        }
+
+                        AssemblyName candidateAssemblyName = AssemblyLoadContext.GetAssemblyName(candidatePath);
+                        if (candidateAssemblyName.Version != assemblyName.Version)
+                        {
+                            continue;
+                        }
+
+                        return LoadAndCache(context, candidatePath);
+                    }
+                }
+            }
+
+            return null;
+        }
+
+        /// <remarks>
+        /// Assumes we have a lock on _guard
+        /// </remarks>
+        private Assembly LoadAndCache(AssemblyLoadContext context, string fullPath)
+        {
+            var assembly = context.LoadFromAssemblyPath(fullPath);
+            var name = assembly.FullName;
+
+            _pathsToAssemblies[fullPath] = assembly;
+            _namesToAssemblies[name] = assembly;
+
+            return assembly;
+        }
+
+        private bool IsAssemblyAlreadyLoaded(string path)
+        {
+            return _pathsToAssemblies.ContainsKey(path);
+        }
     }
 }

src/Shared/MSBuildLoadContext.cs

@@ -19,9 +19,8 @@ namespace Microsoft.Build.Shared
     internal class MSBuildLoadContext : AssemblyLoadContext
     {
         private readonly string _directory;
-        private readonly object _guard = new object();
 
-        private static readonly ImmutableHashSet<string> _wellKnownAssemblyNames =
+        internal static readonly ImmutableHashSet<string> WellKnownAssemblyNames =
             new[]
             {
                 "MSBuild",
@@ -31,7 +30,7 @@ internal class MSBuildLoadContext : AssemblyLoadContext
                 "Microsoft.Build.Utilities.Core",
             }.ToImmutableHashSet();
 
-        private static readonly string[] _extensions = new[] { "ni.dll", "ni.exe", "dll", "exe" };
+        internal static readonly string[] Extensions = new[] { "ni.dll", "ni.exe", "dll", "exe" };
 
 
         public MSBuildLoadContext(string assemblyPath)
@@ -41,7 +40,7 @@ public MSBuildLoadContext(string assemblyPath)
 
         protected override Assembly? Load(AssemblyName assemblyName)
         {
-            if (_wellKnownAssemblyNames.Contains(assemblyName.Name!))
+            if (WellKnownAssemblyNames.Contains(assemblyName.Name!))
             {
                 // Force MSBuild assemblies to be loaded in the default ALC
                 // and unify to the current version.
@@ -57,7 +56,7 @@ public MSBuildLoadContext(string assemblyPath)
                 // bare search directory if that fails.
                 : new[] { assemblyName.CultureName, string.Empty })
             {
-                foreach (var extension in _extensions)
+                foreach (var extension in Extensions)
                 {
                     var candidatePath = Path.Combine(_directory,
                         cultureSubfolder,

src/Shared/Traits.cs

@@ -179,6 +179,11 @@ public bool LogProjectImports
         /// </summary>
         public readonly bool DisableNuGetSdkResolver = Environment.GetEnvironmentVariable("MSBUILDDISABLENUGETSDKRESOLVER") == "1";
 
+        /// <summary>
+        /// Disable AssemblyLoadContext isolation for plugins.
+        /// </summary>
+        public readonly bool UseSingleLoadContext = Environment.GetEnvironmentVariable("MSBUILDSINGLELOADCONTEXT") == "1";
+
         /// <summary>
         /// Enables the user of autorun functionality in CMD.exe on Windows which is disabled by default in MSBuild.
         /// </summary>

src/Shared/TypeLoader.cs

@@ -163,6 +163,8 @@ private static Assembly LoadAssembly(AssemblyLoadInfo assemblyLoadInfo)
 #if !FEATURE_ASSEMBLYLOADCONTEXT
                     loadedAssembly = Assembly.UnsafeLoadFrom(assemblyLoadInfo.AssemblyFile);
 #else
+                    var baseDir = Path.GetDirectoryName(assemblyLoadInfo.AssemblyFile);
+                    s_coreClrAssemblyLoader.AddDependencyLocation(baseDir);
                     loadedAssembly = s_coreClrAssemblyLoader.LoadFromPath(assemblyLoadInfo.AssemblyFile);
 #endif
                 }