🐛 fix(core): dispatch cancellable requests asynchronously

2026-07-14 08:56:30 +08:00
parent efb6eae796
commit 101b69e84f
4 changed files with 387 additions and 11 deletions
@@ -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<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);
    }

@@ -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)