perf(lsp): cancellation checks in blocking code

cli/lsp/language_server.rs

@@ -43,6 +43,7 @@ use serde_json::from_value;
 use tokio::sync::mpsc::unbounded_channel;
 use tokio::sync::mpsc::UnboundedReceiver;
 use tokio::sync::mpsc::UnboundedSender;
+use tokio_util::sync::CancellationToken;
 use tower_lsp::jsonrpc::Error as LspError;
 use tower_lsp::jsonrpc::Result as LspResult;
 use tower_lsp::lsp_types::request::*;
@@ -137,6 +138,7 @@ pub struct LanguageServer {
   init_flag: AsyncFlag,
   performance: Arc<Performance>,
   shutdown_flag: AsyncFlag,
+  worker_runtime: Arc<tokio::runtime::Runtime>,
 }
 
 /// Snapshot of the state used by TSC.
@@ -239,6 +241,12 @@ impl LanguageServer {
       init_flag: Default::default(),
       performance,
       shutdown_flag,
+      worker_runtime: Arc::new(
+        tokio::runtime::Builder::new_multi_thread()
+          .worker_threads(1)
+          .build()
+          .unwrap(),
+      ),
     }
   }
 
@@ -2219,6 +2227,7 @@ impl Inner {
   async fn completion(
     &self,
     params: CompletionParams,
+    token: CancellationToken,
   ) -> LspResult<Option<CompletionResponse>> {
     let specifier = self.url_map.uri_to_specifier(
       &params.text_document_position.text_document.uri,
@@ -2300,6 +2309,7 @@ impl Inner {
             .options)
             .into(),
           scope.cloned(),
+          token.clone(),
         )
         .await
         .unwrap_or_else(|err| {
@@ -2320,6 +2330,7 @@ impl Inner {
             &specifier,
             position,
             self,
+            token,
           ),
         );
       }
@@ -3039,6 +3050,43 @@ impl Inner {
   }
 }
 
+impl LanguageServer {
+  async fn spawn(
+    &self,
+    fut: impl std::future::Future<Output = ()> + Send + 'static,
+  ) {
+    self
+      .worker_runtime
+      .spawn(fut)
+      .await
+      .inspect_err(|err| lsp_warn!("Handler task error: {err}"))
+      .ok();
+  }
+
+  async fn spawn_with_cancellation<
+    T: Send + 'static,
+    TFut: std::future::Future<Output = LspResult<T>> + Send,
+  >(
+    &self,
+    get_fut: impl FnOnce(CancellationToken) -> TFut + Send + 'static,
+  ) -> LspResult<T> {
+    let token = CancellationToken::new();
+    let _drop_guard = token.clone().drop_guard();
+    self
+      .worker_runtime
+      .spawn(async move {
+        tokio::select! {
+          biased;
+          result = get_fut(token.clone()) => result,
+          _ = token.cancelled() => Err(LspError::request_cancelled()),
+        }
+      })
+      .await
+      .inspect_err(|err| lsp_warn!("Handler task error: {err}"))
+      .unwrap_or_else(|_| Err(LspError::internal_error()))
+  }
+}
+
 #[tower_lsp::async_trait]
 impl tower_lsp::LanguageServer for LanguageServer {
   async fn execute_command(
@@ -3133,28 +3181,48 @@ impl tower_lsp::LanguageServer for LanguageServer {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
-    self.inner.write().await.did_open(params).await;
+    let ls = self.clone();
+    self
+      .spawn(async move {
+        ls.inner.write().await.did_open(params).await;
+      })
+      .await;
   }
 
   async fn did_change(&self, params: DidChangeTextDocumentParams) {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
-    self.inner.write().await.did_change(params).await
+    let ls = self.clone();
+    self
+      .spawn(async move {
+        ls.inner.write().await.did_change(params).await;
+      })
+      .await;
   }
 
   async fn did_save(&self, params: DidSaveTextDocumentParams) {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
-    self.inner.write().await.did_save(params);
+    let ls = self.clone();
+    self
+      .spawn(async move {
+        ls.inner.write().await.did_save(params);
+      })
+      .await;
   }
 
   async fn did_close(&self, params: DidCloseTextDocumentParams) {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
-    self.inner.write().await.did_close(params).await
+    let ls = self.clone();
+    self
+      .spawn(async move {
+        ls.inner.write().await.did_close(params).await;
+      })
+      .await;
   }
 
   async fn did_change_configuration(
@@ -3164,17 +3232,21 @@ impl tower_lsp::LanguageServer for LanguageServer {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
-    let mark = self
-      .performance
-      .mark_with_args("lsp.did_change_configuration", &params);
-    self.refresh_configuration().await;
+    let ls = self.clone();
     self
-      .inner
-      .write()
-      .await
-      .did_change_configuration(params)
+      .spawn(async move {
+        let mark = ls
+          .performance
+          .mark_with_args("lsp.did_change_configuration", &params);
+        ls.refresh_configuration().await;
+        ls.inner
+          .write()
+          .await
+          .did_change_configuration(params)
+          .await;
+        ls.performance.measure(mark);
+      })
       .await;
-    self.performance.measure(mark);
   }
 
   async fn did_change_watched_files(
@@ -3184,12 +3256,16 @@ impl tower_lsp::LanguageServer for LanguageServer {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
+    let ls = self.clone();
     self
-      .inner
-      .write()
-      .await
-      .did_change_watched_files(params)
-      .await
+      .spawn(async move {
+        ls.inner
+          .write()
+          .await
+          .did_change_watched_files(params)
+          .await;
+      })
+      .await;
   }
 
   async fn did_change_workspace_folders(
@@ -3199,22 +3275,25 @@ impl tower_lsp::LanguageServer for LanguageServer {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
-    let mark = self
-      .performance
-      .mark_with_args("lsp.did_change_workspace_folders", &params);
-    self
-      .inner
-      .write()
-      .await
-      .pre_did_change_workspace_folders(params);
-    self.refresh_configuration().await;
+    let ls = self.clone();
     self
-      .inner
-      .write()
-      .await
-      .post_did_change_workspace_folders()
+      .spawn(async move {
+        let mark = ls
+          .performance
+          .mark_with_args("lsp.did_change_workspace_folders", &params);
+        ls.inner
+          .write()
+          .await
+          .pre_did_change_workspace_folders(params);
+        ls.refresh_configuration().await;
+        ls.inner
+          .write()
+          .await
+          .post_did_change_workspace_folders()
+          .await;
+        ls.performance.measure(mark);
+      })
       .await;
-    self.performance.measure(mark);
   }
 
   async fn document_symbol(
@@ -3338,7 +3417,12 @@ impl tower_lsp::LanguageServer for LanguageServer {
     if !self.init_flag.is_raised() {
       self.init_flag.wait_raised().await;
     }
-    self.inner.read().await.completion(params).await
+    let ls = self.clone();
+    self
+      .spawn_with_cancellation(move |token| async move {
+        ls.inner.read().await.completion(params, token).await
+      })
+      .await
   }
 
   async fn completion_resolve(

cli/lsp/mod.rs

@@ -82,7 +82,7 @@ pub async fn start() -> Result<(), AnyError> {
   // Force end the server 8 seconds after receiving a shutdown request.
   tokio::select! {
     biased;
-    _ = Server::new(stdin, stdout, socket).serve(service) => {}
+    _ = Server::new(stdin, stdout, socket).concurrency_level(32).serve(service) => {}
     _ = spawn(async move {
       shutdown_flag.wait_raised().await;
       tokio::time::sleep(std::time::Duration::from_secs(8)).await;

cli/lsp/tsc.rs

@@ -894,6 +894,7 @@ impl TsServer {
       })
   }
 
+  #[allow(clippy::too_many_arguments)]
   pub async fn get_completions(
     &self,
     snapshot: Arc<StateSnapshot>,
@@ -902,6 +903,7 @@ impl TsServer {
     options: GetCompletionsAtPositionOptions,
     format_code_settings: FormatCodeSettings,
     scope: Option<ModuleSpecifier>,
+    token: CancellationToken,
   ) -> Result<Option<CompletionInfo>, AnyError> {
     let req = TscRequest::GetCompletionsAtPosition(Box::new((
       self.specifier_map.denormalize(&specifier),
@@ -910,7 +912,9 @@ impl TsServer {
       format_code_settings,
     )));
     self
-      .request::<Option<CompletionInfo>>(snapshot, req, scope)
+      .request_with_cancellation::<Option<CompletionInfo>>(
+        snapshot, req, scope, token,
+      )
       .await
       .map(|mut info| {
         if let Some(info) = &mut info {
@@ -3699,26 +3703,29 @@ impl CompletionInfo {
     specifier: &ModuleSpecifier,
     position: u32,
     language_server: &language_server::Inner,
+    token: CancellationToken,
   ) -> lsp::CompletionResponse {
     // A cache for costly resolution computations.
     // On a test project, it was found to speed up completion requests
     // by 10-20x and contained ~300 entries for 8000 completion items.
     let mut cache = HashMap::with_capacity(512);
-    let items = self
-      .entries
-      .iter()
-      .flat_map(|entry| {
-        entry.as_completion_item(
-          line_index.clone(),
-          self,
-          settings,
-          specifier,
-          position,
-          language_server,
-          &mut cache,
-        )
-      })
-      .collect();
+    let mut items = Vec::with_capacity(self.entries.len());
+    for entry in &self.entries {
+      if token.is_cancelled() {
+        break;
+      }
+      if let Some(item) = entry.as_completion_item(
+        line_index.clone(),
+        self,
+        settings,
+        specifier,
+        position,
+        language_server,
+        &mut cache,
+      ) {
+        items.push(item);
+      }
+    }
     let is_incomplete = self
       .metadata
       .clone()
@@ -6208,6 +6215,7 @@ mod tests {
         },
         Default::default(),
         Some(temp_dir.url()),
+        Default::default(),
       )
       .await
       .unwrap()
@@ -6401,6 +6409,7 @@ mod tests {
         },
         FormatCodeSettings::from(&fmt_options_config),
         Some(temp_dir.url()),
+        Default::default(),
       )
       .await
       .unwrap()