Correctly implement client side request cancellation support

Cargo.toml

@@ -44,9 +44,9 @@ debug = 2
 
 # ungrammar = { path = "lin/ungrammar" }
 
-# salsa = { path = "../salsa" }
-# salsa-macros = { path = "../salsa/components/salsa-macros" }
-# salsa-macro-rules = { path = "../salsa/components/salsa-macro-rules" }
+salsa = { path = "../salsa" }
+salsa-macros = { path = "../salsa/components/salsa-macros" }
+salsa-macro-rules = { path = "../salsa/components/salsa-macro-rules" }
 
 [workspace.dependencies]
 # local crates

crates/ide/src/lib.rs

@@ -67,7 +67,7 @@ use ide_db::{
     FxHashMap, FxIndexSet, LineIndexDatabase,
     base_db::{
         CrateOrigin, CrateWorkspaceData, Env, FileSet, RootQueryDb, SourceDatabase, VfsPath,
-        salsa::{Cancelled, Database},
+        salsa::{CancellationToken, Cancelled, Database},
     },
     prime_caches, symbol_index,
 };
@@ -947,6 +947,10 @@ impl Analysis {
         // We use `attach_db_allow_change()` and not `attach_db()` because fixture injection can change the database.
         hir::attach_db_allow_change(&self.db, || Cancelled::catch(|| f(&self.db)))
     }
+
+    pub fn cancellation_token(&self) -> CancellationToken {
+        self.db.cancellation_token()
+    }
 }
 
 #[test]

crates/rust-analyzer/src/global_state.rs

@@ -15,7 +15,7 @@ use hir::ChangeWithProcMacros;
 use ide::{Analysis, AnalysisHost, Cancellable, FileId, SourceRootId};
 use ide_db::{
     MiniCore,
-    base_db::{Crate, ProcMacroPaths, SourceDatabase, salsa::Revision},
+    base_db::{Crate, ProcMacroPaths, SourceDatabase, salsa::CancellationToken, salsa::Revision},
 };
 use itertools::Itertools;
 use load_cargo::SourceRootConfig;
@@ -89,6 +89,7 @@ pub(crate) struct GlobalState {
     pub(crate) task_pool: Handle<TaskPool<Task>, Receiver<Task>>,
     pub(crate) fmt_pool: Handle<TaskPool<Task>, Receiver<Task>>,
     pub(crate) cancellation_pool: thread::Pool,
+    pub(crate) cancellation_tokens: FxHashMap<lsp_server::RequestId, CancellationToken>,
 
     pub(crate) config: Arc<Config>,
     pub(crate) config_errors: Option<ConfigErrors>,
@@ -265,6 +266,7 @@ impl GlobalState {
             task_pool,
             fmt_pool,
             cancellation_pool,
+            cancellation_tokens: Default::default(),
             loader,
             config: Arc::new(config.clone()),
             analysis_host,
@@ -616,6 +618,7 @@ impl GlobalState {
     }
 
     pub(crate) fn respond(&mut self, response: lsp_server::Response) {
+        self.cancellation_tokens.remove(&response.id);
         if let Some((method, start)) = self.req_queue.incoming.complete(&response.id) {
             if let Some(err) = &response.error
                 && err.message.starts_with("server panicked")
@@ -630,6 +633,9 @@ impl GlobalState {
     }
 
     pub(crate) fn cancel(&mut self, request_id: lsp_server::RequestId) {
+        if let Some(token) = self.cancellation_tokens.remove(&request_id) {
+            token.cancel();
+        }
         if let Some(response) = self.req_queue.incoming.cancel(request_id) {
             self.send(response.into());
         }

crates/rust-analyzer/src/handlers/dispatch.rs

@@ -253,6 +253,9 @@ impl RequestDispatcher<'_> {
         tracing::debug!(?params);
 
         let world = self.global_state.snapshot();
+        self.global_state
+            .cancellation_tokens
+            .insert(req.id.clone(), world.analysis.cancellation_token());
         if RUSTFMT {
             &mut self.global_state.fmt_pool.handle
         } else {
@@ -265,7 +268,19 @@ impl RequestDispatcher<'_> {
             });
             match thread_result_to_response::<R>(req.id.clone(), result) {
                 Ok(response) => Task::Response(response),
-                Err(_cancelled) if ALLOW_RETRYING => Task::Retry(req),
+                Err(HandlerCancelledError::Inner(
+                    Cancelled::PendingWrite | Cancelled::PropagatedPanic,
+                )) if ALLOW_RETRYING => Task::Retry(req),
+                // Note: Technically the return value here does not matter as we have already responded to the client with this error.
+                Err(HandlerCancelledError::Inner(Cancelled::Local)) => Task::Response(Response {
+                    id: req.id,
+                    result: None,
+                    error: Some(ResponseError {
+                        code: ErrorCode::RequestCanceled as i32,
+                        message: "canceled by client".to_owned(),
+                        data: None,
+                    }),
+                }),
                 Err(_cancelled) => {
                     let error = on_cancelled();
                     Task::Response(Response { id: req.id, result: None, error: Some(error) })