Merge pull request #31 from qualified/ws-keep-alive

Cargo.toml

@@ -21,7 +21,7 @@ serde = { version = "1.0.126", features = ["derive"] }
 serde_json = "1.0.64"
 url = "2.2.2"
 
-tokio = { version = "1.6.1", features = ["fs", "process", "macros", "rt", "rt-multi-thread"] }
+tokio = { version = "1.6.1", features = ["fs", "process", "macros", "rt", "rt-multi-thread", "time"] }
 tokio-util = { version = "0.6.7", features = ["codec"] }
 warp = { git = "https://github.com/kazk/warp", branch = "permessage-deflate", default-features = false, features = ["websocket"] }
 

src/api/proxy.rs

@@ -2,7 +2,7 @@ use std::{convert::Infallible, process::Stdio, str::FromStr};
 
 use futures_util::{
     future::{select, Either},
-    SinkExt, StreamExt,
+    stream, SinkExt, StreamExt,
 };
 use tokio::{fs, process::Command};
 use url::Url;
@@ -108,7 +108,19 @@ async fn connected(
     let mut server_send = lsp::framed::writer(server.stdin.take().unwrap());
     let mut server_recv = lsp::framed::reader(server.stdout.take().unwrap());
     let (mut client_send, client_recv) = ws.split();
-    let mut client_recv = client_recv.filter_map(filter_map_warp_ws_message).boxed();
+    let client_recv = client_recv
+        .filter_map(filter_map_warp_ws_message)
+        // Chain this with `Done` so we know when the client disconnects
+        .chain(stream::once(async { Ok(Message::Done) }));
+    // Tick every 30s so we can ping the client to keep the connection alive
+    let ticks = stream::unfold(
+        tokio::time::interval(std::time::Duration::from_secs(30)),
+        |mut interval| async move {
+            interval.tick().await;
+            Some((Ok(Message::Tick), interval))
+        },
+    );
+    let mut client_recv = stream::select(client_recv, ticks).boxed();
 
     let mut client_msg = client_recv.next();
     let mut server_msg = server_recv.next();
@@ -145,15 +157,26 @@ async fn connected(
                         tracing::info!("received Close message");
                     }
 
+                    // Ping the client to keep the connection alive
+                    Some(Ok(Message::Tick)) => {
+                        tracing::debug!("pinging the client");
+                        client_send.send(warp::ws::Message::ping(vec![])).await?;
+                    }
+
+                    // Connection closed
+                    Some(Ok(Message::Done)) => {
+                        tracing::info!("connection closed");
+                        break;
+                    }
+
                     // WebSocket Error
                     Some(Err(err)) => {
                         tracing::error!("websocket error: {}", err);
                     }
 
-                    // Connection closed
                     None => {
-                        tracing::info!("connection closed");
-                        break;
+                        // Unreachable because of the interval stream
+                        unreachable!("should never yield None");
                     }
                 }
 
@@ -206,13 +229,20 @@ async fn connected(
 }
 
 // Type to describe a message from the client conveniently.
+#[allow(clippy::large_enum_variant)]
+#[allow(clippy::enum_variant_names)]
 enum Message {
     // Valid LSP message
     Message(lsp::Message),
     // Invalid JSON
     Invalid(String),
     // Close message
     Close,
+    // Ping the client to keep the connection alive.
+    // Note that this is from the interval stream and not actually from client.
+    Tick,
+    // Client disconnected. Necessary because the combined stream is infinite.
+    Done,
 }
 
 // Parse the message and ignore anything we don't care.

src/lsp/framed/parser.rs

@@ -24,7 +24,7 @@ pub fn parse_message(input: &[u8]) -> IResult<&[u8], &[u8]> {
 
     let header = terminated(terminated(content_len, opt(content_type)), crlf);
 
-    let header = map_res(header, |s: &[u8]| str::from_utf8(s));
+    let header = map_res(header, str::from_utf8);
     let length = map_res(header, |s: &str| s.parse::<usize>());
     let mut message = length_data(length);
 

src/lsp/mod.rs

@@ -18,6 +18,7 @@ use types::Unknown;
 
 #[derive(Clone, Debug, PartialEq, Deserialize)]
 #[serde(untagged)]
+#[allow(clippy::large_enum_variant)]
 pub enum Message {
     Request(Request),
 

src/lsp/request.rs

@@ -10,6 +10,7 @@ use super::types::Id;
 /// [Request message]: https://microsoft.github.io/language-server-protocol/specifications/specification-current/#requestMessage
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 #[serde(tag = "method")]
+#[allow(clippy::large_enum_variant)]
 pub enum Request {
     // To Server
     /// > The [initialize] request is sent as the first request from the client