Merge pull request #64752 from dibarbet/fix_double_dispose_race

Refactor LSP server exit / shutdown / dispose logic for resiliency and clarity

src/EditorFeatures/Core/LanguageServer/AbstractInProcLanguageClient.cs

@@ -29,7 +29,7 @@ internal abstract partial class AbstractInProcLanguageClient : ILanguageClient,
         private readonly ILanguageClientMiddleLayer? _middleLayer;
         private readonly ILspServiceLoggerFactory _lspLoggerFactory;
 
-        private readonly AbstractLspServiceProvider _lspServiceProvider;
+        protected readonly AbstractLspServiceProvider LspServiceProvider;
 
         protected readonly IGlobalOptionService GlobalOptions;
 
@@ -105,7 +105,7 @@ public AbstractInProcLanguageClient(
             IThreadingContext threadingContext,
             AbstractLanguageClientMiddleLayer? middleLayer = null)
         {
-            _lspServiceProvider = lspServiceProvider;
+            LspServiceProvider = lspServiceProvider;
             GlobalOptions = globalOptions;
             _lspLoggerFactory = lspLoggerFactory;
             _threadingContext = threadingContext;
@@ -148,9 +148,7 @@ public AbstractInProcLanguageClient(
 
             if (_languageServer is not null)
             {
-                Contract.ThrowIfFalse(_languageServer.HasShutdownStarted, "The language server has not yet been asked to shutdown.");
-
-                await _languageServer.DisposeAsync().ConfigureAwait(false);
+                await _languageServer.WaitForExitAsync().WithCancellation(cancellationToken).ConfigureAwait(false);
             }
 
             var (clientStream, serverStream) = FullDuplexStream.CreatePair();
@@ -219,14 +217,14 @@ internal static async Task<AbstractLanguageServer<RequestContext>> CreateAsync<T
             return server;
         }
 
-        public AbstractLanguageServer<RequestContext> Create(
+        public virtual AbstractLanguageServer<RequestContext> Create(
             JsonRpc jsonRpc,
             ICapabilitiesProvider capabilitiesProvider,
             WellKnownLspServerKinds serverKind,
             ILspServiceLogger logger)
         {
             var server = new RoslynLanguageServer(
-                _lspServiceProvider,
+                LspServiceProvider,
                 jsonRpc,
                 capabilitiesProvider,
                 logger,

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -649,6 +649,20 @@ public Task CloseDocumentAsync(Uri documentUri)
                 return ExecuteRequestAsync<LSP.DidCloseTextDocumentParams, object>(LSP.Methods.TextDocumentDidCloseName, didCloseParams, CancellationToken.None);
             }
 
+            public async Task ShutdownTestServerAsync()
+            {
+                await _clientRpc.InvokeAsync(LSP.Methods.ShutdownName).ConfigureAwait(false);
+            }
+
+            public async Task ExitTestServerAsync()
+            {
+                // Since exit is a notification that disposes of the json rpc stream we cannot wait on the result
+                // of the request itself since it will throw a ConnectionLostException.
+                // Instead we wait for the server's exit task to be completed.
+                await _clientRpc.NotifyAsync(LSP.Methods.ExitName).ConfigureAwait(false);
+                await LanguageServer.WaitForExitAsync().ConfigureAwait(false);
+            }
+
             public IList<LSP.Location> GetLocations(string locationName) => _locations[locationName];
 
             public Solution GetCurrentSolution() => TestWorkspace.CurrentSolution;
@@ -681,13 +695,16 @@ public async ValueTask DisposeAsync()
                 var solutionCrawlerRegistrationService = (SolutionCrawlerRegistrationService)TestWorkspace.Services.GetRequiredService<ISolutionCrawlerRegistrationService>();
                 solutionCrawlerRegistrationService.Unregister(TestWorkspace);
 
-                // Some tests manually call shutdown, so avoid calling shutdown twice if already called.
-                if (!LanguageServer.HasShutdownStarted)
+                // Some tests will manually call shutdown and exit, so attempting to call this during dispose
+                // will fail as the server's jsonrpc instance will be disposed of.
+                if (!LanguageServer.GetTestAccessor().HasShutdownStarted())
                 {
-                    await LanguageServer.GetTestAccessor().ShutdownServerAsync("Disposing of test lsp server");
+                    await ShutdownTestServerAsync();
+                    await ExitTestServerAsync();
                 }
 
-                await LanguageServer.GetTestAccessor().ExitServerAsync();
+                // Wait for all the exit notifications to run to completion.
+                await LanguageServer.WaitForExitAsync();
 
                 TestWorkspace.Dispose();
                 _clientRpc.Dispose();

src/Features/LanguageServer/Microsoft.CommonLanguageServerProtocol.Framework.Example/ExampleLanguageServer.cs

@@ -21,29 +21,33 @@ protected override ILspServices ConstructLspServices()
     {
         var serviceCollection = new ServiceCollection();
 
-        _ = AddHandlers(serviceCollection)
+        var _ = AddHandlers(serviceCollection)
             .AddSingleton<ILspLogger>(_logger)
             .AddSingleton<IRequestContextFactory<ExampleRequestContext>, ExampleRequestContextFactory>()
             .AddSingleton<IInitializeManager<InitializeParams, InitializeResult>, CapabilitiesManager>()
-            .AddSingleton<ILifeCycleManager>(GetLifeCycleManager())
             .AddSingleton(this);
 
+        var lifeCycleManager = GetLifeCycleManager();
+        if (lifeCycleManager != null)
+        {
+            serviceCollection.AddSingleton(lifeCycleManager);
+        }
+
         var lspServices = new ExampleLspServices(serviceCollection);
 
         return lspServices;
     }
 
-    protected virtual ILifeCycleManager GetLifeCycleManager()
+    protected virtual ILifeCycleManager? GetLifeCycleManager()
     {
-        return this;
+        return null;
     }
 
     private static IServiceCollection AddHandlers(IServiceCollection serviceCollection)
     {
         _ = serviceCollection
             .AddSingleton<IMethodHandler, InitializeHandler<InitializeParams, InitializeResult, ExampleRequestContext>>()
-            .AddSingleton<IMethodHandler, InitializedHandler<InitializedParams, ExampleRequestContext>>()
-            .AddSingleton<IMethodHandler, ShutdownHandler<ExampleRequestContext>>();
+            .AddSingleton<IMethodHandler, InitializedHandler<InitializedParams, ExampleRequestContext>>();
         return serviceCollection;
     }
 }

src/Features/LanguageServer/Microsoft.CommonLanguageServerProtocol.Framework.UnitTests/RequestExecutionQueueTests.cs

@@ -6,24 +6,39 @@
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Elfie.Diagnostics;
 using Microsoft.CommonLanguageServerProtocol.Framework;
 using Moq;
+using Nerdbank.Streams;
+using StreamJsonRpc;
 using Xunit;
 using static Microsoft.CommonLanguageServerProtocol.Framework.UnitTests.HandlerProviderTests;
 
 namespace Microsoft.CommonLanguageServerProtocol.Framework.UnitTests;
 
 public class RequestExecutionQueueTests
 {
+    private class MockServer : AbstractLanguageServer<TestRequestContext>
+    {
+        public MockServer() : base(new JsonRpc(new HeaderDelimitedMessageHandler(FullDuplexStream.CreatePair().Item1)), NoOpLspLogger.Instance)
+        {
+        }
+
+        protected override ILspServices ConstructLspServices()
+        {
+            throw new NotImplementedException();
+        }
+    }
+
     private const string MethodName = "SomeMethod";
 
     private static RequestExecutionQueue<TestRequestContext> GetRequestExecutionQueue(IMethodHandler? methodHandler = null)
     {
         var handlerProvider = new Mock<IHandlerProvider>(MockBehavior.Strict);
         var handler = methodHandler ?? GetTestMethodHandler();
-
         handlerProvider.Setup(h => h.GetMethodHandler(MethodName, TestMethodHandler.RequestType, TestMethodHandler.ResponseType)).Returns(handler);
-        var executionQueue = new RequestExecutionQueue<TestRequestContext>(NoOpLspLogger.Instance, handlerProvider.Object);
+
+        var executionQueue = new RequestExecutionQueue<TestRequestContext>(new MockServer(), NoOpLspLogger.Instance, handlerProvider.Object);
         executionQueue.Start();
 
         return executionQueue;

src/Features/LanguageServer/Microsoft.CommonLanguageServerProtocol.Framework.UnitTests/TestExampleLanguageServer.cs

@@ -45,32 +45,30 @@ internal async Task ExecuteNotificationAsync(string methodName, CancellationToke
 
     protected override ILifeCycleManager GetLifeCycleManager()
     {
-        return new TestLifeCycleManager(this, _shuttingDown, _exiting);
+        return new TestLifeCycleManager(_shuttingDown, _exiting);
     }
 
     private class TestLifeCycleManager : ILifeCycleManager
     {
-        private readonly ILifeCycleManager _lifeCycleManager;
         private readonly TaskCompletionSource<int> _shuttingDownSource;
         private readonly TaskCompletionSource<int> _exitingSource;
 
-        public TestLifeCycleManager(ILifeCycleManager lifeCycleManager, TaskCompletionSource<int> shuttingDownSource, TaskCompletionSource<int> exitingSource)
+        public TestLifeCycleManager(TaskCompletionSource<int> shuttingDownSource, TaskCompletionSource<int> exitingSource)
         {
-            _lifeCycleManager = lifeCycleManager;
             _shuttingDownSource = shuttingDownSource;
             _exitingSource = exitingSource;
         }
 
-        public async Task ExitAsync()
+        public Task ExitAsync()
         {
-            await _lifeCycleManager.ExitAsync();
             _exitingSource.SetResult(0);
+            return Task.CompletedTask;
         }
 
-        public async Task ShutdownAsync(string message = "Shutting down")
+        public Task ShutdownAsync(string message = "Shutting down")
         {
-            await _lifeCycleManager.ShutdownAsync(message);
             _shuttingDownSource.SetResult(0);
+            return Task.CompletedTask;
         }
     }
 
@@ -102,12 +100,6 @@ internal async Task<int> WaitForExit()
         return await _exiting.Task;
     }
 
-    public new ValueTask DisposeAsync()
-    {
-        _clientRpc.Dispose();
-        return base.DisposeAsync();
-    }
-
     private static JsonMessageFormatter CreateJsonMessageFormatter()
     {
         var messageFormatter = new JsonMessageFormatter();