From 101b69e84f7c51864c2fa4a00511c167ca770116 Mon Sep 17 00:00:00 2001 From: csh Date: Tue, 14 Jul 2026 08:56:30 +0800 Subject: [PATCH] :bug: fix(core): dispatch cancellable requests asynchronously --- lsp-server/src/core/dispatcher.cppm | 21 +- lsp-server/src/core/server.cppm | 209 +++++++++++++++++- lsp-server/test/test_core_server.py | 93 +++++++- .../test_provider/core_server_fixture.cppm | 75 +++++++ 4 files changed, 387 insertions(+), 11 deletions(-) diff --git a/lsp-server/src/core/dispatcher.cppm b/lsp-server/src/core/dispatcher.cppm index d6c653c..8baac86 100644 --- a/lsp-server/src/core/dispatcher.cppm +++ b/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) - : async_executor_(scheduler), manager_hub_(manager_hub) + manager::ManagerHub& 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 id, @@ -81,7 +90,8 @@ export namespace lsp::core void RegisterRequestProvider(std::shared_ptr provider); void RegisterNotificationProvider(std::shared_ptr 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 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); } diff --git a/lsp-server/src/core/server.cppm b/lsp-server/src/core/server.cppm index 6b7d06f..ab2fbcc 100644 --- a/lsp-server/src/core/server.cppm +++ b/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 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& 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 fatal_io_error_ = false; std::mutex output_mutex_; + std::mutex requests_mutex_; + std::unordered_map> active_requests_; }; } @@ -109,6 +124,19 @@ namespace lsp::core return protocol::RequestId{ id->second.Get() }; return std::nullopt; } + + std::string RequestKey(const protocol::RequestId& id) + { + return std::visit( + [](const auto& value) { + using Value = std::decay_t; + if constexpr (std::is_same_v) + 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(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(); + 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 { + 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& 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()( + *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 handles; + { + std::lock_guard lock(requests_mutex_); + handles.reserve(active_requests_.size()); + for (const auto& [key, request] : active_requests_) + { + static_cast(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) diff --git a/lsp-server/test/test_core_server.py b/lsp-server/test/test_core_server.py index 7fcd2ee..7bcb4a5 100644 --- a/lsp-server/test/test_core_server.py +++ b/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: - raise RuntimeError("Unknown request method should return MethodNotFound") + client = LspClient(server) + 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 diff --git a/lsp-server/test/test_provider/core_server_fixture.cppm b/lsp-server/test/test_provider/core_server_fixture.cppm index cfbaeb5..9b39b87 100644 --- a/lsp-server/test/test_provider/core_server_fixture.cppm +++ b/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()); + dispatcher.RegisterRequestProvider(std::make_shared()); + dispatcher.RegisterRequestProvider(std::make_shared()); + dispatcher.RegisterNotificationProvider( + std::make_shared()); }, 2, "");