From be561817022969f13d93e8051e007fd7eced92b7 Mon Sep 17 00:00:00 2001
From: Nayeem Rahman <nayeemrmn99@gmail.com>
Date: Wed, 22 Nov 2023 04:08:48 +0000
Subject: [PATCH] fix(lsp): force shutdown after a timeout (#21251)

(cherry picked from commit 64997cce6a9e83dcc57f5be9056eca16708db096)
---
 cli/lsp/language_server.rs | 11 ++++++++---
 cli/lsp/mod.rs             | 21 +++++++++++++++++++--
 cli/lsp/repl.rs            |  6 ++++--
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/cli/lsp/language_server.rs b/cli/lsp/language_server.rs
index 6f0c96746da1d5..060be423b1ffe7 100644
--- a/cli/lsp/language_server.rs
+++ b/cli/lsp/language_server.rs
@@ -30,6 +30,7 @@ use std::env;
 use std::fmt::Write as _;
 use std::path::PathBuf;
 use std::sync::Arc;
+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::*;
@@ -156,7 +157,7 @@ impl LspNpmConfigHash {
 }
 
 #[derive(Debug, Clone)]
-pub struct LanguageServer(Arc<tokio::sync::RwLock<Inner>>);
+pub struct LanguageServer(Arc<tokio::sync::RwLock<Inner>>, CancellationToken);
 
 #[derive(Debug)]
 pub struct StateNpmSnapshot {
@@ -226,8 +227,11 @@ pub struct Inner {
 }
 
 impl LanguageServer {
-  pub fn new(client: Client) -> Self {
-    Self(Arc::new(tokio::sync::RwLock::new(Inner::new(client))))
+  pub fn new(client: Client, token: CancellationToken) -> Self {
+    Self(
+      Arc::new(tokio::sync::RwLock::new(Inner::new(client))),
+      token,
+    )
   }
 
   /// Similar to `deno cache` on the command line, where modules will be cached
@@ -3213,6 +3217,7 @@ impl tower_lsp::LanguageServer for LanguageServer {
   }
 
   async fn shutdown(&self) -> LspResult<()> {
+    self.1.cancel();
     self.0.write().await.shutdown().await
   }
 
diff --git a/cli/lsp/mod.rs b/cli/lsp/mod.rs
index af6ef88c91b55a..2f9cfb0bff1041 100644
--- a/cli/lsp/mod.rs
+++ b/cli/lsp/mod.rs
@@ -1,6 +1,8 @@
 // Copyright 2018-2023 the Deno authors. All rights reserved. MIT license.
 
 use deno_core::error::AnyError;
+use deno_core::unsync::spawn;
+use tokio_util::sync::CancellationToken;
 use tower_lsp::LspService;
 use tower_lsp::Server;
 
@@ -39,8 +41,12 @@ pub async fn start() -> Result<(), AnyError> {
   let stdin = tokio::io::stdin();
   let stdout = tokio::io::stdout();
 
+  let token = CancellationToken::new();
   let builder = LspService::build(|client| {
-    language_server::LanguageServer::new(client::Client::from_tower(client))
+    language_server::LanguageServer::new(
+      client::Client::from_tower(client),
+      token.clone(),
+    )
   })
   // TODO(nayeemrmn): The extension has replaced this with the `deno.cache`
   // command as of vscode_deno 3.21.0 / 2023.09.05. Remove this eventually.
@@ -81,7 +87,18 @@ pub async fn start() -> Result<(), AnyError> {
 
   let (service, socket) = builder.finish();
 
-  Server::new(stdin, stdout, socket).serve(service).await;
+  // TODO(nayeemrmn): This cancellation token is a workaround for
+  // https://github.com/denoland/deno/issues/20700. Remove when
+  // https://github.com/ebkalderon/tower-lsp/issues/399 is fixed.
+  // Force end the server 8 seconds after receiving a shutdown request.
+  tokio::select! {
+    biased;
+    _ = Server::new(stdin, stdout, socket).serve(service) => {}
+    _ = spawn(async move {
+      token.cancelled().await;
+      tokio::time::sleep(std::time::Duration::from_secs(8)).await;
+    }) => {}
+  }
 
   Ok(())
 }
diff --git a/cli/lsp/repl.rs b/cli/lsp/repl.rs
index 836d75a5bbe612..2f023197d1abb8 100644
--- a/cli/lsp/repl.rs
+++ b/cli/lsp/repl.rs
@@ -61,8 +61,10 @@ impl ReplLanguageServer {
     super::logging::set_lsp_log_level(log::Level::Debug);
     super::logging::set_lsp_warn_level(log::Level::Debug);
 
-    let language_server =
-      super::language_server::LanguageServer::new(Client::new_for_repl());
+    let language_server = super::language_server::LanguageServer::new(
+      Client::new_for_repl(),
+      Default::default(),
+    );
 
     let cwd_uri = get_cwd_uri()?;