send keepalive pings on user websockets to survive idle proxy timeouts (#1486)

bin/core/src/api/ws/mod.rs

@@ -24,6 +24,9 @@ use crate::{
 mod terminal;
 mod update;
 
+pub(super) const WS_KEEP_ALIVE_INTERVAL: std::time::Duration =
+  std::time::Duration::from_secs(30);
+
 pub fn router() -> Router {
   Router::new()
     // Periphery facing

bin/core/src/api/ws/terminal.rs

@@ -117,9 +117,25 @@ async fn forward_ws_channel(
   };
 
   let periphery_to_core = async {
+    let mut keep_alive =
+      tokio::time::interval(super::WS_KEEP_ALIVE_INTERVAL);
     loop {
       // Already adheres to cancellation token
-      match periphery_receiver.recv().await {
+      let recv = tokio::select! {
+        // Websocket keep-alive ping
+        _ = keep_alive.tick() => {
+          if client_send
+            .send(ws::Message::Ping(Bytes::new()))
+            .await
+            .is_err()
+          {
+            break;
+          }
+          continue;
+        }
+        recv = periphery_receiver.recv() => recv,
+      };
+      match recv {
         Ok(Ok(bytes)) => {
           if let Err(e) =
             client_send.send(ws::Message::Binary(bytes.into())).await

bin/core/src/api/ws/update.rs

@@ -3,6 +3,7 @@ use axum::{
   extract::{WebSocketUpgrade, ws::Message},
   response::IntoResponse,
 };
+use bytes::Bytes;
 use futures_util::{SinkExt, StreamExt};
 use komodo_client::entities::{
   ResourceTarget, permission::PermissionLevel, user::User,
@@ -38,10 +39,19 @@ pub async fn handler(
     let cancel_clone = cancel.clone();
 
     tokio::spawn(async move {
+      let mut keep_alive =
+        tokio::time::interval(super::WS_KEEP_ALIVE_INTERVAL);
       loop {
-        // poll for updates off the receiver / await cancel.
+        // poll for updates off the receiver / await cancel / send keep-alive ping.
         let update = select! {
           _ = cancel_clone.cancelled() => break,
+          _ = keep_alive.tick() => {
+            if ws_sender.send(Message::Ping(Bytes::new())).await.is_err() {
+              cancel_clone.cancel();
+              break;
+            }
+            continue;
+          }
           update = receiver.recv() => {update.expect("failed to recv update msg")}
         };