Implemented automatic websocket reloading, acking messages and updating state,

and some styles

Author: NicholasLYang

crates/turborepo-ui/src/wui/mod.rs

@@ -1,7 +1,7 @@
 //! Web UI for Turborepo. Creates a WebSocket server that can be subscribed to
 //! by a web client to display the status of tasks.
 
-use std::io::Write;
+use std::{collections::HashSet, io::Write};
 
 use axum::{
     extract::{
@@ -13,8 +13,9 @@ use axum::{
     routing::get,
     Router,
 };
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
+use tokio::select;
 use tower_http::cors::{Any, CorsLayer};
 use tracing::log::warn;
 
@@ -29,6 +30,10 @@ pub enum Error {
     Server(#[from] std::io::Error),
     #[error("failed to start websocket server: {0}")]
     WebSocket(#[source] axum::Error),
+    #[error("failed to serialize message: {0}")]
+    Serde(#[from] serde_json::Error),
+    #[error("failed to send message")]
+    Send(#[from] axum::Error),
 }
 
 #[derive(Clone)]
@@ -98,6 +103,7 @@ impl UISender for WebUISender {
 // Specific events that the websocket server can send to the client,
 // not all the `Event` types from the TUI
 #[derive(Debug, Clone, Serialize)]
+#[serde(tag = "type", content = "payload")]
 pub enum WebUIEvent {
     StartTask {
         task: String,
@@ -122,14 +128,37 @@ pub enum WebUIEvent {
     Stop,
 }
 
+#[derive(Debug, Clone, Serialize)]
+pub struct ServerMessage<'a> {
+    pub id: u32,
+    #[serde(flatten)]
+    pub payload: &'a WebUIEvent,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type", content = "payload")]
+pub enum ClientMessage {
+    /// Acknowledges the receipt of a message.
+    /// If we don't receive an ack, we will resend the message
+    Ack { id: u32 },
+    /// Asks for all messages from the given id onwards
+    CatchUp { start_id: u32 },
+}
+
 struct AppState {
     rx: tokio::sync::broadcast::Receiver<WebUIEvent>,
+    acks: HashSet<u32>,
+    messages: Vec<(WebUIEvent, u32)>,
+    current_id: u32,
 }
 
 impl Clone for AppState {
     fn clone(&self) -> Self {
         Self {
             rx: self.rx.resubscribe(),
+            acks: self.acks.clone(),
+            messages: self.messages.clone(),
+            current_id: self.current_id,
         }
     }
 }
@@ -138,13 +167,60 @@ async fn handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl In
     ws.on_upgrade(|socket| handle_socket(socket, state))
 }
 
-async fn handle_socket(mut socket: WebSocket, state: AppState) {
+async fn handle_socket(socket: WebSocket, state: AppState) {
+    if let Err(e) = handle_socket_inner(socket, state).await {
+        warn!("error handling socket: {e}");
+    }
+}
+
+async fn handle_socket_inner(mut socket: WebSocket, state: AppState) -> Result<(), Error> {
     let mut state = state.clone();
-    while let Ok(event) = state.rx.recv().await {
-        let message_payload = serde_json::to_string(&event).unwrap();
-        if socket.send(Message::Text(message_payload)).await.is_err() {
-            // client disconnected
-            return;
+    let mut interval = tokio::time::interval(std::time::Duration::from_millis(100));
+    loop {
+        select! {
+            biased;
+            Ok(event) = state.rx.recv() => {
+                let id = state.current_id;
+                state.current_id += 1;
+                let message_payload = serde_json::to_string(&ServerMessage {
+                    id,
+                    payload: &event
+                })?;
+
+                state.messages.push((event, id));
+                println!("1");
+                socket.send(Message::Text(message_payload)).await?;
+            }
+            // Every 100ms, check if we need to resend any messages
+            _ = interval.tick() => {
+                for (event, id) in &state.messages {
+                    if !state.acks.contains(&id) {
+                        let message_payload = serde_json::to_string(event).unwrap();
+                        println!("2");
+                        socket.send(Message::Text(message_payload)).await?;
+                    }
+                };
+            }
+            message = socket.recv() => {
+                if let Some(Ok(message)) = message {
+                    let message_payload = message.into_text()?;
+                    if message_payload.is_empty() {
+                        continue;
+                    }
+                    if let Ok(event) = serde_json::from_str::<ClientMessage>(&message_payload) {
+                        match event {
+                            ClientMessage::Ack { id } => {
+                                state.acks.insert(id);
+                            }
+                            ClientMessage::CatchUp { start_id } => {
+                                // TODO: implement
+                            }
+                        }
+                    } else {
+                        warn!("failed to deserialize message from client: {message_payload}");
+                    }
+                }
+            },
         }
     }
 }
@@ -161,7 +237,12 @@ pub async fn start_ws_server(
     let app = Router::new()
         .route("/ws", get(handler))
         .layer(cors)
-        .with_state(AppState { rx });
+        .with_state(AppState {
+            rx,
+            acks: HashSet::new(),
+            messages: Vec::new(),
+            current_id: 0,
+        });
 
     let listener = tokio::net::TcpListener::bind("127.0.0.1:1337").await?;
     println!("Web UI listening on port 1337");