Refactors LSP server extension assembly loading

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

@@ -2,98 +2,48 @@
 // 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.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) : IAssemblyLoader
 {
     /// <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 CustomExportAssemblyLoader(string baseDirectory)
-    {
-        _baseDirectory = baseDirectory;
-    }
-
     public Assembly LoadAssembly(AssemblyName assemblyName)
     {
-        Assembly? value;
-        lock (_loadedAssemblies)
+        // First attempt to load the assembly from the default context.
+        try
         {
-            _loadedAssemblies.TryGetValue(assemblyName, out value);
+            return AssemblyLoadContext.Default.LoadFromAssemblyName(assemblyName);
         }
-
-        if (value == null)
+        catch (FileNotFoundException) 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;
-                }
-            }
+            // continue checking the extension contexts.
+        }
 
-            lock (_loadedAssemblies)
-            {
-                _loadedAssemblies[assemblyName] = value;
-                return value;
-            }
+        var extensionAssembly = extensionAssemblyManager.TryLoadAssemblyInExtensionContext(assemblyName);
+        if (extensionAssembly != null)
+        {
+            return extensionAssembly;
         }
 
-        return value;
+        throw new FileNotFoundException($"Could not find assembly {assemblyName.Name} in any host or extension context.");
     }
 
     public Assembly LoadAssembly(string assemblyFullName, string? codeBasePath)
     {
         var assemblyName = new AssemblyName(assemblyFullName);
         return LoadAssembly(assemblyName);
     }
-
-    private class AssemblyNameComparer : IEqualityComparer<AssemblyName>
-    {
-        public static AssemblyNameComparer Instance = new AssemblyNameComparer();
-
-        public bool Equals(AssemblyName? x, AssemblyName? y)
-        {
-            if (x == null && y == null)
-            {
-                return true;
-            }
-
-            if (x == null || y == null)
-            {
-                return false;
-            }
-
-            return x.Name == y.Name;
-        }
-
-        public int GetHashCode([DisallowNull] AssemblyName obj)
-        {
-            return obj.Name?.GetHashCode(StringComparison.Ordinal) ?? 0;
-        }
-    }
 }

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

@@ -17,14 +17,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(IEnumerable<string> extensionAssemblyPaths, string? devKitDependencyPath, 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 +36,26 @@ public static async Task<ExportProvider> CreateExportProviderAsync(IEnumerable<s
             Assembly.LoadFrom(path);
         }
 
-        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
-        };
-
+        // Add the DevKit assembly to the MEF catalog
         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));
+            logger.LogTrace("loading roslyn devkit deps for MEF");
+            assemblyPaths = assemblyPaths.Concat(devKitDependencyPath);
         }
 
-        foreach (var extensionAssemblyPath in extensionAssemblyPaths)
-        {
-            if (AssemblyLoadContextWrapper.TryLoadExtension(extensionAssemblyPath, logger, out var extensionAssembly))
-            {
-                assemblies.Add(extensionAssembly);
-            }
-        }
+        // Add the extension assemblies to the MEF catalog
+        assemblyPaths = assemblyPaths.Concat(extensionAssemblyPaths);
+
+        // Create a MEF resolver that can resolve assemblies in the extension contexts.
+        var resolver = new Resolver(new CustomExportAssemblyLoader(extensionAssemblyManager));
+
+        var discovery = PartDiscovery.Combine(
+            resolver,
+            new AttributedPartDiscovery(resolver, isNonPublicSupported: true), // "NuGet MEF" attributes (Microsoft.Composition)
+            new AttributedPartDiscoveryV1(resolver));
 
         // 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 +78,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.

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,34 @@ 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>
+    /// <param name="extensionAssemblyManager"></param>
+    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);
+        }
     }
 }

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

@@ -13,6 +13,7 @@
 using Microsoft.CodeAnalysis.LanguageServer.HostWorkspace;
 using Microsoft.CodeAnalysis.LanguageServer.LanguageServer;
 using Microsoft.CodeAnalysis.LanguageServer.Logging;
+using Microsoft.CodeAnalysis.LanguageServer.Services;
 using Microsoft.CodeAnalysis.LanguageServer.StarredSuggestions;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Console;
@@ -76,7 +77,17 @@ static async Task RunAsync(ServerConfiguration serverConfiguration, Cancellation
 
     logger.LogTrace($".NET Runtime Version: {RuntimeInformation.FrameworkDescription}");
 
-    using var exportProvider = await ExportProviderBuilder.CreateExportProviderAsync(serverConfiguration.ExtensionAssemblyPaths, serverConfiguration.DevKitDependencyPath, loggerFactory);
+    var extensionAssemblyPaths = serverConfiguration.ExtensionAssemblyPaths.ToImmutableArray();
+    if (serverConfiguration.StarredCompletionsPath != null)
+    {
+        // TODO - the starred completion component should be passed as a regular extension.
+        // Currently we get passed the path to the directory, but ideally we should get passed the path to the dll as part of the --extension argument.
+        extensionAssemblyPaths = extensionAssemblyPaths.Add(StarredCompletionAssemblyHelper.GetStarredCompletionAssemblyPath(serverConfiguration.StarredCompletionsPath));
+    }
+
+    var extensionManager = ExtensionAssemblyManager.Create(extensionAssemblyPaths, serverConfiguration.DevKitDependencyPath, loggerFactory);
+
+    using var exportProvider = await ExportProviderBuilder.CreateExportProviderAsync(serverConfiguration.ExtensionAssemblyPaths, serverConfiguration.DevKitDependencyPath, extensionManager, loggerFactory);
 
     // The log file directory passed to us by VSCode might not exist yet, though its parent directory is guaranteed to exist.
     Directory.CreateDirectory(serverConfiguration.ExtensionLogDirectory);
@@ -96,10 +107,13 @@ static async Task RunAsync(ServerConfiguration serverConfiguration, Cancellation
         .Select(f => f.FullName)
         .ToImmutableArray();
 
-    await workspaceFactory.InitializeSolutionLevelAnalyzersAsync(analyzerPaths);
+    // Include analyzers from extension assemblies.
+    analyzerPaths = analyzerPaths.AddRange(serverConfiguration.ExtensionAssemblyPaths);
+
+    await workspaceFactory.InitializeSolutionLevelAnalyzersAsync(analyzerPaths, extensionManager);
 
     var serviceBrokerFactory = exportProvider.GetExportedValue<ServiceBrokerFactory>();
-    StarredCompletionAssemblyHelper.InitializeInstance(serverConfiguration.StarredCompletionsPath, loggerFactory, serviceBrokerFactory);
+    StarredCompletionAssemblyHelper.InitializeInstance(serverConfiguration.StarredCompletionsPath, extensionManager, loggerFactory, serviceBrokerFactory);
     // TODO: Remove, the path should match exactly. Workaround for https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1830914.
     Microsoft.CodeAnalysis.EditAndContinue.EditAndContinueMethodDebugInfoReader.IgnoreCaseWhenComparingDocumentNames = Path.DirectorySeparatorChar == '\\';