[Add] AI Keep alive

+ Human Bugfix

src/server/socket.rs

@@ -4,8 +4,9 @@ use futures::StreamExt;
 use hyper::upgrade::OnUpgrade;
 use hyper_util::rt::TokioIo;
 use tokio_tungstenite::WebSocketStream;
-use tungstenite::Message;
+use tungstenite::{Message, Utf8Bytes};
 
+use crate::data::communication::{CommunicationType, CommunicationValue};
 use crate::log;
 use crate::server::omikron_connection::OmikronConnection;
 
@@ -40,18 +41,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 +65,39 @@ 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::Text(Utf8Bytes::from(
+                        CommunicationValue::new(CommunicationType::ping)
+                            .to_json()
+                            .to_string(),
+                    )))
+                    .await
+                {
+                    log!("Failed to send ping: {}. Closing connection.", e);
+                    break;
+                }
+            }
+        }
+    }
     connection.handle_close().await;
 }