Move to immutable dictionary (#77981)

Author: matteo-prosperi

src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.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
     /// <param name="assembly">The assembly to scan for handlers.</param>
     /// <param name="extensionIdentifier">Unique identifier of the extension owning this handler.</param>
     /// <returns>The handlers.</returns>
-    ImmutableArray<IExtensionMessageHandlerWrapper<Solution>> CreateWorkspaceMessageHandlers(Assembly assembly, string extensionIdentifier);
+    ImmutableArray<IExtensionMessageHandlerWrapper<Solution>> CreateWorkspaceMessageHandlers(
+        Assembly assembly, string extensionIdentifier, CancellationToken cancellationToken);
 
     /// <summary>
     /// Creates <see cref="IExtensionMessageHandlerWrapper{Document}"/> instances for each
@@ -28,5 +30,6 @@ internal interface IExtensionMessageHandlerFactory
     /// <param name="assembly">The assembly to scan for handlers.</param>
     /// <param name="extensionIdentifier">Unique identifier of the extension owning this handler.</param>
     /// <returns>The handlers.</returns>
-    ImmutableArray<IExtensionMessageHandlerWrapper<Document>> CreateDocumentMessageHandlers(Assembly assembly, string extensionIdentifier);
+    ImmutableArray<IExtensionMessageHandlerWrapper<Document>> CreateDocumentMessageHandlers(
+        Assembly assembly, string extensionIdentifier, CancellationToken cancellationToken);
 }

src/Features/Core/Portable/Extensions/IExtensionMessageHandlerFactory.cs

@@ -13,14 +13,38 @@ namespace Microsoft.CodeAnalysis.Extensions;
 /// </summary>
 internal interface IExtensionMessageHandlerService : IWorkspaceService
 {
+    /// <summary>
+    /// Registers extension message handlers from the specified assembly.
+    /// </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>
+    /// Unregisters all extension message handlers.
+    /// </summary>
+    /// <remarks>Should be called serially with other <see cref="RegisterExtensionAsync"/>, <see
+    /// cref="UnregisterExtensionAsync"/>, or <see cref="ResetAsync"/> calls.</remarks>
+    ValueTask ResetAsync(CancellationToken cancellationToken);
+
     /// <summary>
     /// Executes a non-document-specific extension message handler with the given message and solution.
     /// </summary>
     /// <param name="solution">The solution the message refers to.</param>
     /// <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>
-    /// <param name="cancellationToken">Cancellation token to cancel the async operation.</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,
@@ -33,30 +57,11 @@ ValueTask<string> HandleExtensionWorkspaceMessageAsync(
     /// <param name="documentId">The document the message refers to.</param>
     /// <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>
-    /// <param name="cancellationToken">Cancellation token to cancel the async operation.</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,
         string jsonMessage,
         CancellationToken cancellationToken);
-
-    /// <summary>
-    /// Registers extension message handlers from the specified assembly.
-    /// </summary>
-    /// <param name="assemblyFilePath">The assembly to register and create message handlers from.</param>
-    /// <returns>The names of the registered handlers.</returns>
-    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>
-    /// <returns>A task representing the async operation.</returns>
-    ValueTask UnregisterExtensionAsync(string assemblyFilePath, CancellationToken cancellationToken);
-
-    /// <summary>
-    /// Unregisters all extension message handlers.
-    /// </summary>
-    ValueTask ResetAsync(CancellationToken cancellationToken);
 }

src/Features/Core/Portable/Extensions/IExtensionMessageHandlerService.cs

@@ -8,13 +8,17 @@
 
 namespace Microsoft.CodeAnalysis.Extensions;
 
+internal interface IExtensionMessageHandlerWrapper
+{
+}
+
 /// <summary>
 /// Wrapper for an <c>IExtensionWorkspaceMessageHandler</c> or <c>IExtensionDocumentMessageHandler</c>
 /// as returned by <see cref="IExtensionMessageHandlerFactory"/>. 
 /// </summary>
 /// <typeparam name="TArgument">The type of object received as parameter by the extension message
 /// handler.</typeparam>
-internal interface IExtensionMessageHandlerWrapper<TArgument>
+internal interface IExtensionMessageHandlerWrapper<TArgument> : IExtensionMessageHandlerWrapper
 {
     /// <summary>
     /// The type of object received as parameter by the extension message handler.

src/Features/Core/Portable/Extensions/IExtensionMessageHandlerWrapper.cs

@@ -20,7 +20,12 @@ internal sealed class ExtensionRegisterHandler()
 {
     private const string MethodName = "roslyn/extensionRegister";
 
-    public bool MutatesSolutionState => false;
+    /// <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 RequiresLSPSolution => true;
 

src/LanguageServer/Protocol/Handler/Extensions/ExtensionRegisterHandler.cs

@@ -20,7 +20,12 @@ internal sealed class ExtensionUnregisterHandler()
 {
     private const string MethodName = "roslyn/extensionUnregister";
 
-    public bool MutatesSolutionState => false;
+    /// <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 RequiresLSPSolution => true;
 

src/LanguageServer/Protocol/Handler/Extensions/ExtensionUnregisterHandler.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<IRemoteExtensionMessageHandlerService>(
-                    (service, cancellationToken) => service.ResetAsync(cancellationToken),
-                    CancellationToken.None).ConfigureAwait(false);
-            }
-            else
-            {
-                var service = hostWorkspace.Services.GetRequiredService<IExtensionMessageHandlerService>();
-                service.Reset();
-            }
+            var service = hostWorkspace.Services.GetRequiredService<IExtensionMessageHandlerService>();
+            await service.ResetAsync(cancellationToken).ConfigureAwait(false);
         }
 
         try
@@ -65,7 +57,7 @@ await client.TryInvokeAsync<IRemoteExtensionMessageHandlerService>(
                 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)
         {

src/LanguageServer/Protocol/Handler/ServerLifetime/LspServiceLifeCycleManager.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<IExtensionMessageHandlerWrapper<Document>> CreateDocumentMessageHandlers(Assembly assembly, string extensionIdentifier)
+    public ImmutableArray<IExtensionMessageHandlerWrapper<Document>> 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<IExtensionMessageHandlerWrapper<Solution>> CreateWorkspaceMessageHandlers(Assembly assembly, string extensionIdentifier)
+    public ImmutableArray<IExtensionMessageHandlerWrapper<Solution>> 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<IExtensionMessageHandlerWrapper<TArgument>> CreateWorkspaceHandlers<TArgument>(
         Assembly assembly,
         Type unboundInterfaceType,
-        Func<object, Type, IExtensionMessageHandlerWrapper<TArgument>> wrapperCreator)
+        Func<object, Type, IExtensionMessageHandlerWrapper<TArgument>> wrapperCreator,
+        CancellationToken cancellationToken)
     {
         var resultBuilder = ImmutableArray.CreateBuilder<IExtensionMessageHandlerWrapper<TArgument>>();
 
         foreach (var candidateType in assembly.GetTypes())
         {
+            cancellationToken.ThrowIfCancellationRequested();
+
             if (candidateType.IsAbstract || candidateType.IsGenericType)
             {
                 continue;

src/Tools/ExternalAccess/Extensions/Internal/ExtensionMessageHandlerFactory.cs

@@ -27,8 +27,8 @@ Microsoft.CodeAnalysis.Extensions.ExtensionHandlerWrapper<TArgument>.Name.get ->
 Microsoft.CodeAnalysis.Extensions.ExtensionHandlerWrapper<TArgument>.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.IExtensionMessageHandlerWrapper<Microsoft.CodeAnalysis.Document!>!>
-Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.CreateWorkspaceMessageHandlers(System.Reflection.Assembly! assembly, string! extensionIdentifier) -> System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.Extensions.IExtensionMessageHandlerWrapper<Microsoft.CodeAnalysis.Solution!>!>
+Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.CreateDocumentMessageHandlers(System.Reflection.Assembly! assembly, string! extensionIdentifier, System.Threading.CancellationToken cancellationToken) -> System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.Extensions.IExtensionMessageHandlerWrapper<Microsoft.CodeAnalysis.Document!>!>
+Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.CreateWorkspaceMessageHandlers(System.Reflection.Assembly! assembly, string! extensionIdentifier, System.Threading.CancellationToken cancellationToken) -> System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.Extensions.IExtensionMessageHandlerWrapper<Microsoft.CodeAnalysis.Solution!>!>
 Microsoft.CodeAnalysis.Extensions.ExtensionMessageHandlerFactory.ExtensionMessageHandlerFactory() -> void
 Microsoft.CodeAnalysis.Extensions.ExtensionWorkspaceMessageHandlerWrapper
 Microsoft.CodeAnalysis.Extensions.ExtensionWorkspaceMessageHandlerWrapper.ExtensionWorkspaceMessageHandlerWrapper(object! handler, System.Type! customMessageHandlerInterface, string! extensionIdentifier) -> void