Migrated to global websocket state

application/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-axum = { version = "0.7", optional = true }
+axum = { version = "0.7", optional = true, features = [ "ws", "macros" ] }
 console_error_panic_hook = "0.1"
 leptos = { version = "0.6", features = [] }
 leptos_axum = { version = "0.6", optional = true }
@@ -26,6 +26,7 @@ serde_json = "1.0.115"
 cfg-if = "1.0.0"
 once_cell = "1.19.0"
 futures = "0.3.30"
+uuid = "1.8.0"
 
 [features]
 hydrate = ["leptos/hydrate", "leptos_meta/hydrate", "leptos_router/hydrate"]

application/src/main.rs

@@ -1,9 +1,9 @@
 #[cfg(feature = "ssr")]
 #[tokio::main]
 async fn main() {
-    use application::app::*;
     use application::fileserv::file_and_error_handler;
-    use axum::Router;
+    use application::{app::*, util::websocket::server};
+    use axum::{routing::get, Router};
     use leptos::*;
     use leptos_axum::{generate_route_list, LeptosRoutes};
 
@@ -23,6 +23,7 @@ async fn main() {
 
     // build our application with a route
     let app = Router::new()
+        .route("/ws", get(server::websocket_handler))
         .leptos_routes(&leptos_options, routes, App)
         .fallback(file_and_error_handler)
         .with_state(leptos_options);

application/src/pages/add_participant.rs

@@ -4,12 +4,17 @@ use leptos_router::ActionForm;
 cfg_if::cfg_if! {
     if #[cfg(feature = "ssr")] {
         use crate::util::surrealdb::{DB, schemas};
+        use crate::util::websocket::server;
         use leptos::logging;
     }
 }
 
 #[server(AddParticipant)]
 async fn add_participant(name: String, group: String) -> Result<(), ServerFnError> {
+    let websocket_sate = &server::WEBSOCKET_STATE;
+
+    websocket_sate.tx.send("BOE".to_string()).unwrap();
+
     let created: Vec<schemas::ParticipantRecord> = DB
         .create("participant")
         .content(schemas::NewParticipant { name, group })

application/src/pages/index.rs

@@ -1,8 +1,6 @@
 use leptos::*;
 use leptos_router::*;
 
-use crate::util::surrealdb::schemas;
-
 /// Renders the home page of your application.
 #[component]
 pub fn HomePage() -> impl IntoView {

application/src/util.rs

@@ -1 +1,2 @@
 pub mod surrealdb;
+pub mod websocket;

application/src/util/websocket.rs

@@ -0,0 +1,2 @@
+#[cfg(feature = "ssr")]
+pub mod server;

application/src/util/websocket/server.rs

@@ -0,0 +1,82 @@
+use axum::{
+    extract::ws::{Message, WebSocket, WebSocketUpgrade},
+    response::IntoResponse,
+};
+use futures::{sink::SinkExt, stream::StreamExt};
+use leptos::LeptosOptions;
+use leptos_router::RouteListing;
+use once_cell::sync::Lazy;
+use tokio::sync::{broadcast, Mutex};
+
+use std::{collections::HashSet, sync::Arc};
+
+pub static WEBSOCKET_STATE: Lazy<WebSocketState> = Lazy::new(|| {
+    let client_set = Arc::new(Mutex::new(HashSet::<uuid::Uuid>::new()));
+    let (tx, _rx) = broadcast::channel(100);
+
+    WebSocketState { client_set, tx }
+});
+
+#[derive(Clone)]
+pub struct WebSocketState {
+    pub client_set: Arc<Mutex<HashSet<uuid::Uuid>>>,
+    pub tx: broadcast::Sender<String>,
+}
+
+#[derive(Clone, axum::extract::FromRef)]
+pub struct AppState {
+    pub leptos_options: LeptosOptions,
+    pub websocket_state: Arc<WebSocketState>,
+    pub routes: Vec<RouteListing>,
+}
+
+pub async fn websocket_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
+    ws.on_upgrade(|socket| websocket(socket))
+}
+
+async fn websocket(stream: WebSocket) {
+    let state = &WEBSOCKET_STATE;
+
+    let (mut sender, mut receiver) = stream.split();
+
+    let mut client_set = state.client_set.lock().await;
+
+    let uuid = uuid::Uuid::new_v4();
+
+    client_set.insert(uuid);
+    drop(client_set);
+
+    let mut rx = state.tx.subscribe();
+
+    let msg = format!("{uuid} joined");
+    println!("{uuid} joined");
+    let _ = state.tx.send(msg);
+
+    let mut send_task = tokio::spawn(async move {
+        while let Ok(msg) = rx.recv().await {
+            if sender.send(Message::Text(msg)).await.is_err() {
+                break;
+            }
+        }
+    });
+
+    let tx = state.tx.clone();
+    let uuid_clone = uuid.clone();
+
+    let mut recv_task = tokio::spawn(async move {
+        while let Some(Ok(Message::Text(text))) = receiver.next().await {
+            let _ = tx.send(format!("{uuid_clone}: {text}"));
+        }
+    });
+
+    tokio::select! {
+        _ = (&mut send_task) => recv_task.abort(),
+        _ = (&mut recv_task) => send_task.abort(),
+    };
+
+    let msg = format!("{uuid} left");
+    println!("{uuid} left");
+    let _ = state.tx.send(msg);
+
+    state.client_set.lock().await.remove(&uuid);
+}