Added cancellation to all async methods that didn't already have it. (#225)

* Added cancellation tokens everywhere +semver:major

* Added new methods for sending requests when using Mediatr typed request objects.  Simplified extenison methods to use these new methods.  Added method name to each params type.  Added unit tests to validate that all param classes and handler interfaces have all the appropriate pieces (method, abstract handler, delegating handler, etc)

* Added small logging change to show more information over state

* Added xml confguration support

Author: david-driscoll

sample/SampleServer/Program.cs

@@ -76,7 +76,7 @@ static async Task MainAsync(string[] args)
                             Section = "terminal",
                         });
                     })
-                    .OnInitialize(new InitializeDelegate(async (server, request) => {
+                    .OnInitialize(async (server, request, token) => {
                         var manager = server.ProgressManager.WorkDone(request, new WorkDoneProgressBegin() {
                             Title = "Server is starting...",
                             Percentage = 10,
@@ -89,8 +89,8 @@ static async Task MainAsync(string[] args)
                             Percentage = 20,
                             Message = "loading in progress"
                         });
-                    }))
-                    .OnInitialized(new InitializedDelegate(async (server, request, response) => {
+                    })
+                    .OnInitialized(async (server, request, response, token) => {
                         workDone.OnNext(new WorkDoneProgressReport() {
                             Percentage = 40,
                             Message = "loading almost done",
@@ -102,9 +102,10 @@ static async Task MainAsync(string[] args)
                             Message = "loading done",
                             Percentage = 100,
                         });
-                    }))
-                    .OnStarted(async (languageServer, result) => {
-                        using var manager = await languageServer.ProgressManager.Create(new WorkDoneProgressBegin() { Title = "Doing some work..." });
+                        workDone.OnCompleted();
+                    })
+                    .OnStarted(async (languageServer, result, token) => {
+                        using var manager = languageServer.ProgressManager.Create(new WorkDoneProgressBegin() { Title = "Doing some work..." });
 
                         manager.OnNext(new WorkDoneProgressReport() { Message = "doing things..." });
                         await Task.Delay(10000);

src/Dap.Protocol/Requests/RunInTerminalExtensions.cs

@@ -1,12 +1,13 @@
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace OmniSharp.Extensions.DebugAdapter.Protocol.Requests
 {
     public static class RunInTerminalExtensions
     {
-        public static Task<RunInTerminalResponse> RunInTerminal(this IDebugAdapterClient mediator, RunInTerminalArguments @params)
+        public static Task<RunInTerminalResponse> RunInTerminal(this IDebugAdapterClient mediator, RunInTerminalArguments @params, CancellationToken cancellationToken = default)
         {
-            return mediator.SendRequest<RunInTerminalArguments, RunInTerminalResponse>(RequestNames.RunInTerminal, @params);
+            return mediator.SendRequest(@params, cancellationToken);
         }
     }
 

src/JsonRpc/CancelParams.cs

@@ -4,6 +4,7 @@
 
 namespace OmniSharp.Extensions.JsonRpc
 {
+    [Method(JsonRpcNames.CancelRequest)]
     public class CancelParams : IRequest
     {
         /// <summary>

src/JsonRpc/IOutputHandler.cs

@@ -1,12 +1,13 @@
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace OmniSharp.Extensions.JsonRpc
 {
     public interface IOutputHandler : IDisposable
     {
         void Start();
-        void Send(object value);
+        void Send(object value, CancellationToken cancellationToken);
         Task WaitForShutdown();
     }
 }

src/JsonRpc/IResponseRouter.cs

@@ -1,4 +1,6 @@
+using System.Threading;
 using System.Threading.Tasks;
+using MediatR;
 using Newtonsoft.Json.Linq;
 
 namespace OmniSharp.Extensions.JsonRpc
@@ -7,9 +9,12 @@ public interface IResponseRouter
     {
         void SendNotification(string method);
         void SendNotification<T>(string method, T @params);
-        Task<TResponse> SendRequest<T, TResponse>(string method, T @params);
-        Task<TResponse> SendRequest<TResponse>(string method);
-        Task SendRequest<T>(string method, T @params);
+        void SendNotification(IRequest @params);
+        Task<TResponse> SendRequest<T, TResponse>(string method, T @params, CancellationToken cancellationToken);
+        Task<TResponse> SendRequest<TResponse>(IRequest<TResponse> @params, CancellationToken cancellationToken);
+        Task SendRequest(IRequest @params, CancellationToken cancellationToken);
+        Task<TResponse> SendRequest<TResponse>(string method, CancellationToken cancellationToken);
+        Task SendRequest<T>(string method, T @params, CancellationToken cancellationToken);
         TaskCompletionSource<JToken> GetRequest(long id);
     }
 }

src/JsonRpc/InputHandler.cs

@@ -104,7 +104,7 @@ private void ProcessInputStream()
 
                     if (length == 0 || length >= int.MaxValue)
                     {
-                        HandleRequest(string.Empty);
+                        HandleRequest(string.Empty, CancellationToken.None);
                     }
                     else
                     {
@@ -118,7 +118,7 @@ private void ProcessInputStream()
                         }
                         // TODO sometimes: encoding should be based on the respective header (including the wrong "utf8" value)
                         var payload = System.Text.Encoding.UTF8.GetString(requestBuffer);
-                        HandleRequest(payload);
+                        HandleRequest(payload, CancellationToken.None);
                     }
                 }
                 catch (IOException)
@@ -129,7 +129,7 @@ private void ProcessInputStream()
             }
         }
 
-        private void HandleRequest(string request)
+        private void HandleRequest(string request, CancellationToken cancellationToken)
         {
             JToken payload;
             try
@@ -138,13 +138,13 @@ private void HandleRequest(string request)
             }
             catch
             {
-                _outputHandler.Send(new ParseError());
+                _outputHandler.Send(new ParseError(), cancellationToken);
                 return;
             }
 
             if (!_receiver.IsValid(payload))
             {
-                _outputHandler.Send(new InvalidRequest());
+                _outputHandler.Send(new InvalidRequest(), cancellationToken);
                 return;
             }
 
@@ -186,12 +186,12 @@ private void HandleRequest(string request)
                         async () => {
                             try
                             {
-                                var result = await _requestRouter.RouteRequest(descriptor, item.Request, CancellationToken.None);
+                                var result = await _requestRouter.RouteRequest(descriptor, item.Request, cancellationToken);
                                 if (result.IsError && result.Error is RequestCancelled)
                                 {
                                     return;
                                 }
-                                _outputHandler.Send(result.Value);
+                                _outputHandler.Send(result.Value, cancellationToken);
                             }
                             catch (Exception e)
                             {
@@ -230,7 +230,7 @@ private void HandleRequest(string request)
                 if (item.IsError)
                 {
                     // TODO:
-                    _outputHandler.Send(item.Error);
+                    _outputHandler.Send(item.Error, cancellationToken);
                 }
             }
 

src/JsonRpc/JsonRpcServer.cs

@@ -180,19 +180,34 @@ public void SendNotification<T>(string method, T @params)
             _responseRouter.SendNotification(method, @params);
         }
 
-        public Task<TResponse> SendRequest<T, TResponse>(string method, T @params)
+        public void SendNotification(IRequest @params)
         {
-            return _responseRouter.SendRequest<T, TResponse>(method, @params);
+            _responseRouter.SendNotification(@params);
         }
 
-        public Task<TResponse> SendRequest<TResponse>(string method)
+        public Task<TResponse> SendRequest<T, TResponse>(string method, T @params, CancellationToken cancellationToken)
         {
-            return _responseRouter.SendRequest<TResponse>(method);
+            return _responseRouter.SendRequest<T, TResponse>(method, @params, cancellationToken);
         }
 
-        public Task SendRequest<T>(string method, T @params)
+        public Task<TResponse> SendRequest<TResponse>(IRequest<TResponse> @params, CancellationToken cancellationToken)
         {
-            return _responseRouter.SendRequest(method, @params);
+            return _responseRouter.SendRequest(@params, cancellationToken);
+        }
+
+        public Task SendRequest(IRequest @params, CancellationToken cancellationToken)
+        {
+            return _responseRouter.SendRequest(@params, cancellationToken);
+        }
+
+        public Task<TResponse> SendRequest<TResponse>(string method, CancellationToken cancellationToken)
+        {
+            return _responseRouter.SendRequest<TResponse>(method, cancellationToken);
+        }
+
+        public Task SendRequest<T>(string method, T @params, CancellationToken cancellationToken)
+        {
+            return _responseRouter.SendRequest(method, @params, cancellationToken);
         }
 
         public TaskCompletionSource<JToken> GetRequest(long id)

src/JsonRpc/OutputHandler.cs

@@ -18,23 +18,25 @@ public class OutputHandler : IOutputHandler
 
         public OutputHandler(Stream output, ISerializer serializer)
         {
-            if (!output.CanWrite) throw new ArgumentException($"must provide a writable stream for {nameof(output)}", nameof(output));
+            if (!output.CanWrite)
+                throw new ArgumentException($"must provide a writable stream for {nameof(output)}", nameof(output));
             _output = output;
             _serializer = serializer;
             _queue = new BlockingCollection<object>();
             _cancel = new CancellationTokenSource();
             _outputIsFinished = new TaskCompletionSource<object>();
-            _thread = new Thread(ProcessOutputQueue) { IsBackground = true, Name = "ProcessOutputQueue" };
+            _thread = new Thread(ProcessOutputQueue) {IsBackground = true, Name = "ProcessOutputQueue"};
         }
 
         public void Start()
         {
             _thread.Start();
         }
 
-        public void Send(object value)
+        public void Send(object value, CancellationToken cancellationToken)
         {
-            _queue.Add(value);
+            if (!cancellationToken.IsCancellationRequested)
+                _queue.Add(value);
         }
 
         private void ProcessOutputQueue()
@@ -62,7 +64,7 @@ private void ProcessOutputQueue()
                             ms.Write(contentBytes, 0, contentBytes.Length);
                             if (!token.IsCancellationRequested)
                             {
-                                _output.Write(ms.ToArray(), 0, (int)ms.Position);
+                                _output.Write(ms.ToArray(), 0, (int) ms.Position);
                             }
                         }
 

src/JsonRpc/ResponseRouter.cs

@@ -1,5 +1,9 @@
+using System;
 using System.Collections.Concurrent;
+using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
+using MediatR;
 using Newtonsoft.Json.Linq;
 
 namespace OmniSharp.Extensions.JsonRpc
@@ -10,6 +14,7 @@ public class ResponseRouter : IResponseRouter
         private readonly ISerializer _serializer;
         private readonly object _lock = new object();
         private readonly ConcurrentDictionary<long, TaskCompletionSource<JToken>> _requests = new ConcurrentDictionary<long, TaskCompletionSource<JToken>>();
+        private static readonly ConcurrentDictionary<Type, string> _methodCache = new ConcurrentDictionary<Type, string>();
 
         public ResponseRouter(IOutputHandler outputHandler, ISerializer serializer)
         {
@@ -21,41 +26,61 @@ public void SendNotification(string method)
         {
             _outputHandler.Send(new Client.Notification() {
                 Method = method
-            });
+            }, CancellationToken.None);
         }
 
         public void SendNotification<T>(string method, T @params)
         {
             _outputHandler.Send(new Client.Notification() {
                 Method = method,
                 Params = @params
-            });
+            }, CancellationToken.None);
         }
 
-        public async Task<TResponse> SendRequest<T, TResponse>(string method, T @params)
+        public void SendNotification(IRequest @params)
+        {
+            SendNotification(GetMethodName(@params.GetType()), @params);
+        }
+
+        public async Task<TResponse> SendRequest<T, TResponse>(string method, T @params, CancellationToken cancellationToken)
         {
             var tcs = new TaskCompletionSource<JToken>();
+
             var nextId = _serializer.GetNextId();
             _requests.TryAdd(nextId, tcs);
 
             _outputHandler.Send(new Client.Request() {
                 Method = method,
                 Params = @params,
                 Id = nextId
-            });
+            }, cancellationToken);
 
             try
             {
                 var result = await tcs.Task;
+                if (typeof(TResponse) == typeof(Unit))
+                {
+                    return (TResponse)(object)Unit.Value;
+                }
                 return result.ToObject<TResponse>(_serializer.JsonSerializer);
             }
             finally
             {
-                _requests.TryRemove(nextId, out var _);
+                _requests.TryRemove(nextId, out _);
             }
         }
 
-        public async Task<TResponse> SendRequest<TResponse>(string method)
+        public Task<TResponse> SendRequest<TResponse>(IRequest<TResponse> @params, CancellationToken cancellationToken)
+        {
+            return SendRequest<IRequest<TResponse>, TResponse>(GetMethodName(@params.GetType()), @params, cancellationToken);
+        }
+
+        public Task SendRequest(IRequest @params, CancellationToken cancellationToken)
+        {
+            return SendRequest(GetMethodName(@params.GetType()), @params, cancellationToken);
+        }
+
+        public async Task<TResponse> SendRequest<TResponse>(string method, CancellationToken cancellationToken)
         {
             var nextId = _serializer.GetNextId();
 
@@ -66,7 +91,7 @@ public async Task<TResponse> SendRequest<TResponse>(string method)
                 Method = method,
                 Params = null,
                 Id = nextId
-            });
+            }, cancellationToken);
 
             try
             {
@@ -79,7 +104,7 @@ public async Task<TResponse> SendRequest<TResponse>(string method)
             }
         }
 
-        public async Task SendRequest<T>(string method, T @params)
+        public async Task SendRequest<T>(string method, T @params, CancellationToken cancellationToken)
         {
             var nextId = _serializer.GetNextId();
 
@@ -90,7 +115,7 @@ public async Task SendRequest<T>(string method, T @params)
                 Method = method,
                 Params = @params,
                 Id = nextId
-            });
+            }, cancellationToken);
 
             try
             {
@@ -107,5 +132,22 @@ public TaskCompletionSource<JToken> GetRequest(long id)
             _requests.TryGetValue(id, out var source);
             return source;
         }
+
+        private string GetMethodName(Type type)
+        {
+            if (!_methodCache.TryGetValue(type, out var methodName))
+            {
+                var attribute = type.GetCustomAttribute<MethodAttribute>(true);
+                if (attribute == null)
+                {
+                    throw new NotSupportedException($"Unable to infer method name for type {type.FullName}");
+                }
+
+                methodName = attribute.Method;
+                _methodCache.TryAdd(type, methodName);
+            }
+
+            return methodName;
+        }
     }
 }

src/Protocol/Client/Server/RegisterCapabilityExtensions.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
@@ -9,9 +10,9 @@ namespace OmniSharp.Extensions.LanguageServer.Protocol.Server
 {
     public static class RegisterCapabilityExtensions
     {
-        public static async Task RegisterCapability(this ILanguageServerClient mediator, RegistrationParams @params)
+        public static Task RegisterCapability(this ILanguageServerClient mediator, RegistrationParams @params, CancellationToken cancellationToken = default)
         {
-            await mediator.SendRequest(ClientNames.RegisterCapability, @params);
+            return mediator.SendRequest(@params, cancellationToken);
         }
     }
 }