Shutdown

src/Draco.JsonRpc/IJsonRpcConnection.cs

@@ -23,6 +23,11 @@ public interface IJsonRpcConnection
     /// <returns>The task that completes when the connection closes.</returns>
     public Task ListenAsync();
 
+    /// <summary>
+    /// Shuts down this connection.
+    /// </summary>
+    public void Shutdown();
+
     // TODO: Doc
     public Task<TResponse?> SendRequestAsync<TResponse>(string method, object? @params);
 

src/Draco.JsonRpc/JsonRpcConnection.cs

@@ -5,6 +5,7 @@
 using System.Text.Json;
 using System.Threading.Channels;
 using System.Collections.Concurrent;
+using static System.Runtime.InteropServices.JavaScript.JSType;
 
 namespace Draco.JsonRpc;
 
@@ -85,6 +86,9 @@ private sealed class OutgoingRequest<TResponse> : IOutgoingRequest
     private readonly ConcurrentDictionary<object, CancellationTokenSource> pendingIncomingRequests = new();
     private readonly ConcurrentDictionary<int, IOutgoingRequest> pendingOutgoingRequests = new();
 
+    // Shutdown
+    private readonly CancellationTokenSource shutdownTokenSource = new();
+
     // Communication state
     private int lastMessageId = 0;
 
@@ -95,30 +99,47 @@ public Task ListenAsync() => Task.WhenAll(
         this.WriterLoopAsync(),
         this.ProcessorLoopAsync());
 
+    public void Shutdown() => this.shutdownTokenSource.Cancel();
+
     protected int NextMessageId() => Interlocked.Increment(ref this.lastMessageId);
 
     #region Message Loops
     private async Task ReaderLoopAsync()
     {
         while (true)
         {
-            var (message, foundMessage) = await this.ReadMessageAsync();
-            if (!foundMessage) break;
+            try
+            {
+                var (message, foundMessage) = await this
+                    .ReadMessageAsync()
+                    .WaitAsync(this.shutdownTokenSource.Token);
+                if (!foundMessage) break;
 
-            if (TMessageAdapter.IsResponse(message!))
+                if (TMessageAdapter.IsResponse(message!))
+                {
+                    this.ProcessIncomingResponse(message!);
+                }
+                else
+                {
+                    await this.incomingMessages.Writer.WriteAsync(message!);
+                }
+            }
+            catch (OperationCanceledException oce) when (oce.CancellationToken == this.shutdownTokenSource.Token)
             {
-                this.ProcessIncomingResponse(message!);
+                break;
             }
-            else
+            catch (JsonException ex)
             {
-                await this.incomingMessages.Writer.WriteAsync(message!);
+                var error = TMessageAdapter.CreateJsonExceptionError(ex);
+                await this.outgoingMessages.Writer.WriteAsync(TMessageAdapter.CreateErrorResponse(null!, error));
+                continue;
             }
         }
     }
 
     private async Task WriterLoopAsync()
     {
-        await foreach (var message in this.outgoingMessages.Reader.ReadAllAsync())
+        await foreach (var message in this.outgoingMessages.Reader.ReadAllAsync(this.shutdownTokenSource.Token))
         {
             await this.WriteMessageAsync(message);
         }
@@ -134,7 +155,7 @@ bool IsMutating(TMessage message)
         }
 
         var currentTasks = new List<Task>();
-        await foreach (var message in this.incomingMessages.Reader.ReadAllAsync())
+        await foreach (var message in this.incomingMessages.Reader.ReadAllAsync(this.shutdownTokenSource.Token))
         {
             if (IsMutating(message))
             {

src/Draco.Lsp/Server/LanguageServerLifecycle.cs

@@ -55,8 +55,7 @@ public async Task InitializedAsync(InitializedParams param)
 
     public Task ExitAsync()
     {
-        // TODO
-        // this.connection.Shutdown();
+        this.connection.Shutdown();
         return Task.CompletedTask;
     }