Cache the MEF composition in the Roslyn LSP.

Author: JoeRobich

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

@@ -39,6 +39,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: null, loggerFactory);
     }
 }

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.Security.Cryptography;
+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);
 
+        string? compositionCacheFile = cacheDirectory is not null
+            ? GetCompositionCacheFilePath(cacheDirectory, assemblyPaths)
+            : null;
+
+        // Try to load a cached composition.
+        try
+        {
+            if (compositionCacheFile is not null && File.Exists(compositionCacheFile))
+            {
+                logger.LogTrace($"Loading cached MEF catalog: {compositionCacheFile}");
+
+                CachedComposition cachedComposition = new();
+                using FileStream cacheStream = new(compositionCacheFile, FileMode.Open, FileAccess.Read, FileShare.Read, 4096, useAsync: true);
+                return await cachedComposition.LoadExportProviderFactoryAsync(cacheStream, resolver);
+            }
+        }
+        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,52 @@ public static async Task<ExportProvider> CreateExportProviderAsync(
         // Verify we only have expected errors.
         ThrowOnUnexpectedErrors(config, catalog, logger);
 
-        // Prepare an ExportProvider factory based on this graph.
-        var exportProviderFactory = config.CreateExportProviderFactory();
+        // Try to cache the composition.
+        if (compositionCacheFile is not null)
+        {
+            if (Path.GetDirectoryName(compositionCacheFile) is string directory)
+            {
+                Directory.CreateDirectory(directory);
+            }
 
-        // 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();
+            CachedComposition cachedComposition = new();
+            var tempFilePath = Path.Combine(Path.GetTempPath(), Path.GetTempFileName());
+            using (FileStream cacheStream = new(tempFilePath, FileMode.Create, FileAccess.Write, FileShare.None, 4096, useAsync: true))
+            {
+                await cachedComposition.SaveAsync(config, cacheStream);
+            }
 
-        // Immediately set the logger factory, so that way it'll be available for the rest of the composition
-        exportProvider.GetExportedValue<ServerLoggerFactory>().SetFactory(loggerFactory);
+            File.Move(tempFilePath, compositionCacheFile, overwrite: true);
+        }
 
-        // Also add the ExtensionAssemblyManager so it will be available for the rest of the composition.
-        exportProvider.GetExportedValue<ExtensionAssemblyManagerMefProvider>().SetMefExtensionAssemblyManager(extensionManager);
+        // Prepare an ExportProvider factory based on this graph.
+        return config.CreateExportProviderFactory();
+    }
 
-        return exportProvider;
+    private static string GetCompositionCacheFilePath(string cacheDirectory, ImmutableArray<string> assemblyPaths)
+    {
+        // Include the .NET runtime version in the cache path so that running on a newer
+        // runtime causes the cache to be rebuilt.
+        var cacheSubdirectory = $".NET {Environment.Version.Major}";
+        return Path.Combine(cacheDirectory, cacheSubdirectory, $"c#-languageserver.{ComputeAssemblyHash(assemblyPaths)}.mef-composition");
+
+        static string ComputeAssemblyHash(ImmutableArray<string> assemblyPaths)
+        {
+            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 = SHA256.HashData(Encoding.UTF8.GetBytes(assemblies.ToString()));
+            // Convert to filename safe base64 string.
+            return Convert.ToBase64String(hash).Replace('+', '-').Replace('/', '_').TrimEnd('=');
+        }
     }
 
     private static void ThrowOnUnexpectedErrors(CompositionConfiguration configuration, ComposableCatalog catalog, ILogger logger)

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

@@ -86,7 +86,9 @@ static async Task RunAsync(ServerConfiguration serverConfiguration, Cancellation
     var assemblyLoader = new CustomExportAssemblyLoader(extensionManager, loggerFactory);
     var typeRefResolver = new ExtensionTypeRefResolver(assemblyLoader, loggerFactory);
 
-    using var exportProvider = await ExportProviderBuilder.CreateExportProviderAsync(extensionManager, assemblyLoader, serverConfiguration.DevKitDependencyPath, loggerFactory);
+    var cacheDirectory = Path.Combine(Path.GetDirectoryName(typeof(Program).Assembly.Location)!, "cache");
+
+    using var exportProvider = await ExportProviderBuilder.CreateExportProviderAsync(extensionManager, assemblyLoader, serverConfiguration.DevKitDependencyPath, cacheDirectory, loggerFactory);
 
     // LSP server doesn't have the pieces yet to support 'balanced' mode for source-generators.  Hardcode us to
     // 'automatic' for now.