From 58101fa9048c2df19ee528432b5c304d187dafb7 Mon Sep 17 00:00:00 2001
From: Nayeem Rahman <nayeemrmn99@gmail.com>
Date: Thu, 30 Nov 2023 03:35:39 +0000
Subject: [PATCH] perf(lsp): cancel ts requests on future drop (#21387)

When an old request is obsoleted while the user is typing, the client
will say so to the server and tower-lsp will drop the future associated
with that request.

This wires that up to the ts server by having any request's token be
cancelled when the surrounding state is dropped.
---
 cli/lsp/tsc.rs | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/cli/lsp/tsc.rs b/cli/lsp/tsc.rs
index 94770603c8..996389cbe3 100644
--- a/cli/lsp/tsc.rs
+++ b/cli/lsp/tsc.rs
@@ -230,8 +230,12 @@ impl TsServer {
             start(&mut ts_runtime, false).unwrap();
             started = true;
           }
-          let value = request(&mut ts_runtime, state_snapshot, req, token);
-          if tx.send(value).is_err() {
+          let value =
+            request(&mut ts_runtime, state_snapshot, req, token.clone());
+          let was_sent = tx.send(value).is_ok();
+          // Don't print the send error if the token is cancelled, it's expected
+          // to fail in that case and this commonly occurs.
+          if !was_sent && !token.is_cancelled() {
             lsp_warn!("Unable to send result to client.");
           }
         }
@@ -961,11 +965,24 @@ impl TsServer {
   where
     R: de::DeserializeOwned,
   {
+    // When an LSP request is cancelled by the client, the future this is being
+    // executed under and any local variables here will be dropped at the next
+    // await point. To pass on that cancellation to the TS thread, we make this
+    // wrapper which cancels the request's token on drop.
+    struct DroppableToken(CancellationToken);
+    impl Drop for DroppableToken {
+      fn drop(&mut self) {
+        self.0.cancel();
+      }
+    }
+    let token = token.child_token();
+    let droppable_token = DroppableToken(token.clone());
     let (tx, rx) = oneshot::channel::<Result<Value, AnyError>>();
     if self.sender.send((req, snapshot, tx, token)).is_err() {
       return Err(anyhow!("failed to send request to tsc thread"));
     }
     let value = rx.await??;
+    drop(droppable_token);
     Ok(serde_json::from_value::<R>(value)?)
   }
 }