Ping the client every 30s to keep the connection alive

kazk · kazk · commit 676032b12063 · 2022-03-16T00:07:25.000-07:00
diff --git a/Cargo.toml b/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"] }
 
diff --git a/src/api/proxy.rs b/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,19 @@ async fn connected(
 }
 
 // Type to describe a message from the client conveniently.
+#[allow(clippy::large_enum_variant)]
 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.
