Cache the MEF composition in the Roslyn LSP.

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/ExportProviderBuilderTests.cs

@@ -0,0 +1,48 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Xunit.Abstractions;
+
+namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests
+{
+    public sealed class ExportProviderBuilderTests(ITestOutputHelper testOutputHelper)
+        : AbstractLanguageServerHostTests(testOutputHelper)
+    {
+        [Fact]
+        public async Task MefCompositionIsCached()
+        {
+            await using var testServer = await CreateLanguageServerAsync(includeDevKitComponents: false);
+
+            var mefCompositions = Directory.EnumerateFiles(MefCacheDirectory.Path, "*.mef-composition", SearchOption.AllDirectories);
+
+            Assert.Single(mefCompositions);
+        }
+
+        [Fact]
+        public async Task MefCompositionIsReused()
+        {
+            await using var testServer = await CreateLanguageServerAsync(includeDevKitComponents: false);
+
+            // Second test server with the same set of assemblies.
+            await using var testServer2 = await CreateLanguageServerAsync(includeDevKitComponents: false);
+
+            var mefCompositions = Directory.EnumerateFiles(MefCacheDirectory.Path, "*.mef-composition", SearchOption.AllDirectories);
+
+            Assert.Single(mefCompositions);
+        }
+
+        [Fact]
+        public async Task MultipleMefCompositionsAreCached()
+        {
+            await using var testServer = await CreateLanguageServerAsync(includeDevKitComponents: false);
+
+            // Second test server with a different set of assemblies.
+            await using var testServer2 = await CreateLanguageServerAsync(includeDevKitComponents: true);
+
+            var mefCompositions = Directory.EnumerateFiles(MefCacheDirectory.Path, "*.mef-composition", SearchOption.AllDirectories);
+
+            Assert.Equal(2, mefCompositions.Count());
+        }
+    }
+}

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/LspFileChangeWatcherTests.cs

@@ -14,7 +14,8 @@
 
 namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests;
 
-public class LspFileChangeWatcherTests : AbstractLanguageServerHostTests
+public class LspFileChangeWatcherTests(ITestOutputHelper testOutputHelper)
+    : AbstractLanguageServerHostTests(testOutputHelper)
 {
     private readonly ClientCapabilities _clientCapabilitiesWithFileWatcherSupport = new ClientCapabilities
     {
@@ -24,22 +25,18 @@ public class LspFileChangeWatcherTests : AbstractLanguageServerHostTests
         }
     };
 
-    public LspFileChangeWatcherTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper)
-    {
-    }
-
     [Fact]
     public async Task LspFileWatcherNotSupportedWithoutClientSupport()
     {
-        await using var testLspServer = await TestLspServer.CreateAsync(new ClientCapabilities(), TestOutputLogger);
+        await using var testLspServer = await TestLspServer.CreateAsync(new ClientCapabilities(), TestOutputLogger, MefCacheDirectory.Path);
 
         Assert.False(LspFileChangeWatcher.SupportsLanguageServerHost(testLspServer.LanguageServerHost));
     }
 
     [Fact]
     public async Task LspFileWatcherSupportedWithClientSupport()
     {
-        await using var testLspServer = await TestLspServer.CreateAsync(_clientCapabilitiesWithFileWatcherSupport, TestOutputLogger);
+        await using var testLspServer = await TestLspServer.CreateAsync(_clientCapabilitiesWithFileWatcherSupport, TestOutputLogger, MefCacheDirectory.Path);
 
         Assert.True(LspFileChangeWatcher.SupportsLanguageServerHost(testLspServer.LanguageServerHost));
     }
@@ -49,16 +46,15 @@ public async Task CreatingDirectoryWatchRequestsDirectoryWatch()
     {
         AsynchronousOperationListenerProvider.Enable(enable: true);
 
-        await using var testLspServer = await TestLspServer.CreateAsync(_clientCapabilitiesWithFileWatcherSupport, TestOutputLogger);
+        await using var testLspServer = await TestLspServer.CreateAsync(_clientCapabilitiesWithFileWatcherSupport, TestOutputLogger, MefCacheDirectory.Path);
         var lspFileChangeWatcher = new LspFileChangeWatcher(
             testLspServer.LanguageServerHost,
             testLspServer.ExportProvider.GetExportedValue<IAsynchronousOperationListenerProvider>());
 
         var dynamicCapabilitiesRpcTarget = new DynamicCapabilitiesRpcTarget();
         testLspServer.AddClientLocalRpcTarget(dynamicCapabilitiesRpcTarget);
 
-        using var tempRoot = new TempRoot();
-        var tempDirectory = tempRoot.CreateDirectory();
+        var tempDirectory = TempRoot.CreateDirectory();
 
         // Try creating a context and ensure we created the registration
         var context = lspFileChangeWatcher.CreateContext([new ProjectSystem.WatchedDirectory(tempDirectory.Path, extensionFilters: [])]);
@@ -80,16 +76,15 @@ public async Task CreatingFileWatchRequestsFileWatch()
     {
         AsynchronousOperationListenerProvider.Enable(enable: true);
 
-        await using var testLspServer = await TestLspServer.CreateAsync(_clientCapabilitiesWithFileWatcherSupport, TestOutputLogger);
+        await using var testLspServer = await TestLspServer.CreateAsync(_clientCapabilitiesWithFileWatcherSupport, TestOutputLogger, MefCacheDirectory.Path);
         var lspFileChangeWatcher = new LspFileChangeWatcher(
             testLspServer.LanguageServerHost,
             testLspServer.ExportProvider.GetExportedValue<IAsynchronousOperationListenerProvider>());
 
         var dynamicCapabilitiesRpcTarget = new DynamicCapabilitiesRpcTarget();
         testLspServer.AddClientLocalRpcTarget(dynamicCapabilitiesRpcTarget);
 
-        using var tempRoot = new TempRoot();
-        var tempDirectory = tempRoot.CreateDirectory();
+        var tempDirectory = TempRoot.CreateDirectory();
 
         // Try creating a single file watch and ensure we created the registration
         var context = lspFileChangeWatcher.CreateContext([]);

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/TelemetryReporterTests.cs

@@ -9,16 +9,17 @@
 
 namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests;
 
-public sealed class TelemetryReporterTests : AbstractLanguageServerHostTests
+public sealed class TelemetryReporterTests(ITestOutputHelper testOutputHelper)
+    : AbstractLanguageServerHostTests(testOutputHelper)
 {
-    public TelemetryReporterTests(ITestOutputHelper testOutputHelper)
-        : base(testOutputHelper)
-    {
-    }
-
     private async Task<ITelemetryReporter> CreateReporterAsync()
     {
-        var exportProvider = await LanguageServerTestComposition.CreateExportProviderAsync(TestOutputLogger.Factory, includeDevKitComponents: true, out var _, out var _);
+        var exportProvider = await LanguageServerTestComposition.CreateExportProviderAsync(
+            TestOutputLogger.Factory,
+            includeDevKitComponents: true,
+            MefCacheDirectory.Path,
+            out var _,
+            out var _);
 
         // VS Telemetry requires this environment variable to be set.
         Environment.SetEnvironmentVariable("CommonPropertyBagPath", Path.GetTempFileName());

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/Utilities/AbstractLanguageServerHostTests.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using Microsoft.CodeAnalysis.LanguageServer.LanguageServer;
+using Microsoft.CodeAnalysis.Test.Utilities;
 using Microsoft.VisualStudio.Composition;
 using Nerdbank.Streams;
 using Roslyn.LanguageServer.Protocol;
@@ -11,18 +12,27 @@
 
 namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests;
 
-public abstract class AbstractLanguageServerHostTests
+public abstract class AbstractLanguageServerHostTests : IDisposable
 {
     protected TestOutputLogger TestOutputLogger { get; }
+    protected TempRoot TempRoot { get; }
+    protected TempDirectory MefCacheDirectory { get; }
 
     protected AbstractLanguageServerHostTests(ITestOutputHelper testOutputHelper)
     {
         TestOutputLogger = new TestOutputLogger(testOutputHelper);
+        TempRoot = new();
+        MefCacheDirectory = TempRoot.CreateDirectory();
     }
 
     protected Task<TestLspServer> CreateLanguageServerAsync(bool includeDevKitComponents = true)
     {
-        return TestLspServer.CreateAsync(new ClientCapabilities(), TestOutputLogger, includeDevKitComponents);
+        return TestLspServer.CreateAsync(new ClientCapabilities(), TestOutputLogger, MefCacheDirectory.Path, includeDevKitComponents);
+    }
+
+    public void Dispose()
+    {
+        TempRoot.Dispose();
     }
 
     protected sealed class TestLspServer : IAsyncDisposable
@@ -32,10 +42,10 @@ protected sealed class TestLspServer : IAsyncDisposable
 
         private ServerCapabilities? _serverCapabilities;
 
-        internal static async Task<TestLspServer> CreateAsync(ClientCapabilities clientCapabilities, TestOutputLogger logger, bool includeDevKitComponents = true)
+        internal static async Task<TestLspServer> CreateAsync(ClientCapabilities clientCapabilities, TestOutputLogger logger, string cacheDirectory, bool includeDevKitComponents = true)
         {
             var exportProvider = await LanguageServerTestComposition.CreateExportProviderAsync(
-                logger.Factory, includeDevKitComponents, out var _, out var assemblyLoader);
+                logger.Factory, includeDevKitComponents, cacheDirectory, out var _, out var assemblyLoader);
             var testLspServer = new TestLspServer(exportProvider, logger, assemblyLoader);
             var initializeResponse = await testLspServer.ExecuteRequestAsync<InitializeParams, InitializeResult>(Methods.InitializeName, new InitializeParams { Capabilities = clientCapabilities }, CancellationToken.None);
             Assert.NotNull(initializeResponse?.Capabilities);

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

@@ -23,6 +23,7 @@ private static string GetDevKitExtensionPath()
     public static Task<ExportProvider> CreateExportProviderAsync(
         ILoggerFactory loggerFactory,
         bool includeDevKitComponents,
+        string cacheDirectory,
         out ServerConfiguration serverConfiguration,
         out IAssemblyLoader assemblyLoader)
     {
@@ -39,6 +40,6 @@ public static Task<ExportProvider> CreateExportProviderAsync(
             ExtensionLogDirectory: string.Empty);
         var extensionManager = ExtensionAssemblyManager.Create(serverConfiguration, loggerFactory);
         assemblyLoader = new CustomExportAssemblyLoader(extensionManager, loggerFactory);
-        return ExportProviderBuilder.CreateExportProviderAsync(extensionManager, assemblyLoader, devKitDependencyPath, loggerFactory);
+        return ExportProviderBuilder.CreateExportProviderAsync(extensionManager, assemblyLoader, devKitDependencyPath, cacheDirectory, loggerFactory);
     }
 }

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/WorkspaceProjectFactoryServiceTests.cs

@@ -7,17 +7,19 @@
 using Microsoft.CodeAnalysis.Remote.ProjectSystem;
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.Shell.ServiceBroker;
+using Xunit.Abstractions;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests;
 
-public class WorkspaceProjectFactoryServiceTests
+public class WorkspaceProjectFactoryServiceTests(ITestOutputHelper testOutputHelper)
+    : AbstractLanguageServerHostTests(testOutputHelper)
 {
     [Fact]
     public async Task CreateProjectAndBatch()
     {
         var loggerFactory = new LoggerFactory();
         using var exportProvider = await LanguageServerTestComposition.CreateExportProviderAsync(
-            loggerFactory, includeDevKitComponents: false, out var serverConfiguration, out var _);
+            loggerFactory, includeDevKitComponents: false, MefCacheDirectory.Path, out var serverConfiguration, out var _);
 
         exportProvider.GetExportedValue<ServerConfigurationFactory>()
             .InitializeConfiguration(serverConfiguration);

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

@@ -3,6 +3,8 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Immutable;
+using System.IO.Hashing;
+using System.Text;
 using Microsoft.CodeAnalysis.LanguageServer.Logging;
 using Microsoft.CodeAnalysis.LanguageServer.Services;
 using Microsoft.CodeAnalysis.Shared.Collections;
@@ -18,6 +20,7 @@ public static async Task<ExportProvider> CreateExportProviderAsync(
         ExtensionAssemblyManager extensionManager,
         IAssemblyLoader assemblyLoader,
         string? devKitDependencyPath,
+        string cacheDirectory,
         ILoggerFactory loggerFactory)
     {
         var logger = loggerFactory.CreateLogger<ExportProviderBuilder>();
@@ -38,17 +41,60 @@ public static async Task<ExportProvider> CreateExportProviderAsync(
         // Add the extension assemblies to the MEF catalog.
         assemblyPaths = assemblyPaths.Concat(extensionManager.ExtensionAssemblyPaths);
 
-        logger.LogTrace($"Composing MEF catalog using:{Environment.NewLine}{string.Join($"    {Environment.NewLine}", assemblyPaths)}.");
+        // Get the cached MEF composition or create a new one.
+        var exportProviderFactory = await GetCompositionConfigurationAsync(assemblyPaths.ToImmutableArray(), assemblyLoader, cacheDirectory, logger);
+
+        // Create an export provider, which represents a unique container of values.
+        // You can create as many of these as you want, but typically an app needs just one.
+        var exportProvider = exportProviderFactory.CreateExportProvider();
+
+        // Immediately set the logger factory, so that way it'll be available for the rest of the composition
+        exportProvider.GetExportedValue<ServerLoggerFactory>().SetFactory(loggerFactory);
 
+        // Also add the ExtensionAssemblyManager so it will be available for the rest of the composition.
+        exportProvider.GetExportedValue<ExtensionAssemblyManagerMefProvider>().SetMefExtensionAssemblyManager(extensionManager);
+
+        return exportProvider;
+    }
+
+    private static async Task<IExportProviderFactory> GetCompositionConfigurationAsync(
+        ImmutableArray<string> assemblyPaths,
+        IAssemblyLoader assemblyLoader,
+        string cacheDirectory,
+        ILogger logger)
+    {
         // Create a MEF resolver that can resolve assemblies in the extension contexts.
         var resolver = new Resolver(assemblyLoader);
 
+        var compositionCacheFile = GetCompositionCacheFilePath(cacheDirectory, assemblyPaths);
+
+        // Try to load a cached composition.
+        try
+        {
+            if (File.Exists(compositionCacheFile))
+            {
+                logger.LogTrace($"Loading cached MEF catalog: {compositionCacheFile}");
+
+                CachedComposition cachedComposition = new();
+                using FileStream cacheStream = new(compositionCacheFile, FileMode.Open, FileAccess.Read, FileShare.Read, bufferSize: 4096, useAsync: true);
+                var exportProviderFactory = await cachedComposition.LoadExportProviderFactoryAsync(cacheStream, resolver);
+
+                return exportProviderFactory;
+            }
+        }
+        catch (Exception ex)
+        {
+            // Log the error, and move on to recover by recreating the MEF composition.
+            logger.LogError($"Loading cached MEF composition failed: {ex}");
+        }
+
+        logger.LogTrace($"Composing MEF catalog using:{Environment.NewLine}{string.Join($"    {Environment.NewLine}", assemblyPaths)}.");
+
         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(assemblyPaths))
             .WithCompositionService(); // Makes an ICompositionService export available to MEF parts to import
@@ -59,20 +105,69 @@ public static async Task<ExportProvider> CreateExportProviderAsync(
         // Verify we only have expected errors.
         ThrowOnUnexpectedErrors(config, catalog, logger);
 
+        // Try to cache the composition.
+        _ = WriteCompositionCacheAsync(compositionCacheFile, config, logger).ReportNonFatalErrorAsync();
+
         // Prepare an ExportProvider factory based on this graph.
-        var exportProviderFactory = config.CreateExportProviderFactory();
+        return config.CreateExportProviderFactory();
+    }
 
-        // Create an export provider, which represents a unique container of values.
-        // You can create as many of these as you want, but typically an app needs just one.
-        var exportProvider = exportProviderFactory.CreateExportProvider();
+    private static string GetCompositionCacheFilePath(string cacheDirectory, ImmutableArray<string> assemblyPaths)
+    {
+        // This should vary based on .NET runtime major version so that as some of our processes switch between our target
+        // .NET version and the user's selected SDK runtime version (which may be newer), the MEF cache is kept isolated.
+        // This can be important when the MEF catalog records full assembly names such as "System.Runtime, 8.0.0.0" yet
+        // we might be running on .NET 7 or .NET 8, depending on the particular session and user settings.
+        var cacheSubdirectory = $".NET {Environment.Version.Major}";
 
-        // Immediately set the logger factory, so that way it'll be available for the rest of the composition
-        exportProvider.GetExportedValue<ServerLoggerFactory>().SetFactory(loggerFactory);
+        return Path.Combine(cacheDirectory, cacheSubdirectory, $"c#-languageserver.{ComputeAssemblyHash(assemblyPaths)}.mef-composition");
 
-        // Also add the ExtensionAssemblyManager so it will be available for the rest of the composition.
-        exportProvider.GetExportedValue<ExtensionAssemblyManagerMefProvider>().SetMefExtensionAssemblyManager(extensionManager);
+        static string ComputeAssemblyHash(ImmutableArray<string> assemblyPaths)
+        {
+            // Ensure AssemblyPaths are always in the same order.
+            assemblyPaths = assemblyPaths.Sort();
 
-        return exportProvider;
+            var assemblies = new StringBuilder();
+            foreach (var assemblyPath in assemblyPaths)
+            {
+                // Include assembly path in the hash so that changes to the set of included
+                // assemblies cause the composition to be rebuilt.
+                assemblies.Append(assemblyPath);
+                // Include the last write time in the hash so that newer assemblies written
+                // to the same location cause the composition to be rebuilt.
+                assemblies.Append(File.GetLastWriteTimeUtc(assemblyPath).ToString("F"));
+            }
+
+            var hash = XxHash128.Hash(Encoding.UTF8.GetBytes(assemblies.ToString()));
+            // Convert to filename safe base64 string.
+            return Convert.ToBase64String(hash).Replace('+', '-').Replace('/', '_').TrimEnd('=');
+        }
+    }
+
+    private static async Task WriteCompositionCacheAsync(string compositionCacheFile, CompositionConfiguration config, ILogger logger)
+    {
+        try
+        {
+            await Task.Yield();
+
+            if (Path.GetDirectoryName(compositionCacheFile) is string directory)
+            {
+                Directory.CreateDirectory(directory);
+            }
+
+            CachedComposition cachedComposition = new();
+            var tempFilePath = Path.Combine(Path.GetTempPath(), Path.GetTempFileName());
+            using (FileStream cacheStream = new(tempFilePath, FileMode.Create, FileAccess.Write, FileShare.None, bufferSize: 4096, useAsync: true))
+            {
+                await cachedComposition.SaveAsync(config, cacheStream);
+            }
+
+            File.Move(tempFilePath, compositionCacheFile, overwrite: true);
+        }
+        catch (Exception ex)
+        {
+            logger.LogError($"Failed to save MEF cache: {ex}");
+        }
     }
 
     private static void ThrowOnUnexpectedErrors(CompositionConfiguration configuration, ComposableCatalog catalog, ILogger logger)