From de862de4ccba91f6573416ef52cee854d8772778 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 18:55:41 -0700 Subject: [PATCH 01/22] Move to immutable dictionary --- .../ExtensionMessageHandlerService.cs | 39 +++++++++---------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 0d0882e73b60e..3aee1ad54e057 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -46,20 +46,17 @@ internal sealed class ExtensionMessageHandlerService( /// /// Extensions assembly load contexts and loaded handlers, indexed by handler file path. The handlers are indexed by type name. /// - private readonly Dictionary _extensions = new(); + private ImmutableDictionary> _extensions = ImmutableDictionary>.Empty; /// /// Handlers of document-related messages, indexed by handler message name. /// - private readonly Dictionary> _documentHandlers = new(); + private ImmutableDictionary?>> _documentHandlers = ImmutableDictionary?>>.Empty; /// /// Handlers of non-document-related messages, indexed by handler message name. /// - private readonly Dictionary> _workspaceHandlers = new(); - - // Used to protect access to _extensions, _handlers, _documentHandlers and Extension._assemblies. - private readonly SemaphoreSlim _lock = new(initialCount: 1); + private ImmutableDictionary?>> _workspaceHandlers = ImmutableDictionary?>>.Empty; private async ValueTask ExecuteInRemoteOrCurrentProcessAsync( Solution? solution, @@ -111,7 +108,18 @@ public async ValueTask RegisterExtensionInCurrentProc var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); - var analyzerAssemblyLoaderProvider = _solutionServices.GetRequiredService(); + // var analyzerAssemblyLoaderProvider = _solutionServices.GetRequiredService(); + + var lazy = ImmutableInterlocked.GetOrAdd( + ref _extensions, + assemblyFolderPath, + static (assemblyFolderPath, @this) => Extension.Create(@this, assemblyFolderPath), + this); + + var extension = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (extension is null) + throw new InvalidOperationException($"A loading assemblies from {assemblyFolderPath} failed."); + Extension? extension; using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) { @@ -236,20 +244,11 @@ await ExecuteInRemoteOrCurrentProcessAsync( cancellationToken).ConfigureAwait(false); } - private async ValueTask ResetInCurrentProcessAsync(CancellationToken cancellationToken) + private ValueTask ResetInCurrentProcessAsync(CancellationToken cancellationToken) { - List extensions; - using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - { - extensions = [.. _extensions.Values]; - _extensions.Clear(); - _workspaceHandlers.Clear(); - _documentHandlers.Clear(); - } - - foreach (var extension in extensions) - extension.AnalyzerAssemblyLoader.Dispose(); - + _extensions = ImmutableDictionary>.Empty; + _workspaceHandlers = ImmutableDictionary?>>.Empty; + _documentHandlers = ImmutableDictionary?>>.Empty; return default; } From 5ea19b84467cefd5dd000f228dff2b5495fcd555 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:03:23 -0700 Subject: [PATCH 02/22] in progress --- .../ExtensionMessageHandlerService.cs | 167 ++++++++++-------- 1 file changed, 98 insertions(+), 69 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 3aee1ad54e057..5444e9129d0d8 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -113,71 +113,72 @@ public async ValueTask RegisterExtensionInCurrentProc var lazy = ImmutableInterlocked.GetOrAdd( ref _extensions, assemblyFolderPath, - static (assemblyFolderPath, @this) => Extension.Create(@this, assemblyFolderPath), + static (assemblyFolderPath, @this) => AsyncLazy.Create( + cancellationToken => Extension.CreateAsync(@this, assemblyFolderPath, cancellationToken)), this); var extension = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); if (extension is null) throw new InvalidOperationException($"A loading assemblies from {assemblyFolderPath} failed."); - Extension? extension; - using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - { - // Check if the assembly is already loaded. - if (!_extensions.TryGetValue(assemblyFolderPath, out extension)) - { - try - { - var analyzerAssemblyLoader = analyzerAssemblyLoaderProvider.CreateNewShadowCopyLoader(); - - // Allow this assembly loader to load any dll in assemblyFolderPath. - foreach (var dll in Directory.EnumerateFiles(assemblyFolderPath, "*.dll")) - { - try - { - // Check if the file is a valid .NET assembly. - AssemblyName.GetAssemblyName(dll); - } - catch - { - // The file is not a valid .NET assembly, skip it. - continue; - } - - analyzerAssemblyLoader.AddDependencyLocation(dll); - } - - extension = new Extension(this, analyzerAssemblyLoader, assemblyFolderPath); - } - catch - { - _extensions[assemblyFolderPath] = null; - throw; - } - - _extensions[assemblyFolderPath] = extension; - } - - if (extension is null) - { - throw new InvalidOperationException($"A previous attempt to load assemblies from {assemblyFolderPath} failed."); - } - - if (extension.TryGetAssemblyHandlers(assemblyFileName, out var assemblyHandlers)) - { - if (assemblyHandlers is null) - { - throw new InvalidOperationException($"A previous attempt to load {assemblyFilePath} failed."); - } - - return new( - [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], - [.. assemblyHandlers.DocumentMessageHandlers.Keys]); - } - } - - // Intentionally call this outside of the lock. - return await extension.LoadAssemblyAsync(assemblyFileName, cancellationToken).ConfigureAwait(false); + //Extension? extension; + //using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + //{ + // // Check if the assembly is already loaded. + // if (!_extensions.TryGetValue(assemblyFolderPath, out extension)) + // { + // try + // { + // var analyzerAssemblyLoader = analyzerAssemblyLoaderProvider.CreateNewShadowCopyLoader(); + + // // Allow this assembly loader to load any dll in assemblyFolderPath. + // foreach (var dll in Directory.EnumerateFiles(assemblyFolderPath, "*.dll")) + // { + // try + // { + // // Check if the file is a valid .NET assembly. + // AssemblyName.GetAssemblyName(dll); + // } + // catch + // { + // // The file is not a valid .NET assembly, skip it. + // continue; + // } + + // analyzerAssemblyLoader.AddDependencyLocation(dll); + // } + + // extension = new Extension(this, analyzerAssemblyLoader, assemblyFolderPath); + // } + // catch + // { + // _extensions[assemblyFolderPath] = null; + // throw; + // } + + // _extensions[assemblyFolderPath] = extension; + // } + + // if (extension is null) + // { + // throw new InvalidOperationException($"A previous attempt to load assemblies from {assemblyFolderPath} failed."); + // } + + // if (extension.TryGetAssemblyHandlers(assemblyFileName, out var assemblyHandlers)) + // { + // if (assemblyHandlers is null) + // { + // throw new InvalidOperationException($"A previous attempt to load {assemblyFilePath} failed."); + // } + + // return new( + // [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], + // [.. assemblyHandlers.DocumentMessageHandlers.Keys]); + // } + //} + + //// Intentionally call this outside of the lock. + //return await extension.LoadAssemblyAsync(assemblyFileName, cancellationToken).ConfigureAwait(false); } public async ValueTask UnregisterExtensionAsync( @@ -362,20 +363,48 @@ private sealed class Extension( string assemblyFolderPath) { private readonly ExtensionMessageHandlerService _extensionMessageHandlerService = extensionMessageHandlerService; + private readonly IAnalyzerAssemblyLoaderInternal _analyzerAssemblyLoader = analyzerAssemblyLoader; - public readonly IAnalyzerAssemblyLoaderInternal AnalyzerAssemblyLoader = analyzerAssemblyLoader; + private readonly string _assemblyFolderPath = assemblyFolderPath; - public readonly string AssemblyFolderPath = assemblyFolderPath; + private readonly Dictionary _assemblies = new(); - /// - /// Gets the object that is used to lock in order to avoid multiple calls from the same extensions to load the - /// same assembly concurrently resulting in the constructors of the same handlers being called more than once. - /// All other concurrent operations, including modifying are protected by . - /// - private readonly SemaphoreSlim _assemblyLoadLock = new(initialCount: 1); + public static async Task CreateAsync( + ExtensionMessageHandlerService extensionMessageHandlerService, + string assemblyFolderPath, + CancellationToken cancellationToken) + { + var analyzerAssemblyLoaderProvider = extensionMessageHandlerService._solutionServices.GetRequiredService(); + var analyzerAssemblyLoader = analyzerAssemblyLoaderProvider.CreateNewShadowCopyLoader(); - private readonly Dictionary _assemblies = new(); + try + { + // Allow this assembly loader to load any dll in assemblyFolderPath. + foreach (var dll in Directory.EnumerateFiles(assemblyFolderPath, "*.dll")) + { + cancellationToken.ThrowIfCancellationRequested(); + try + { + // Check if the file is a valid .NET assembly. + AssemblyName.GetAssemblyName(dll); + } + catch + { + // The file is not a valid .NET assembly, skip it. + continue; + } + + analyzerAssemblyLoader.AddDependencyLocation(dll); + } + + return new Extension(extensionMessageHandlerService, analyzerAssemblyLoader, assemblyFolderPath); + } + catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical)) + { + // If loading the assembly fails, we don't want to cache it. + return null; + } + } public void SetAssemblyHandlers(string assemblyFileName, AssemblyHandlers? value) { From f907906d30d783c415ad62ed068b3e3d5bae24e6 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:18:03 -0700 Subject: [PATCH 03/22] in progress --- .../ExtensionMessageHandlerService.cs | 38 ++++++++++++++++--- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 5444e9129d0d8..f852f171b5dce 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -46,17 +46,17 @@ internal sealed class ExtensionMessageHandlerService( /// /// Extensions assembly load contexts and loaded handlers, indexed by handler file path. The handlers are indexed by type name. /// - private ImmutableDictionary> _extensions = ImmutableDictionary>.Empty; + private ImmutableDictionary> _extensions = ImmutableDictionary>.Empty; /// /// Handlers of document-related messages, indexed by handler message name. /// - private ImmutableDictionary?>> _documentHandlers = ImmutableDictionary?>>.Empty; + private ImmutableDictionary>>> _cachedDocumentHandlers = ImmutableDictionary>>>.Empty; /// /// Handlers of non-document-related messages, indexed by handler message name. /// - private ImmutableDictionary?>> _workspaceHandlers = ImmutableDictionary?>>.Empty; + private ImmutableDictionary?>> _cachedWorkspaceHandlers = ImmutableDictionary?>>.Empty; private async ValueTask ExecuteInRemoteOrCurrentProcessAsync( Solution? solution, @@ -200,6 +200,18 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); + if (_extensions.TryGetValue(assemblyFolderPath, out var extension) && + extension != null) + { + // Unregister this particular assembly file from teh assembly folder. If it was the last extension within this folder, + // we can remove the registeration for the extension entirely. + if (extension.Unregister(assemblyFileName)) + _extensions = _extensions.Remove(assemblyFolderPath); + } + + // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. + ClearCachedHandlers(); + Extension? extension = null; using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) { @@ -248,9 +260,13 @@ await ExecuteInRemoteOrCurrentProcessAsync( private ValueTask ResetInCurrentProcessAsync(CancellationToken cancellationToken) { _extensions = ImmutableDictionary>.Empty; + return default; + } + + private void ClearCachedHandlers() + { _workspaceHandlers = ImmutableDictionary?>>.Empty; _documentHandlers = ImmutableDictionary?>>.Empty; - return default; } public async ValueTask HandleExtensionWorkspaceMessageAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) @@ -271,6 +287,14 @@ public async ValueTask HandleExtensionDocumentMessageAsync(Document docu cancellationToken).ConfigureAwait(false); } + private ValueTask HandleExtensionWorkspaceMessageInCurrentProcess1Async(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) + { + var lazy = _cachedWorkspaceHandlers.GetOrAdd( + messageName, + (messageName, arg) => AsyncLazy.Create(cancellationToken => ComputeHandler(arg.@this), + (@this: this, solution, jsonMessage)) + } + private ValueTask HandleExtensionWorkspaceMessageInCurrentProcessAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) => HandleExtensionMessageAsync(solution, messageName, jsonMessage, _workspaceHandlers, cancellationToken); @@ -281,9 +305,11 @@ private async ValueTask HandleExtensionMessageAsync( TArgument argument, string messageName, string jsonMessage, - Dictionary> handlers, + ImmutableDictionary> handlers, CancellationToken cancellationToken) { + + var lazy = _cachedDocumentHandlers. IExtensionMessageHandlerWrapper? handler; using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) { @@ -447,7 +473,7 @@ public async ValueTask LoadAssemblyAsync( Exception? exception = null; try { - var assembly = AnalyzerAssemblyLoader.LoadFromPath(assemblyFilePath); + var assembly = _analyzerAssemblyLoader.LoadFromPath(assemblyFilePath); var factory = _extensionMessageHandlerService._customMessageHandlerFactory; var messageWorkspaceHandlers = factory.CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath) .ToImmutableDictionary(h => h.Name, h => h); From 31ade893f9dc4a225f45bf7e1f8ca11396a85106 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:29:24 -0700 Subject: [PATCH 04/22] in progress --- .../ExtensionMessageHandlerService.cs | 186 +++++++++++------- 1 file changed, 120 insertions(+), 66 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index f852f171b5dce..b0f9a46b05039 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -17,6 +17,7 @@ using Microsoft.CodeAnalysis.ErrorReporting; using Microsoft.CodeAnalysis.Host; using Microsoft.CodeAnalysis.Host.Mef; +using Microsoft.CodeAnalysis.PooledObjects; using Microsoft.CodeAnalysis.Remote; using Microsoft.CodeAnalysis.Threading; using Roslyn.Utilities; @@ -56,7 +57,7 @@ internal sealed class ExtensionMessageHandlerService( /// /// Handlers of non-document-related messages, indexed by handler message name. /// - private ImmutableDictionary?>> _cachedWorkspaceHandlers = ImmutableDictionary?>>.Empty; + private ImmutableDictionary>>> _cachedWorkspaceHandlers = ImmutableDictionary>>>.Empty; private async ValueTask ExecuteInRemoteOrCurrentProcessAsync( Solution? solution, @@ -203,8 +204,8 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( if (_extensions.TryGetValue(assemblyFolderPath, out var extension) && extension != null) { - // Unregister this particular assembly file from teh assembly folder. If it was the last extension within this folder, - // we can remove the registeration for the extension entirely. + // 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 (extension.Unregister(assemblyFileName)) _extensions = _extensions.Remove(assemblyFolderPath); } @@ -287,44 +288,28 @@ public async ValueTask HandleExtensionDocumentMessageAsync(Document docu cancellationToken).ConfigureAwait(false); } - private ValueTask HandleExtensionWorkspaceMessageInCurrentProcess1Async(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) + private async ValueTask HandleExtensionWorkspaceMessageInCurrentProcessAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) { var lazy = _cachedWorkspaceHandlers.GetOrAdd( messageName, - (messageName, arg) => AsyncLazy.Create(cancellationToken => ComputeHandler(arg.@this), - (@this: this, solution, jsonMessage)) - } + static (messageName, @this) => AsyncLazy.Create( + static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, cancellationToken), + (messageName, @this)), + this); - private ValueTask HandleExtensionWorkspaceMessageInCurrentProcessAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) - => HandleExtensionMessageAsync(solution, messageName, jsonMessage, _workspaceHandlers, cancellationToken); + var handlers = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (handlers.Length == 0) + throw new InvalidOperationException($"No handler found for message {messageName}."); - public ValueTask HandleExtensionDocumentMessageInCurrentProcessAsync(Document document, string messageName, string jsonMessage, CancellationToken cancellationToken) - => HandleExtensionMessageAsync(document, messageName, jsonMessage, _documentHandlers, cancellationToken); + if (handlers.Length > 1) + throw new InvalidOperationException($"Multiple handlers found for message {messageName}."); - private async ValueTask HandleExtensionMessageAsync( - TArgument argument, - string messageName, - string jsonMessage, - ImmutableDictionary> handlers, - CancellationToken cancellationToken) - { - - var lazy = _cachedDocumentHandlers. - IExtensionMessageHandlerWrapper? handler; - using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - { - // handlers here is either _workspaceHandlers or _documentHandlers, so it must be protected - // by _lockObject. - if (!handlers.TryGetValue(messageName, out handler)) - { - throw new InvalidOperationException($"No handler found for message {messageName}."); - } - } + var handler = handlers[0]; try { var message = JsonSerializer.Deserialize(jsonMessage, handler.MessageType); - var result = await handler.ExecuteAsync(message, argument, cancellationToken) + var result = await handler.ExecuteAsync(message, solution, cancellationToken) .ConfigureAwait(false); var responseJson = JsonSerializer.Serialize(result, handler.ResponseType); return responseJson; @@ -338,51 +323,120 @@ private async ValueTask HandleExtensionMessageAsync( } } - private async Task RegisterAssemblyAsync( - Extension extension, - string assemblyFileName, - AssemblyHandlers? assemblyHandlers, - CancellationToken cancellationToken) + private static async Task>> ComputeHandlersAsync( + ExtensionMessageHandlerService @this, string messageName, bool solution, CancellationToken cancellationToken) { - using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + using var _ = ArrayBuilder>.GetInstance(out var result); + foreach (var (_, lazyExtension) in @this._extensions) { - // Make sure a call to UnloadCustomMessageHandlersAsync hasn't happened while we relinquished the lock on _lockObject - if (!_extensions.TryGetValue(extension.AssemblyFolderPath, out var currentExtension) || !extension.Equals(currentExtension)) - { - throw new InvalidOperationException($"Handlers in {extension.AssemblyFolderPath} were unregistered while loading handlers."); - } + cancellationToken.ThrowIfCancellationRequested(); + var extension = await lazyExtension.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (extension is null) + continue; - try + if (solution) { - if (assemblyHandlers is not null) - { - var duplicateHandler = _workspaceHandlers.Keys.Intersect(assemblyHandlers.WorkspaceMessageHandlers.Keys).Concat( - _documentHandlers.Keys.Intersect(assemblyHandlers.DocumentMessageHandlers.Keys)).FirstOrDefault(); - - if (duplicateHandler is not null) - { - assemblyHandlers = null; - throw new InvalidOperationException($"Handler name {duplicateHandler} is already registered."); - } - - foreach (var handler in assemblyHandlers.WorkspaceMessageHandlers) - { - _workspaceHandlers.Add(handler.Key, handler.Value); - } - - foreach (var handler in assemblyHandlers.DocumentMessageHandlers) - { - _documentHandlers.Add(handler.Key, handler.Value); - } - } + if (extension.WorkspaceMessageHandlers.TryGetValue(messageName, out var handler)) + result.Add((IExtensionMessageHandlerWrapper)handler); } - finally + else { - extension.SetAssemblyHandlers(assemblyFileName, assemblyHandlers); + if (extension.DocumentMessageHandlers.TryGetValue(messageName, out var handler)) + result.Add((IExtensionMessageHandlerWrapper)handler); } } + + return result.ToImmutable(); } + //private ValueTask HandleExtensionWorkspaceMessageInCurrentProcessAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) + // => HandleExtensionMessageAsync(solution, messageName, jsonMessage, _workspaceHandlers, cancellationToken); + + //public ValueTask HandleExtensionDocumentMessageInCurrentProcessAsync(Document document, string messageName, string jsonMessage, CancellationToken cancellationToken) + // => HandleExtensionMessageAsync(document, messageName, jsonMessage, _documentHandlers, cancellationToken); + + //private async ValueTask HandleExtensionMessageAsync( + // TArgument argument, + // string messageName, + // string jsonMessage, + // ImmutableDictionary> handlers, + // CancellationToken cancellationToken) + //{ + + // var lazy = _cachedDocumentHandlers. + // IExtensionMessageHandlerWrapper? handler; + // using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + // { + // // handlers here is either _workspaceHandlers or _documentHandlers, so it must be protected + // // by _lockObject. + // if (!handlers.TryGetValue(messageName, out handler)) + // { + // throw new InvalidOperationException($"No handler found for message {messageName}."); + // } + // } + + // try + // { + // var message = JsonSerializer.Deserialize(jsonMessage, handler.MessageType); + // var result = await handler.ExecuteAsync(message, argument, cancellationToken) + // .ConfigureAwait(false); + // var responseJson = JsonSerializer.Serialize(result, handler.ResponseType); + // return responseJson; + // } + // catch + // { + // // Any exception thrown in this method is left to bubble up to the extension. + // // But we unregister all handlers from that assembly to minimize the impact of a bad extension. + // await UnregisterExtensionAsync(assemblyFilePath: handler.ExtensionIdentifier, cancellationToken).ConfigureAwait(false); + // throw; + // } + //} + + //private async Task RegisterAssemblyAsync( + // Extension extension, + // string assemblyFileName, + // AssemblyHandlers? assemblyHandlers, + // CancellationToken cancellationToken) + //{ + // using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + // { + // // Make sure a call to UnloadCustomMessageHandlersAsync hasn't happened while we relinquished the lock on _lockObject + // if (!_extensions.TryGetValue(extension.AssemblyFolderPath, out var currentExtension) || !extension.Equals(currentExtension)) + // { + // throw new InvalidOperationException($"Handlers in {extension.AssemblyFolderPath} were unregistered while loading handlers."); + // } + + // try + // { + // if (assemblyHandlers is not null) + // { + // var duplicateHandler = _workspaceHandlers.Keys.Intersect(assemblyHandlers.WorkspaceMessageHandlers.Keys).Concat( + // _documentHandlers.Keys.Intersect(assemblyHandlers.DocumentMessageHandlers.Keys)).FirstOrDefault(); + + // if (duplicateHandler is not null) + // { + // assemblyHandlers = null; + // throw new InvalidOperationException($"Handler name {duplicateHandler} is already registered."); + // } + + // foreach (var handler in assemblyHandlers.WorkspaceMessageHandlers) + // { + // _workspaceHandlers.Add(handler.Key, handler.Value); + // } + + // foreach (var handler in assemblyHandlers.DocumentMessageHandlers) + // { + // _documentHandlers.Add(handler.Key, handler.Value); + // } + // } + // } + // finally + // { + // extension.SetAssemblyHandlers(assemblyFileName, assemblyHandlers); + // } + // } + //} + private sealed class Extension( ExtensionMessageHandlerService extensionMessageHandlerService, IAnalyzerAssemblyLoaderInternal analyzerAssemblyLoader, From 508f0920ef79ef18cb264555a88176f77b6698de Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:42:14 -0700 Subject: [PATCH 05/22] in progress --- .../ExtensionMessageHandlerService.cs | 226 +++++++++--------- 1 file changed, 110 insertions(+), 116 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index b0f9a46b05039..d4277d82ef9ca 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -45,9 +45,9 @@ internal sealed class ExtensionMessageHandlerService( private readonly IExtensionMessageHandlerFactory _customMessageHandlerFactory = customMessageHandlerFactory; /// - /// Extensions assembly load contexts and loaded handlers, indexed by handler file path. The handlers are indexed by type name. + /// Extensions assembly load contexts and loaded handlers, indexed by extension folder path. /// - private ImmutableDictionary> _extensions = ImmutableDictionary>.Empty; + private ImmutableDictionary> _folderPathToExtension = ImmutableDictionary>.Empty; /// /// Handlers of document-related messages, indexed by handler message name. @@ -105,22 +105,30 @@ public async ValueTask RegisterExtensionInCurrentProc string assemblyFilePath, CancellationToken cancellationToken) { - var assemblyFileName = Path.GetFileName(assemblyFilePath); + // var assemblyFileName = Path.GetFileName(assemblyFilePath); var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); // var analyzerAssemblyLoaderProvider = _solutionServices.GetRequiredService(); var lazy = ImmutableInterlocked.GetOrAdd( - ref _extensions, + ref _folderPathToExtension, assemblyFolderPath, static (assemblyFolderPath, @this) => AsyncLazy.Create( - cancellationToken => Extension.CreateAsync(@this, assemblyFolderPath, cancellationToken)), + cancellationToken => ExtensionFolder.CreateAsync(@this, assemblyFolderPath, cancellationToken)), this); - var extension = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - if (extension is null) - throw new InvalidOperationException($"A loading assemblies from {assemblyFolderPath} failed."); + var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (extensionFolder is null) + throw new InvalidOperationException($"Loading extensions from {assemblyFolderPath} failed."); + + var assemblyHandlers = await extensionFolder.GetAssemblyHandlersAsync(assemblyFilePath, cancellationToken).ConfigureAwait(false); + if (assemblyHandlers is null) + throw new InvalidOperationException($"Loading extensions from {assemblyFilePath} failed."); + + return new( + [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], + [.. assemblyHandlers.DocumentMessageHandlers.Keys]); //Extension? extension; //using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) @@ -197,56 +205,56 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( string assemblyFilePath, CancellationToken cancellationToken) { - var assemblyFileName = Path.GetFileName(assemblyFilePath); + // var assemblyFileName = Path.GetFileName(assemblyFilePath); var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); - if (_extensions.TryGetValue(assemblyFolderPath, out var extension) && + if (_folderPathToExtension.TryGetValue(assemblyFolderPath, out var extension) && extension != null) { // 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 (extension.Unregister(assemblyFileName)) - _extensions = _extensions.Remove(assemblyFolderPath); + _folderPathToExtension = _folderPathToExtension.RemoveEtension(assemblyFilePath); } // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. ClearCachedHandlers(); - Extension? extension = null; - using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - { - if (_extensions.TryGetValue(assemblyFolderPath, out extension)) - { - // If loading assemblies from this folder failed earlier, don't do anything. - if (extension is null) - return default; + //Extension? extension = null; + //using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + //{ + // if (_folderPathToExtension.TryGetValue(assemblyFolderPath, out extension)) + // { + // // If loading assemblies from this folder failed earlier, don't do anything. + // if (extension is null) + // return default; - if (extension.RemoveAssemblyHandlers(assemblyFileName, out var assemblyHandlers)) - { - if (assemblyHandlers is not null) - { - foreach (var workspaceHandler in assemblyHandlers.WorkspaceMessageHandlers.Keys) - { - _workspaceHandlers.Remove(workspaceHandler); - } - - foreach (var documentHandler in assemblyHandlers.DocumentMessageHandlers.Keys) - { - _documentHandlers.Remove(documentHandler); - } - } - } + // if (extension.RemoveAssemblyHandlers(assemblyFileName, out var assemblyHandlers)) + // { + // if (assemblyHandlers is not null) + // { + // foreach (var workspaceHandler in assemblyHandlers.WorkspaceMessageHandlers.Keys) + // { + // _workspaceHandlers.Remove(workspaceHandler); + // } - if (extension.AssemblyHandlersCount > 0) - return default; + // foreach (var documentHandler in assemblyHandlers.DocumentMessageHandlers.Keys) + // { + // _documentHandlers.Remove(documentHandler); + // } + // } + // } - _extensions.Remove(assemblyFolderPath); - } - } + // if (extension.AssemblyHandlersCount > 0) + // return default; - extension?.AnalyzerAssemblyLoader.Dispose(); - return default; + // _folderPathToExtension.Remove(assemblyFolderPath); + // } + //} + + //extension?.AnalyzerAssemblyLoader.Dispose(); + //return default; } public async ValueTask ResetAsync(CancellationToken cancellationToken) @@ -260,7 +268,7 @@ await ExecuteInRemoteOrCurrentProcessAsync( private ValueTask ResetInCurrentProcessAsync(CancellationToken cancellationToken) { - _extensions = ImmutableDictionary>.Empty; + _folderPathToExtension = ImmutableDictionary>.Empty; return default; } @@ -327,7 +335,7 @@ private static async Task>.GetInstance(out var result); - foreach (var (_, lazyExtension) in @this._extensions) + foreach (var (_, lazyExtension) in @this._folderPathToExtension) { cancellationToken.ThrowIfCancellationRequested(); var extension = await lazyExtension.GetValueAsync(cancellationToken).ConfigureAwait(false); @@ -437,7 +445,7 @@ private static async Task _assemblies = new(); + private ImmutableDictionary> _assemblyFilePathToHandlers = ImmutableDictionary>.Empty; - public static async Task CreateAsync( + public static async Task CreateAsync( ExtensionMessageHandlerService extensionMessageHandlerService, string assemblyFolderPath, CancellationToken cancellationToken) @@ -477,7 +485,7 @@ private sealed class Extension( analyzerAssemblyLoader.AddDependencyLocation(dll); } - return new Extension(extensionMessageHandlerService, analyzerAssemblyLoader, assemblyFolderPath); + return new ExtensionFolder(extensionMessageHandlerService, analyzerAssemblyLoader, assemblyFolderPath); } catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical)) { @@ -486,85 +494,71 @@ private sealed class Extension( } } - public void SetAssemblyHandlers(string assemblyFileName, AssemblyHandlers? value) - { - EnsureGlobalLockIsOwned(); - _assemblies[assemblyFileName] = value; - } + //public void SetAssemblyHandlers(string assemblyFileName, AssemblyHandlers? value) + //{ + // EnsureGlobalLockIsOwned(); + // _assemblies[assemblyFileName] = value; + //} - public bool TryGetAssemblyHandlers(string assemblyFileName, out AssemblyHandlers? value) - { - EnsureGlobalLockIsOwned(); - return _assemblies.TryGetValue(assemblyFileName, out value); - } + //public bool TryGetAssemblyHandlers(string assemblyFileName, out AssemblyHandlers? value) + //{ + // EnsureGlobalLockIsOwned(); + // return _assemblies.TryGetValue(assemblyFileName, out value); + //} - public bool RemoveAssemblyHandlers(string assemblyFileName, out AssemblyHandlers? value) - { - EnsureGlobalLockIsOwned(); - return _assemblies.Remove(assemblyFileName, out value); - } + //public bool RemoveAssemblyHandlers(string assemblyFileName, out AssemblyHandlers? value) + //{ + // EnsureGlobalLockIsOwned(); + // return _assemblies.Remove(assemblyFileName, out value); + //} - public int AssemblyHandlersCount + //public int AssemblyHandlersCount + //{ + // get + // { + // EnsureGlobalLockIsOwned(); + // return _assemblies.Count; + // } + //} + + public async ValueTask GetAssemblyHandlersAsync( + string assemblyFilePath, CancellationToken cancellationToken) { - get - { - EnsureGlobalLockIsOwned(); - return _assemblies.Count; - } + var lazy = _assemblyFilePathToHandlers.GetOrAdd( + assemblyFilePath, + static (@this, assemblyFilePath) => AsyncLazy.Create( + static (args, cancellationToken) => CreateAssemblyHandlersAsync(args.@this, args.assemblyFilePath, cancellationToken), + (@this, assemblyFilePath))); + + return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); } - public async ValueTask LoadAssemblyAsync( - string assemblyFileName, CancellationToken cancellationToken) + private static async Task CreateAssemblyHandlersAsync( + ExtensionFolder @this, string assemblyFilePath, CancellationToken cancellationToken) { - var assemblyFilePath = Path.Combine(AssemblyFolderPath, assemblyFileName); - - // AssemblyLoadLockObject is only used to avoid multiple calls from the same extensions to load the same assembly concurrently - // resulting in the constructors of the same handlers being called more than once. - // All other concurrent operations, including modifying extension.Assemblies are protected by _lockObject. - using (await _assemblyLoadLock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) + try { - AssemblyHandlers? assemblyHandlers = null; - Exception? exception = null; - try + var assembly = @this._analyzerAssemblyLoader.LoadFromPath(assemblyFilePath); + var factory = @this._extensionMessageHandlerService._customMessageHandlerFactory; + var messageWorkspaceHandlers = factory.CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath) + .ToImmutableDictionary(h => h.Name, h => h); + var messageDocumentHandlers = factory.CreateDocumentMessageHandlers(assembly, extensionIdentifier: assemblyFilePath) + .ToImmutableDictionary(h => h.Name, h => h); + + return new AssemblyHandlers() { - var assembly = _analyzerAssemblyLoader.LoadFromPath(assemblyFilePath); - var factory = _extensionMessageHandlerService._customMessageHandlerFactory; - var messageWorkspaceHandlers = factory.CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath) - .ToImmutableDictionary(h => h.Name, h => h); - var messageDocumentHandlers = factory.CreateDocumentMessageHandlers(assembly, extensionIdentifier: assemblyFilePath) - .ToImmutableDictionary(h => h.Name, h => h); - - assemblyHandlers = new AssemblyHandlers() - { - WorkspaceMessageHandlers = messageWorkspaceHandlers, - DocumentMessageHandlers = messageDocumentHandlers, - }; + WorkspaceMessageHandlers = messageWorkspaceHandlers, + DocumentMessageHandlers = messageDocumentHandlers, + }; - // We don't add assemblyHandlers to _assemblies here and instead let _extensionMessageHandlerService.RegisterAssembly do it - // since RegisterAssembly can still fail if there are duplicated handler names. - } - catch (Exception e) when (FatalError.ReportAndPropagate(exception = e, ErrorSeverity.General)) - { - throw ExceptionUtilities.Unreachable(); - } - finally - { - // In case of exception, we cache null so that we don't try to load the same assembly again. - await _extensionMessageHandlerService.RegisterAssemblyAsync( - this, assemblyFileName, exception is null ? assemblyHandlers : null, cancellationToken).ConfigureAwait(false); - } - - // The return is here, after RegisterAssembly, since RegisterAssembly can also throw an exception: the registration is not - // completed until RegisterAssembly returns. - return new( - [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], - [.. assemblyHandlers.DocumentMessageHandlers.Keys]); + // We don't add assemblyHandlers to _assemblies here and instead let _extensionMessageHandlerService.RegisterAssembly do it + // since RegisterAssembly can still fail if there are duplicated handler names. + } + catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.General)) + { + // TODO: Log error so it is visible to user. + return null; } - } - - private void EnsureGlobalLockIsOwned() - { - Contract.ThrowIfTrue(_extensionMessageHandlerService._lock.CurrentCount != 0, "Global lock should be owned"); } } From 5ef943b6165e6a1dc7dceec6aebeab712afdefcb Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:44:03 -0700 Subject: [PATCH 06/22] in progress --- .../Portable/Extensions/ExtensionMessageHandlerService.cs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index d4277d82ef9ca..c34336b087ef8 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -524,11 +524,13 @@ private sealed class ExtensionFolder( public async ValueTask GetAssemblyHandlersAsync( string assemblyFilePath, CancellationToken cancellationToken) { - var lazy = _assemblyFilePathToHandlers.GetOrAdd( + var lazy = ImmutableInterlocked.GetOrAdd( + ref _assemblyFilePathToHandlers, assemblyFilePath, - static (@this, assemblyFilePath) => AsyncLazy.Create( + static (assemblyFilePath, @this) => AsyncLazy.Create( static (args, cancellationToken) => CreateAssemblyHandlersAsync(args.@this, args.assemblyFilePath, cancellationToken), - (@this, assemblyFilePath))); + (assemblyFilePath, @this)), + this); return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); } From 9745b90c42026a396ddb81b899f3a032ced5e214 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:47:08 -0700 Subject: [PATCH 07/22] in progress --- .../ExtensionMessageHandlerService.cs | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index c34336b087ef8..469fd53fc5df1 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -342,16 +342,7 @@ private static async Task)handler); - } - else - { - if (extension.DocumentMessageHandlers.TryGetValue(messageName, out var handler)) - result.Add((IExtensionMessageHandlerWrapper)handler); - } + await extension.AddHandlersAsync(messageName, solution, result, cancellationToken).ConfigureAwait(false); } return result.ToImmutable(); @@ -562,6 +553,29 @@ private sealed class ExtensionFolder( return null; } } + + public async Task AddHandlersAsync(string messageName, bool solution, ArrayBuilder> result, CancellationToken cancellationToken) + { + foreach (var (_, lazy) in _assemblyFilePathToHandlers) + { + cancellationToken.ThrowIfCancellationRequested(); + + var handlers = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (handlers is null) + continue; + + if (solution) + { + 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); + } + } + } } private sealed class AssemblyHandlers From bd715126d9e78d0d7e16d7ff3509dae88e919838 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:50:37 -0700 Subject: [PATCH 08/22] Simplify --- .../ExtensionMessageHandlerService.cs | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 469fd53fc5df1..585bc2cf3e0db 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -47,7 +47,7 @@ internal sealed class ExtensionMessageHandlerService( /// /// Extensions assembly load contexts and loaded handlers, indexed by extension folder path. /// - private ImmutableDictionary> _folderPathToExtension = ImmutableDictionary>.Empty; + private ImmutableDictionary> _folderPathToExtensionFolder = ImmutableDictionary>.Empty; /// /// Handlers of document-related messages, indexed by handler message name. @@ -112,7 +112,7 @@ public async ValueTask RegisterExtensionInCurrentProc // var analyzerAssemblyLoaderProvider = _solutionServices.GetRequiredService(); var lazy = ImmutableInterlocked.GetOrAdd( - ref _folderPathToExtension, + ref _folderPathToExtensionFolder, assemblyFolderPath, static (assemblyFolderPath, @this) => AsyncLazy.Create( cancellationToken => ExtensionFolder.CreateAsync(@this, assemblyFolderPath, cancellationToken)), @@ -209,13 +209,13 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); - if (_folderPathToExtension.TryGetValue(assemblyFolderPath, out var extension) && + if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var extension) && extension != null) { // 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 (extension.Unregister(assemblyFileName)) - _folderPathToExtension = _folderPathToExtension.RemoveEtension(assemblyFilePath); + _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); } // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. @@ -268,14 +268,15 @@ await ExecuteInRemoteOrCurrentProcessAsync( private ValueTask ResetInCurrentProcessAsync(CancellationToken cancellationToken) { - _folderPathToExtension = ImmutableDictionary>.Empty; + _folderPathToExtensionFolder = ImmutableDictionary>.Empty; + ClearCachedHandlers(); return default; } private void ClearCachedHandlers() { - _workspaceHandlers = ImmutableDictionary?>>.Empty; - _documentHandlers = ImmutableDictionary?>>.Empty; + _cachedWorkspaceHandlers = ImmutableDictionary>>>.Empty; + _cachedDocumentHandlers = ImmutableDictionary>>>.Empty; } public async ValueTask HandleExtensionWorkspaceMessageAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) @@ -335,7 +336,7 @@ private static async Task>.GetInstance(out var result); - foreach (var (_, lazyExtension) in @this._folderPathToExtension) + foreach (var (_, lazyExtension) in @this._folderPathToExtensionFolder) { cancellationToken.ThrowIfCancellationRequested(); var extension = await lazyExtension.GetValueAsync(cancellationToken).ConfigureAwait(false); From f01331e7b6090341118ae9df7f8e1053af2c17ad Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 19:57:39 -0700 Subject: [PATCH 09/22] Serialize --- .../ExtensionMessageHandlerService.cs | 72 +++---------------- .../Extensions/ExtensionRegisterHandler.cs | 7 +- .../Extensions/ExtensionUnregisterHandler.cs | 7 +- 3 files changed, 20 insertions(+), 66 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 585bc2cf3e0db..fa5cda3ccd556 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -44,18 +44,22 @@ 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. This allows for a safe approach to + /// /// Extensions assembly load contexts and loaded handlers, indexed by extension folder path. /// - private ImmutableDictionary> _folderPathToExtensionFolder = ImmutableDictionary>.Empty; + private ImmutableDictionary> _folderPathToExtensionFolder = ImmutableDictionary>.Empty; /// - /// Handlers of document-related messages, indexed by handler message name. + /// Cached handlers of document-related messages, indexed by handler message name. /// private ImmutableDictionary>>> _cachedDocumentHandlers = ImmutableDictionary>>>.Empty; /// - /// Handlers of non-document-related messages, indexed by handler message name. + /// Cached handlers of non-document-related messages, indexed by handler message name. /// private ImmutableDictionary>>> _cachedWorkspaceHandlers = ImmutableDictionary>>>.Empty; @@ -129,65 +133,6 @@ public async ValueTask RegisterExtensionInCurrentProc return new( [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], [.. assemblyHandlers.DocumentMessageHandlers.Keys]); - - //Extension? extension; - //using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - //{ - // // Check if the assembly is already loaded. - // if (!_extensions.TryGetValue(assemblyFolderPath, out extension)) - // { - // try - // { - // var analyzerAssemblyLoader = analyzerAssemblyLoaderProvider.CreateNewShadowCopyLoader(); - - // // Allow this assembly loader to load any dll in assemblyFolderPath. - // foreach (var dll in Directory.EnumerateFiles(assemblyFolderPath, "*.dll")) - // { - // try - // { - // // Check if the file is a valid .NET assembly. - // AssemblyName.GetAssemblyName(dll); - // } - // catch - // { - // // The file is not a valid .NET assembly, skip it. - // continue; - // } - - // analyzerAssemblyLoader.AddDependencyLocation(dll); - // } - - // extension = new Extension(this, analyzerAssemblyLoader, assemblyFolderPath); - // } - // catch - // { - // _extensions[assemblyFolderPath] = null; - // throw; - // } - - // _extensions[assemblyFolderPath] = extension; - // } - - // if (extension is null) - // { - // throw new InvalidOperationException($"A previous attempt to load assemblies from {assemblyFolderPath} failed."); - // } - - // if (extension.TryGetAssemblyHandlers(assemblyFileName, out var assemblyHandlers)) - // { - // if (assemblyHandlers is null) - // { - // throw new InvalidOperationException($"A previous attempt to load {assemblyFilePath} failed."); - // } - - // return new( - // [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], - // [.. assemblyHandlers.DocumentMessageHandlers.Keys]); - // } - //} - - //// Intentionally call this outside of the lock. - //return await extension.LoadAssemblyAsync(assemblyFileName, cancellationToken).ConfigureAwait(false); } public async ValueTask UnregisterExtensionAsync( @@ -205,7 +150,6 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( string assemblyFilePath, CancellationToken cancellationToken) { - // var assemblyFileName = Path.GetFileName(assemblyFilePath); var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); @@ -214,7 +158,7 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( { // 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 (extension.Unregister(assemblyFileName)) + if (extension.UnregisterHandlers(assemblyFilePath)) _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); } diff --git a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs index f72272f641816..dafdb4c7379ca 100644 --- a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs +++ b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs @@ -20,7 +20,12 @@ internal sealed class ExtensionRegisterHandler() { private const string MethodName = "roslyn/extensionRegister"; - public bool MutatesSolutionState => false; + /// + /// 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 RequiresLSPSolution => true; diff --git a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs index 42ca41e4d1e71..c9d1ce218310d 100644 --- a/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs +++ b/src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.cs @@ -20,7 +20,12 @@ internal sealed class ExtensionUnregisterHandler() { private const string MethodName = "roslyn/extensionUnregister"; - public bool MutatesSolutionState => false; + /// + /// 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 RequiresLSPSolution => true; From 05020453e32aa8c3300a41a5338b42cbf2885a63 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 20:01:35 -0700 Subject: [PATCH 10/22] simplify --- .../ExtensionMessageHandlerService.cs | 125 +----------------- .../IExtensionMessageHandlerService.cs | 47 ++++--- 2 files changed, 27 insertions(+), 145 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index fa5cda3ccd556..cd8596b874974 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -46,7 +46,7 @@ internal sealed class ExtensionMessageHandlerService( // 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. This allows for a safe approach to + // safely if missing. /// /// Extensions assembly load contexts and loaded handlers, indexed by extension folder path. @@ -164,41 +164,6 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. ClearCachedHandlers(); - - //Extension? extension = null; - //using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - //{ - // if (_folderPathToExtension.TryGetValue(assemblyFolderPath, out extension)) - // { - // // If loading assemblies from this folder failed earlier, don't do anything. - // if (extension is null) - // return default; - - // if (extension.RemoveAssemblyHandlers(assemblyFileName, out var assemblyHandlers)) - // { - // if (assemblyHandlers is not null) - // { - // foreach (var workspaceHandler in assemblyHandlers.WorkspaceMessageHandlers.Keys) - // { - // _workspaceHandlers.Remove(workspaceHandler); - // } - - // foreach (var documentHandler in assemblyHandlers.DocumentMessageHandlers.Keys) - // { - // _documentHandlers.Remove(documentHandler); - // } - // } - // } - - // if (extension.AssemblyHandlersCount > 0) - // return default; - - // _folderPathToExtension.Remove(assemblyFolderPath); - // } - //} - - //extension?.AnalyzerAssemblyLoader.Dispose(); - //return default; } public async ValueTask ResetAsync(CancellationToken cancellationToken) @@ -293,94 +258,6 @@ private static async Task HandleExtensionWorkspaceMessageInCurrentProcessAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) - // => HandleExtensionMessageAsync(solution, messageName, jsonMessage, _workspaceHandlers, cancellationToken); - - //public ValueTask HandleExtensionDocumentMessageInCurrentProcessAsync(Document document, string messageName, string jsonMessage, CancellationToken cancellationToken) - // => HandleExtensionMessageAsync(document, messageName, jsonMessage, _documentHandlers, cancellationToken); - - //private async ValueTask HandleExtensionMessageAsync( - // TArgument argument, - // string messageName, - // string jsonMessage, - // ImmutableDictionary> handlers, - // CancellationToken cancellationToken) - //{ - - // var lazy = _cachedDocumentHandlers. - // IExtensionMessageHandlerWrapper? handler; - // using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - // { - // // handlers here is either _workspaceHandlers or _documentHandlers, so it must be protected - // // by _lockObject. - // if (!handlers.TryGetValue(messageName, out handler)) - // { - // throw new InvalidOperationException($"No handler found for message {messageName}."); - // } - // } - - // try - // { - // var message = JsonSerializer.Deserialize(jsonMessage, handler.MessageType); - // var result = await handler.ExecuteAsync(message, argument, cancellationToken) - // .ConfigureAwait(false); - // var responseJson = JsonSerializer.Serialize(result, handler.ResponseType); - // return responseJson; - // } - // catch - // { - // // Any exception thrown in this method is left to bubble up to the extension. - // // But we unregister all handlers from that assembly to minimize the impact of a bad extension. - // await UnregisterExtensionAsync(assemblyFilePath: handler.ExtensionIdentifier, cancellationToken).ConfigureAwait(false); - // throw; - // } - //} - - //private async Task RegisterAssemblyAsync( - // Extension extension, - // string assemblyFileName, - // AssemblyHandlers? assemblyHandlers, - // CancellationToken cancellationToken) - //{ - // using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false)) - // { - // // Make sure a call to UnloadCustomMessageHandlersAsync hasn't happened while we relinquished the lock on _lockObject - // if (!_extensions.TryGetValue(extension.AssemblyFolderPath, out var currentExtension) || !extension.Equals(currentExtension)) - // { - // throw new InvalidOperationException($"Handlers in {extension.AssemblyFolderPath} were unregistered while loading handlers."); - // } - - // try - // { - // if (assemblyHandlers is not null) - // { - // var duplicateHandler = _workspaceHandlers.Keys.Intersect(assemblyHandlers.WorkspaceMessageHandlers.Keys).Concat( - // _documentHandlers.Keys.Intersect(assemblyHandlers.DocumentMessageHandlers.Keys)).FirstOrDefault(); - - // if (duplicateHandler is not null) - // { - // assemblyHandlers = null; - // throw new InvalidOperationException($"Handler name {duplicateHandler} is already registered."); - // } - - // foreach (var handler in assemblyHandlers.WorkspaceMessageHandlers) - // { - // _workspaceHandlers.Add(handler.Key, handler.Value); - // } - - // foreach (var handler in assemblyHandlers.DocumentMessageHandlers) - // { - // _documentHandlers.Add(handler.Key, handler.Value); - // } - // } - // } - // finally - // { - // extension.SetAssemblyHandlers(assemblyFileName, assemblyHandlers); - // } - // } - //} - private sealed class ExtensionFolder( ExtensionMessageHandlerService extensionMessageHandlerService, IAnalyzerAssemblyLoaderInternal analyzerAssemblyLoader, diff --git a/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs index 0692a71cb759d..af16a65172367 100644 --- a/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs @@ -13,14 +13,38 @@ namespace Microsoft.CodeAnalysis.Extensions; /// internal interface IExtensionMessageHandlerService : IWorkspaceService { + /// + /// Registers extension message handlers from the specified assembly. + /// + /// 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); + + /// + /// Unregisters all extension message handlers. + /// + /// Should be called serially with other , , or calls. + ValueTask ResetAsync(CancellationToken cancellationToken); + /// /// Executes a non-document-specific extension message handler with the given message and solution. /// /// The solution the message refers to. /// 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. - /// Cancellation token to cancel the async operation. /// The json message returned by the handler. + /// Can be called concurrently with other message requests. ValueTask HandleExtensionWorkspaceMessageAsync( Solution solution, string messageName, @@ -33,30 +57,11 @@ ValueTask HandleExtensionWorkspaceMessageAsync( /// The document the message refers to. /// 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. - /// Cancellation token to cancel the async operation. /// The json message returned by the handler. + /// Can be called concurrently with other message requests. ValueTask HandleExtensionDocumentMessageAsync( Document documentId, string messageName, string jsonMessage, CancellationToken cancellationToken); - - /// - /// Registers extension message handlers from the specified assembly. - /// - /// The assembly to register and create message handlers from. - /// The names of the registered handlers. - ValueTask RegisterExtensionAsync(string assemblyFilePath, CancellationToken cancellationToken); - - /// - /// Unregisters extension message handlers previously registered from . - /// - /// The assembly for which handlers should be unregistered. - /// A task representing the async operation. - ValueTask UnregisterExtensionAsync(string assemblyFilePath, CancellationToken cancellationToken); - - /// - /// Unregisters all extension message handlers. - /// - ValueTask ResetAsync(CancellationToken cancellationToken); } From e22cdc986831628ebe3be65c5b6171626f49dd0d Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 20:06:58 -0700 Subject: [PATCH 11/22] simplify --- .../ExtensionMessageHandlerService.cs | 34 ++++++++++++++----- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index cd8596b874974..fd2b11aa889fa 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -119,7 +119,7 @@ public async ValueTask RegisterExtensionInCurrentProc ref _folderPathToExtensionFolder, assemblyFolderPath, static (assemblyFolderPath, @this) => AsyncLazy.Create( - cancellationToken => ExtensionFolder.CreateAsync(@this, assemblyFolderPath, cancellationToken)), + cancellationToken => ExtensionFolder.Create(@this, assemblyFolderPath, cancellationToken)), this); var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); @@ -153,17 +153,26 @@ 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 extension) && - extension != null) + if (!_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazy)) + return default; + + var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (extensionFolder is null) + { + // was an extension that failed to load. Remove the entry marking that so wee can try loading it again in the future. + _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); + } + else { // 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 (extension.UnregisterHandlers(assemblyFilePath)) + if (extensionFolder.UnregisterHandlers(assemblyFilePath)) _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); } // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. ClearCachedHandlers(); + return default; } public async ValueTask ResetAsync(CancellationToken cancellationToken) @@ -260,8 +269,7 @@ private static async Task> _assemblyFilePathToHandlers = ImmutableDictionary>.Empty; - public static async Task CreateAsync( + public static ExtensionFolder? Create( ExtensionMessageHandlerService extensionMessageHandlerService, string assemblyFolderPath, CancellationToken cancellationToken) @@ -298,7 +306,7 @@ private sealed class ExtensionFolder( analyzerAssemblyLoader.AddDependencyLocation(dll); } - return new ExtensionFolder(extensionMessageHandlerService, analyzerAssemblyLoader, assemblyFolderPath); + return new ExtensionFolder(extensionMessageHandlerService, analyzerAssemblyLoader); } catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical)) { @@ -398,6 +406,16 @@ public async Task AddHandlersAsync(string messageName, bool solution, A } } } + + /// + /// 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. + /// + public bool UnregisterHandlers(string assemblyFilePath) + { + _assemblyFilePathToHandlers = _assemblyFilePathToHandlers.Remove(assemblyFilePath); + return _assemblyFilePathToHandlers.IsEmpty; + } } private sealed class AssemblyHandlers From 0fe035253bb7e7d0634d8e1359511180ace1f930 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 20:10:01 -0700 Subject: [PATCH 12/22] simplify --- .../ExtensionMessageHandlerService.cs | 42 ++++--------------- .../IExtensionMessageHandlerFactory.cs | 7 +++- .../ExtensionMessageHandlerFactory.cs | 18 +++++--- .../Extensions/InternalAPI.Unshipped.txt | 4 +- 4 files changed, 29 insertions(+), 42 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index fd2b11aa889fa..54fda5a5aa502 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -179,12 +179,12 @@ public async ValueTask ResetAsync(CancellationToken cancellationToken) { await ExecuteInRemoteOrCurrentProcessAsync( solution: null, - cancellationToken => ResetInCurrentProcessAsync(cancellationToken), + _ => ResetInCurrentProcessAsync(), (service, _, cancellationToken) => service.ResetAsync(cancellationToken), cancellationToken).ConfigureAwait(false); } - private ValueTask ResetInCurrentProcessAsync(CancellationToken cancellationToken) + private ValueTask ResetInCurrentProcessAsync() { _folderPathToExtensionFolder = ImmutableDictionary>.Empty; ClearCachedHandlers(); @@ -315,33 +315,6 @@ private sealed class ExtensionFolder( } } - //public void SetAssemblyHandlers(string assemblyFileName, AssemblyHandlers? value) - //{ - // EnsureGlobalLockIsOwned(); - // _assemblies[assemblyFileName] = value; - //} - - //public bool TryGetAssemblyHandlers(string assemblyFileName, out AssemblyHandlers? value) - //{ - // EnsureGlobalLockIsOwned(); - // return _assemblies.TryGetValue(assemblyFileName, out value); - //} - - //public bool RemoveAssemblyHandlers(string assemblyFileName, out AssemblyHandlers? value) - //{ - // EnsureGlobalLockIsOwned(); - // return _assemblies.Remove(assemblyFileName, out value); - //} - - //public int AssemblyHandlersCount - //{ - // get - // { - // EnsureGlobalLockIsOwned(); - // return _assemblies.Count; - // } - //} - public async ValueTask GetAssemblyHandlersAsync( string assemblyFilePath, CancellationToken cancellationToken) { @@ -349,23 +322,26 @@ private sealed class ExtensionFolder( ref _assemblyFilePathToHandlers, assemblyFilePath, static (assemblyFilePath, @this) => AsyncLazy.Create( - static (args, cancellationToken) => CreateAssemblyHandlersAsync(args.@this, args.assemblyFilePath, cancellationToken), + static (args, cancellationToken) => CreateAssemblyHandlers(args.@this, args.assemblyFilePath, cancellationToken), (assemblyFilePath, @this)), this); return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); } - private static async Task CreateAssemblyHandlersAsync( + private static AssemblyHandlers? CreateAssemblyHandlers( ExtensionFolder @this, string assemblyFilePath, CancellationToken cancellationToken) { try { var assembly = @this._analyzerAssemblyLoader.LoadFromPath(assemblyFilePath); var factory = @this._extensionMessageHandlerService._customMessageHandlerFactory; - var messageWorkspaceHandlers = factory.CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath) + + var messageWorkspaceHandlers = factory + .CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath, cancellationToken) .ToImmutableDictionary(h => h.Name, h => h); - var messageDocumentHandlers = factory.CreateDocumentMessageHandlers(assembly, extensionIdentifier: assemblyFilePath) + var messageDocumentHandlers = factory + .CreateDocumentMessageHandlers(assembly, extensionIdentifier: assemblyFilePath, cancellationToken) .ToImmutableDictionary(h => h.Name, h => h); return new AssemblyHandlers() diff --git a/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerFactory.cs b/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerFactory.cs index 8e2b084368495..7f19b218d6c9c 100644 --- a/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerFactory.cs +++ b/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerFactory.cs @@ -4,6 +4,7 @@ using System.Collections.Immutable; using System.Reflection; +using System.Threading; namespace Microsoft.CodeAnalysis.Extensions; @@ -19,7 +20,8 @@ internal interface IExtensionMessageHandlerFactory /// The assembly to scan for handlers. /// Unique identifier of the extension owning this handler. /// The handlers. - ImmutableArray> CreateWorkspaceMessageHandlers(Assembly assembly, string extensionIdentifier); + ImmutableArray> CreateWorkspaceMessageHandlers( + Assembly assembly, string extensionIdentifier, CancellationToken cancellationToken); /// /// Creates instances for each @@ -28,5 +30,6 @@ internal interface IExtensionMessageHandlerFactory /// The assembly to scan for handlers. /// Unique identifier of the extension owning this handler. /// The handlers. - ImmutableArray> CreateDocumentMessageHandlers(Assembly assembly, string extensionIdentifier); + ImmutableArray> CreateDocumentMessageHandlers( + Assembly assembly, string extensionIdentifier, CancellationToken cancellationToken); } diff --git a/src/Tools/ExternalAccess/Extensions/Internal/ExtensionMessageHandlerFactory.cs b/src/Tools/ExternalAccess/Extensions/Internal/ExtensionMessageHandlerFactory.cs index ddd0a3b317c44..9f8bc626bd28d 100644 --- a/src/Tools/ExternalAccess/Extensions/Internal/ExtensionMessageHandlerFactory.cs +++ b/src/Tools/ExternalAccess/Extensions/Internal/ExtensionMessageHandlerFactory.cs @@ -6,6 +6,7 @@ using System.Collections.Immutable; using System.Composition; using System.Reflection; +using System.Threading; using Microsoft.CodeAnalysis.Host.Mef; namespace Microsoft.CodeAnalysis.Extensions; @@ -15,27 +16,34 @@ namespace Microsoft.CodeAnalysis.Extensions; [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)] internal sealed class ExtensionMessageHandlerFactory() : IExtensionMessageHandlerFactory { - public ImmutableArray> CreateDocumentMessageHandlers(Assembly assembly, string extensionIdentifier) + public ImmutableArray> CreateDocumentMessageHandlers( + Assembly assembly, string extensionIdentifier, CancellationToken cancellationToken) => CreateWorkspaceHandlers( assembly, typeof(IExtensionDocumentMessageHandler<,>), - (handler, handlerInterface) => new ExtensionDocumentMessageHandlerWrapper(handler, handlerInterface, extensionIdentifier)); + (handler, handlerInterface) => new ExtensionDocumentMessageHandlerWrapper(handler, handlerInterface, extensionIdentifier), + cancellationToken); - public ImmutableArray> CreateWorkspaceMessageHandlers(Assembly assembly, string extensionIdentifier) + public ImmutableArray> CreateWorkspaceMessageHandlers( + Assembly assembly, string extensionIdentifier, CancellationToken cancellationToken) => CreateWorkspaceHandlers( assembly, typeof(IExtensionWorkspaceMessageHandler<,>), - (handler, handlerInterface) => new ExtensionWorkspaceMessageHandlerWrapper(handler, handlerInterface, extensionIdentifier)); + (handler, handlerInterface) => new ExtensionWorkspaceMessageHandlerWrapper(handler, handlerInterface, extensionIdentifier), + cancellationToken); private static ImmutableArray> CreateWorkspaceHandlers( Assembly assembly, Type unboundInterfaceType, - Func> wrapperCreator) + Func> wrapperCreator, + CancellationToken cancellationToken) { var resultBuilder = ImmutableArray.CreateBuilder>(); foreach (var candidateType in assembly.GetTypes()) { + cancellationToken.ThrowIfCancellationRequested(); + if (candidateType.IsAbstract || candidateType.IsGenericType) { continue; diff --git a/src/Tools/ExternalAccess/Extensions/InternalAPI.Unshipped.txt b/src/Tools/ExternalAccess/Extensions/InternalAPI.Unshipped.txt index 585364db866e3..d556753bb5665 100644 --- a/src/Tools/ExternalAccess/Extensions/InternalAPI.Unshipped.txt +++ b/src/Tools/ExternalAccess/Extensions/InternalAPI.Unshipped.txt @@ -27,8 +27,8 @@ Microsoft.CodeAnalysis.Extensions.ExtensionHandlerWrapper.Name.get -> Microsoft.CodeAnalysis.Extensions.ExtensionHandlerWrapper.ResponseType.get -> System.Type! Microsoft.CodeAnalysis.Extensions.ExtensionMessageContext.ExtensionMessageContext(Microsoft.CodeAnalysis.Solution! solution) -> void Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory -Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.CreateDocumentMessageHandlers(System.Reflection.Assembly! assembly, string! extensionIdentifier) -> System.Collections.Immutable.ImmutableArray!> -Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.CreateWorkspaceMessageHandlers(System.Reflection.Assembly! assembly, string! extensionIdentifier) -> System.Collections.Immutable.ImmutableArray!> +Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.CreateDocumentMessageHandlers(System.Reflection.Assembly! assembly, string! extensionIdentifier, System.Threading.CancellationToken cancellationToken) -> System.Collections.Immutable.ImmutableArray!> +Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.CreateWorkspaceMessageHandlers(System.Reflection.Assembly! assembly, string! extensionIdentifier, System.Threading.CancellationToken cancellationToken) -> System.Collections.Immutable.ImmutableArray!> Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.ExtensionMessageHandlerFactory() -> void Microsoft.CodeAnalysis.Extensions.ExtensionWorkspaceMessageHandlerWrapper Microsoft.CodeAnalysis.Extensions.ExtensionWorkspaceMessageHandlerWrapper.ExtensionWorkspaceMessageHandlerWrapper(object! handler, System.Type! customMessageHandlerInterface, string! extensionIdentifier) -> void \ No newline at end of file From eb6c6f6c32edd6e02f49b79c7253646329152aa3 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 20:12:02 -0700 Subject: [PATCH 13/22] in progress --- .../Core/Portable/Extensions/ExtensionMessageHandlerService.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 54fda5a5aa502..80a3d03d5f001 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -220,7 +220,7 @@ private async ValueTask HandleExtensionWorkspaceMessageInCurrentProcessA var lazy = _cachedWorkspaceHandlers.GetOrAdd( messageName, static (messageName, @this) => AsyncLazy.Create( - static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, cancellationToken), + static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, true, cancellationToken), (messageName, @this)), this); From cda45a38eccec3c0e4a5c08518cd6625d4fcbf36 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 20:15:29 -0700 Subject: [PATCH 14/22] in progress --- .../ExtensionMessageHandlerService.cs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 80a3d03d5f001..0c66eff429054 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -215,14 +215,18 @@ public async ValueTask HandleExtensionDocumentMessageAsync(Document docu cancellationToken).ConfigureAwait(false); } - private async ValueTask HandleExtensionWorkspaceMessageInCurrentProcessAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken) + private async ValueTask HandleExtensionMessageInCurrentProcessAsync( + TArgument executeArgument, bool solution, string messageName, string jsonMessage, + ImmutableDictionary>>> cachedHandlers, + CancellationToken cancellationToken) { - var lazy = _cachedWorkspaceHandlers.GetOrAdd( + var lazy = ImmutableInterlocked.GetOrAdd( + ref cachedHandlers, messageName, - static (messageName, @this) => AsyncLazy.Create( - static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, true, cancellationToken), - (messageName, @this)), - this); + static (messageName, arg) => AsyncLazy.Create( + static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, arg.solution, cancellationToken), + (messageName, arg.@this, arg.solution)), + (@this: this, executeArgument, solution)); var handlers = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); if (handlers.Length == 0) @@ -236,7 +240,7 @@ private async ValueTask HandleExtensionWorkspaceMessageInCurrentProcessA try { var message = JsonSerializer.Deserialize(jsonMessage, handler.MessageType); - var result = await handler.ExecuteAsync(message, solution, cancellationToken) + var result = await handler.ExecuteAsync(message, executeArgument, cancellationToken) .ConfigureAwait(false); var responseJson = JsonSerializer.Serialize(result, handler.ResponseType); return responseJson; From a20078d563d345e1ae43ffa54b48dc0f75b4e52f Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 20:18:14 -0700 Subject: [PATCH 15/22] Reset code --- .../ExtensionMessageHandlerService.cs | 6 ++++-- .../LspServiceLifeCycleManager.cs | 20 ++++++------------- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 0c66eff429054..019689e466d6a 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -201,7 +201,8 @@ public async ValueTask HandleExtensionWorkspaceMessageAsync(Solution sol { return await ExecuteInRemoteOrCurrentProcessAsync( solution, - cancellationToken => HandleExtensionWorkspaceMessageInCurrentProcessAsync(solution, messageName, jsonMessage, cancellationToken), + cancellationToken => HandleExtensionMessageInCurrentProcessAsync( + executeArgument: solution, solution: true, messageName, jsonMessage, _cachedWorkspaceHandlers, cancellationToken), (service, checksum, cancellationToken) => service.HandleExtensionWorkspaceMessageAsync(checksum!.Value, messageName, jsonMessage, cancellationToken), cancellationToken).ConfigureAwait(false); } @@ -210,7 +211,8 @@ public async ValueTask HandleExtensionDocumentMessageAsync(Document docu { return await ExecuteInRemoteOrCurrentProcessAsync( document.Project.Solution, - cancellationToken => HandleExtensionDocumentMessageInCurrentProcessAsync(document, messageName, jsonMessage, cancellationToken), + cancellationToken => HandleExtensionMessageInCurrentProcessAsync( + executeArgument: document, solution: false, messageName, jsonMessage, _cachedDocumentHandlers, cancellationToken), (service, checksum, cancellationToken) => service.HandleExtensionDocumentMessageAsync(checksum!.Value, messageName, jsonMessage, document.Id, cancellationToken), cancellationToken).ConfigureAwait(false); } diff --git a/src/LanguageServer/Protocol/Handler/ServerLifetime/LspServiceLifeCycleManager.cs b/src/LanguageServer/Protocol/Handler/ServerLifetime/LspServiceLifeCycleManager.cs index 94cea1ac624d1..94ca5ca32f332 100644 --- a/src/LanguageServer/Protocol/Handler/ServerLifetime/LspServiceLifeCycleManager.cs +++ b/src/LanguageServer/Protocol/Handler/ServerLifetime/LspServiceLifeCycleManager.cs @@ -9,7 +9,6 @@ using System.Threading.Tasks; using Microsoft.CodeAnalysis.Extensions; using Microsoft.CodeAnalysis.Host.Mef; -using Microsoft.CodeAnalysis.Remote; using Microsoft.CommonLanguageServerProtocol.Framework; using Roslyn.LanguageServer.Protocol; using StreamJsonRpc; @@ -41,21 +40,14 @@ private LspServiceLifeCycleManager(IClientLanguageServerManager clientLanguageSe public async Task ShutdownAsync(string message = "Shutting down") { + // Shutting down is not cancellable. + var cancellationToken = CancellationToken.None; + var hostWorkspace = _lspWorkspaceRegistrationService.GetAllRegistrations().SingleOrDefault(w => w.Kind == WorkspaceKind.Host); if (hostWorkspace is not null) { - var client = await RemoteHostClient.TryGetClientAsync(hostWorkspace, CancellationToken.None).ConfigureAwait(false); - if (client is not null) - { - await client.TryInvokeAsync( - (service, cancellationToken) => service.ResetAsync(cancellationToken), - CancellationToken.None).ConfigureAwait(false); - } - else - { - var service = hostWorkspace.Services.GetRequiredService(); - service.Reset(); - } + var service = hostWorkspace.Services.GetRequiredService(); + await service.ResetAsync(cancellationToken).ConfigureAwait(false); } try @@ -65,7 +57,7 @@ await client.TryInvokeAsync( MessageType = MessageType.Info, Message = message }; - await _clientLanguageServerManager.SendNotificationAsync("window/logMessage", messageParams, CancellationToken.None).ConfigureAwait(false); + await _clientLanguageServerManager.SendNotificationAsync("window/logMessage", messageParams, cancellationToken).ConfigureAwait(false); } catch (Exception ex) when (ex is ObjectDisposedException or ConnectionLostException) { From 2a1de8856ed4a0d50d2fe3d2f5d98fe39fa15829 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 23:48:30 -0700 Subject: [PATCH 16/22] Clelar --- .../ExtensionMessageHandlerService.cs | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 019689e466d6a..9460dc1d353a6 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -113,8 +113,6 @@ public async ValueTask RegisterExtensionInCurrentProc var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath) ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}."); - // var analyzerAssemblyLoaderProvider = _solutionServices.GetRequiredService(); - var lazy = ImmutableInterlocked.GetOrAdd( ref _folderPathToExtensionFolder, assemblyFolderPath, @@ -130,6 +128,9 @@ public async ValueTask RegisterExtensionInCurrentProc if (assemblyHandlers is null) throw new InvalidOperationException($"Loading extensions from {assemblyFilePath} failed."); + // After registering, clear out the cached handler names. They will be recomputed the next time we need them. + ClearCachedHandlers(); + return new( [.. assemblyHandlers.WorkspaceMessageHandlers.Keys], [.. assemblyHandlers.DocumentMessageHandlers.Keys]); @@ -153,21 +154,21 @@ 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)) - return default; - - var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - if (extensionFolder is null) + if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazy)) { - // was an extension that failed to load. Remove the entry marking that so wee can try loading it again in the future. - _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); - } - else - { - // 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.UnregisterHandlers(assemblyFilePath)) + var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); + if (extensionFolder is null) + { + // was an extension that failed to load. Remove the entry marking that so wee can try loading it again in the future. _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); + } + else + { + // 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.UnregisterHandlers(assemblyFilePath)) + _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); + } } // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. From 81c07cc4dcde84df1a596e1b3f604d265d3732fc Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 23:50:04 -0700 Subject: [PATCH 17/22] renames --- .../Extensions/ExtensionMessageHandlerService.cs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 9460dc1d353a6..b727443cbb631 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -203,7 +203,7 @@ public async ValueTask HandleExtensionWorkspaceMessageAsync(Solution sol return await ExecuteInRemoteOrCurrentProcessAsync( solution, cancellationToken => HandleExtensionMessageInCurrentProcessAsync( - executeArgument: solution, solution: true, messageName, jsonMessage, _cachedWorkspaceHandlers, cancellationToken), + executeArgument: solution, isSolution: true, messageName, jsonMessage, _cachedWorkspaceHandlers, cancellationToken), (service, checksum, cancellationToken) => service.HandleExtensionWorkspaceMessageAsync(checksum!.Value, messageName, jsonMessage, cancellationToken), cancellationToken).ConfigureAwait(false); } @@ -213,13 +213,13 @@ public async ValueTask HandleExtensionDocumentMessageAsync(Document docu return await ExecuteInRemoteOrCurrentProcessAsync( document.Project.Solution, cancellationToken => HandleExtensionMessageInCurrentProcessAsync( - executeArgument: document, solution: false, messageName, jsonMessage, _cachedDocumentHandlers, cancellationToken), + executeArgument: document, isSolution: false, messageName, jsonMessage, _cachedDocumentHandlers, cancellationToken), (service, checksum, cancellationToken) => service.HandleExtensionDocumentMessageAsync(checksum!.Value, messageName, jsonMessage, document.Id, cancellationToken), cancellationToken).ConfigureAwait(false); } private async ValueTask HandleExtensionMessageInCurrentProcessAsync( - TArgument executeArgument, bool solution, string messageName, string jsonMessage, + TArgument executeArgument, bool isSolution, string messageName, string jsonMessage, ImmutableDictionary>>> cachedHandlers, CancellationToken cancellationToken) { @@ -227,9 +227,9 @@ private async ValueTask HandleExtensionMessageInCurrentProcessAsync AsyncLazy.Create( - static (arg, cancellationToken) => ComputeHandlersAsync(arg.@this, arg.messageName, arg.solution, cancellationToken), - (messageName, arg.@this, arg.solution)), - (@this: this, executeArgument, solution)); + 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); if (handlers.Length == 0) @@ -258,7 +258,7 @@ private async ValueTask HandleExtensionMessageInCurrentProcessAsync>> ComputeHandlersAsync( - ExtensionMessageHandlerService @this, string messageName, bool solution, CancellationToken cancellationToken) + ExtensionMessageHandlerService @this, string messageName, bool isSolution, CancellationToken cancellationToken) { using var _ = ArrayBuilder>.GetInstance(out var result); foreach (var (_, lazyExtension) in @this._folderPathToExtensionFolder) @@ -268,7 +268,7 @@ private static async Task Date: Wed, 2 Apr 2025 23:56:51 -0700 Subject: [PATCH 18/22] Simplify --- .../ExtensionMessageHandlerService.cs | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index b727443cbb631..a4e3ddba60313 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -76,7 +76,7 @@ private async ValueTask ExecuteInRemoteOrCurrentProcessAsync( if (solution is null) { var result = await client.TryInvokeAsync( - (service, cancellationToken) => executeOutOfProcessAsync(service, null, cancellationToken), + (remoteService, cancellationToken) => executeOutOfProcessAsync(remoteService, null, cancellationToken), cancellationToken).ConfigureAwait(false); // If the remote call succeeded, this will have a valid valid in it and can be returned. If it did not @@ -88,7 +88,7 @@ private async ValueTask ExecuteInRemoteOrCurrentProcessAsync( { var result = await client.TryInvokeAsync( solution, - (service, checksum, cancellationToken) => executeOutOfProcessAsync(service, checksum, cancellationToken), + (remoteService, checksum, cancellationToken) => executeOutOfProcessAsync(remoteService, checksum, cancellationToken), cancellationToken).ConfigureAwait(false); return result.Value; } @@ -101,7 +101,7 @@ public async ValueTask RegisterExtensionAsync( return await ExecuteInRemoteOrCurrentProcessAsync( solution: null, cancellationToken => RegisterExtensionInCurrentProcessAsync(assemblyFilePath, cancellationToken), - (service, _, cancellationToken) => service.RegisterExtensionAsync(assemblyFilePath, cancellationToken), + (remoteService, _, cancellationToken) => remoteService.RegisterExtensionAsync(assemblyFilePath, cancellationToken), cancellationToken).ConfigureAwait(false); } @@ -143,7 +143,7 @@ public async ValueTask UnregisterExtensionAsync( await ExecuteInRemoteOrCurrentProcessAsync( solution: null, cancellationToken => UnregisterExtensionInCurrentProcessAsync(assemblyFilePath, cancellationToken), - (service, _, cancellationToken) => service.UnregisterExtensionAsync(assemblyFilePath, cancellationToken), + (remoteService, _, cancellationToken) => remoteService.UnregisterExtensionAsync(assemblyFilePath, cancellationToken), cancellationToken).ConfigureAwait(false); } @@ -181,7 +181,7 @@ public async ValueTask ResetAsync(CancellationToken cancellationToken) await ExecuteInRemoteOrCurrentProcessAsync( solution: null, _ => ResetInCurrentProcessAsync(), - (service, _, cancellationToken) => service.ResetAsync(cancellationToken), + (remoteService, _, cancellationToken) => remoteService.ResetAsync(cancellationToken), cancellationToken).ConfigureAwait(false); } @@ -204,7 +204,7 @@ public async ValueTask HandleExtensionWorkspaceMessageAsync(Solution sol solution, cancellationToken => HandleExtensionMessageInCurrentProcessAsync( executeArgument: solution, isSolution: true, messageName, jsonMessage, _cachedWorkspaceHandlers, cancellationToken), - (service, checksum, cancellationToken) => service.HandleExtensionWorkspaceMessageAsync(checksum!.Value, messageName, jsonMessage, cancellationToken), + (remoteService, checksum, cancellationToken) => remoteService.HandleExtensionWorkspaceMessageAsync(checksum!.Value, messageName, jsonMessage, cancellationToken), cancellationToken).ConfigureAwait(false); } @@ -214,7 +214,7 @@ public async ValueTask HandleExtensionDocumentMessageAsync(Document docu document.Project.Solution, cancellationToken => HandleExtensionMessageInCurrentProcessAsync( executeArgument: document, isSolution: false, messageName, jsonMessage, _cachedDocumentHandlers, cancellationToken), - (service, checksum, cancellationToken) => service.HandleExtensionDocumentMessageAsync(checksum!.Value, messageName, jsonMessage, document.Id, cancellationToken), + (remoteService, checksum, cancellationToken) => remoteService.HandleExtensionDocumentMessageAsync(checksum!.Value, messageName, jsonMessage, document.Id, cancellationToken), cancellationToken).ConfigureAwait(false); } @@ -243,10 +243,8 @@ private async ValueTask HandleExtensionMessageInCurrentProcessAsync> _assemblyFilePathToHandlers = ImmutableDictionary>.Empty; public static ExtensionFolder? Create( From ddb3e235f889604a0868e01de2fa73ed074841cd Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Wed, 2 Apr 2025 23:57:07 -0700 Subject: [PATCH 19/22] rename --- .../Portable/Extensions/ExtensionMessageHandlerService.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index a4e3ddba60313..3485cb1318351 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -363,7 +363,7 @@ private sealed class ExtensionFolder( } } - public async Task AddHandlersAsync(string messageName, bool solution, ArrayBuilder> result, CancellationToken cancellationToken) + public async Task AddHandlersAsync(string messageName, bool isSolution, ArrayBuilder> result, CancellationToken cancellationToken) { foreach (var (_, lazy) in _assemblyFilePathToHandlers) { @@ -373,7 +373,7 @@ public async Task AddHandlersAsync(string messageName, bool solution, A if (handlers is null) continue; - if (solution) + if (isSolution) { if (handlers.WorkspaceMessageHandlers.TryGetValue(messageName, out var handler)) result.Add((IExtensionMessageHandlerWrapper)handler); From 2e02b4836ae15239cf3a37aab9844fd2a4bb8b64 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Thu, 3 Apr 2025 00:12:36 -0700 Subject: [PATCH 20/22] No more null --- .../ExtensionMessageHandlerService.cs | 85 +++++++++++++------ 1 file changed, 57 insertions(+), 28 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 3485cb1318351..b2c125ec99052 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -51,7 +51,7 @@ internal sealed class ExtensionMessageHandlerService( /// /// Extensions assembly load contexts and loaded handlers, indexed by extension folder path. /// - private ImmutableDictionary> _folderPathToExtensionFolder = ImmutableDictionary>.Empty; + private ImmutableDictionary> _folderPathToExtensionFolder = ImmutableDictionary>.Empty; /// /// Cached handlers of document-related messages, indexed by handler message name. @@ -121,12 +121,7 @@ public async ValueTask RegisterExtensionInCurrentProc this); var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - if (extensionFolder is null) - throw new InvalidOperationException($"Loading extensions from {assemblyFolderPath} failed."); - - var assemblyHandlers = await extensionFolder.GetAssemblyHandlersAsync(assemblyFilePath, cancellationToken).ConfigureAwait(false); - if (assemblyHandlers is null) - throw new InvalidOperationException($"Loading extensions from {assemblyFilePath} failed."); + 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(); @@ -166,7 +161,7 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( { // 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.UnregisterHandlers(assemblyFilePath)) + if (await extensionFolder.UnregisterAssemblyAsync(assemblyFilePath, cancellationToken).ConfigureAwait(false)) _folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath); } } @@ -187,7 +182,7 @@ await ExecuteInRemoteOrCurrentProcessAsync( private ValueTask ResetInCurrentProcessAsync() { - _folderPathToExtensionFolder = ImmutableDictionary>.Empty; + _folderPathToExtensionFolder = ImmutableDictionary>.Empty; ClearCachedHandlers(); return default; } @@ -272,16 +267,51 @@ 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. + /// + ValueTask UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken); + + ValueTask AddHandlersAsync(string messageName, bool isSolution, ArrayBuilder> result, CancellationToken cancellationToken); + } + + private sealed class TrivialExtensionFolder : IExtensionFolder + { + public static readonly TrivialExtensionFolder Instance = new(); + + private readonly List _registeredFilePaths = []; + + public ValueTask RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken) + { + _registeredFilePaths.Add(assemblyFilePath); + return new(AssemblyHandlers.Empty); + } + + public ValueTask UnregisterAssemblyAsync(string assemblyFilePath, 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) + IAnalyzerAssemblyLoaderInternal analyzerAssemblyLoader) : IExtensionFolder { private readonly ExtensionMessageHandlerService _extensionMessageHandlerService = extensionMessageHandlerService; private readonly IAnalyzerAssemblyLoaderInternal _analyzerAssemblyLoader = analyzerAssemblyLoader; - private ImmutableDictionary> _assemblyFilePathToHandlers = ImmutableDictionary>.Empty; + private ImmutableDictionary> _assemblyFilePathToHandlers = ImmutableDictionary>.Empty; - public static ExtensionFolder? Create( + public static IExtensionFolder Create( ExtensionMessageHandlerService extensionMessageHandlerService, string assemblyFolderPath, CancellationToken cancellationToken) @@ -313,12 +343,12 @@ private sealed class ExtensionFolder( } catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical)) { - // If loading the assembly fails, we don't want to cache it. - return null; + // TODO: Log this exception so the client knows something went wrong. + return new TrivialExtensionFolder(); } } - public async ValueTask GetAssemblyHandlersAsync( + public async ValueTask RegisterAssemblyAsync( string assemblyFilePath, CancellationToken cancellationToken) { var lazy = ImmutableInterlocked.GetOrAdd( @@ -332,7 +362,7 @@ private sealed class ExtensionFolder( return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); } - private static AssemblyHandlers? CreateAssemblyHandlers( + private static AssemblyHandlers CreateAssemblyHandlers( ExtensionFolder @this, string assemblyFilePath, CancellationToken cancellationToken) { try @@ -358,21 +388,18 @@ private sealed class ExtensionFolder( } catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.General)) { - // TODO: Log error so it is visible to user. - return null; + // TODO: Log this exception so the client knows something went wrong. + return AssemblyHandlers.Empty; } } - public async Task AddHandlersAsync(string messageName, bool isSolution, ArrayBuilder> result, CancellationToken cancellationToken) + 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 (handlers is null) - continue; - if (isSolution) { if (handlers.WorkspaceMessageHandlers.TryGetValue(messageName, out var handler)) @@ -386,19 +413,21 @@ public async Task AddHandlersAsync(string messageName, bool isSolution, } } - /// - /// 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. - /// - public bool UnregisterHandlers(string assemblyFilePath) + public ValueTask UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken) { _assemblyFilePathToHandlers = _assemblyFilePathToHandlers.Remove(assemblyFilePath); - return _assemblyFilePathToHandlers.IsEmpty; + return new(_assemblyFilePathToHandlers.IsEmpty); } } private sealed class AssemblyHandlers { + public static readonly AssemblyHandlers Empty = new() + { + DocumentMessageHandlers = ImmutableDictionary>.Empty, + WorkspaceMessageHandlers = ImmutableDictionary>.Empty, + }; + /// /// Gets the document-specific handlers that can be passed to , indexed by their name. /// From a8ab20420bc735b412fb5bd8e821c5e08d3aa40d Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Thu, 3 Apr 2025 00:19:04 -0700 Subject: [PATCH 21/22] Simplify --- .../ExtensionMessageHandlerService.cs | 29 +++++++++---------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index b2c125ec99052..696a3fcfb637e 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -152,18 +152,10 @@ private async ValueTask UnregisterExtensionInCurrentProcessAsync( if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazy)) { var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false); - if (extensionFolder is null) - { - // was an extension that failed to load. Remove the entry marking that so wee can try loading it again in the future. + // 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); - } - else - { - // 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); - } } // After unregistering, clear out the cached handler names. They will be recomputed the next time we need them. @@ -257,10 +249,8 @@ private static async Task(string messageName, bool isSolution, ArrayBuilder> result, 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(); + /// + /// No lock needed as registratin/unregistration must happen serially. + /// private readonly List _registeredFilePaths = []; public ValueTask RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken) @@ -372,10 +369,10 @@ private static AssemblyHandlers CreateAssemblyHandlers( var messageWorkspaceHandlers = factory .CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath, cancellationToken) - .ToImmutableDictionary(h => h.Name, h => h); + .ToImmutableDictionary(h => h.Name); var messageDocumentHandlers = factory .CreateDocumentMessageHandlers(assembly, extensionIdentifier: assemblyFilePath, cancellationToken) - .ToImmutableDictionary(h => h.Name, h => h); + .ToImmutableDictionary(h => h.Name); return new AssemblyHandlers() { From 2185e399070b7fdaf78b29c79884bf012ceb6809 Mon Sep 17 00:00:00 2001 From: Cyrus Najmabadi Date: Thu, 3 Apr 2025 00:27:17 -0700 Subject: [PATCH 22/22] Exception throwing --- .../ExtensionMessageHandlerService.cs | 22 ++++++++++++++----- .../IExtensionMessageHandlerWrapper.cs | 6 ++++- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs index 696a3fcfb637e..f13a988e017e1 100644 --- a/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs +++ b/src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs @@ -11,6 +11,7 @@ using System.IO; using System.Linq; using System.Reflection; +using System.Runtime.CompilerServices; using System.Text.Json; using System.Threading; using System.Threading.Tasks; @@ -41,6 +42,8 @@ internal sealed class ExtensionMessageHandlerService( IExtensionMessageHandlerFactory customMessageHandlerFactory) : IExtensionMessageHandlerService { + private static readonly ConditionalWeakTable s_disabledExtensionHandlers = new(); + private readonly SolutionServices _solutionServices = solutionServices; private readonly IExtensionMessageHandlerFactory _customMessageHandlerFactory = customMessageHandlerFactory; @@ -226,6 +229,8 @@ private async ValueTask HandleExtensionMessageInCurrentProcessAsync HandleExtensionMessageInCurrentProcessAsync /// Wrapper for an IExtensionWorkspaceMessageHandler or IExtensionDocumentMessageHandler /// as returned by . /// /// The type of object received as parameter by the extension message /// handler. -internal interface IExtensionMessageHandlerWrapper +internal interface IExtensionMessageHandlerWrapper : IExtensionMessageHandlerWrapper { /// /// The type of object received as parameter by the extension message handler.