wsio_server/namespace/
mod.rs

1use std::sync::Arc;
2
3use anyhow::Result;
4use arc_swap::ArcSwap;
5use futures_util::{
6    SinkExt,
7    StreamExt,
8};
9use http::{
10    HeaderMap,
11    Uri,
12};
13use hyper::upgrade::{
14    OnUpgrade,
15    Upgraded,
16};
17use hyper_util::rt::TokioIo;
18use kikiutils::{
19    atomic::enum_cell::AtomicEnumCell,
20    types::fx_collections::FxDashMap,
21};
22use num_enum::{
23    IntoPrimitive,
24    TryFromPrimitive,
25};
26use roaring::RoaringTreemap;
27use serde::Serialize;
28use tokio::{
29    join,
30    select,
31    spawn,
32    sync::Mutex,
33    task::JoinSet,
34    time::timeout,
35};
36use tokio_tungstenite::{
37    WebSocketStream,
38    tungstenite::{
39        Message,
40        protocol::Role,
41    },
42};
43
44pub(crate) mod builder;
45mod config;
46pub mod operators;
47
48use self::{
49    config::WsIoServerNamespaceConfig,
50    operators::broadcast::WsIoServerNamespaceBroadcastOperator,
51};
52use crate::{
53    WsIoServer,
54    connection::WsIoServerConnection,
55    core::packet::WsIoPacket,
56    runtime::{
57        WsIoServerRuntime,
58        WsIoServerRuntimeStatus,
59    },
60};
61
62// Enums
63#[repr(u8)]
64#[derive(Debug, Eq, IntoPrimitive, PartialEq, TryFromPrimitive)]
65enum NamespaceStatus {
66    Running,
67    Stopped,
68    Stopping,
69}
70
71// Structs
72#[derive(Debug)]
73pub struct WsIoServerNamespace {
74    pub(crate) config: WsIoServerNamespaceConfig,
75    connection_ids: ArcSwap<RoaringTreemap>,
76    connections: FxDashMap<u64, Arc<WsIoServerConnection>>,
77    connection_task_set: Mutex<JoinSet<()>>,
78    rooms: FxDashMap<String, RoaringTreemap>,
79    runtime: Arc<WsIoServerRuntime>,
80    status: AtomicEnumCell<NamespaceStatus>,
81}
82
83impl WsIoServerNamespace {
84    fn new(config: WsIoServerNamespaceConfig, runtime: Arc<WsIoServerRuntime>) -> Arc<Self> {
85        Arc::new(Self {
86            config,
87            connection_ids: ArcSwap::new(Arc::new(RoaringTreemap::new())),
88            connections: FxDashMap::default(),
89            connection_task_set: Mutex::new(JoinSet::new()),
90            rooms: FxDashMap::default(),
91            runtime,
92            status: AtomicEnumCell::new(NamespaceStatus::Running),
93        })
94    }
95
96    // Private methods
97    async fn handle_upgraded_request(
98        self: &Arc<Self>,
99        headers: HeaderMap,
100        request_uri: Uri,
101        upgraded: Upgraded,
102    ) -> Result<()> {
103        // Create ws stream
104        let mut ws_stream =
105            WebSocketStream::from_raw_socket(TokioIo::new(upgraded), Role::Server, Some(self.config.websocket_config))
106                .await;
107
108        // Check runtime and namespace status
109        if !self.runtime.status.is(WsIoServerRuntimeStatus::Running) || !self.status.is(NamespaceStatus::Running) {
110            ws_stream
111                .send((*self.encode_packet_to_message(&WsIoPacket::new_disconnect())?).clone())
112                .await?;
113
114            let _ = ws_stream.close(None).await;
115            return Ok(());
116        }
117
118        // Create connection
119        let (connection, mut message_rx) = WsIoServerConnection::new(headers, self.clone(), request_uri);
120
121        // Split ws stream and spawn read and write tasks
122        let (mut ws_stream_writer, mut ws_stream_reader) = ws_stream.split();
123        let connection_clone = connection.clone();
124        let mut read_ws_stream_task = spawn(async move {
125            while let Some(message) = ws_stream_reader.next().await {
126                if match message {
127                    Ok(Message::Binary(bytes)) => {
128                        // Treat any single-byte binary frame as a client heartbeat and ignore it
129                        if bytes.len() == 1 {
130                            continue;
131                        }
132
133                        connection_clone.handle_incoming_packet(&bytes).await
134                    },
135                    Ok(Message::Close(_)) | Err(_) => break,
136                    Ok(Message::Text(text)) => connection_clone.handle_incoming_packet(text.as_bytes()).await,
137                    _ => Ok(()),
138                }
139                .is_err()
140                {
141                    break;
142                }
143            }
144        });
145
146        let mut write_ws_stream_task = spawn(async move {
147            while let Some(message) = message_rx.recv().await {
148                let message = (*message).clone();
149                let is_close = matches!(message, Message::Close(_));
150                if ws_stream_writer.send(message).await.is_err() {
151                    break;
152                }
153
154                if is_close {
155                    let _ = ws_stream_writer.close().await;
156                    break;
157                }
158            }
159        });
160
161        // Try to init connection
162        match connection.init().await {
163            Ok(_) => {
164                // Wait for either read or write task to finish
165                select! {
166                    _ = &mut read_ws_stream_task => {
167                        write_ws_stream_task.abort();
168                    },
169                    _ = &mut write_ws_stream_task => {
170                        read_ws_stream_task.abort();
171                    },
172                }
173            },
174            Err(_) => {
175                // Close connection
176                read_ws_stream_task.abort();
177                connection.close();
178                let _ = join!(read_ws_stream_task, write_ws_stream_task);
179            },
180        }
181
182        // Cleanup connection
183        connection.cleanup().await;
184        Ok(())
185    }
186
187    // Protected methods
188    #[inline]
189    pub(crate) fn add_connection_id_to_room(&self, room_name: &str, connection_id: u64) {
190        self.rooms.entry(room_name.into()).or_default().insert(connection_id);
191    }
192
193    #[inline]
194    pub(crate) fn encode_packet_to_message(&self, packet: &WsIoPacket) -> Result<Arc<Message>> {
195        let bytes = self.config.packet_codec.encode(packet)?;
196        Ok(Arc::new(match self.config.packet_codec.is_text() {
197            // SAFETY: text packet codecs only produce valid UTF-8 payloads.
198            true => Message::Text(unsafe { String::from_utf8_unchecked(bytes).into() }),
199            false => Message::Binary(bytes.into()),
200        }))
201    }
202
203    pub(crate) async fn handle_on_upgrade_request(
204        self: &Arc<Self>,
205        headers: HeaderMap,
206        on_upgrade: OnUpgrade,
207        request_uri: Uri,
208    ) {
209        let namespace = self.clone();
210        self.connection_task_set.lock().await.spawn(async move {
211            if let Ok(Ok(upgraded)) = timeout(namespace.config.http_request_upgrade_timeout, on_upgrade).await {
212                let _ = namespace.handle_upgraded_request(headers, request_uri, upgraded).await;
213            }
214        });
215    }
216
217    #[inline]
218    pub(crate) fn insert_connection(&self, connection: Arc<WsIoServerConnection>) {
219        self.connections.insert(connection.id(), connection.clone());
220        self.runtime.insert_connection_id(connection.id());
221        self.connection_ids.rcu(|old_connection_ids| {
222            let mut new_connection_ids = (**old_connection_ids).clone();
223            new_connection_ids.insert(connection.id());
224            new_connection_ids
225        });
226    }
227
228    #[inline]
229    pub(crate) fn remove_connection(&self, id: u64) {
230        self.connections.remove(&id);
231        self.runtime.remove_connection_id(id);
232        self.connection_ids.rcu(|old_connection_ids| {
233            let mut new_connection_ids = (**old_connection_ids).clone();
234            new_connection_ids.remove(id);
235            new_connection_ids
236        });
237    }
238
239    #[inline]
240    pub(crate) fn remove_connection_id_from_room(&self, room_name: &str, connection_id: u64) {
241        if let Some(mut entry) = self.rooms.get_mut(room_name) {
242            entry.remove(connection_id);
243        }
244
245        self.rooms.remove_if(room_name, |_, entry| entry.is_empty());
246    }
247
248    // Public methods
249    pub async fn close_all(self: &Arc<Self>) {
250        WsIoServerNamespaceBroadcastOperator::new(self.clone()).close().await;
251    }
252
253    #[inline]
254    pub fn connection_count(&self) -> usize {
255        self.connections.len()
256    }
257
258    pub async fn disconnect_all(self: &Arc<Self>) -> Result<()> {
259        WsIoServerNamespaceBroadcastOperator::new(self.clone())
260            .disconnect()
261            .await
262    }
263
264    pub async fn emit<D: Serialize>(self: &Arc<Self>, event: impl AsRef<str>, data: Option<&D>) -> Result<()> {
265        WsIoServerNamespaceBroadcastOperator::new(self.clone())
266            .emit(event, data)
267            .await
268    }
269
270    #[inline]
271    pub fn except(
272        self: &Arc<Self>,
273        room_names: impl IntoIterator<Item = impl Into<String>>,
274    ) -> WsIoServerNamespaceBroadcastOperator {
275        WsIoServerNamespaceBroadcastOperator::new(self.clone()).except(room_names)
276    }
277
278    #[inline]
279    pub fn path(&self) -> &str {
280        &self.config.path
281    }
282
283    #[inline]
284    pub fn server(&self) -> WsIoServer {
285        WsIoServer(self.runtime.clone())
286    }
287
288    pub async fn shutdown(self: &Arc<Self>) {
289        match self.status.get() {
290            NamespaceStatus::Stopped => return,
291            NamespaceStatus::Running => self.status.store(NamespaceStatus::Stopping),
292            _ => unreachable!(),
293        }
294
295        self.close_all().await;
296        let mut connection_task_set = self.connection_task_set.lock().await;
297        while connection_task_set.join_next().await.is_some() {}
298
299        self.status.store(NamespaceStatus::Stopped);
300    }
301
302    #[inline]
303    pub fn to(
304        self: &Arc<Self>,
305        room_names: impl IntoIterator<Item = impl Into<String>>,
306    ) -> WsIoServerNamespaceBroadcastOperator {
307        WsIoServerNamespaceBroadcastOperator::new(self.clone()).to(room_names)
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use std::time::Duration;
314
315    use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
316
317    use super::*;
318    use crate::{
319        config::WsIoServerConfig,
320        core::packet::codecs::WsIoPacketCodec,
321    };
322
323    fn create_test_namespace() -> Arc<WsIoServerNamespace> {
324        let runtime = WsIoServerRuntime::new(WsIoServerConfig {
325            broadcast_concurrency_limit: 16,
326            http_request_upgrade_timeout: Duration::from_secs(3),
327            init_request_handler_timeout: Duration::from_secs(3),
328            init_response_handler_timeout: Duration::from_secs(3),
329            init_response_timeout: Duration::from_secs(3),
330            middleware_execution_timeout: Duration::from_secs(3),
331            on_close_handler_timeout: Duration::from_secs(3),
332            on_connect_handler_timeout: Duration::from_secs(3),
333            packet_codec: WsIoPacketCodec::SerdeJson,
334            request_path: "/socket".into(),
335            websocket_config: WebSocketConfig::default(),
336        });
337        runtime.new_namespace_builder("/test").register().unwrap()
338    }
339
340    #[tokio::test]
341    async fn test_namespace_new() {
342        let namespace = create_test_namespace();
343        assert_eq!(namespace.path(), "/test");
344        assert_eq!(namespace.connection_count(), 0);
345    }
346
347    #[tokio::test]
348    async fn test_namespace_connection_count() {
349        let namespace = create_test_namespace();
350        assert_eq!(namespace.connection_count(), 0);
351    }
352
353    #[tokio::test]
354    async fn test_namespace_server() {
355        let namespace = create_test_namespace();
356        namespace.server();
357    }
358
359    #[tokio::test]
360    async fn test_namespace_to_broadcast_operator() {
361        let namespace = create_test_namespace();
362        namespace.to(["room1", "room2"]);
363    }
364
365    #[tokio::test]
366    async fn test_namespace_except_broadcast_operator() {
367        let namespace = create_test_namespace();
368        namespace.except(["room1", "room2"]);
369    }
370
371    #[tokio::test]
372    async fn test_namespace_add_remove_connection_id_to_room() {
373        let namespace = create_test_namespace();
374        namespace.add_connection_id_to_room("room1", 1);
375        namespace.add_connection_id_to_room("room1", 2);
376        namespace.add_connection_id_to_room("room2", 3);
377
378        assert_eq!(namespace.rooms.get("room1").unwrap().len(), 2);
379        assert!(namespace.rooms.get("room1").unwrap().contains(1));
380        assert!(namespace.rooms.get("room1").unwrap().contains(2));
381        assert_eq!(namespace.rooms.get("room2").unwrap().len(), 1);
382
383        namespace.remove_connection_id_from_room("room1", 1);
384        assert_eq!(namespace.rooms.get("room1").unwrap().len(), 1);
385        assert!(namespace.rooms.get("room1").unwrap().contains(2));
386
387        namespace.remove_connection_id_from_room("room1", 2);
388        namespace.remove_connection_id_from_room("room2", 3);
389
390        assert!(!namespace.rooms.contains_key("room1"));
391        assert!(!namespace.rooms.contains_key("room2"));
392    }
393
394    #[tokio::test]
395    async fn test_namespace_remove_connection_id_from_empty_room() {
396        let namespace = create_test_namespace();
397        // Removing from non-existent room should not panic
398        namespace.remove_connection_id_from_room("nonexistent", 1);
399    }
400
401    #[tokio::test]
402    async fn test_namespace_encode_packet_to_message() {
403        let namespace = create_test_namespace();
404        let packet = WsIoPacket::new_disconnect();
405        let message = namespace.encode_packet_to_message(&packet).unwrap();
406
407        assert!(matches!(&*message, Message::Text(_)));
408    }
409
410    #[tokio::test]
411    async fn test_namespace_shutdown_idempotent() {
412        let namespace = create_test_namespace();
413        namespace.clone().shutdown().await;
414        // Shutting down again should be safe
415        namespace.shutdown().await;
416    }
417
418    #[tokio::test]
419    async fn test_broadcast_operator_disconnect_with_no_connections() {
420        let namespace = create_test_namespace();
421        let op = namespace.to(["room1"]);
422        let result = op.clone().disconnect().await;
423        assert!(result.is_ok());
424    }
425
426    #[tokio::test]
427    async fn test_broadcast_operator_emit_requires_running() {
428        let namespace = create_test_namespace();
429        // Shutdown to make status invalid
430        namespace.clone().shutdown().await;
431
432        let op = namespace.to(["room1"]);
433        let result = op.emit("event", Option::<&()>::None).await;
434        assert!(result.is_err());
435        let err_msg = result.unwrap_err().to_string();
436        assert!(err_msg.contains("invalid status"));
437    }
438}
wsio_server/namespace/mod.rs

wsio_server/namespace/
mod.rs
