attempt to fix issue where cancellation didn't actually do anything useful

Author: david-driscoll

src/JsonRpc/InputHandler.cs

@@ -5,7 +5,6 @@
 using System.Threading.Tasks;
 using Newtonsoft.Json.Linq;
 using Microsoft.Extensions.Logging;
-using Newtonsoft.Json;
 using OmniSharp.Extensions.JsonRpc.Server;
 using OmniSharp.Extensions.JsonRpc.Server.Messages;
 
@@ -72,7 +71,8 @@ private void ProcessInputStream()
             // content is encoded in UTF-8
             while (true)
             {
-                try {
+                try
+                {
                     if (_inputThread == null) return;
 
                     var buffer = new byte[300];
@@ -182,8 +182,7 @@ private void HandleRequest(string request)
                     _scheduler.Add(
                         type,
                         item.Request.Method,
-                        async () =>
-                        {
+                        async () => {
                             try
                             {
                                 var result = await _requestRouter.RouteRequest(descriptor, item.Request, CancellationToken.None);
@@ -202,26 +201,24 @@ private void HandleRequest(string request)
 
                 if (item.IsNotification)
                 {
+
                     var descriptor = _requestRouter.GetDescriptor(item.Notification);
                     if (descriptor is null) continue;
+
+                    // We need to special case cancellation so that we can cancel any request that is currently in flight.
+                    if (descriptor.Method == JsonRpcNames.CancelRequest)
+                    {
+                        var cancelParams = item.Notification.Params?.ToObject<CancelParams>();
+                        if (cancelParams == null) { continue; }
+                        _requestRouter.CancelRequest(cancelParams.Id);
+                        continue;
+                    }
+
                     var type = _requestProcessIdentifier.Identify(descriptor);
                     _scheduler.Add(
                         type,
                         item.Notification.Method,
-                        async () =>
-                        {
-                            try
-                            {
-                                await _requestRouter.RouteNotification(descriptor, item.Notification, CancellationToken.None);
-                            }
-                            catch (Exception e)
-                            {
-                                _logger.LogCritical(Events.UnhandledNotification, e, "Unhandled exception executing notification {Method}", item.Notification.Method);
-                                // TODO: Should we rethrow or swallow?
-                                // If an exception happens... the whole system could be in a bad state, hence this throwing currently.
-                                throw;
-                            }
-                        }
+                        DoNotification(descriptor, item.Notification)
                     );
                 }
 
@@ -231,6 +228,23 @@ private void HandleRequest(string request)
                     _outputHandler.Send(item.Error);
                 }
             }
+
+            Func<Task> DoNotification(IHandlerDescriptor descriptor, Notification notification)
+            {
+                return async () => {
+                    try
+                    {
+                        await _requestRouter.RouteNotification(descriptor, notification, CancellationToken.None);
+                    }
+                    catch (Exception e)
+                    {
+                        _logger.LogCritical(Events.UnhandledNotification, e, "Unhandled exception executing notification {Method}", notification.Method);
+                        // TODO: Should we rethrow or swallow?
+                        // If an exception happens... the whole system could be in a bad state, hence this throwing currently.
+                        throw;
+                    }
+                };
+            }
         }
 
 

test/JsonRpc.Tests/InputHandlerTests.cs

@@ -7,6 +7,7 @@
 using System.Threading.Tasks;
 using FluentAssertions;
 using Microsoft.Extensions.Logging;
+using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
 using NSubstitute;
 using OmniSharp.Extensions.JsonRpc;
@@ -262,5 +263,52 @@ public void ShouldHandleResponse()
                 tcs.Task.Result.ToString().Should().Be("{}");
             }
         }
+
+        [Fact]
+        public async Task ShouldCancelRequest()
+        {
+            var inputStream = new MemoryStream(Encoding.ASCII.GetBytes("Content-Length: 2\r\n\r\n{}"));
+            var outputHandler = Substitute.For<IOutputHandler>();
+            var reciever = Substitute.For<IReciever>();
+            var incomingRequestRouter = Substitute.For<IRequestRouter<IHandlerDescriptor>>();
+            var requestDescription = Substitute.For<IHandlerDescriptor>();
+            requestDescription.Method.Returns("abc");
+            var cancelDescription = Substitute.For<IHandlerDescriptor>();
+            cancelDescription.Method.Returns(JsonRpcNames.CancelRequest);
+
+            var req = new Request(1, "abc", null);
+            var cancel = new Notification(JsonRpcNames.CancelRequest, "{ \"id\": 1 }");
+            reciever.IsValid(Arg.Any<JToken>()).Returns(true);
+            reciever.GetRequests(Arg.Any<JToken>())
+                .Returns(c => (new Renor[] { req, cancel }, false));
+
+            incomingRequestRouter.When(z => z.CancelRequest(Arg.Any<object>()));
+            incomingRequestRouter.GetDescriptor(cancel).Returns(cancelDescription);
+            incomingRequestRouter.GetDescriptor(req).Returns(requestDescription);
+
+            incomingRequestRouter.RouteRequest(requestDescription, req, CancellationToken.None)
+                .Returns(new Response(1, req));
+
+            incomingRequestRouter.RouteNotification(cancelDescription, cancel, CancellationToken.None)
+                .Returns(Task.CompletedTask);
+
+            using (NewHandler(
+                inputStream,
+                outputHandler,
+                reciever,
+                Substitute.For<IRequestProcessIdentifier>(),
+                incomingRequestRouter,
+                Substitute.For<IResponseRouter>(),
+                cts => {
+                    outputHandler.When(x => x.Send(Arg.Any<object>()))
+                        .Do(x => {
+                            cts.Cancel();
+                        });
+                }))
+            {
+                await incomingRequestRouter.Received().RouteNotification(cancelDescription, cancel, CancellationToken.None);
+                incomingRequestRouter.Received().CancelRequest(1L);
+            }
+        }
     }
 }