CHange how locking works

CyrusNajmabadi · CyrusNajmabadi · commit 16a0b29c1072 · 2025-04-03T10:58:30.000-07:00
diff --git 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.
+    /// <summary>
+    /// Lock for <see cref="_folderPathToExtensionFolder"/>, <see cref="_cachedDocumentHandlers"/>, and <see
+    /// cref="_cachedWorkspaceHandlers"/>.
+    /// </summary>
+    private readonly SemaphoreSlim _gate = new(initialCount: 1);
 
     /// <summary>
     /// Extensions assembly load contexts and loaded handlers, indexed by extension folder path.
     /// </summary>
-    private ImmutableDictionary<string, AsyncLazy<IExtensionFolder>> _folderPathToExtensionFolder = ImmutableDictionary<string, AsyncLazy<IExtensionFolder>>.Empty;
+    private readonly Dictionary<string, AsyncLazy<IExtensionFolder>> _folderPathToExtensionFolder = new();
 
     /// <summary>
     /// Cached handlers of document-related messages, indexed by handler message name.
     /// </summary>
-    private ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>> _cachedDocumentHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>>.Empty;
+    private readonly Dictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>> _cachedDocumentHandlers = new();
 
     /// <summary>
     /// Cached handlers of non-document-related messages, indexed by handler message name.
     /// </summary>
-    private ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>> _cachedWorkspaceHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>>.Empty;
+    private readonly Dictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>> _cachedWorkspaceHandlers = new();
 
     private async ValueTask<TResult> ExecuteInRemoteOrCurrentProcessAsync<TResult>(
         Solution? solution,
@@ -116,18 +118,22 @@ public async ValueTask<RegisterExtensionResponse> 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<IExtensionFolder> 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,21 @@ private async ValueTask<VoidResult> UnregisterExtensionInCurrentProcessAsync(
         var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath)
             ?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}.");
 
-        if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazy))
+        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 +187,19 @@ await ExecuteInRemoteOrCurrentProcessAsync(
 
     private ValueTask<VoidResult> ResetInCurrentProcessAsync()
     {
-        _folderPathToExtensionFolder = ImmutableDictionary<string, AsyncLazy<IExtensionFolder>>.Empty;
-        ClearCachedHandlers();
-        return default;
+        lock (_gate)
+        {
+            _folderPathToExtensionFolder.Clear();
+            ClearCachedHandlers();
+            return default;
+        }
     }
 
     private void ClearCachedHandlers()
     {
-        _cachedWorkspaceHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>>.Empty;
-        _cachedDocumentHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>>.Empty;
+        Contract.ThrowIfTrue(!Monitor.IsEntered(_gate));
+        _cachedWorkspaceHandlers.Clear();
+        _cachedDocumentHandlers.Clear();
     }
 
     public async ValueTask<string> HandleExtensionWorkspaceMessageAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken)
@@ -210,18 +224,21 @@ public async ValueTask<string> HandleExtensionDocumentMessageAsync(Document docu
 
     private async ValueTask<string> HandleExtensionMessageInCurrentProcessAsync<TArgument>(
         TArgument executeArgument, bool isSolution, string messageName, string jsonMessage,
-        ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<TArgument>>>> cachedHandlers,
+        Dictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<TArgument>>>> cachedHandlers,
         CancellationToken cancellationToken)
     {
-        var lazy = ImmutableInterlocked.GetOrAdd(
-            ref cachedHandlers,
-            messageName,
-            static (messageName, arg) => AsyncLazy.Create(
-                static (arg, cancellationToken) => ComputeHandlersAsync<TArgument>(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<ImmutableArray<IExtensionMessageHandlerWrapper<TArgument>>> lazyHandlers;
+        using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
+        {
+            lazyHandlers = cachedHandlers.GetOrAdd(
+                messageName,
+                static (messageName, arg) => AsyncLazy.Create(
+                    static (arg, cancellationToken) => ComputeHandlersAsync<TArgument>(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}.");
 
@@ -271,13 +288,13 @@ private static async Task<ImmutableArray<IExtensionMessageHandlerWrapper<TResult
 
     private interface IExtensionFolder
     {
-        ValueTask<AssemblyHandlers> RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken);
+        AsyncLazy<AssemblyHandlers> RegisterAssembly(string assemblyFilePath);
 
         /// <summary>
         /// 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.
         /// </summary>
-        ValueTask<bool> UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken);
+        bool UnregisterAssembly(string assemblyFilePath);
 
         ValueTask AddHandlersAsync<TResult>(string messageName, bool isSolution, ArrayBuilder<IExtensionMessageHandlerWrapper<TResult>> result, CancellationToken cancellationToken);
     }
@@ -291,20 +308,20 @@ private sealed class TrivialExtensionFolder : IExtensionFolder
         public static readonly TrivialExtensionFolder Instance = new();
 
         /// <summary>
-        /// No lock needed as registratin/unregistration must happen serially.
+        /// No lock needed as registration/unregistration must happen serially.
         /// </summary>
         private readonly List<string> _registeredFilePaths = [];
 
-        public ValueTask<AssemblyHandlers> RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken)
+        public AsyncLazy<AssemblyHandlers> RegisterAssembly(string assemblyFilePath)
         {
             _registeredFilePaths.Add(assemblyFilePath);
-            return new(AssemblyHandlers.Empty);
+            return AsyncLazy.Create(AssemblyHandlers.Empty);
         }
 
-        public ValueTask<bool> UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken)
+        public bool UnregisterAssembly(string assemblyFilePath)
         {
             _registeredFilePaths.Remove(assemblyFilePath);
-            return new(_registeredFilePaths.Count == 0);
+            return _registeredFilePaths.Count == 0;
         }
 
         public ValueTask AddHandlersAsync<TResult>(string messageName, bool isSolution, ArrayBuilder<IExtensionMessageHandlerWrapper<TResult>> result, CancellationToken cancellationToken)
@@ -357,18 +374,15 @@ public static IExtensionFolder Create(
             }
         }
 
-        public async ValueTask<AssemblyHandlers> RegisterAssemblyAsync(
-            string assemblyFilePath, CancellationToken cancellationToken)
+        public AsyncLazy<AssemblyHandlers> RegisterAssembly(string assemblyFilePath)
         {
-            var lazy = ImmutableInterlocked.GetOrAdd(
+            return ImmutableInterlocked.GetOrAdd(
                 ref _assemblyFilePathToHandlers,
                 assemblyFilePath,
                 static (assemblyFilePath, @this) => AsyncLazy.Create(
                     static (args, cancellationToken) => CreateAssemblyHandlers(args.@this, args.assemblyFilePath, cancellationToken),
                     (assemblyFilePath, @this)),
                 this);
-
-            return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
         }
 
         private static AssemblyHandlers CreateAssemblyHandlers(
@@ -422,10 +436,10 @@ public async ValueTask AddHandlersAsync<TResult>(string messageName, bool isSolu
             }
         }
 
-        public ValueTask<bool> UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken)
+        public bool UnregisterAssembly(string assemblyFilePath)
         {
-            _assemblyFilePathToHandlers = _assemblyFilePathToHandlers.Remove(assemblyFilePath);
-            return new(_assemblyFilePathToHandlers.IsEmpty);
+            _assemblyFilePathToHandlers.Remove(assemblyFilePath);
+            return _assemblyFilePathToHandlers.IsEmpty;
         }
     }
 
diff --git a/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs b/src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs
@@ -18,16 +18,12 @@ internal interface IExtensionMessageHandlerService : IWorkspaceService
     /// </summary>
     /// <param name="assemblyFilePath">The assembly to register and create message handlers from.</param>
     /// <returns>The names of the registered handlers.</returns>
-    /// <remarks>Should be called serially with other <see cref="RegisterExtensionAsync"/>, <see
-    /// cref="UnregisterExtensionAsync"/>, or <see cref="ResetAsync"/> calls.</remarks>
     ValueTask<RegisterExtensionResponse> RegisterExtensionAsync(string assemblyFilePath, CancellationToken cancellationToken);
 
     /// <summary>
     /// Unregisters extension message handlers previously registered from <paramref name="assemblyFilePath"/>.
     /// </summary>
     /// <param name="assemblyFilePath">The assembly for which handlers should be unregistered.</param>
-    /// <remarks>Should be called serially with other <see cref="RegisterExtensionAsync"/>, <see
-    /// cref="UnregisterExtensionAsync"/>, or <see cref="ResetAsync"/> calls.</remarks>
     ValueTask UnregisterExtensionAsync(string assemblyFilePath, CancellationToken cancellationToken);
 
     /// <summary>
@@ -44,7 +40,6 @@ internal interface IExtensionMessageHandlerService : IWorkspaceService
     /// <param name="messageName">The name of the handler to execute. This is generally the full name of the type implementing the handler.</param>
     /// <param name="jsonMessage">The json message to be passed to the handler.</param>
     /// <returns>The json message returned by the handler.</returns>
-    /// <remarks>Can be called concurrently with other message requests.</remarks>
     ValueTask<string> HandleExtensionWorkspaceMessageAsync(
         Solution solution,
         string messageName,
@@ -58,7 +53,6 @@ ValueTask<string> HandleExtensionWorkspaceMessageAsync(
     /// <param name="messageName">The name of the handler to execute. This is generally the full name of the type implementing the handler.</param>
     /// <param name="jsonMessage">The json message to be passed to the handler.</param>
     /// <returns>The json message returned by the handler.</returns>
-    /// <remarks>Can be called concurrently with other message requests.</remarks>
     ValueTask<string> HandleExtensionDocumentMessageAsync(
         Document documentId,
         string messageName,
diff --git 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";
 
-    /// <summary>
-    /// 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.
-    /// </summary>
-    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
@@ -20,12 +20,7 @@ internal sealed class ExtensionUnregisterHandler()
 {
     private const string MethodName = "roslyn/extensionUnregister";
 
-    /// <summary>
-    /// 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.
-    /// </summary>
-    public bool MutatesSolutionState => true;
+    public bool MutatesSolutionState => false;
 
     public bool RequiresLSPSolution => true;
 
