sova_ws/
hub.rs

1//! In-memory pub/sub rooms for WebSocket sessions.
2
3use std::collections::HashMap;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::{Arc, Mutex};
6
7use tokio::sync::mpsc;
8use tokio_tungstenite::tungstenite::Message;
9
10type ClientId = u64;
11
12type Sender = mpsc::UnboundedSender<Message>;
13type RoomClients = HashMap<ClientId, Sender>;
14type Rooms = HashMap<String, RoomClients>;
15
16#[derive(Clone)]
17struct HubInner {
18    rooms: Arc<Mutex<Rooms>>,
19    next_id: Arc<AtomicU64>,
20}
21
22/// Shared room hub installed by [`crate::Ws`].
23#[derive(Clone)]
24pub struct Hub {
25    inner: HubInner,
26}
27
28impl Hub {
29    pub(crate) fn new() -> Self {
30        Self {
31            inner: HubInner {
32                rooms: Arc::new(Mutex::new(HashMap::new())),
33                next_id: Arc::new(AtomicU64::new(1)),
34            },
35        }
36    }
37
38    pub(crate) fn register(
39        &self,
40        room: &str,
41        tx: mpsc::UnboundedSender<Message>,
42    ) -> (ClientId, RoomHandle) {
43        let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed);
44        self.inner
45            .rooms
46            .lock()
47            .unwrap()
48            .entry(room.to_string())
49            .or_default()
50            .insert(id, tx);
51        (
52            id,
53            RoomHandle {
54                inner: Arc::new(RoomHandleInner {
55                    hub: self.inner.clone(),
56                    room: room.to_string(),
57                    id,
58                }),
59            },
60        )
61    }
62
63    fn unregister(&self, room: &str, id: ClientId) {
64        let mut rooms = self.inner.rooms.lock().unwrap();
65        if let Some(clients) = rooms.get_mut(room) {
66            clients.remove(&id);
67            if clients.is_empty() {
68                rooms.remove(room);
69            }
70        }
71    }
72
73    /// Send a message to every client in `room`.
74    pub async fn broadcast(&self, room: &str, msg: Message) {
75        let senders: Vec<mpsc::UnboundedSender<Message>> = {
76            let rooms = self.inner.rooms.lock().unwrap();
77            rooms
78                .get(room)
79                .map(|clients| clients.values().cloned().collect())
80                .unwrap_or_default()
81        };
82        for tx in senders {
83            let _ = tx.send(msg.clone());
84        }
85    }
86}
87
88struct RoomHandleInner {
89    hub: HubInner,
90    room: String,
91    id: ClientId,
92}
93
94/// Handle for a room membership; dropped automatically on leave / session end.
95#[derive(Clone)]
96pub struct RoomHandle {
97    inner: Arc<RoomHandleInner>,
98}
99
100impl RoomHandle {
101    pub fn room(&self) -> &str {
102        &self.inner.room
103    }
104}
105
106impl Drop for RoomHandleInner {
107    fn drop(&mut self) {
108        Hub {
109            inner: self.hub.clone(),
110        }
111        .unregister(&self.room, self.id);
112    }
113}
sova_ws/hub.rs

sova_ws/
hub.rs
