heartbeat

crates/backend/Cargo.toml

@@ -18,3 +18,4 @@ uuid.workspace = true
 tower-http = { version = "0.4.0", features = ["fs", "trace", "cors"] }
 tokio-stream = { version = "0.1.14", features = ["sync"] }
 clap = { version = "4.3.3", features = ["derive"] }
+futures-util = "0.3.28"

crates/backend/src/main.rs

@@ -20,13 +20,15 @@ use common::DeleteMilestone;
 use common::Milestone;
 use common::RestResponse;
 use common::ToggleAchievement;
+use futures_util::SinkExt;
+use futures_util::StreamExt;
 use serde::Deserialize;
 use serde::Serialize;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::fs;
-use tokio_stream::StreamExt;
+use tokio::select;
 use tower::ServiceBuilder;
 use tower_http::cors::CorsLayer;
 use tower_http::trace::DefaultMakeSpan;
@@ -85,7 +87,7 @@ async fn main() {
             );
             AppState::default()
         }
-        Err(e) => panic!("Unexpected error: {:?}", e),
+        Err(e) => panic!("Failed to read app state. Unexpected error: {:?}", e),
     };
     let (app_state_watch_tx, app_state_watch_rx) = tokio::sync::watch::channel(init_app_state.state.clone());
     let app_state: SharedState = Arc::new(tokio::sync::RwLock::new(SharedStateParts {
@@ -184,14 +186,54 @@ async fn ws_handler(
 }
 
 /// Websocket statemachine (one will be spawned per connection)
-async fn handle_socket(mut socket: WebSocket, state_watch_rx: tokio::sync::watch::Receiver<common::State>) {
-    let mut stream = tokio_stream::wrappers::WatchStream::new(state_watch_rx);
-    while let Some(state) = stream.next().await {
-        let state: common::WebSocketMessage = state;
-        let serialized = serde_json::to_string(&state).expect("Failed to serialize app state to JSON");
-        if socket.send(Message::Text(serialized)).await.is_err() {
-            tracing::debug!("Websocket client disconnected");
-            break;
+async fn handle_socket(socket: WebSocket, state_watch_rx: tokio::sync::watch::Receiver<common::State>) {
+    let (mut send, mut recv) = socket.split();
+
+    let mut outgoing = tokio_stream::wrappers::WatchStream::new(state_watch_rx).map(|state| {
+        let msg = common::WebSocketMessageServerToApp::State(state);
+        let msg = serde_json::to_string(&msg).expect("Failed to serialize app state to JSON");
+        Message::Text(msg)
+    });
+
+    let incoming = &mut recv;
+
+    loop {
+        select! {
+            out = outgoing.next() => match out {
+                Some(msg) => {
+                    if send.send(msg).await.is_err() {
+                        tracing::debug!("Websocket client disconnected");
+                        break;
+                    }
+                },
+                None => break,
+            },
+            inc = incoming.next() => match inc {
+                Some(Ok(Message::Text(msg))) => {
+                    match serde_json::from_str::<common::WebSocketMessageAppToServer>(&msg) {
+                        Ok(msg) => match msg {
+                            common::WebSocketMessageAppToServer::HeartBeat => {
+                                tracing::info!("ws recv heartbeat");
+                            },
+                        },
+                        Err(err) => {
+                            tracing::error!("ws deserialization error: {err}");
+                        },
+                    }
+                },
+                Some(Ok(Message::Binary(_))) => {},
+                Some(Ok(Message::Ping(_))) => {
+                    tracing::info!("ws recv ping");
+                },
+                Some(Ok(Message::Pong(_))) => {
+                    tracing::info!("ws recv pong");
+                },
+                Some(Ok(Message::Close(_))) => {},
+                Some(Err(err)) => {
+                    tracing::error!("ws error: {err}");
+                },
+                None => break,
+            }
         }
     }