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        #[cfg(feature = "tracing")]
104        tracing::debug!(
105            namespace = self.config.path,
106            request_path = request_uri.path(),
107            "handling upgraded WebSocket request"
108        );
109
110        // Create ws stream
111        let mut ws_stream =
112            WebSocketStream::from_raw_socket(TokioIo::new(upgraded), Role::Server, Some(self.config.websocket_config))
113                .await;
114
115        // Check runtime and namespace status
116        if !self.runtime.status.is(WsIoServerRuntimeStatus::Running) || !self.status.is(NamespaceStatus::Running) {
117            #[cfg(feature = "tracing")]
118            tracing::debug!(
119                namespace = self.config.path,
120                runtime_status = ?self.runtime.status.get(),
121                namespace_status = ?self.status.get(),
122                "rejecting upgraded request because server or namespace is not running"
123            );
124
125            ws_stream
126                .send((*self.encode_packet_to_message(&WsIoPacket::new_disconnect())?).clone())
127                .await?;
128
129            let _ = ws_stream.close(None).await;
130            return Ok(());
131        }
132
133        // Create connection
134        let (connection, mut message_rx) = WsIoServerConnection::new(headers, self.clone(), request_uri);
135
136        #[cfg(feature = "tracing")]
137        tracing::debug!(
138            namespace = self.config.path,
139            connection_id = connection.id(),
140            "accepted WebSocket connection"
141        );
142
143        // Split ws stream and spawn read and write tasks
144        let (mut ws_stream_writer, mut ws_stream_reader) = ws_stream.split();
145        let connection_clone = connection.clone();
146        let mut read_ws_stream_task = spawn(async move {
147            while let Some(message) = ws_stream_reader.next().await {
148                if match message {
149                    Ok(Message::Binary(bytes)) => {
150                        // Treat any single-byte binary frame as a client heartbeat and ignore it
151                        if bytes.len() == 1 {
152                            continue;
153                        }
154
155                        connection_clone.handle_incoming_packet(&bytes).await
156                    },
157                    Ok(Message::Close(_)) => {
158                        #[cfg(feature = "tracing")]
159                        tracing::debug!(connection_id = connection_clone.id(), "server read task received close frame");
160                        break;
161                    },
162                    Err(_err) => {
163                        #[cfg(feature = "tracing")]
164                        tracing::debug!(connection_id = connection_clone.id(), error = %_err, "server read task failed");
165                        break;
166                    },
167                    Ok(Message::Text(text)) => connection_clone.handle_incoming_packet(text.as_bytes()).await,
168                    _ => Ok(()),
169                }
170                .is_err()
171                {
172                    #[cfg(feature = "tracing")]
173                    tracing::debug!(
174                        connection_id = connection_clone.id(),
175                        "server read task stopped after packet handling error"
176                    );
177
178                    break;
179                }
180            }
181        });
182
183        let mut write_ws_stream_task = spawn(async move {
184            while let Some(message) = message_rx.recv().await {
185                let message = (*message).clone();
186                let is_close = matches!(message, Message::Close(_));
187                if ws_stream_writer.send(message).await.is_err() {
188                    #[cfg(feature = "tracing")]
189                    tracing::debug!("server write task failed to send message");
190                    break;
191                }
192
193                if is_close {
194                    #[cfg(feature = "tracing")]
195                    tracing::debug!("server write task sent close frame");
196                    let _ = ws_stream_writer.close().await;
197                    break;
198                }
199            }
200        });
201
202        // Try to init connection
203        match connection.init().await {
204            Ok(_) => {
205                #[cfg(feature = "tracing")]
206                tracing::debug!(connection_id = connection.id(), "server connection initialized");
207                // Wait for either read or write task to finish
208                select! {
209                    _ = &mut read_ws_stream_task => {
210                        #[cfg(feature = "tracing")]
211                        tracing::debug!(connection_id = connection.id(), "server read task finished; aborting write task");
212                        write_ws_stream_task.abort();
213                    },
214                    _ = &mut write_ws_stream_task => {
215                        #[cfg(feature = "tracing")]
216                        tracing::debug!(connection_id = connection.id(), "server write task finished; aborting read task");
217                        read_ws_stream_task.abort();
218                    },
219                }
220            },
221            Err(_err) => {
222                #[cfg(feature = "tracing")]
223                tracing::debug!(connection_id = connection.id(), error = %_err, "server connection initialization failed");
224                // Close connection
225                read_ws_stream_task.abort();
226                connection.close();
227                let _ = join!(read_ws_stream_task, write_ws_stream_task);
228            },
229        }
230
231        // Cleanup connection
232        connection.cleanup().await;
233
234        #[cfg(feature = "tracing")]
235        tracing::debug!(connection_id = connection.id(), "server connection stopped");
236        Ok(())
237    }
238
239    // Protected methods
240    #[inline]
241    pub(crate) fn add_connection_id_to_room(&self, room_name: &str, connection_id: u64) {
242        self.rooms.entry(room_name.into()).or_default().insert(connection_id);
243    }
244
245    #[inline]
246    pub(crate) fn encode_packet_to_message(&self, packet: &WsIoPacket) -> Result<Arc<Message>> {
247        let bytes = self.config.packet_codec.encode(packet)?;
248        Ok(Arc::new(match self.config.packet_codec.is_text() {
249            // SAFETY: text packet codecs only produce valid UTF-8 payloads.
250            true => Message::Text(unsafe { String::from_utf8_unchecked(bytes).into() }),
251            false => Message::Binary(bytes.into()),
252        }))
253    }
254
255    pub(crate) async fn handle_on_upgrade_request(
256        self: &Arc<Self>,
257        headers: HeaderMap,
258        on_upgrade: OnUpgrade,
259        request_uri: Uri,
260    ) {
261        #[cfg(feature = "tracing")]
262        tracing::trace!(
263            namespace = self.config.path,
264            request_path = request_uri.path(),
265            "spawning WebSocket upgrade task"
266        );
267
268        let namespace = self.clone();
269        self.connection_task_set.lock().await.spawn(async move {
270            match timeout(namespace.config.http_request_upgrade_timeout, on_upgrade).await {
271                Ok(Ok(upgraded)) => {
272                    if let Err(_err) = namespace.handle_upgraded_request(headers, request_uri, upgraded).await {
273                        #[cfg(feature = "tracing")]
274                        tracing::debug!(namespace = namespace.config.path, error = %_err, "upgraded request handling failed");
275                    }
276                },
277                Ok(Err(_err)) => {
278                    #[cfg(feature = "tracing")]
279                    tracing::warn!(namespace = namespace.config.path, error = %_err, "HTTP upgrade failed");
280                },
281                Err(_err) => {
282                    #[cfg(feature = "tracing")]
283                    tracing::warn!(
284                        namespace = namespace.config.path,
285                        error = %_err,
286                        timeout_ms = namespace.config.http_request_upgrade_timeout.as_millis() as u64,
287                        "HTTP upgrade timed out"
288                    );
289                },
290            }
291        });
292    }
293
294    #[inline]
295    pub(crate) fn insert_connection(&self, connection: Arc<WsIoServerConnection>) {
296        self.connections.insert(connection.id(), connection.clone());
297        self.runtime.insert_connection_id(connection.id());
298        self.connection_ids.rcu(|old_connection_ids| {
299            let mut new_connection_ids = (**old_connection_ids).clone();
300            new_connection_ids.insert(connection.id());
301            new_connection_ids
302        });
303    }
304
305    #[inline]
306    pub(crate) fn remove_connection(&self, id: u64) {
307        self.connections.remove(&id);
308        self.runtime.remove_connection_id(id);
309        self.connection_ids.rcu(|old_connection_ids| {
310            let mut new_connection_ids = (**old_connection_ids).clone();
311            new_connection_ids.remove(id);
312            new_connection_ids
313        });
314    }
315
316    #[inline]
317    pub(crate) fn remove_connection_id_from_room(&self, room_name: &str, connection_id: u64) {
318        if let Some(mut entry) = self.rooms.get_mut(room_name) {
319            entry.remove(connection_id);
320        }
321
322        self.rooms.remove_if(room_name, |_, entry| entry.is_empty());
323    }
324
325    // Public methods
326    pub async fn close_all(self: &Arc<Self>) {
327        WsIoServerNamespaceBroadcastOperator::new(self.clone()).close().await;
328    }
329
330    #[inline]
331    pub fn connection_count(&self) -> usize {
332        self.connections.len()
333    }
334
335    pub async fn disconnect_all(self: &Arc<Self>) -> Result<()> {
336        WsIoServerNamespaceBroadcastOperator::new(self.clone())
337            .disconnect()
338            .await
339    }
340
341    pub async fn emit<D: Serialize>(self: &Arc<Self>, event: impl AsRef<str>, data: Option<&D>) -> Result<()> {
342        WsIoServerNamespaceBroadcastOperator::new(self.clone())
343            .emit(event, data)
344            .await
345    }
346
347    #[inline]
348    pub fn except(
349        self: &Arc<Self>,
350        room_names: impl IntoIterator<Item = impl Into<String>>,
351    ) -> WsIoServerNamespaceBroadcastOperator {
352        WsIoServerNamespaceBroadcastOperator::new(self.clone()).except(room_names)
353    }
354
355    #[inline]
356    pub fn path(&self) -> &str {
357        &self.config.path
358    }
359
360    #[inline]
361    pub fn server(&self) -> WsIoServer {
362        WsIoServer(self.runtime.clone())
363    }
364
365    pub async fn shutdown(self: &Arc<Self>) {
366        match self.status.get() {
367            NamespaceStatus::Stopped => return,
368            NamespaceStatus::Running => {
369                #[cfg(feature = "tracing")]
370                tracing::info!(namespace = self.config.path, "shutting down namespace");
371                self.status.store(NamespaceStatus::Stopping)
372            },
373            _ => unreachable!(),
374        }
375
376        self.close_all().await;
377        let mut connection_task_set = self.connection_task_set.lock().await;
378        while connection_task_set.join_next().await.is_some() {}
379
380        self.status.store(NamespaceStatus::Stopped);
381
382        #[cfg(feature = "tracing")]
383        tracing::info!(namespace = self.config.path, "namespace stopped");
384    }
385
386    #[inline]
387    pub fn to(
388        self: &Arc<Self>,
389        room_names: impl IntoIterator<Item = impl Into<String>>,
390    ) -> WsIoServerNamespaceBroadcastOperator {
391        WsIoServerNamespaceBroadcastOperator::new(self.clone()).to(room_names)
392    }
393}
394
395#[cfg(test)]
396mod tests {
397    use std::time::Duration;
398
399    use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
400
401    use super::*;
402    use crate::{
403        config::WsIoServerConfig,
404        core::packet::codecs::WsIoPacketCodec,
405    };
406
407    fn create_test_namespace() -> Arc<WsIoServerNamespace> {
408        let runtime = WsIoServerRuntime::new(WsIoServerConfig {
409            broadcast_concurrency_limit: 16,
410            http_request_upgrade_timeout: Duration::from_secs(3),
411            init_request_handler_timeout: Duration::from_secs(3),
412            init_response_handler_timeout: Duration::from_secs(3),
413            init_response_timeout: Duration::from_secs(3),
414            middleware_execution_timeout: Duration::from_secs(3),
415            on_close_handler_timeout: Duration::from_secs(3),
416            on_connect_handler_timeout: Duration::from_secs(3),
417            packet_codec: WsIoPacketCodec::SerdeJson,
418            request_path: "/socket".into(),
419            websocket_config: WebSocketConfig::default(),
420        });
421        runtime.new_namespace_builder("/test").register().unwrap()
422    }
423
424    #[tokio::test]
425    async fn test_namespace_new() {
426        let namespace = create_test_namespace();
427        assert_eq!(namespace.path(), "/test");
428        assert_eq!(namespace.connection_count(), 0);
429    }
430
431    #[tokio::test]
432    async fn test_namespace_connection_count() {
433        let namespace = create_test_namespace();
434        assert_eq!(namespace.connection_count(), 0);
435    }
436
437    #[tokio::test]
438    async fn test_namespace_server() {
439        let namespace = create_test_namespace();
440        namespace.server();
441    }
442
443    #[tokio::test]
444    async fn test_namespace_to_broadcast_operator() {
445        let namespace = create_test_namespace();
446        namespace.to(["room1", "room2"]);
447    }
448
449    #[tokio::test]
450    async fn test_namespace_except_broadcast_operator() {
451        let namespace = create_test_namespace();
452        namespace.except(["room1", "room2"]);
453    }
454
455    #[tokio::test]
456    async fn test_namespace_add_remove_connection_id_to_room() {
457        let namespace = create_test_namespace();
458        namespace.add_connection_id_to_room("room1", 1);
459        namespace.add_connection_id_to_room("room1", 2);
460        namespace.add_connection_id_to_room("room2", 3);
461
462        assert_eq!(namespace.rooms.get("room1").unwrap().len(), 2);
463        assert!(namespace.rooms.get("room1").unwrap().contains(1));
464        assert!(namespace.rooms.get("room1").unwrap().contains(2));
465        assert_eq!(namespace.rooms.get("room2").unwrap().len(), 1);
466
467        namespace.remove_connection_id_from_room("room1", 1);
468        assert_eq!(namespace.rooms.get("room1").unwrap().len(), 1);
469        assert!(namespace.rooms.get("room1").unwrap().contains(2));
470
471        namespace.remove_connection_id_from_room("room1", 2);
472        namespace.remove_connection_id_from_room("room2", 3);
473
474        assert!(!namespace.rooms.contains_key("room1"));
475        assert!(!namespace.rooms.contains_key("room2"));
476    }
477
478    #[tokio::test]
479    async fn test_namespace_remove_connection_id_from_empty_room() {
480        let namespace = create_test_namespace();
481        // Removing from non-existent room should not panic
482        namespace.remove_connection_id_from_room("nonexistent", 1);
483    }
484
485    #[tokio::test]
486    async fn test_namespace_encode_packet_to_message() {
487        let namespace = create_test_namespace();
488        let packet = WsIoPacket::new_disconnect();
489        let message = namespace.encode_packet_to_message(&packet).unwrap();
490
491        assert!(matches!(&*message, Message::Text(_)));
492    }
493
494    #[tokio::test]
495    async fn test_namespace_shutdown_idempotent() {
496        let namespace = create_test_namespace();
497        namespace.clone().shutdown().await;
498        // Shutting down again should be safe
499        namespace.shutdown().await;
500    }
501
502    #[tokio::test]
503    async fn test_broadcast_operator_disconnect_with_no_connections() {
504        let namespace = create_test_namespace();
505        let op = namespace.to(["room1"]);
506        let result = op.clone().disconnect().await;
507        assert!(result.is_ok());
508    }
509
510    #[tokio::test]
511    async fn test_broadcast_operator_emit_requires_running() {
512        let namespace = create_test_namespace();
513        // Shutdown to make status invalid
514        namespace.clone().shutdown().await;
515
516        let op = namespace.to(["room1"]);
517        let result = op.emit("event", Option::<&()>::None).await;
518        assert!(result.is_err());
519        let err_msg = result.unwrap_err().to_string();
520        assert!(err_msg.contains("invalid status"));
521    }
522}
wsio_server/namespace/mod.rs

wsio_server/namespace/
mod.rs
