wsio_server/namespace/
mod.rs

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

wsio_server/namespace/
mod.rs
