lsp: Gracefully shutdown language servers, fix unwraps in transport

helix-lsp/src/client.rs

@@ -272,6 +272,21 @@ impl Client {
         self.notify::<lsp::notification::Exit>(())
     }
 
+    /// Tries to shut down the language server but returns
+    /// early if server responds with an error.
+    pub async fn shutdown_and_exit(&self) -> Result<()> {
+        self.shutdown().await?;
+        self.exit().await
+    }
+
+    /// Forcefully shuts down the language server ignoring any errors.
+    pub async fn force_shutdown(&self) -> Result<()> {
+        if let Err(e) = self.shutdown().await {
+            log::warn!("language server failed to terminate gracefully - {}", e);
+        }
+        self.exit().await
+    }
+
     // -------------------------------------------------------------------------------------------
     // Text document
     // -------------------------------------------------------------------------------------------

helix-lsp/src/lib.rs

@@ -310,6 +310,10 @@ impl Registry {
             Err(Error::LspNotDefined)
         }
     }
+
+    pub fn iter_clients(&self) -> impl Iterator<Item = &Arc<Client>> {
+        self.inner.values().map(|(_, client)| client)
+    }
 }
 
 #[derive(Debug)]

helix-lsp/src/transport.rs

@@ -1,4 +1,5 @@
 use crate::Result;
+use anyhow::Context;
 use jsonrpc_core as jsonrpc;
 use log::{error, info};
 use serde::{Deserialize, Serialize};
@@ -90,7 +91,7 @@ impl Transport {
 
             match (parts.next(), parts.next(), parts.next()) {
                 (Some("Content-Length"), Some(value), None) => {
-                    content_length = Some(value.parse().unwrap());
+                    content_length = Some(value.parse().context("invalid content length")?);
                 }
                 (Some(_), Some(_), None) => {}
                 _ => {
@@ -103,12 +104,12 @@ impl Transport {
             }
         }
 
-        let content_length = content_length.unwrap();
+        let content_length = content_length.context("missing content length")?;
 
         //TODO: reuse vector
         let mut content = vec![0; content_length];
         reader.read_exact(&mut content).await?;
-        let msg = String::from_utf8(content).unwrap();
+        let msg = String::from_utf8(content).context("invalid utf8 from server")?;
 
         info!("<- {}", msg);
 
@@ -162,7 +163,9 @@ impl Transport {
         match msg {
             ServerMessage::Output(output) => self.process_request_response(output).await?,
             ServerMessage::Call(call) => {
-                self.client_tx.send((self.id, call)).unwrap();
+                self.client_tx
+                    .send((self.id, call))
+                    .context("failed to send a message to server")?;
                 // let notification = Notification::parse(&method, params);
             }
         };

helix-term/src/application.rs

@@ -23,7 +23,7 @@ use crossterm::{
 
 use tui::layout::Rect;
 
-use futures_util::stream::FuturesUnordered;
+use futures_util::{future, stream::FuturesUnordered};
 
 type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
 pub type LspCallback =
@@ -406,6 +406,8 @@ impl Application {
 
         self.event_loop().await;
 
+        self.editor.close_language_servers(None).await;
+
         // reset cursor shape
         write!(stdout, "\x1B[2 q");
 

helix-view/Cargo.toml

@@ -27,6 +27,7 @@ once_cell = "1.8"
 url = "2"
 
 tokio = { version = "1", features = ["full"] }
+futures-util = { version = "0.3", features = ["std", "async-await"], default-features = false }
 
 slotmap = "1"
 

helix-view/src/editor.rs

@@ -2,7 +2,9 @@ use crate::{theme::Theme, tree::Tree, Document, DocumentId, RegisterSelection, V
 use tui::layout::Rect;
 use tui::terminal::CursorKind;
 
+use futures_util::future;
 use std::path::PathBuf;
+use std::time::Duration;
 
 use slotmap::SlotMap;
 
@@ -270,4 +272,22 @@ impl Editor {
             (None, CursorKind::Hidden)
         }
     }
+
+    /// Closes language servers with timeout. The default timeout is 500 ms, use
+    /// `timeout` parameter to override this.
+    pub async fn close_language_servers(
+        &self,
+        timeout: Option<u64>,
+    ) -> Result<(), tokio::time::error::Elapsed> {
+        tokio::time::timeout(
+            Duration::from_millis(timeout.unwrap_or(500)),
+            future::join_all(
+                self.language_servers
+                    .iter_clients()
+                    .map(|client| client.force_shutdown()),
+            ),
+        )
+        .await
+        .map(|_| ())
+    }
 }