Skip to main content

howler_core/
websocket.rs

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}