🐛 fix(core): dispatch cancellable requests asynchronously

lsp-server/src/core/dispatcher.cppm

@@ -44,8 +44,11 @@ export namespace lsp::core
     {
     public:
         ExecutionContext(scheduler::async_executor::AsyncExecutor& scheduler,
-                         manager::ManagerHub& manager_hub)
+                         manager::ManagerHub& manager_hub,
-            : async_executor_(scheduler), manager_hub_(manager_hub)
+                         std::stop_token stop_token = {})
+            : async_executor_(scheduler),
+              manager_hub_(manager_hub),
+              stop_token_(stop_token)
         {
         }
 
@@ -59,9 +62,15 @@ export namespace lsp::core
             return manager_hub_;
         }
 
+        std::stop_token GetStopToken() const
+        {
+            return stop_token_;
+        }
+
     private:
         scheduler::async_executor::AsyncExecutor& async_executor_;
         manager::ManagerHub& manager_hub_;
+        std::stop_token stop_token_;
     };
 
     std::string BuildErrorResponseMessage(std::optional<protocol::RequestId> id,
@@ -81,7 +90,8 @@ export namespace lsp::core
         void RegisterRequestProvider(std::shared_ptr<IRequestProvider> provider);
         void RegisterNotificationProvider(std::shared_ptr<INotificationProvider> provider);
 
-        std::string Dispatch(const protocol::RequestMessage& request);
+        std::string Dispatch(const protocol::RequestMessage& request,
+                             std::stop_token stop_token = {});
         void Dispatch(const protocol::NotificationMessage& notification);
 
         bool SupportsRequest(const std::string& method) const;
@@ -131,7 +141,8 @@ namespace lsp::core
         notification_providers_[method] = std::move(provider);
     }
 
-    std::string RequestDispatcher::Dispatch(const protocol::RequestMessage& request)
+    std::string RequestDispatcher::Dispatch(const protocol::RequestMessage& request,
+                                            std::stop_token stop_token)
     {
         std::shared_ptr<IRequestProvider> provider;
         {
@@ -144,7 +155,7 @@ namespace lsp::core
         if (!provider)
             return HandleUnknownRequest(request);
 
-        ExecutionContext context(async_executor_, manager_hub_);
+        ExecutionContext context(async_executor_, manager_hub_, stop_token);
         return provider->ProvideResponse(request, context);
     }
 

lsp-server/src/core/server.cppm

@@ -44,6 +44,11 @@ export namespace lsp::core
         int Run();
 
     private:
+        struct ActiveRequest
+        {
+            scheduler::async_executor::TaskHandle handle;
+        };
+
         std::optional<std::string> ReadMessage();
         void HandleMessage(const std::string& raw_message);
         void SendMessage(const std::string& message);
@@ -51,6 +56,14 @@ export namespace lsp::core
         void HandleRequest(const protocol::RequestMessage& request);
         void HandleNotification(const protocol::NotificationMessage& notification);
         void HandleResponse(const protocol::ResponseMessage& response);
+        void SubmitRequest(const protocol::RequestMessage& request);
+        void FinishRequest(
+            const protocol::RequestMessage& request,
+            const std::string& key,
+            const std::shared_ptr<ActiveRequest>& active_request,
+            const scheduler::async_executor::TaskResult& result);
+        void HandleCancelRequest(const protocol::NotificationMessage& notification);
+        void DrainActiveRequests();
 
         void InitializeManagerHub();
         void RegisterProviders(ProviderRegistrar registrar);
@@ -80,6 +93,8 @@ export namespace lsp::core
         int exit_code_ = 1;
         std::atomic<bool> fatal_io_error_ = false;
         std::mutex output_mutex_;
+        std::mutex requests_mutex_;
+        std::unordered_map<std::string, std::shared_ptr<ActiveRequest>> active_requests_;
     };
 }
 
@@ -109,6 +124,19 @@ namespace lsp::core
                 return protocol::RequestId{ id->second.Get<protocol::string>() };
             return std::nullopt;
         }
+
+        std::string RequestKey(const protocol::RequestId& id)
+        {
+            return std::visit(
+                [](const auto& value) {
+                    using Value = std::decay_t<decltype(value)>;
+                    if constexpr (std::is_same_v<Value, protocol::integer>)
+                        return "i:" + std::to_string(value);
+                    else
+                        return "s:" + value;
+                },
+                id);
+        }
     }
 
     LspServer::LspServer(std::istream& input,
@@ -176,6 +204,7 @@ namespace lsp::core
             }
         }
 
+        DrainActiveRequests();
         spdlog::info("LSP server main loop ended");
         return exit_code_;
     }
@@ -335,7 +364,23 @@ namespace lsp::core
                 return;
             }
 
-            const auto response = dispatcher_.Dispatch(request);
+            std::string response;
+            try
+            {
+                response = dispatcher_.Dispatch(request);
+            }
+            catch (const std::exception& error)
+            {
+                spdlog::error("Initialize request failed: {}", error.what());
+                SendError(request, protocol::ErrorCodes::InternalError, "Internal error");
+                return;
+            }
+            catch (...)
+            {
+                spdlog::error("Initialize request failed with unknown exception");
+                SendError(request, protocol::ErrorCodes::InternalError, "Internal error");
+                return;
+            }
             SendMessage(response);
             const auto parsed = transform::Deserialize<protocol::ResponseMessage>(response);
             if (!parsed)
@@ -362,6 +407,7 @@ namespace lsp::core
                 return;
             }
 
+            DrainActiveRequests();
             async_executor_.WaitAll();
             manager_hub_.Shutdown();
 
@@ -388,7 +434,7 @@ namespace lsp::core
             return;
         }
 
-        SendMessage(dispatcher_.Dispatch(request));
+        SubmitRequest(request);
     }
 
     void LspServer::HandleNotification(const protocol::NotificationMessage& notification)
@@ -411,7 +457,164 @@ namespace lsp::core
             return;
         }
 
-        dispatcher_.Dispatch(notification);
+        if (notification.method == "$/cancelRequest")
+        {
+            HandleCancelRequest(notification);
+            return;
+        }
+
+        try
+        {
+            dispatcher_.Dispatch(notification);
+        }
+        catch (const std::exception& error)
+        {
+            spdlog::error("Notification {} failed: {}",
+                          notification.method,
+                          error.what());
+        }
+        catch (...)
+        {
+            spdlog::error("Notification {} failed with unknown exception",
+                          notification.method);
+        }
+    }
+
+    void LspServer::SubmitRequest(const protocol::RequestMessage& request)
+    {
+        const std::string key = RequestKey(request.id);
+        auto active_request = std::make_shared<ActiveRequest>();
+        bool duplicate = false;
+        {
+            std::lock_guard lock(requests_mutex_);
+            if (active_requests_.contains(key))
+                duplicate = true;
+            else
+                active_requests_.emplace(key, active_request);
+        }
+
+        if (duplicate)
+        {
+            SendError(request,
+                      protocol::ErrorCodes::InvalidRequest,
+                      "Duplicate active request id");
+            return;
+        }
+
+        active_request->handle = async_executor_.Submit(
+            "lsp-request:" + key,
+            [this, request](std::stop_token stop_token) -> std::optional<std::string> {
+                return dispatcher_.Dispatch(request, stop_token);
+            },
+            [this, request, key, active_request](
+                const scheduler::async_executor::TaskResult& result) {
+                FinishRequest(request, key, active_request, result);
+            });
+    }
+
+    void LspServer::FinishRequest(
+        const protocol::RequestMessage& request,
+        const std::string& key,
+        const std::shared_ptr<ActiveRequest>& active_request,
+        const scheduler::async_executor::TaskResult& result)
+    {
+        {
+            std::lock_guard lock(requests_mutex_);
+            const auto current = active_requests_.find(key);
+            if (current != active_requests_.end() && current->second == active_request)
+                active_requests_.erase(current);
+        }
+
+        switch (result.status)
+        {
+        case scheduler::async_executor::TaskStatus::kCompleted:
+            if (result.value)
+                SendMessage(*result.value);
+            else
+                SendError(request, protocol::ErrorCodes::InternalError, "Internal error");
+            break;
+        case scheduler::async_executor::TaskStatus::kCancelled:
+            SendError(request,
+                      protocol::ErrorCodes::RequestCancelled,
+                      "Request cancelled");
+            break;
+        case scheduler::async_executor::TaskStatus::kFailed:
+            if (result.error)
+            {
+                try
+                {
+                    std::rethrow_exception(result.error);
+                }
+                catch (const std::exception& error)
+                {
+                    spdlog::error("Request {} failed: {}", request.method, error.what());
+                }
+                catch (...)
+                {
+                    spdlog::error("Request {} failed with unknown exception",
+                                  request.method);
+                }
+            }
+            SendError(request, protocol::ErrorCodes::InternalError, "Internal error");
+            break;
+        }
+    }
+
+    void LspServer::HandleCancelRequest(
+        const protocol::NotificationMessage& notification)
+    {
+        if (!notification.params)
+        {
+            spdlog::warn("Ignoring cancel request without params");
+            return;
+        }
+
+        protocol::CancelParams params;
+        try
+        {
+            params = transform::FromLSPAny.template operator()<protocol::CancelParams>(
+                *notification.params);
+        }
+        catch (const std::exception& error)
+        {
+            spdlog::warn("Ignoring invalid cancel request: {}", error.what());
+            return;
+        }
+
+        const std::string key = RequestKey(params.id);
+        scheduler::async_executor::TaskHandle handle;
+        {
+            std::lock_guard lock(requests_mutex_);
+            const auto request = active_requests_.find(key);
+            if (request == active_requests_.end())
+            {
+                spdlog::debug("Cancel request did not match an active request: {}", key);
+                return;
+            }
+            handle = request->second->handle;
+        }
+
+        if (!handle.Cancel())
+            spdlog::debug("Request was already completed or cancelled: {}", key);
+    }
+
+    void LspServer::DrainActiveRequests()
+    {
+        std::vector<scheduler::async_executor::TaskHandle> handles;
+        {
+            std::lock_guard lock(requests_mutex_);
+            handles.reserve(active_requests_.size());
+            for (const auto& [key, request] : active_requests_)
+            {
+                static_cast<void>(key);
+                handles.push_back(request->handle);
+            }
+        }
+
+        for (const auto& handle : handles)
+            handle.Cancel();
+        for (const auto& handle : handles)
+            handle.Wait();
     }
 
     void LspServer::HandleResponse(const protocol::ResponseMessage& response)

lsp-server/test/test_core_server.py

@@ -184,7 +184,6 @@ def assert_json_rpc_errors(server: Path) -> None:
             "method": "initialize",
             "params": {},
         }),
-        frame({"jsonrpc": "2.0", "id": 9, "method": "missing/method"}),
         frame({"jsonrpc": "2.0", "id": 10, "method": "shutdown"}),
         frame({"jsonrpc": "2.0", "method": "exit"}),
     ])
@@ -204,8 +203,95 @@ def assert_json_rpc_errors(server: Path) -> None:
         raise RuntimeError("A non-object JSON value should return InvalidRequest")
     if response_by_id(result.stdout, 7).get("error", {}).get("code") != -32600:
         raise RuntimeError("jsonrpc other than 2.0 should return InvalidRequest")
-    if response_by_id(result.stdout, 9).get("error", {}).get("code") != -32601:
+    client = LspClient(server)
-        raise RuntimeError("Unknown request method should return MethodNotFound")
+    try:
+        client.send({
+            "jsonrpc": "2.0",
+            "id": 11,
+            "method": "initialize",
+            "params": {},
+        })
+        client.read()
+        client.send({"jsonrpc": "2.0", "id": 9, "method": "missing/method"})
+        if client.read().get("error", {}).get("code") != -32601:
+            raise RuntimeError("Unknown request method should return MethodNotFound")
+        client.send({"jsonrpc": "2.0", "id": 12, "method": "shutdown"})
+        client.read()
+        client.send({"jsonrpc": "2.0", "method": "exit"})
+        if client.close_input() != 0:
+            raise RuntimeError("unknown method sequence should shut down cleanly")
+    finally:
+        client.kill()
+
+
+def assert_cancellation_and_failures(server: Path) -> None:
+    client = LspClient(server)
+    try:
+        client.send({
+            "jsonrpc": "2.0",
+            "id": 1,
+            "method": "initialize",
+            "params": {},
+        })
+        if client.read().get("id") != 1:
+            raise RuntimeError("missing initialize response")
+
+        client.send({"jsonrpc": "2.0", "id": 2, "method": "test/block"})
+        client.send({"jsonrpc": "2.0", "id": "2", "method": "test/block"})
+        client.send({
+            "jsonrpc": "2.0",
+            "method": "$/cancelRequest",
+            "params": {"id": 2},
+        })
+        cancelled = client.read(timeout=1)
+        if cancelled.get("id") != 2 or \
+                cancelled.get("error", {}).get("code") != -32800:
+            raise RuntimeError("integer request id should be cancelled independently")
+
+        client.send({
+            "jsonrpc": "2.0",
+            "method": "$/cancelRequest",
+            "params": {"id": "2"},
+        })
+        string_cancelled = client.read(timeout=1)
+        if string_cancelled.get("id") != "2" or \
+                string_cancelled.get("error", {}).get("code") != -32800:
+            raise RuntimeError("string request id should be cancelled independently")
+
+        client.send({"jsonrpc": "2.0", "id": 3, "method": "test/throw"})
+        failed = client.read(timeout=1)
+        if failed.get("id") != 3 or \
+                failed.get("error", {}).get("code") != -32603:
+            raise RuntimeError("request exceptions should return InternalError")
+
+        client.send({"jsonrpc": "2.0", "id": 4, "method": "test/block"})
+        client.send({"jsonrpc": "2.0", "id": 4, "method": "test/block"})
+        duplicate = client.read(timeout=1)
+        if duplicate.get("id") != 4 or \
+                duplicate.get("error", {}).get("code") != -32600:
+            raise RuntimeError("duplicate active request id should return InvalidRequest")
+
+        client.send({
+            "jsonrpc": "2.0",
+            "method": "$/cancelRequest",
+            "params": {"id": 4},
+        })
+        cancelled_original = client.read(timeout=1)
+        if cancelled_original.get("id") != 4 or \
+                cancelled_original.get("error", {}).get("code") != -32800:
+            raise RuntimeError("duplicate id must not replace the original request")
+
+        client.send({"jsonrpc": "2.0", "method": "test/throwNotification"})
+        client.send({"jsonrpc": "2.0", "id": 5, "method": "shutdown"})
+        shutdown = client.read(timeout=1)
+        if shutdown.get("id") != 5 or shutdown.get("result", "missing") is not None:
+            raise RuntimeError("notification exception should not stop the server")
+
+        client.send({"jsonrpc": "2.0", "method": "exit"})
+        if client.close_input() != 0:
+            raise RuntimeError("fixture should exit cleanly after cancellation tests")
+    finally:
+        client.kill()
 
 
 def main() -> int:
@@ -215,6 +301,7 @@ def main() -> int:
 
     assert_lifecycle(args.server)
     assert_json_rpc_errors(args.server)
+    assert_cancellation_and_failures(args.server)
     return 0
 
 

lsp-server/test/test_provider/core_server_fixture.cppm

@@ -40,6 +40,77 @@ namespace lsp::test::provider
         }
     };
 
+    class FixtureBlock final : public core::IRequestProvider
+    {
+    public:
+        std::string GetMethod() const override
+        {
+            return "test/block";
+        }
+
+        std::string GetProviderName() const override
+        {
+            return "FixtureBlock";
+        }
+
+        std::string ProvideResponse(const protocol::RequestMessage& request,
+                                    core::ExecutionContext& context) override
+        {
+            const auto deadline =
+                std::chrono::steady_clock::now() + std::chrono::seconds(2);
+            while (!context.GetStopToken().stop_requested() &&
+                   std::chrono::steady_clock::now() < deadline)
+            {
+                std::this_thread::sleep_for(std::chrono::milliseconds(5));
+            }
+
+            protocol::ResponseMessage response;
+            response.id = request.id;
+            response.result = protocol::LSPAny(protocol::string("completed"));
+            return codec::Serialize(response).value();
+        }
+    };
+
+    class FixtureThrow final : public core::IRequestProvider
+    {
+    public:
+        std::string GetMethod() const override
+        {
+            return "test/throw";
+        }
+
+        std::string GetProviderName() const override
+        {
+            return "FixtureThrow";
+        }
+
+        std::string ProvideResponse(const protocol::RequestMessage&,
+                                    core::ExecutionContext&) override
+        {
+            throw std::runtime_error("fixture request failure");
+        }
+    };
+
+    class FixtureThrowNotification final : public core::INotificationProvider
+    {
+    public:
+        std::string GetMethod() const override
+        {
+            return "test/throwNotification";
+        }
+
+        std::string GetProviderName() const override
+        {
+            return "FixtureThrowNotification";
+        }
+
+        void HandleNotification(const protocol::NotificationMessage&,
+                                core::ExecutionContext&) override
+        {
+            throw std::runtime_error("fixture notification failure");
+        }
+    };
+
     int RunCoreServerFixture()
     {
         spdlog::set_level(spdlog::level::off);
@@ -48,6 +119,10 @@ namespace lsp::test::provider
             std::cout,
             [](core::RequestDispatcher& dispatcher) {
                 dispatcher.RegisterRequestProvider(std::make_shared<FixtureInitialize>());
+                dispatcher.RegisterRequestProvider(std::make_shared<FixtureBlock>());
+                dispatcher.RegisterRequestProvider(std::make_shared<FixtureThrow>());
+                dispatcher.RegisterNotificationProvider(
+                    std::make_shared<FixtureThrowNotification>());
             },
             2,
             "");