websockets + send messages as json

src/websockets.rs

@@ -1,8 +1,6 @@
-use std::time::Duration;
-
 use axum::{
     extract::{
-        WebSocketUpgrade,
+        State, WebSocketUpgrade,
         ws::{Message, WebSocket},
     },
     response::Response,
@@ -12,22 +10,45 @@ use futures_util::{
     SinkExt as _, StreamExt,
     stream::{SplitSink, SplitStream},
 };
+use tokio::sync::broadcast::Receiver;
 
-pub async fn websocket_handler(upgrade: WebSocketUpgrade) -> Response {
-    upgrade.on_upgrade(handler)
+use crate::{message::ChatMessage, state::AppState};
+
+pub async fn websocket_handler(
+    upgrade: WebSocketUpgrade,
+    State(state): State<AppState>,
+) -> Response {
+    upgrade.on_upgrade(move |websocket| handler(websocket, state))
 }
 
-async fn handler(socket: WebSocket) {
+async fn handler(mut socket: WebSocket, state: AppState) {
+    let mut next_client_id = state.next_client_id.lock().await;
+
+    *next_client_id += 1;
+    let _ = socket.send((*next_client_id).to_string().into()).await;
+
+    drop(next_client_id);
+
     let (sender, receiver) = socket.split();
 
+    let broadcast_receiver = state.broadcast_sender.subscribe();
+
     tokio::spawn(receive(receiver));
-    tokio::spawn(send(sender));
+    tokio::spawn(send(sender, broadcast_receiver));
 }
 
-async fn send(mut sender: SplitSink<WebSocket, Message>) {
-    loop {
-        let _ = sender.send("Hello client!".into());
-        tokio::time::sleep(Duration::from_secs(5)).await;
+async fn send(
+    mut sender: SplitSink<WebSocket, Message>,
+    mut broadcast_receiver: Receiver<ChatMessage>,
+) {
+    while let Ok(message) = broadcast_receiver.recv().await {
+        let Ok(json_message) = serde_json::to_string(&message) else {
+            continue;
+        };
+
+        if sender.send(json_message.into()).await.is_err() {
+            break;
+        }
     }
 }