perf(lsp): cancellable TS diagnostics (#13565)

cli/lsp/diagnostics.rs

@@ -218,6 +218,7 @@ impl DiagnosticsServer {
                     snapshot.clone(),
                     &config,
                     &ts_server,
+                    token.clone(),
                   )
                   .await
                   .map_err(|err| {
@@ -495,6 +496,7 @@ async fn generate_ts_diagnostics(
   snapshot: Arc<language_server::StateSnapshot>,
   config: &ConfigSnapshot,
   ts_server: &tsc::TsServer,
+  token: CancellationToken,
 ) -> Result<DiagnosticVec, AnyError> {
   let mut diagnostics_vec = Vec::new();
   let specifiers = snapshot
@@ -509,7 +511,9 @@ async fn generate_ts_diagnostics(
     .partition::<Vec<_>, _>(|s| config.specifier_enabled(s));
   let ts_diagnostics_map: TsDiagnosticsMap = if !enabled_specifiers.is_empty() {
     let req = tsc::RequestMethod::GetDiagnostics(enabled_specifiers);
-    ts_server.request(snapshot.clone(), req).await?
+    ts_server
+      .request_with_cancellation(snapshot.clone(), req, token)
+      .await?
   } else {
     Default::default()
   };
@@ -781,8 +785,12 @@ let c: number = "a";
       )
       .await;
       assert_eq!(get_diagnostics_for_single(diagnostics).len(), 6);
-      let diagnostics =
+      let diagnostics = generate_ts_diagnostics(
-        generate_ts_diagnostics(snapshot.clone(), &enabled_config, &ts_server)
+        snapshot.clone(),
+        &enabled_config,
+        &ts_server,
+        Default::default(),
+      )
       .await
       .unwrap();
       assert_eq!(get_diagnostics_for_single(diagnostics).len(), 4);
@@ -817,8 +825,12 @@ let c: number = "a";
       )
       .await;
       assert_eq!(get_diagnostics_for_single(diagnostics).len(), 0);
-      let diagnostics =
+      let diagnostics = generate_ts_diagnostics(
-        generate_ts_diagnostics(snapshot.clone(), &disabled_config, &ts_server)
+        snapshot.clone(),
+        &disabled_config,
+        &ts_server,
+        Default::default(),
+      )
       .await
       .unwrap();
       assert_eq!(get_diagnostics_for_single(diagnostics).len(), 0);
@@ -839,4 +851,27 @@ let c: number = "a";
     let (_, _, diagnostics) = diagnostic_vec.into_iter().next().unwrap();
     diagnostics
   }
+
+  #[tokio::test]
+  async fn test_cancelled_ts_diagnostics_request() {
+    let specifier = ModuleSpecifier::parse("file:///a.ts").unwrap();
+    let (snapshot, _) = setup(&[(
+      "file:///a.ts",
+      r#"export let a: string = 5;"#,
+      1,
+      LanguageId::TypeScript,
+    )]);
+    let snapshot = Arc::new(snapshot);
+    let ts_server = TsServer::new(Default::default());
+
+    let config = mock_config();
+    let token = CancellationToken::new();
+    token.cancel();
+    let diagnostics =
+      generate_ts_diagnostics(snapshot.clone(), &config, &ts_server, token)
+        .await
+        .unwrap();
+    // should be none because it's cancelled
+    assert_eq!(diagnostics.len(), 0);
+  }
 }

cli/lsp/tsc.rs

@@ -61,6 +61,7 @@ use text_size::TextRange;
 use text_size::TextSize;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 
 static BRACKET_ACCESSOR_RE: Lazy<Regex> =
   Lazy::new(|| Regex::new(r#"^\[['"](.+)[\['"]\]$"#).unwrap());
@@ -87,6 +88,7 @@ type Request = (
   RequestMethod,
   Arc<StateSnapshot>,
   oneshot::Sender<Result<Value, AnyError>>,
+  CancellationToken,
 );
 
 #[derive(Clone, Debug)]
@@ -101,14 +103,14 @@ impl TsServer {
       let runtime = create_basic_runtime();
       runtime.block_on(async {
         let mut started = false;
-        while let Some((req, state_snapshot, tx)) = rx.recv().await {
+        while let Some((req, state_snapshot, tx, token)) = rx.recv().await {
           if !started {
             // TODO(@kitsonk) need to reflect the debug state of the lsp here
             start(&mut ts_runtime, false, &state_snapshot)
               .expect("could not start tsc");
             started = true;
           }
-          let value = request(&mut ts_runtime, state_snapshot, req);
+          let value = request(&mut ts_runtime, state_snapshot, req, token);
           if tx.send(value).is_err() {
             warn!("Unable to send result to client.");
           }
@@ -124,11 +126,25 @@ impl TsServer {
     snapshot: Arc<StateSnapshot>,
     req: RequestMethod,
   ) -> Result<R, AnyError>
+  where
+    R: de::DeserializeOwned,
+  {
+    self
+      .request_with_cancellation(snapshot, req, Default::default())
+      .await
+  }
+
+  pub(crate) async fn request_with_cancellation<R>(
+    &self,
+    snapshot: Arc<StateSnapshot>,
+    req: RequestMethod,
+    token: CancellationToken,
+  ) -> Result<R, AnyError>
   where
     R: de::DeserializeOwned,
   {
     let (tx, rx) = oneshot::channel::<Result<Value, AnyError>>();
-    if self.0.send((req, snapshot, tx)).is_err() {
+    if self.0.send((req, snapshot, tx, token)).is_err() {
       return Err(anyhow!("failed to send request to tsc thread"));
     }
     rx.await?.map(|v| serde_json::from_value::<R>(v).unwrap())
@@ -2256,6 +2272,7 @@ struct State<'a> {
   state_snapshot: Arc<StateSnapshot>,
   snapshots: HashMap<(ModuleSpecifier, Cow<'a, str>), String>,
   specifiers: HashMap<String, String>,
+  token: CancellationToken,
 }
 
 impl<'a> State<'a> {
@@ -2270,6 +2287,7 @@ impl<'a> State<'a> {
       state_snapshot,
       snapshots: HashMap::default(),
       specifiers: HashMap::default(),
+      token: Default::default(),
     }
   }
 
@@ -2489,6 +2507,10 @@ fn op_get_text(
   Ok(text::slice(content, args.start..args.end).to_string())
 }
 
+fn op_is_cancelled(state: &mut State, _: ()) -> Result<bool, AnyError> {
+  Ok(state.token.is_cancelled())
+}
+
 fn op_load(
   state: &mut State,
   args: SpecifierArgs,
@@ -2602,6 +2624,7 @@ fn init_extension(performance: Arc<Performance>) -> Extension {
       ("op_get_change_range", op_lsp(op_get_change_range)),
       ("op_get_length", op_lsp(op_get_length)),
       ("op_get_text", op_lsp(op_get_text)),
+      ("op_is_cancelled", op_lsp(op_is_cancelled)),
       ("op_load", op_lsp(op_load)),
       ("op_resolve", op_lsp(op_resolve)),
       ("op_respond", op_lsp(op_respond)),
@@ -3069,12 +3092,14 @@ pub(crate) fn request(
   runtime: &mut JsRuntime,
   state_snapshot: Arc<StateSnapshot>,
   method: RequestMethod,
+  token: CancellationToken,
 ) -> Result<Value, AnyError> {
   let (performance, request_params) = {
     let op_state = runtime.op_state();
     let mut op_state = op_state.borrow_mut();
     let state = op_state.borrow_mut::<State>();
     state.state_snapshot = state_snapshot;
+    state.token = token;
     state.last_id += 1;
     let id = state.last_id;
     (state.performance.clone(), method.to_value(state, id))
@@ -3148,7 +3173,8 @@ mod tests {
       request(
         &mut runtime,
         state_snapshot.clone(),
-        RequestMethod::Configure(ts_config)
+        RequestMethod::Configure(ts_config),
+        Default::default(),
       )
       .expect("failed request"),
       json!(true)
@@ -3205,6 +3231,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::Configure(ts_config),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3232,6 +3259,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3282,6 +3310,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3316,6 +3345,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3346,6 +3376,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3400,6 +3431,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3437,6 +3469,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3499,6 +3532,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3544,6 +3578,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetAsset(specifier),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response: Option<String> =
@@ -3588,6 +3623,7 @@ mod tests {
       &mut runtime,
       state_snapshot.clone(),
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3625,6 +3661,7 @@ mod tests {
       &mut runtime,
       state_snapshot,
       RequestMethod::GetDiagnostics(vec![specifier]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();
@@ -3696,6 +3733,7 @@ mod tests {
       &mut runtime,
       state_snapshot.clone(),
       RequestMethod::GetDiagnostics(vec![specifier.clone()]),
+      Default::default(),
     );
     assert!(result.is_ok());
     let result = request(
@@ -3712,6 +3750,7 @@ mod tests {
           trigger_character: Some(".".to_string()),
         },
       )),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response: CompletionInfo =
@@ -3727,6 +3766,7 @@ mod tests {
         source: None,
         data: None,
       }),
+      Default::default(),
     );
     assert!(result.is_ok());
     let response = result.unwrap();

cli/tsc/99_main_compiler.js

@@ -242,6 +242,39 @@ delete Object.prototype.__proto__;
     }
   }
 
+  /** Error thrown on cancellation. */
+  class OperationCanceledError extends Error {
+  }
+
+  /**
+   * Inspired by ThrottledCancellationToken in ts server.
+   *
+   * We don't want to continually call back into Rust and so
+   * we throttle cancellation checks to only occur once
+   * in a while.
+   * @implements {ts.CancellationToken}
+   */
+  class ThrottledCancellationToken {
+    #lastCheckTimeMs = 0;
+
+    isCancellationRequested() {
+      const timeMs = Date.now();
+      // TypeScript uses 20ms
+      if ((timeMs - this.#lastCheckTimeMs) < 20) {
+        return false;
+      }
+
+      this.#lastCheckTimeMs = timeMs;
+      return core.opSync("op_is_cancelled", {});
+    }
+
+    throwIfCancellationRequested() {
+      if (this.isCancellationRequested()) {
+        throw new OperationCanceledError();
+      }
+    }
+  }
+
   /** @type {ts.CompilerOptions} */
   let compilationSettings = {};
 
@@ -262,6 +295,10 @@ delete Object.prototype.__proto__;
       debug(`host.readFile("${specifier}")`);
       return core.opSync("op_load", { specifier }).data;
     },
+    getCancellationToken() {
+      // createLanguageService will call this immediately and cache it
+      return new ThrottledCancellationToken();
+    },
     getSourceFile(
       specifier,
       languageVersion,
@@ -706,11 +743,13 @@ delete Object.prototype.__proto__;
           }
           return respond(id, diagnosticMap);
         } catch (e) {
+          if (!(e instanceof OperationCanceledError)) {
             if ("stack" in e) {
               error(e.stack);
             } else {
               error(e);
             }
+          }
           return respond(id, {});
         }
       }