Cache MEF composition in OOP and VS Code (#12041)

Fixes #11830

Code mostly stolen from Todd who stole it from Joey. Mistakes are all me
though 😛

Author: davidwengier

eng/targets/Services.props

@@ -19,6 +19,7 @@
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.SemanticTokens" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteSemanticTokensService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.HtmlDocument" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteHtmlDocumentService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.TagHelperProvider" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteTagHelperProviderService+Factory"/>
+    <ServiceHubService Include="Microsoft.VisualStudio.Razor.MEFInitialization" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteMEFInitializationService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.ClientInitialization" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteClientInitializationService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.UriPresentation" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteUriPresentationService+Factory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.FoldingRange" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteFoldingRangeService+Factory" />

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteMEFInitializationService.cs

@@ -0,0 +1,12 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.CodeAnalysis.Razor.Remote;
+
+internal interface IRemoteMEFInitializationService : IRemoteJsonService
+{
+    ValueTask InitializeAsync(string cacheDirectory, CancellationToken cancellationToken);
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/RazorServices.cs

@@ -43,6 +43,7 @@ internal static class RazorServices
             (typeof(IRemoteCompletionService), null),
             (typeof(IRemoteCodeActionsService), null),
             (typeof(IRemoteFindAllReferencesService), null),
+            (typeof(IRemoteMEFInitializationService), null),
         ];
 
     private const string ComponentName = "Razor";

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Initialization/RemoteMEFInitializationService.cs

@@ -0,0 +1,53 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Diagnostics;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Api;
+using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.ServiceHub.Framework;
+using static Microsoft.CodeAnalysis.Remote.Razor.RazorBrokeredServiceBase;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+/// <summary>
+/// A special service that is used to initialize the MEF composition for Razor in the remote host.
+/// </summary>
+/// <remarks>
+/// It's special because it doesn't use MEF. Nor can it use anything else really.
+/// </remarks>
+internal sealed class RemoteMEFInitializationService : IRemoteMEFInitializationService
+{
+    internal sealed class Factory : FactoryBase<IRemoteMEFInitializationService>
+    {
+        protected override Task<object> CreateInternalAsync(Stream? stream, IServiceProvider hostProvidedServices, IServiceBroker? serviceBroker)
+        {
+            var traceSource = (TraceSource?)hostProvidedServices.GetService(typeof(TraceSource));
+
+            var service = new RemoteMEFInitializationService();
+            if (stream is not null)
+            {
+                var serverConnection = CreateServerConnection(stream, traceSource);
+                ConnectService(serverConnection, service);
+            }
+
+            return Task.FromResult<object>(service);
+        }
+
+        protected override IRemoteMEFInitializationService CreateService(in ServiceArgs args)
+            => Assumed.Unreachable<IRemoteMEFInitializationService>("This service overrides CreateInternalAsync to avoid MEF instatiation, so the CreateService method should never be called.");
+    }
+
+    public ValueTask InitializeAsync(string cacheDirectory, CancellationToken cancellationToken)
+    {
+        return RazorBrokeredServiceImplementation.RunServiceAsync(_ =>
+          {
+              RemoteMefComposition.CacheDirectory = cacheDirectory;
+              return new();
+          }, cancellationToken);
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RazorBrokeredServiceBase.FactoryBase`1.cs

@@ -60,7 +60,7 @@ private Task<object> CreateAsync(Stream stream, IServiceProvider hostProvidedSer
 #endif
         }
 
-        protected async Task<object> CreateInternalAsync(
+        protected virtual async Task<object> CreateInternalAsync(
             Stream? stream,
             IServiceProvider hostProvidedServices,
             IServiceBroker? serviceBroker)
@@ -103,21 +103,31 @@ protected async Task<object> CreateInternalAsync(
                 return CreateService(in inProcArgs);
             }
 
+            var serverConnection = CreateServerConnection(stream, traceSource);
+
+            var args = new ServiceArgs(serviceBroker.AssumeNotNull(), exportProvider, targetLoggerFactory, workspaceProvider, serverConnection, brokeredServiceData?.Interceptor);
+            var service = CreateService(in args);
+
+            ConnectService(serverConnection, service);
+
+            return service;
+        }
+
+        protected static ServiceRpcDescriptor.RpcConnection CreateServerConnection(Stream stream, TraceSource? traceSource)
+        {
             var pipe = stream.UsePipe();
 
             var descriptor = typeof(IRemoteJsonService).IsAssignableFrom(typeof(TService))
                 ? RazorServices.JsonDescriptors.GetDescriptorForServiceFactory(typeof(TService))
                 : RazorServices.Descriptors.GetDescriptorForServiceFactory(typeof(TService));
             var serverConnection = descriptor.WithTraceSource(traceSource).ConstructRpcConnection(pipe);
+            return serverConnection;
+        }
 
-            var args = new ServiceArgs(serviceBroker.AssumeNotNull(), exportProvider, targetLoggerFactory, workspaceProvider, serverConnection, brokeredServiceData?.Interceptor);
-
-            var service = CreateService(in args);
-
+        protected static void ConnectService(ServiceRpcDescriptor.RpcConnection serverConnection, TService service)
+        {
             serverConnection.AddLocalRpcTarget(service);
             serverConnection.StartListening();
-
-            return service;
         }
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RemoteMefComposition.cs

@@ -1,10 +1,15 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
+using System.IO;
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
+using Microsoft.AspNetCore.Razor.Utilities;
 using Microsoft.VisualStudio.Composition;
 using Microsoft.VisualStudio.Threading;
 
@@ -21,9 +26,13 @@ internal sealed class RemoteMefComposition
         joinableTaskFactory: null);
 
     private static readonly AsyncLazy<ExportProvider> s_lazyExportProvider = new(
-        static () => CreateExportProviderAsync(CancellationToken.None),
+        static () => CreateExportProviderAsync(CacheDirectory, CancellationToken.None),
         joinableTaskFactory: null);
 
+    private static Task? s_saveCacheFileTask;
+
+    public static string? CacheDirectory { get; set; }
+
     /// <summary>
     ///  Gets a <see cref="CompositionConfiguration"/> built from <see cref="Assemblies"/>. Note that the
     ///  same <see cref="CompositionConfiguration"/> instance is returned for subsequent calls to this method.
@@ -52,17 +61,125 @@ private static async Task<CompositionConfiguration> CreateConfigurationAsync(Can
     ///  Creates a new MEF composition and returns an <see cref="ExportProvider"/>. The catalog and configuration
     ///  are reused for subsequent calls to this method.
     /// </summary>
-    public static async Task<ExportProvider> CreateExportProviderAsync(CancellationToken cancellationToken)
+    public static async Task<ExportProvider> CreateExportProviderAsync(string? cacheDirectory, CancellationToken cancellationToken)
     {
+        var cache = new CachedComposition();
+        var compositionCacheFile = GetCompositionCacheFile(cacheDirectory);
+        if (await TryLoadCachedExportProviderAsync(cache, compositionCacheFile, cancellationToken).ConfigureAwait(false) is { } cachedProvider)
+        {
+            return cachedProvider;
+        }
+
         var configuration = await s_lazyConfiguration.GetValueAsync(cancellationToken).ConfigureAwait(false);
         cancellationToken.ThrowIfCancellationRequested();
 
         var runtimeComposition = RuntimeComposition.CreateRuntimeComposition(configuration);
         var exportProviderFactory = runtimeComposition.CreateExportProviderFactory();
 
+        // We don't need to block on saving the cache, because if it fails or is corrupt, we'll just try again next time, but
+        // we capture the task just so that tests can verify things.
+        s_saveCacheFileTask = TrySaveCachedExportProviderAsync(cache, compositionCacheFile, runtimeComposition, cancellationToken);
+
         return exportProviderFactory.CreateExportProvider();
     }
 
+    private static async Task<ExportProvider?> TryLoadCachedExportProviderAsync(CachedComposition cache, string? compositionCacheFile, CancellationToken cancellationToken)
+    {
+        if (compositionCacheFile is null)
+        {
+            return null;
+        }
+
+        try
+        {
+            if (File.Exists(compositionCacheFile))
+            {
+                var resolver = new Resolver(SimpleAssemblyLoader.Instance);
+                using var cacheStream = new FileStream(compositionCacheFile, FileMode.Open, FileAccess.Read, FileShare.Read, bufferSize: 4096, useAsync: true);
+                var cachedFactory = await cache.LoadExportProviderFactoryAsync(cacheStream, resolver, cancellationToken).ConfigureAwait(false);
+                return cachedFactory.CreateExportProvider();
+            }
+        }
+        catch (Exception)
+        {
+            // We ignore all errors when loading the cache, because if the cache is corrupt we will just create a new export provider.
+        }
+
+        return null;
+    }
+
+    private static async Task TrySaveCachedExportProviderAsync(CachedComposition cache, string? compositionCacheFile, RuntimeComposition runtimeComposition, CancellationToken cancellationToken)
+    {
+        if (compositionCacheFile is null)
+        {
+            return;
+        }
+
+        try
+        {
+            var cacheDirectory = Path.GetDirectoryName(compositionCacheFile).AssumeNotNull();
+            var directoryInfo = Directory.CreateDirectory(cacheDirectory);
+
+            CleanCacheDirectory(directoryInfo, cancellationToken);
+
+            var tempFilePath = Path.Combine(cacheDirectory, Path.GetRandomFileName());
+            using (var cacheStream = new FileStream(compositionCacheFile, FileMode.Create, FileAccess.Write, FileShare.None, bufferSize: 4096, useAsync: true))
+            {
+                await cache.SaveAsync(runtimeComposition, cacheStream, cancellationToken).ConfigureAwait(false);
+            }
+
+            File.Move(tempFilePath, compositionCacheFile);
+        }
+        catch (Exception)
+        {
+            // We ignore all errors when saving the cache, because if something goes wrong, the next run will just create a new export provider.
+        }
+    }
+
+    private static void CleanCacheDirectory(DirectoryInfo directoryInfo, CancellationToken cancellationToken)
+    {
+        try
+        {
+            // Delete any existing cached files.
+            foreach (var fileInfo in directoryInfo.EnumerateFiles())
+            {
+                // Failing to delete any file is fine, we'll just try again the next VS session in which we attempt
+                // to write a new cache
+                fileInfo.Delete();
+                cancellationToken.ThrowIfCancellationRequested();
+            }
+        }
+        catch (Exception)
+        {
+            // We ignore all errors when cleaning the cache directory, because we'll try again if the cache is corrupt.
+        }
+    }
+
+    [return: NotNullIfNotNull(nameof(cacheDirectory))]
+    private static string? GetCompositionCacheFile(string? cacheDirectory)
+    {
+        if (cacheDirectory is null)
+        {
+            return null;
+        }
+
+        var checksum = new Checksum.Builder();
+        foreach (var assembly in Assemblies)
+        {
+            var assemblyPath = assembly.Location.AssumeNotNull();
+            checksum.AppendData(Path.GetFileName(assemblyPath));
+            checksum.AppendData(File.GetLastWriteTimeUtc(assemblyPath).ToString("F"));
+        }
+
+        // Create base64 string of the hash.
+        var hashAsBase64String = checksum.FreeAndGetChecksum().ToBase64String();
+
+        // Convert to filename safe base64 string.
+        hashAsBase64String = hashAsBase64String.Replace('+', '-').Replace('/', '_').TrimEnd('=');
+
+        return Path.Combine(cacheDirectory, $"razor.mef.{hashAsBase64String}.cache");
+    }
+
     private sealed class SimpleAssemblyLoader : IAssemblyLoader
     {
         public static readonly IAssemblyLoader Instance = new SimpleAssemblyLoader();
@@ -83,4 +200,19 @@ public Assembly LoadAssembly(string assemblyFullName, string? codeBasePath)
             return LoadAssembly(assemblyName);
         }
     }
+
+    public static class TestAccessor
+    {
+        public static Task? SaveCacheFileTask => s_saveCacheFileTask;
+
+        public static void ClearSaveCacheFileTask()
+        {
+            s_saveCacheFileTask = null;
+        }
+
+        public static string GetCacheCompositionFile(string cacheDirectory)
+        {
+            return GetCompositionCacheFile(cacheDirectory);
+        }
+    }
 }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/Remote/RemoteServiceInvoker.cs

@@ -16,6 +16,9 @@
 using Microsoft.CodeAnalysis.Razor.SemanticTokens;
 using Microsoft.CodeAnalysis.Razor.Telemetry;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.VisualStudio.Settings;
+using Microsoft.VisualStudio.Shell;
+using Microsoft.VisualStudio.Shell.Settings;
 
 namespace Microsoft.VisualStudio.Razor.Remote;
 
@@ -26,12 +29,14 @@ internal sealed class RemoteServiceInvoker(
     LanguageServerFeatureOptions languageServerFeatureOptions,
     IClientCapabilitiesService clientCapabilitiesService,
     ISemanticTokensLegendService semanticTokensLegendService,
+    SVsServiceProvider serviceProvider,
     ITelemetryReporter telemetryReporter,
     ILoggerFactory loggerFactory) : IRemoteServiceInvoker, IDisposable
 {
     private readonly LanguageServerFeatureOptions _languageServerFeatureOptions = languageServerFeatureOptions;
     private readonly IClientCapabilitiesService _clientCapabilitiesService = clientCapabilitiesService;
     private readonly ISemanticTokensLegendService _semanticTokensLegendService = semanticTokensLegendService;
+    private readonly IServiceProvider _serviceProvider = serviceProvider;
     private readonly ITelemetryReporter _telemetryReporter = telemetryReporter;
     private readonly ILogger _logger = loggerFactory.GetOrCreateLogger<RemoteServiceInvoker>();
 
@@ -163,8 +168,15 @@ async Task InitializeCoreAsync(bool oopInitialized, bool lspInitialized)
                 await _initializeLspTask.ConfigureAwait(false);
             }
 
-            Task InitializeOOPAsync(RazorRemoteHostClient remoteClient)
+            async Task InitializeOOPAsync(RazorRemoteHostClient remoteClient)
             {
+                // The first call to OOP must be to initialize the MEF services, because everything after that relies on MEF.
+                var localSettingsDirectory = new ShellSettingsManager(_serviceProvider).GetApplicationDataFolder(ApplicationDataFolder.LocalSettings);
+                var cacheDirectory = Path.Combine(localSettingsDirectory, "Razor", "RemoteMEFCache");
+                await remoteClient.TryInvokeAsync<IRemoteMEFInitializationService>(
+                    (s, ct) => s.InitializeAsync(cacheDirectory, ct),
+                    _disposeTokenSource.Token).ConfigureAwait(false);
+
                 var initParams = new RemoteClientInitializationOptions
                 {
                     UseRazorCohostServer = _languageServerFeatureOptions.UseRazorCohostServer,
@@ -179,11 +191,10 @@ Task InitializeOOPAsync(RazorRemoteHostClient remoteClient)
 
                 _logger.LogDebug($"First OOP call, so initializing OOP service.");
 
-                return remoteClient
+                await remoteClient
                     .TryInvokeAsync<IRemoteClientInitializationService>(
                         (s, ct) => s.InitializeAsync(initParams, ct),
-                        _disposeTokenSource.Token)
-                    .AsTask();
+                        _disposeTokenSource.Token).ConfigureAwait(false);
             }
 
             Task InitializeLspAsync(RazorRemoteHostClient remoteClient)

src/Razor/src/Microsoft.VisualStudioCode.RazorExtension/Services/VSCodeRemoteServicesInitializer.cs

@@ -2,13 +2,15 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Composition;
+using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Razor.SemanticTokens;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Remote.Razor;
 using Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
 namespace Microsoft.VisualStudioCode.RazorExtension.Services;
@@ -41,6 +43,9 @@ public async Task StartupAsync(VSInternalClientCapabilities clientCapabilities,
         // we know this is VS Code specific, its all just smoke and mirrors anyway. We can avoid the smoke :)
         var serviceInterceptor = new VSCodeBrokeredServiceInterceptor();
 
+        // First things first, set the cache directory for the MEF composition.
+        RemoteMefComposition.CacheDirectory = Path.Combine(Path.GetDirectoryName(this.GetType().Assembly.Location)!, "cache");
+
         var logger = _loggerFactory.GetOrCreateLogger<VSCodeRemoteServicesInitializer>();
         logger.LogDebug("Initializing remote services.");
         var service = await InProcServiceFactory.CreateServiceAsync<IRemoteClientInitializationService>(serviceInterceptor, _workspaceProvider, _loggerFactory).ConfigureAwait(false);

src/Razor/test/Microsoft.CodeAnalysis.Remote.Razor.Test/RazorServicesTest.cs

@@ -54,7 +54,6 @@ public void JsonServicesHaveTheRightParameters(Type serviceType, Type? _)
                 {
                     Assert.Fail($"Method {method.Name} in a Json service has a pinned solution info wrapper parameter that isn't Json serializable");
                 }
-
             }
         }
     }