Review feedback

src/EditorFeatures/Core/LanguageServer/AbstractInProcLanguageClient.cs

@@ -148,7 +148,7 @@ public AbstractInProcLanguageClient(
 
             if (_languageServer is not null)
             {
-                await _languageServer.WaitForExitAsync().ConfigureAwait(false);
+                await _languageServer.WaitForExitAsync().WithCancellation(cancellationToken).ConfigureAwait(false);
             }
 
             var (clientStream, serverStream) = FullDuplexStream.CreatePair();

src/Features/LanguageServer/Microsoft.CommonLanguageServerProtocol.Framework/AbstractLanguageServer.cs

@@ -49,7 +49,7 @@ public abstract class AbstractLanguageServer<TRequestContext>
     /// Task completion source that is started when the server starts and completes when the server exits.
     /// Used when callers need to wait for the server to cleanup.
     /// </summary>
-    private readonly TaskCompletionSource<bool> _serverExitedSource = new(false);
+    private readonly TaskCompletionSource<object?> _serverExitedSource = new();
 
     protected AbstractLanguageServer(
         JsonRpc jsonRpc,
@@ -218,7 +218,7 @@ private void CheckServerState()
         }
     }
 
-    public async Task WaitForExitAsync()
+    public Task WaitForExitAsync()
     {
         lock (_lifeCycleLock)
         {
@@ -229,10 +229,10 @@ public async Task WaitForExitAsync()
             }
         }
 
-        // Note - we wait for the _serverExitedSource task here instead of the _exitNotification task as we may not have
+        // Note - we return the _serverExitedSource task here instead of the _exitNotification task as we may not have
         // finished processing the exit notification before a client calls into us asking to restart.
         // This is because unlike shutdown, exit is a notification where clients do not need to wait for a response.
-        await _serverExitedSource.Task.ConfigureAwait(false);
+        return _serverExitedSource.Task;
     }
 
     /// <summary>
@@ -247,78 +247,76 @@ public Task ShutdownAsync(string message = "Shutting down")
             // Run shutdown or return the already running shutdown request.
             _shutdownRequestTask ??= Shutdown_NoLockAsync(message);
             shutdownTask = _shutdownRequestTask;
+            return shutdownTask;
         }
 
-        return shutdownTask;
+        // Runs the actual shutdown outside of the lock - guaranteed to be only called once by the above code.
+        async Task Shutdown_NoLockAsync(string message)
+        {
+            // Immediately yield so that this does not run under the lock.
+            await Task.Yield();
+
+            _logger.LogInformation(message);
+
+            // Allow implementations to do any additional cleanup on shutdown.
+            var lifeCycleManager = GetLspServices().GetRequiredService<ILifeCycleManager>();
+            await lifeCycleManager.ShutdownAsync(message).ConfigureAwait(false);
+
+            await ShutdownRequestExecutionQueueAsync().ConfigureAwait(false);
+        }
     }
 
     /// <summary>
     /// Tells the LSP server to exit.  Requires that <see cref="ShutdownAsync(string)"/> was called first.
     /// Typically called from an LSP exit notification.
     /// </summary>
-    public async Task ExitAsync()
+    public Task ExitAsync()
     {
         Task exitTask;
         lock (_lifeCycleLock)
         {
             if (_shutdownRequestTask?.IsCompleted != true)
             {
-                throw new InvalidOperationException("The language server has not yet been asked to shutdown.");
+                throw new InvalidOperationException("The language server has not yet been asked to shutdown or has not finished shutting down.");
             }
 
             // Run exit or return the already running exit request.
             _exitNotificationTask ??= Exit_NoLockAsync();
             exitTask = _exitNotificationTask;
+            return exitTask;
         }
 
-        await exitTask.ConfigureAwait(false);
-    }
-
-    /// <summary>
-    /// Performs the actual shutdown work outside of the lock.
-    /// Guaranteed to only be called once by <see cref="ShutdownAsync(string)"/>
-    /// </summary>
-    private async Task Shutdown_NoLockAsync(string message = "Shutting down")
-    {
-        _logger.LogInformation(message);
-
-        // Allow implementations to do any additional cleanup on shutdown.
-        var lifeCycleManager = GetLspServices().GetRequiredService<ILifeCycleManager>();
-        await lifeCycleManager.ShutdownAsync(message).ConfigureAwait(false);
-
-        await ShutdownRequestExecutionQueueAsync().ConfigureAwait(false);
-    }
-
-    /// <summary>
-    /// Performs the actual shutdown work outside of the lock.
-    /// Guaranteed to only be called once by <see cref="ExitAsync"/>
-    /// </summary>
-    private async Task Exit_NoLockAsync()
-    {
-        try
+        // Runs the actual exit outside of the lock - guaranteed to be only called once by the above code.
+        async Task Exit_NoLockAsync()
         {
-            var lspServices = GetLspServices();
+            // Immediately yield so that this does not run under the lock.
+            await Task.Yield();
 
-            // Allow implementations to do any additional cleanup on exit.
-            var lifeCycleManager = lspServices.GetRequiredService<ILifeCycleManager>();
-            await lifeCycleManager.ExitAsync().ConfigureAwait(false);
+            try
+            {
+                var lspServices = GetLspServices();
 
-            await ShutdownRequestExecutionQueueAsync().ConfigureAwait(false);
+                // Allow implementations to do any additional cleanup on exit.
+                var lifeCycleManager = lspServices.GetRequiredService<ILifeCycleManager>();
+                await lifeCycleManager.ExitAsync().ConfigureAwait(false);
 
-            lspServices.Dispose();
+                await ShutdownRequestExecutionQueueAsync().ConfigureAwait(false);
 
-            _jsonRpc.Disconnected -= JsonRpc_Disconnected;
-            _jsonRpc.Dispose();
-        }
-        catch (Exception)
-        {
-            // Swallow exceptions thrown by disposing our JsonRpc object. Disconnected events can potentially throw their own exceptions so
-            // we purposefully ignore all of those exceptions in an effort to shutdown gracefully.
-        }
-        finally
-        {
-            _logger.LogInformation("Exiting server");
-            _serverExitedSource.TrySetResult(true);
+                lspServices.Dispose();
+
+                _jsonRpc.Disconnected -= JsonRpc_Disconnected;
+                _jsonRpc.Dispose();
+            }
+            catch (Exception)
+            {
+                // Swallow exceptions thrown by disposing our JsonRpc object. Disconnected events can potentially throw their own exceptions so
+                // we purposefully ignore all of those exceptions in an effort to shutdown gracefully.
+            }
+            finally
+            {
+                _logger.LogInformation("Exiting server");
+                _serverExitedSource.TrySetResult(null);
+            }
         }
     }
 

src/Features/LanguageServer/Microsoft.CommonLanguageServerProtocol.Framework/ILifeCycleManager.cs

@@ -17,13 +17,14 @@ public interface ILifeCycleManager
     /// Called when the server recieves the LSP exit notification.
     /// </summary>
     /// <remarks>
-    /// This is called before LSP services and the JsonRpc connection is disposed of.
+    /// This is always called after the LSP shutdown request and <see cref="ShutdownAsync(string)"/> runs
+    /// but before LSP services and the JsonRpc connection is disposed of in LSP exit.
     /// Implementations are not expected to be threadsafe.
     /// </remarks>
     Task ExitAsync();
 
     /// <summary>
-    /// Called when the server recieves the LSP shutdown request.
+    /// Called when the server receives the LSP shutdown request.
     /// </summary>
     /// <remarks>
     /// This is called before the request execution is closed.

src/Features/LanguageServer/Protocol/Handler/ServerLifetime/LspServiceLifeCycleManager.cs

@@ -39,6 +39,7 @@ public async Task ShutdownAsync(string message = "Shutting down")
 
     public Task ExitAsync()
     {
+        // We don't need any custom logic to run on exit.
         return Task.CompletedTask;
     }
 }