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