[Add] AI Keep alive

src/server/socket.rs

@@ -40,18 +40,19 @@ 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
+
-        } {
+        match tokio::time::timeout(IDLE_TIMEOUT, receiver.next()).await {
-            Some(Ok(msg)) => msg,
+            Ok(Some(Ok(msg))) => {
-            Some(Err(e)) => {
+                // Drop the lock so other tasks can use the receiver if needed,
-                log!("WS Error: {}", e);
+                // and so we can handle the message without holding the lock.
-                break;
+                drop(receiver);
-            }
-            _ => break,
-        };
 
                 match msg {
                     Message::Text(text) => {
@@ -63,8 +64,32 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
                     Message::Close(_) => {
                         break;
                     }
+                    Message::Pong(_) => {
+                        // Received a pong, connection is alive.
+                    }
                     _ => {}
                 }
             }
+            Ok(Some(Err(e))) => {
+                log!("WS Error: {}", e);
+                break;
+            }
+            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;
 }