wsio_server/namespace/operators/
broadcast.rs

1use std::{
2    collections::HashSet,
3    sync::Arc,
4};
5
6use anyhow::Result;
7use futures_util::{
8    StreamExt,
9    future::ready,
10    stream::iter,
11};
12use roaring::RoaringTreemap;
13use serde::Serialize;
14
15use super::super::{
16    NamespaceStatus,
17    WsIoServerNamespace,
18};
19use crate::{
20    connection::WsIoServerConnection,
21    core::packet::WsIoPacket,
22};
23
24// Structs
25#[derive(Clone)]
26pub struct WsIoServerNamespaceBroadcastOperator {
27    exclude_connection_ids: HashSet<u64>,
28    exclude_rooms: HashSet<String>,
29    include_rooms: HashSet<String>,
30    namespace: Arc<WsIoServerNamespace>,
31}
32
33impl WsIoServerNamespaceBroadcastOperator {
34    #[inline]
35    pub(in super::super) fn new(namespace: Arc<WsIoServerNamespace>) -> Self {
36        Self {
37            exclude_connection_ids: HashSet::new(),
38            exclude_rooms: HashSet::new(),
39            include_rooms: HashSet::new(),
40            namespace,
41        }
42    }
43
44    // Private methods
45    async fn for_each_target_connections<F, Fut>(&self, f: F)
46    where
47        F: Fn(Arc<WsIoServerConnection>) -> Fut + Send + Sync + 'static,
48        Fut: Future<Output = Result<()>> + Send + 'static,
49    {
50        let mut target_connection_ids = if self.include_rooms.is_empty() {
51            (**self.namespace.connection_ids.load()).clone()
52        } else {
53            let mut connection_ids = RoaringTreemap::new();
54            for room_name in &self.include_rooms {
55                if let Some(room) = self.namespace.rooms.get(room_name) {
56                    connection_ids |= room.value();
57                }
58            }
59
60            connection_ids
61        };
62
63        for room_name in &self.exclude_rooms {
64            if let Some(room) = self.namespace.rooms.get(room_name) {
65                target_connection_ids -= room.value();
66                if target_connection_ids.is_empty() {
67                    break;
68                }
69            }
70        }
71
72        for exclude_connection_id in &self.exclude_connection_ids {
73            target_connection_ids.remove(*exclude_connection_id);
74        }
75
76        if target_connection_ids.is_empty() {
77            return;
78        }
79
80        iter(target_connection_ids)
81            .filter_map(|target_connection_id| {
82                ready(
83                    self.namespace
84                        .connections
85                        .get(&target_connection_id)
86                        .map(|entry| entry.value().clone()),
87                )
88            })
89            .for_each_concurrent(self.namespace.config.broadcast_concurrency_limit, |connection| async {
90                let _ = f(connection).await;
91            })
92            .await;
93    }
94
95    // Public methods
96    pub async fn close(self) {
97        self.for_each_target_connections(|connection| async move {
98            connection.close();
99            Ok(())
100        })
101        .await;
102    }
103
104    pub async fn disconnect(self) -> Result<()> {
105        let message = self.namespace.encode_packet_to_message(&WsIoPacket::new_disconnect())?;
106        self.for_each_target_connections(move |connection| {
107            let message = message.clone();
108            async move { connection.send_message(message).await }
109        })
110        .await;
111
112        Ok(())
113    }
114
115    pub async fn emit<D: Serialize>(self, event: impl AsRef<str>, data: Option<&D>) -> Result<()> {
116        self.namespace.status.ensure(NamespaceStatus::Running, |status| {
117            format!("Cannot emit in invalid status: {status:?}")
118        })?;
119
120        let message = self.namespace.encode_packet_to_message(&WsIoPacket::new_event(
121            event.as_ref(),
122            data.map(|data| self.namespace.config.packet_codec.encode_data(data))
123                .transpose()?,
124        ))?;
125
126        self.for_each_target_connections(move |connection| {
127            let message = message.clone();
128            async move { connection.emit_event_message(message).await }
129        })
130        .await;
131
132        Ok(())
133    }
134
135    #[inline]
136    pub fn except(mut self, room_names: impl IntoIterator<Item = impl Into<String>>) -> Self {
137        self.exclude_rooms.extend(room_names.into_iter().map(Into::into));
138        self
139    }
140
141    pub fn except_connection_ids(mut self, connection_ids: impl IntoIterator<Item = u64>) -> Self {
142        self.exclude_connection_ids.extend(connection_ids);
143        self
144    }
145
146    #[inline]
147    pub fn to(mut self, room_names: impl IntoIterator<Item = impl Into<String>>) -> Self {
148        self.include_rooms.extend(room_names.into_iter().map(Into::into));
149        self
150    }
151}
wsio_server/namespace/operators/broadcast.rs

wsio_server/namespace/operators/
broadcast.rs
