Refactors LSP server extension assembly loading

src/Features/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/Utilities/LanguageServerTestComposition.cs

@@ -2,6 +2,9 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Immutable;
+using System.Runtime.Loader;
+using Microsoft.CodeAnalysis.LanguageServer.Services;
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.Composition;
 
@@ -20,5 +23,16 @@ private static string GetDevKitExtensionPath()
         => Path.Combine(AppContext.BaseDirectory, DevKitExtensionSubdirectory, DevKitAssemblyFileName);
 
     public static Task<ExportProvider> CreateExportProviderAsync(ILoggerFactory loggerFactory, bool includeDevKitComponents)
-        => ExportProviderBuilder.CreateExportProviderAsync(extensionAssemblyPaths: [], devKitDependencyPath: includeDevKitComponents ? GetDevKitExtensionPath() : null, loggerFactory: loggerFactory);
+    {
+        var serverConfiguration = new ServerConfiguration(LaunchDebugger: false,
+            MinimumLogLevel: LogLevel.Trace,
+            StarredCompletionsPath: null,
+            TelemetryLevel: null,
+            SessionId: null,
+            ExtensionAssemblyPaths: [],
+            DevKitDependencyPath: includeDevKitComponents ? GetDevKitExtensionPath() : null,
+            ExtensionLogDirectory: string.Empty);
+        var extensionAssemblyManager = ExtensionAssemblyManager.Create(serverConfiguration, loggerFactory);
+        return ExportProviderBuilder.CreateExportProviderAsync(extensionAssemblyManager, loggerFactory: loggerFactory);
+    }
 }

src/Features/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/CustomExportAssemblyLoader.cs

@@ -2,98 +2,95 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System.Diagnostics.CodeAnalysis;
 using System.Reflection;
+using System.Runtime.Loader;
+using Microsoft.CodeAnalysis.LanguageServer.Services;
+using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.Composition;
 
 namespace Microsoft.CodeAnalysis.LanguageServer;
 
-internal class CustomExportAssemblyLoader : IAssemblyLoader
+/// <summary>
+/// Defines a MEF assembly loader that knows how to load assemblies from both the default assembly load context
+/// and from the assembly load contexts for any of our extensions.
+/// </summary>
+internal class CustomExportAssemblyLoader(ExtensionAssemblyManager extensionAssemblyManager, ILoggerFactory loggerFactory) : IAssemblyLoader
 {
+    private readonly ILogger _logger = loggerFactory.CreateLogger("MEF Assembly Loader");
     /// <summary>
-    /// Cache assemblies that are already loaded by AssemblyName comparison
+    /// Loads assemblies from either the host or from our extensions.
+    /// If an assembly exists in both the host and an extension, we will use the host assembly for the MEF catalog.
+    /// If an assembly exists in two extensions, we use the first one we find for the MEF catalog.
     /// </summary>
-    private readonly Dictionary<AssemblyName, Assembly> _loadedAssemblies = new Dictionary<AssemblyName, Assembly>(AssemblyNameComparer.Instance);
-
-    /// <summary>
-    /// Base directory to search for <see cref="Assembly.LoadFrom(string)"/> if initial load fails
-    /// </summary>
-    private readonly string _baseDirectory;
+    public Assembly LoadAssembly(AssemblyName assemblyName)
+    {
+        // VS-MEF generally tries to populate AssemblyName.CodeBase with the path to the assembly being loaded.
+        // We need to read this in order to figure out which ALC we should load the assembly into.
+#pragma warning disable SYSLIB0044 // Type or member is obsolete
+        var codeBasePath = assemblyName.CodeBase;
+#pragma warning restore SYSLIB0044 // Type or member is obsolete
+        return LoadAssembly(assemblyName, codeBasePath);
+    }
 
-    public CustomExportAssemblyLoader(string baseDirectory)
+    public Assembly LoadAssembly(string assemblyFullName, string? codeBasePath)
     {
-        _baseDirectory = baseDirectory;
+        var assemblyName = new AssemblyName(assemblyFullName);
+        return LoadAssembly(assemblyName, codeBasePath);
     }
 
-    public Assembly LoadAssembly(AssemblyName assemblyName)
+    private Assembly LoadAssembly(AssemblyName assemblyName, string? codeBasePath)
     {
-        Assembly? value;
-        lock (_loadedAssemblies)
+        _logger.LogTrace($"Loading assembly {assemblyName}");
+        // First attempt to load the assembly from the default context.
+        Exception loadException;
+        try
         {
-            _loadedAssemblies.TryGetValue(assemblyName, out value);
+            return AssemblyLoadContext.Default.LoadFromAssemblyName(assemblyName);
         }
-
-        if (value == null)
+        catch (FileNotFoundException ex) when (assemblyName.Name is not null)
         {
-            // Attempt to load the assembly normally, but fall back to Assembly.LoadFrom in the base
-            // directory if the assembly load fails
-            try
-            {
-                value = Assembly.Load(assemblyName);
-            }
-            catch (FileNotFoundException) when (assemblyName.Name is not null)
-            {
-                var filePath = Path.Combine(_baseDirectory, assemblyName.Name)
-                    + (assemblyName.Name.EndsWith(".dll", StringComparison.OrdinalIgnoreCase)
-                        ? ""
-                        : ".dll");
-
-                value = Assembly.LoadFrom(filePath);
-
-                if (value is null)
-                {
-                    throw;
-                }
-            }
+            loadException = ex;
+            // continue checking the extension contexts.
+        }
 
-            lock (_loadedAssemblies)
-            {
-                _loadedAssemblies[assemblyName] = value;
-                return value;
-            }
+        if (codeBasePath is not null)
+        {
+            return LoadAssemblyFromCodeBase(assemblyName, codeBasePath);
         }
 
-        return value;
-    }
+        // We don't have a code base path for this assembly.  We'll search the extension contexts
+        // and load the first one that ships the assembly.
+        var assembly = extensionAssemblyManager.SearchExtensionContextsForAssembly(assemblyName);
+        if (assembly is not null)
+        {
+            _logger.LogTrace("{assemblyName} found in extension context without code base", assemblyName);
+            return assembly;
+        }
 
-    public Assembly LoadAssembly(string assemblyFullName, string? codeBasePath)
-    {
-        var assemblyName = new AssemblyName(assemblyFullName);
-        return LoadAssembly(assemblyName);
+        _logger.LogTrace("{assemblyName} not found in any host or extension context", assemblyName);
+        throw loadException;
     }
 
-    private class AssemblyNameComparer : IEqualityComparer<AssemblyName>
+    private Assembly LoadAssemblyFromCodeBase(AssemblyName assemblyName, string codeBaseUriStr)
     {
-        public static AssemblyNameComparer Instance = new AssemblyNameComparer();
-
-        public bool Equals(AssemblyName? x, AssemblyName? y)
+        // CodeBase is spec'd as being a URL string.
+        var codeBaseUri = ProtocolConversions.CreateAbsoluteUri(codeBaseUriStr);
+        if (!codeBaseUri.IsFile)
         {
-            if (x == null && y == null)
-            {
-                return true;
-            }
-
-            if (x == null || y == null)
-            {
-                return false;
-            }
-
-            return x.Name == y.Name;
+            throw new ArgumentException($"Code base {codeBaseUriStr} is not a file URI.", codeBaseUriStr);
         }
 
-        public int GetHashCode([DisallowNull] AssemblyName obj)
+        var codeBasePath = codeBaseUri.LocalPath;
+
+        var assembly = extensionAssemblyManager.TryLoadAssemblyInExtensionContext(codeBasePath);
+        if (assembly is not null)
         {
-            return obj.Name?.GetHashCode(StringComparison.Ordinal) ?? 0;
+            _logger.LogTrace("{assemblyName} with code base {codeBase} found in extension context", assemblyName, codeBasePath);
+            return assembly;
         }
+
+        // We were given an explicit code base path, but no extension context had the assembly.
+        // This is unexpected, so we'll throw an exception.
+        throw new FileNotFoundException($"Could not find assembly {assemblyName} with code base {codeBasePath} in any extension context.");
     }
 }

src/Features/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/ExportProviderBuilder.cs

@@ -4,10 +4,8 @@
 
 using System.Collections.Immutable;
 using System.Reflection;
-using System.Runtime.Loader;
 using Microsoft.CodeAnalysis.LanguageServer.Logging;
 using Microsoft.CodeAnalysis.LanguageServer.Services;
-using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.Composition;
@@ -17,14 +15,11 @@ namespace Microsoft.CodeAnalysis.LanguageServer;
 
 internal sealed class ExportProviderBuilder
 {
-    public static async Task<ExportProvider> CreateExportProviderAsync(IEnumerable<string> extensionAssemblyPaths, string? devKitDependencyPath, ILoggerFactory loggerFactory)
+    public static async Task<ExportProvider> CreateExportProviderAsync(ExtensionAssemblyManager extensionAssemblyManager, ILoggerFactory loggerFactory)
     {
         var logger = loggerFactory.CreateLogger<ExportProviderBuilder>();
-
         var baseDirectory = AppContext.BaseDirectory;
 
-        var resolver = new Resolver(new CustomExportAssemblyLoader(baseDirectory));
-
         // Load any Roslyn assemblies from the extension directory
         var assemblyPaths = Directory.EnumerateFiles(baseDirectory, "Microsoft.CodeAnalysis*.dll");
         assemblyPaths = assemblyPaths.Concat(Directory.EnumerateFiles(baseDirectory, "Microsoft.ServiceHub*.dll"));
@@ -39,34 +34,31 @@ public static async Task<ExportProvider> CreateExportProviderAsync(IEnumerable<s
             Assembly.LoadFrom(path);
         }
 
+        // DevKit assemblies are not shipped in the main language server folder
+        // and not included in ExtensionAssemblyPaths (they get loaded into the default ALC).
+        // So manually add them to the MEF catalog here.
+        if (extensionAssemblyManager.DevKitDependencyPath != null)
+        {
+            assemblyPaths = assemblyPaths.Concat(extensionAssemblyManager.DevKitDependencyPath);
+        }
+
+        // Add the extension assemblies to the MEF catalog.
+        assemblyPaths = assemblyPaths.Concat(extensionAssemblyManager.ExtensionAssemblyPaths);
+
+        ValidateNoDuplicateAssemblies(assemblyPaths, logger);
+
+        logger.LogTrace($"Composing MEF catalog using:{Environment.NewLine}{string.Join($"    {Environment.NewLine}", assemblyPaths)}.");
+
+        // Create a MEF resolver that can resolve assemblies in the extension contexts.
+        var resolver = new Resolver(new CustomExportAssemblyLoader(extensionAssemblyManager, loggerFactory));
+
         var discovery = PartDiscovery.Combine(
             resolver,
             new AttributedPartDiscovery(resolver, isNonPublicSupported: true), // "NuGet MEF" attributes (Microsoft.Composition)
             new AttributedPartDiscoveryV1(resolver));
 
-        var assemblies = new List<Assembly>()
-        {
-            typeof(ExportProviderBuilder).Assembly
-        };
-
-        if (devKitDependencyPath != null)
-        {
-            // Load devkit dependencies before other extensions to ensure dependencies
-            // like VS Telemetry are available from the host.
-            assemblies.AddRange(LoadDevKitAssemblies(devKitDependencyPath, logger));
-        }
-
-        foreach (var extensionAssemblyPath in extensionAssemblyPaths)
-        {
-            if (AssemblyLoadContextWrapper.TryLoadExtension(extensionAssemblyPath, logger, out var extensionAssembly))
-            {
-                assemblies.Add(extensionAssembly);
-            }
-        }
-
         // TODO - we should likely cache the catalog so we don't have to rebuild it every time.
         var catalog = ComposableCatalog.Create(resolver)
-            .AddParts(await discovery.CreatePartsAsync(assemblies))
             .AddParts(await discovery.CreatePartsAsync(assemblyPaths))
             .WithCompositionService(); // Makes an ICompositionService export available to MEF parts to import
 
@@ -89,25 +81,6 @@ public static async Task<ExportProvider> CreateExportProviderAsync(IEnumerable<s
         return exportProvider;
     }
 
-    private static ImmutableArray<Assembly> LoadDevKitAssemblies(string devKitDependencyPath, ILogger logger)
-    {
-        var directoryName = Path.GetDirectoryName(devKitDependencyPath);
-        Contract.ThrowIfNull(directoryName);
-        logger.LogTrace("Loading DevKit assemblies from {directory}", directoryName);
-
-        var directory = new DirectoryInfo(directoryName);
-        using var _ = ArrayBuilder<Assembly>.GetInstance(out var builder);
-        foreach (var file in directory.GetFiles("*.dll"))
-        {
-            logger.LogTrace("Loading {assemblyName}", file.Name);
-            // DevKit assemblies are loaded into the default load context. This allows extensions
-            // to share the host's instance of these assemblies as long as they do not ship their own copy.
-            builder.Add(AssemblyLoadContext.Default.LoadFromAssemblyPath(file.FullName));
-        }
-
-        return builder.ToImmutable();
-    }
-
     private static void ThrowOnUnexpectedErrors(CompositionConfiguration configuration, ILogger logger)
     {
         // Verify that we have exactly the MEF errors that we expect.  If we have less or more this needs to be updated to assert the expected behavior.
@@ -134,4 +107,22 @@ private static void ThrowOnUnexpectedErrors(CompositionConfiguration configurati
             }
         }
     }
+
+    private static void ValidateNoDuplicateAssemblies(IEnumerable<string> assemblyPaths, ILogger logger)
+    {
+        // MEF relies on type full names (*without* their assembly names) to identify parts.
+        // If an assembly is added to the catalog twice then we will almost certainly get duplicate MEF parts
+        // which breaks consumers who are only expecting one part.
+        //
+        // We validate this constraint here by checking for duplicate assembly full names in the MEF composition.
+        var duplicateAssemblyNames = assemblyPaths
+            .Select(p => AssemblyName.GetAssemblyName(p).FullName)
+            .GroupBy(n => n)
+            .Where(key => key.Count() > 1);
+        if (duplicateAssemblyNames.Any())
+        {
+            logger.LogError("Found duplicate assemblies in the MEF composition:{line}:{assemblies}", Environment.NewLine, string.Join($"    {Environment.NewLine}", duplicateAssemblyNames.Select(g => g.Key)));
+            throw new InvalidOperationException("Found duplicate assemblies in the MEF composition");
+        }
+    }
 }

src/Features/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/HostWorkspace/LanguageServerWorkspaceFactory.cs

@@ -7,8 +7,8 @@
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
-using Microsoft.CodeAnalysis.LanguageServer.ExternalAccess.VSCode.API;
 using Microsoft.CodeAnalysis.LanguageServer.Handler.DebugConfiguration;
+using Microsoft.CodeAnalysis.LanguageServer.Services;
 using Microsoft.CodeAnalysis.ProjectSystem;
 using Microsoft.CodeAnalysis.Workspaces.ProjectSystem;
 using Microsoft.Extensions.Logging;
@@ -54,10 +54,10 @@ public LanguageServerWorkspaceFactory(
     public ProjectSystemHostInfo ProjectSystemHostInfo { get; }
     public ProjectTargetFrameworkManager TargetFrameworkManager { get; }
 
-    public async Task InitializeSolutionLevelAnalyzersAsync(ImmutableArray<string> analyzerPaths)
+    public async Task InitializeSolutionLevelAnalyzersAsync(ImmutableArray<string> analyzerPaths, ExtensionAssemblyManager extensionAssemblyManager)
     {
         var references = new List<AnalyzerFileReference>();
-        var analyzerLoader = VSCodeAnalyzerLoader.CreateAnalyzerAssemblyLoader();
+        var analyzerLoader = VSCodeAnalyzerLoader.CreateAnalyzerAssemblyLoader(extensionAssemblyManager, _logger);
 
         foreach (var analyzerPath in analyzerPaths)
         {

src/Features/LanguageServer/Protocol/ExternalAccess/VSCode/API/VSCodeAnalyzerLoader.cs → src/Features/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/HostWorkspace/VSCodeAnalyzerLoader.cs

@@ -2,14 +2,17 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System;
+using System.Collections.Immutable;
 using System.Composition;
+using System.Reflection;
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.LanguageServer.Services;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.SolutionCrawler;
+using Microsoft.Extensions.Logging;
 
-namespace Microsoft.CodeAnalysis.LanguageServer.ExternalAccess.VSCode.API;
+namespace Microsoft.CodeAnalysis.LanguageServer.HostWorkspace;
 
 [Export(typeof(VSCodeAnalyzerLoader)), Shared]
 internal class VSCodeAnalyzerLoader
@@ -34,8 +37,33 @@ public void InitializeDiagnosticsServices(Workspace workspace)
         _diagnosticService.Register((IDiagnosticUpdateSource)_analyzerService);
     }
 
-    public static IAnalyzerAssemblyLoader CreateAnalyzerAssemblyLoader()
+    public static IAnalyzerAssemblyLoader CreateAnalyzerAssemblyLoader(ExtensionAssemblyManager extensionAssemblyManager, ILogger logger)
     {
-        return new DefaultAnalyzerAssemblyLoader();
+        return new VSCodeExtensionAssemblyAnalyzerLoader(extensionAssemblyManager, logger);
+    }
+
+    /// <summary>
+    /// Analyzer loader that will re-use already loaded assemblies from the extension load context.
+    /// </summary>
+    private class VSCodeExtensionAssemblyAnalyzerLoader(ExtensionAssemblyManager extensionAssemblyManager, ILogger logger) : IAnalyzerAssemblyLoader
+    {
+        private readonly DefaultAnalyzerAssemblyLoader _defaultLoader = new();
+
+        public void AddDependencyLocation(string fullPath)
+        {
+            _defaultLoader.AddDependencyLocation(fullPath);
+        }
+
+        public Assembly LoadFromPath(string fullPath)
+        {
+            var assembly = extensionAssemblyManager.TryLoadAssemblyInExtensionContext(fullPath);
+            if (assembly is not null)
+            {
+                logger.LogTrace("Loaded analyzer {fullPath} from extension context", fullPath);
+                return assembly;
+            }
+
+            return _defaultLoader.LoadFromPath(fullPath);
+        }
     }
 }