[Add] AI Keep alive

src/server/socket.rs

@@ -40,30 +40,55 @@ pub fn handle(path: String, upgrades: OnUpgrade) {
     });
 }
 pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
+    use futures::SinkExt;
+    use tokio::time::Duration;
+
+    const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
+
     loop {
-        let msg = match {
-            let mut receiver = connection.receiver.write().await;
-            receiver.next().await
-        } {
-            Some(Ok(msg)) => msg,
-            Some(Err(e)) => {
+        let mut receiver = connection.receiver.write().await;
+
+        match tokio::time::timeout(IDLE_TIMEOUT, receiver.next()).await {
+            Ok(Some(Ok(msg))) => {
+                // Drop the lock so other tasks can use the receiver if needed,
+                // and so we can handle the message without holding the lock.
+                drop(receiver);
+
+                match msg {
+                    Message::Text(text) => {
+                        let conn_clone = connection.clone();
+                        tokio::spawn(async move {
+                            conn_clone.handle_message(text.to_string()).await;
+                        });
+                    }
+                    Message::Close(_) => {
+                        break;
+                    }
+                    Message::Pong(_) => {
+                        // Received a pong, connection is alive.
+                    }
+                    _ => {}
+                }
+            }
+            Ok(Some(Err(e))) => {
                 log!("WS Error: {}", e);
                 break;
             }
-            _ => break,
-        };
-
-        match msg {
-            Message::Text(text) => {
-                let conn_clone = connection.clone();
-                tokio::spawn(async move {
-                    conn_clone.handle_message(text.to_string()).await;
-                });
-            }
-            Message::Close(_) => {
+            Ok(None) => {
+                // Stream is closed
                 break;
             }
-            _ => {}
+            Err(_) => {
+                // Timeout, we need to send a ping.
+                // Drop receiver lock before acquiring sender lock to avoid deadlock.
+                drop(receiver);
+                log!("WebSocket connection is idle. Sending a ping.");
+                let mut sender = connection.sender.write().await;
+                if let Err(e) = sender.send(Message::Ping(vec![])).await {
+                    log!("Failed to send ping: {}. Closing connection.", e);
+                    break;
+                }
+            }
         }
     }
     connection.handle_close().await;