diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index f13a988e017e1..22d1fb987f524 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -47,24 +47,26 @@ internal sealed class ExtensionMessageHandlerService( private readonly SolutionServices _solutionServices = solutionServices; private readonly IExtensionMessageHandlerFactory _customMessageHandlerFactory = customMessageHandlerFactory; - // Core design: To make things lightweight, and to avoid locking, all work is computed and cached in simple - // immutable dictionaries. These dictionaries are populated on demand, but contain data that can be recomputed - // safely if missing. + /// + /// Lock for , , and . + /// + private readonly SemaphoreSlim _gate = new(initialCount: 1); /// /// Extensions assembly load contexts and loaded handlers, indexed by extension folder path. /// - private ImmutableDictionary> _folderPathToExtensionFolder = ImmutableDictionary>.Empty; + private readonly Dictionary> _folderPathToExtensionFolder = new(); /// /// Cached handlers of document-related messages, indexed by handler message name. /// - private ImmutableDictionary>>> _cachedDocumentHandlers = ImmutableDictionary>>>.Empty; + private readonly Dictionary>>> _cachedDocumentHandlers = new(); /// /// Cached handlers of non-document-related messages, indexed by handler message name. /// - private ImmutableDictionary>>> _cachedWorkspaceHandlers = ImmutableDictionary>>>.Empty; + private readonly Dictionary>>> _cachedWorkspaceHandlers = new(); private async ValueTask ExecuteInRemoteOrCurrentProcessAsync( Solution? solution, @@ -116,18 +118,22 @@ public async ValueTask RegisterExtensionInCurrentProc var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); - var lazy = ImmutableInterlocked.GetOrAdd( - ref _folderPathToExtensionFolder, - assemblyFolderPath, - static (assemblyFolderPath, @this) => AsyncLazy.Create( - cancellationToken => ExtensionFolder.Create(@this, assemblyFolderPath, cancellationToken)), - this); + AsyncLazy lazyExtensionFolder; + using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + { + lazyExtensionFolder = _folderPathToExtensionFolder.GetOrAdd( + assemblyFolderPath, + static (assemblyFolderPath, @this) => AsyncLazy.Create( + cancellationToken => ExtensionFolder.Create(@this, assemblyFolderPath, cancellationToken)), + this); - var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - var assemblyHandlers = await extensionFolder.RegisterAssemblyAsync(assemblyFilePath, cancellationToken).ConfigureAwait(false); + // After registering, clear out the cached handler names. They will be recomputed the next time we need them. + ClearCachedHandlers(); + } - // After registering, clear out the cached handler names. They will be recomputed the next time we need them. - ClearCachedHandlers(); + var extensionFolder = await lazyExtensionFolder.GetValueAsync(cancellationToken).ConfigureAwait(false); + var lazyAssemblyHandlers = extensionFolder.RegisterAssembly(assemblyFilePath); + var assemblyHandlers = await lazyAssemblyHandlers.GetValueAsync(cancellationToken).ConfigureAwait(false); return new( [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], @@ -152,17 +158,23 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); - if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazy)) + // Note: unregistering is slightly expensive as we do everything under a lock, to ensure that we have a + // consistent view of the world. This is fine as we don't expect this to be called very often. + using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) { - var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - // Unregister this particular assembly file from teh assembly folder. If it was the last extension within - // this folder, we can remove the registration for the extension entirely. - if (await extensionFolder.UnregisterAssemblyAsync(assemblyFilePath, cancellationToken).ConfigureAwait(false)) - _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); + if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazyExtensionFolder)) + { + var extensionFolder = await lazyExtensionFolder.GetValueAsync(cancellationToken).ConfigureAwait(false); + // Unregister this particular assembly file from teh assembly folder. If it was the last extension within + // this folder, we can remove the registration for the extension entirely. + if (extensionFolder.UnregisterAssembly(assemblyFilePath)) + _folderPathToExtensionFolder.Remove(assemblyFolderPath); + } + + // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. + ClearCachedHandlers(); } - // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. - ClearCachedHandlers(); return default; } @@ -177,15 +189,19 @@ await ExecuteInRemoteOrCurrentProcessAsync( private ValueTask ResetInCurrentProcessAsync() { - _folderPathToExtensionFolder = ImmutableDictionary>.Empty; - ClearCachedHandlers(); - return default; + lock (_gate) + { + _folderPathToExtensionFolder.Clear(); + ClearCachedHandlers(); + return default; + } } private void ClearCachedHandlers() { - _cachedWorkspaceHandlers = ImmutableDictionary>>>.Empty; - _cachedDocumentHandlers = ImmutableDictionary>>>.Empty; + Contract.ThrowIfTrue(!Monitor.IsEntered(_gate)); + _cachedWorkspaceHandlers.Clear(); + _cachedDocumentHandlers.Clear(); } public async ValueTask HandleExtensionWorkspaceMessageAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) @@ -210,18 +226,21 @@ public async ValueTask HandleExtensionDocumentMessageAsync(Document docu private async ValueTask HandleExtensionMessageInCurrentProcessAsync( TArgument executeArgument, bool isSolution, string messageName, string jsonMessage, - ImmutableDictionary>>> cachedHandlers, + Dictionary>>> cachedHandlers, CancellationToken cancellationToken) { - var lazy = ImmutableInterlocked.GetOrAdd( - ref cachedHandlers, - messageName, - static (messageName, arg) => AsyncLazy.Create( - static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, arg.isSolution, cancellationToken), - (messageName, arg.@this, arg.isSolution)), - (@this: this, executeArgument, isSolution)); - - var handlers = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); + AsyncLazy>> lazyHandlers; + using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + { + lazyHandlers = cachedHandlers.GetOrAdd( + messageName, + static (messageName, arg) => AsyncLazy.Create( + static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, arg.isSolution, cancellationToken), + (messageName, arg.@this, arg.isSolution)), + (messageName, @this: this, isSolution)); + } + + var handlers = await lazyHandlers.GetValueAsync(cancellationToken).ConfigureAwait(false); if (handlers.Length == 0) throw new InvalidOperationException($"No handler found for message {messageName}."); @@ -269,58 +288,62 @@ private static async Task RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken); - /// - /// Unregisters this assembly path from this extension folder. If this was the last registered path, then this - /// will return true so that this folder can be unloaded. + /// Mapping from assembly file path to the handlers it contains. Used as its own lock when mutating. /// - ValueTask UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken); + private readonly Dictionary> _assemblyFilePathToHandlers = new(); - ValueTask AddHandlersAsync(string messageName, bool isSolution, ArrayBuilder> result, CancellationToken cancellationToken); - } + protected abstract AssemblyHandlers CreateAssemblyHandlers(string assemblyFilePath, CancellationToken cancellationToken); - /// - /// Trivial placeholder impl of when we fail for some reason to even process the - /// folder we are told contains extensions. - /// - private sealed class TrivialExtensionFolder : IExtensionFolder - { - public static readonly TrivialExtensionFolder Instance = new(); + public AsyncLazy RegisterAssembly(string assemblyFilePath) + { + lock (_assemblyFilePathToHandlers) + { + return _assemblyFilePathToHandlers.GetOrAdd( + assemblyFilePath, + static (assemblyFilePath, @this) => AsyncLazy.Create( + static (args, cancellationToken) => args.@this.CreateAssemblyHandlers(args.assemblyFilePath, cancellationToken), + (assemblyFilePath, @this)), + this); + } + } /// - /// No lock needed as registratin/unregistration must happen serially. + /// Unregisters this assembly path from this extension folder. If this was the last registered path, then this + /// will return true so that this folder can be unloaded. /// - private readonly List _registeredFilePaths = []; - - public ValueTask RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken) + public bool UnregisterAssembly(string assemblyFilePath) { - _registeredFilePaths.Add(assemblyFilePath); - return new(AssemblyHandlers.Empty); + lock (_assemblyFilePathToHandlers) + { + _assemblyFilePathToHandlers.Remove(assemblyFilePath); + return _assemblyFilePathToHandlers.Count == 0; + } } - public ValueTask UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken) + public async ValueTask AddHandlersAsync(string messageName, bool isSolution, ArrayBuilder> result, CancellationToken cancellationToken) { - _registeredFilePaths.Remove(assemblyFilePath); - return new(_registeredFilePaths.Count == 0); - } - - public ValueTask AddHandlersAsync(string messageName, bool isSolution, ArrayBuilder> result, CancellationToken cancellationToken) - => default; - } - - private sealed class ExtensionFolder( - ExtensionMessageHandlerService extensionMessageHandlerService, - IAnalyzerAssemblyLoaderInternal analyzerAssemblyLoader) : IExtensionFolder - { - private readonly ExtensionMessageHandlerService _extensionMessageHandlerService = extensionMessageHandlerService; - private readonly IAnalyzerAssemblyLoaderInternal _analyzerAssemblyLoader = analyzerAssemblyLoader; + foreach (var (_, lazyHandlers) in _assemblyFilePathToHandlers) + { + cancellationToken.ThrowIfCancellationRequested(); - private ImmutableDictionary> _assemblyFilePathToHandlers = ImmutableDictionary>.Empty; + var handlers = await lazyHandlers.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (isSolution) + { + if (handlers.WorkspaceMessageHandlers.TryGetValue(messageName, out var handler)) + result.Add((IExtensionMessageHandlerWrapper)handler); + } + else + { + if (handlers.DocumentMessageHandlers.TryGetValue(messageName, out var handler)) + result.Add((IExtensionMessageHandlerWrapper)handler); + } + } + } - public static IExtensionFolder Create( + public static ExtensionFolder Create( ExtensionMessageHandlerService extensionMessageHandlerService, string assemblyFolderPath, CancellationToken cancellationToken) @@ -348,7 +371,7 @@ public static IExtensionFolder Create( analyzerAssemblyLoader.AddDependencyLocation(dll); } - return new ExtensionFolder(extensionMessageHandlerService, analyzerAssemblyLoader); + return new ShadowCopyExtensionFolder(extensionMessageHandlerService, analyzerAssemblyLoader); } catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical)) { @@ -356,28 +379,34 @@ public static IExtensionFolder Create( return new TrivialExtensionFolder(); } } + } - public async ValueTask RegisterAssemblyAsync( - string assemblyFilePath, CancellationToken cancellationToken) - { - var lazy = ImmutableInterlocked.GetOrAdd( - ref _assemblyFilePathToHandlers, - assemblyFilePath, - static (assemblyFilePath, @this) => AsyncLazy.Create( - static (args, cancellationToken) => CreateAssemblyHandlers(args.@this, args.assemblyFilePath, cancellationToken), - (assemblyFilePath, @this)), - this); + /// + /// Trivial placeholder impl of when we fail for some reason to even process the + /// folder we are told contains extensions. + /// + private sealed class TrivialExtensionFolder : ExtensionFolder + { + protected override AssemblyHandlers CreateAssemblyHandlers(string assemblyFilePath, CancellationToken cancellationToken) + => AssemblyHandlers.Empty; + } - return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - } + /// + /// Standard impl of that uses a shadow copy loader to load extensions. + /// + private sealed class ShadowCopyExtensionFolder( + ExtensionMessageHandlerService extensionMessageHandlerService, + IAnalyzerAssemblyLoaderInternal analyzerAssemblyLoader) : ExtensionFolder + { + private readonly ExtensionMessageHandlerService _extensionMessageHandlerService = extensionMessageHandlerService; + private readonly IAnalyzerAssemblyLoaderInternal _analyzerAssemblyLoader = analyzerAssemblyLoader; - private static AssemblyHandlers CreateAssemblyHandlers( - ExtensionFolder @this, string assemblyFilePath, CancellationToken cancellationToken) + protected override AssemblyHandlers CreateAssemblyHandlers(string assemblyFilePath, CancellationToken cancellationToken) { try { - var assembly = @this._analyzerAssemblyLoader.LoadFromPath(assemblyFilePath); - var factory = @this._extensionMessageHandlerService._customMessageHandlerFactory; + var assembly = _analyzerAssemblyLoader.LoadFromPath(assemblyFilePath); + var factory = _extensionMessageHandlerService._customMessageHandlerFactory; var messageWorkspaceHandlers = factory .CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath, cancellationToken) @@ -401,32 +430,6 @@ private static AssemblyHandlers CreateAssemblyHandlers( return AssemblyHandlers.Empty; } } - - public async ValueTask AddHandlersAsync(string messageName, bool isSolution, ArrayBuilder> result, CancellationToken cancellationToken) - { - foreach (var (_, lazy) in _assemblyFilePathToHandlers) - { - cancellationToken.ThrowIfCancellationRequested(); - - var handlers = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - if (isSolution) - { - if (handlers.WorkspaceMessageHandlers.TryGetValue(messageName, out var handler)) - result.Add((IExtensionMessageHandlerWrapper)handler); - } - else - { - if (handlers.DocumentMessageHandlers.TryGetValue(messageName, out var handler)) - result.Add((IExtensionMessageHandlerWrapper)handler); - } - } - } - - public ValueTask UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken) - { - _assemblyFilePathToHandlers = _assemblyFilePathToHandlers.Remove(assemblyFilePath); - return new(_assemblyFilePathToHandlers.IsEmpty); - } } private sealed class AssemblyHandlers diff --git a/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs index af16a65172367..c64270cd4026a 100644 --- a/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs @@ -18,16 +18,12 @@ internal interface IExtensionMessageHandlerService : IWorkspaceService /// /// The assembly to register and create message handlers from. /// The names of the registered handlers. - /// Should be called serially with other , , or calls. ValueTask RegisterExtensionAsync(string assemblyFilePath, CancellationToken cancellationToken); /// /// Unregisters extension message handlers previously registered from . /// /// The assembly for which handlers should be unregistered. - /// Should be called serially with other , , or calls. ValueTask UnregisterExtensionAsync(string assemblyFilePath, CancellationToken cancellationToken); /// @@ -44,7 +40,6 @@ internal interface IExtensionMessageHandlerService : IWorkspaceService /// The name of the handler to execute. This is generally the full name of the type implementing the handler. /// The json message to be passed to the handler. /// The json message returned by the handler. - /// Can be called concurrently with other message requests. ValueTask HandleExtensionWorkspaceMessageAsync( Solution solution, string messageName, @@ -58,7 +53,6 @@ ValueTask HandleExtensionWorkspaceMessageAsync( /// The name of the handler to execute. This is generally the full name of the type implementing the handler. /// The json message to be passed to the handler. /// The json message returned by the handler. - /// Can be called concurrently with other message requests. ValueTask HandleExtensionDocumentMessageAsync( Document documentId, string messageName, diff --git a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs index dafdb4c7379ca..f72272f641816 100644 --- a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs +++ b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs @@ -20,12 +20,7 @@ internal sealed class ExtensionRegisterHandler() { private const string MethodName = "roslyn/extensionRegister"; - /// - /// Report that we mutate solution state so that we only attempt to register or unregister one extension at a time. - /// This ensures we don't have to handle any threading concerns while this is happening. As this should be a rare - /// operation, this simplifies things while ideally being low cost. - /// - public bool MutatesSolutionState => true; + public bool MutatesSolutionState => false; public bool RequiresLSPSolution => true; diff --git a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs index c9d1ce218310d..42ca41e4d1e71 100644 --- a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs +++ b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs @@ -20,12 +20,7 @@ internal sealed class ExtensionUnregisterHandler() { private const string MethodName = "roslyn/extensionUnregister"; - /// - /// Report that we mutate solution state so that we only attempt to register or unregister one extension at a time. - /// This ensures we don't have to handle any threading concerns while this is happening. As this should be a rare - /// operation, this simplifies things while ideally being low cost. - /// - public bool MutatesSolutionState => true; + public bool MutatesSolutionState => false; public bool RequiresLSPSolution => true;