1use crate::streaming::Broadcast;
2use axum::{
3 extract::{
4 ws::{Message, WebSocket},
5 State, WebSocketUpgrade,
6 },
7 response::IntoResponse,
8};
9use futures_util::{SinkExt, StreamExt};
10
11pub async fn ws_handler(
12 ws: WebSocketUpgrade,
13 State(state): State<crate::web_server::AppState>,
14) -> impl IntoResponse {
15 ws.on_upgrade(move |socket| handle_socket(socket, state.broadcast))
16}
17
18async fn handle_socket(socket: WebSocket, broadcast: Broadcast) {
19 let (mut sender, mut receiver) = socket.split();
20 let mut rx = broadcast.subscribe();
21
22 let mut send_task = tokio::spawn(async move {
23 while let Ok(event) = rx.recv().await {
24 if let Ok(json) = serde_json::to_string(&event) {
25 if sender.send(Message::Text(json)).await.is_err() {
26 break;
27 }
28 }
29 }
30 });
31
32 let mut recv_task =
33 tokio::spawn(async move { while let Some(Ok(_)) = receiver.next().await {} });
34
35 tokio::select! {
36 _ = &mut send_task => recv_task.abort(),
37 _ = &mut recv_task => send_task.abort(),
38 }
39}