Skip to main content

socketioxide_core/adapter/
mod.rs

1//! The adapter module contains the [`CoreAdapter`] trait and other related types.
2//!
3//! It is used to implement communication between socket.io servers to share messages and state.
4//!
5//! The [`CoreLocalAdapter`] provide a local implementation that will allow any implementors to apply local
6//! operations (`broadcast_with_ack`, `broadcast`, `rooms`, etc...).
7//!
8use std::{
9    borrow::Cow,
10    collections::{HashMap, HashSet, hash_map, hash_set},
11    error::Error as StdError,
12    future::{self, Future},
13    hash::Hash,
14    slice,
15    sync::{Arc, RwLock},
16    time::Duration,
17};
18
19use engineioxide_core::{Sid, Str};
20use futures_core::{FusedStream, Stream};
21use serde::{Deserialize, Serialize, de::DeserializeOwned};
22use smallvec::SmallVec;
23
24use crate::{Uid, Value, packet::Packet, parser::Parse};
25use errors::{AdapterError, BroadcastError, SocketError};
26
27pub mod errors;
28#[cfg(feature = "remote-adapter")]
29pub mod remote_packet;
30
31/// A room identifier
32pub type Room = Cow<'static, str>;
33
34/// Flags that can be used to modify the behavior of the broadcast methods.
35#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
36pub enum BroadcastFlags {
37    /// Broadcast only to the current server
38    Local = 0x01,
39    /// Broadcast to all clients except the sender
40    Broadcast = 0x02,
41    /// The event may be dropped if the client is not ready to receive it
42    /// (e.g. the connection is buffering or not connected).
43    /// This is useful for events that are not critical, like position updates in a game.
44    /// See [socket.io volatile events](https://socket.io/docs/v4/emitting-events/#volatile-events).
45    Volatile = 0x04,
46}
47
48/// Options that can be used to modify the behavior of the broadcast methods.
49#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
50pub struct BroadcastOptions {
51    /// The flags to apply to the broadcast represented as a bitflag.
52    flags: u8,
53    /// The rooms to broadcast to.
54    pub rooms: SmallVec<[Room; 4]>,
55    /// The rooms to exclude from the broadcast.
56    pub except: SmallVec<[Room; 4]>,
57    /// The socket id of the sender.
58    pub sid: Option<Sid>,
59    /// The target server id can be used to optimize the broadcast.
60    /// More specifically when we use broadcasting to apply a single action on a remote socket.
61    /// We now the server_id of the remote socket, so we can send the action directly to the server.
62    pub server_id: Option<Uid>,
63}
64impl BroadcastOptions {
65    /// Add any flags to the options.
66    pub fn add_flag(&mut self, flag: BroadcastFlags) {
67        self.flags |= flag as u8;
68    }
69    /// Check if the options have a flag.
70    pub fn has_flag(&self, flag: BroadcastFlags) -> bool {
71        self.flags & flag as u8 == flag as u8
72    }
73
74    /// get the flags of the options.
75    pub fn flags(&self) -> u8 {
76        self.flags
77    }
78
79    /// Set the socket id of the sender.
80    pub fn new(sid: Sid) -> Self {
81        Self {
82            sid: Some(sid),
83            ..Default::default()
84        }
85    }
86    /// Create a new broadcast options from a remote socket data.
87    pub fn new_remote(data: &RemoteSocketData) -> Self {
88        Self {
89            sid: Some(data.id),
90            server_id: Some(data.server_id),
91            ..Default::default()
92        }
93    }
94
95    /// Check if the selected options are local to the current server.
96    #[inline]
97    pub fn is_local(&self, uid: Uid) -> bool {
98        let target_sock_is_local = !self.has_flag(BroadcastFlags::Broadcast)
99            && self.server_id == Some(uid)
100            && self.rooms.is_empty()
101            && self.sid.is_some();
102        self.has_flag(BroadcastFlags::Local) || target_sock_is_local
103    }
104}
105
106/// A trait for types that can be used as a room parameter.
107///
108/// [`String`], [`Vec<String>`], [`Vec<&str>`], [`&'static str`](str) and const arrays are implemented by default.
109pub trait RoomParam: Send + 'static {
110    /// The type of the iterator returned by `into_room_iter`.
111    type IntoIter: Iterator<Item = Room>;
112
113    /// Convert `self` into an iterator of rooms.
114    fn into_room_iter(self) -> Self::IntoIter;
115}
116
117impl RoomParam for Room {
118    type IntoIter = std::iter::Once<Room>;
119    #[inline(always)]
120    fn into_room_iter(self) -> Self::IntoIter {
121        std::iter::once(self)
122    }
123}
124impl RoomParam for String {
125    type IntoIter = std::iter::Once<Room>;
126    #[inline(always)]
127    fn into_room_iter(self) -> Self::IntoIter {
128        std::iter::once(Cow::Owned(self))
129    }
130}
131impl RoomParam for Vec<String> {
132    type IntoIter = std::iter::Map<std::vec::IntoIter<String>, fn(String) -> Room>;
133    #[inline(always)]
134    fn into_room_iter(self) -> Self::IntoIter {
135        self.into_iter().map(Cow::Owned)
136    }
137}
138impl RoomParam for Vec<&'static str> {
139    type IntoIter = std::iter::Map<std::vec::IntoIter<&'static str>, fn(&'static str) -> Room>;
140    #[inline(always)]
141    fn into_room_iter(self) -> Self::IntoIter {
142        self.into_iter().map(Cow::Borrowed)
143    }
144}
145
146impl RoomParam for Vec<Room> {
147    type IntoIter = std::vec::IntoIter<Room>;
148    #[inline(always)]
149    fn into_room_iter(self) -> Self::IntoIter {
150        self.into_iter()
151    }
152}
153impl RoomParam for &'static str {
154    type IntoIter = std::iter::Once<Room>;
155    #[inline(always)]
156    fn into_room_iter(self) -> Self::IntoIter {
157        std::iter::once(Cow::Borrowed(self))
158    }
159}
160impl<const COUNT: usize> RoomParam for [&'static str; COUNT] {
161    type IntoIter =
162        std::iter::Map<std::array::IntoIter<&'static str, COUNT>, fn(&'static str) -> Room>;
163
164    #[inline(always)]
165    fn into_room_iter(self) -> Self::IntoIter {
166        self.into_iter().map(Cow::Borrowed)
167    }
168}
169impl RoomParam for &'static [&'static str] {
170    type IntoIter =
171        std::iter::Map<std::slice::Iter<'static, &'static str>, fn(&'static &'static str) -> Room>;
172
173    #[inline(always)]
174    fn into_room_iter(self) -> Self::IntoIter {
175        self.iter().map(|i| Cow::Borrowed(*i))
176    }
177}
178impl<const COUNT: usize> RoomParam for [String; COUNT] {
179    type IntoIter = std::iter::Map<std::array::IntoIter<String, COUNT>, fn(String) -> Room>;
180    #[inline(always)]
181    fn into_room_iter(self) -> Self::IntoIter {
182        self.into_iter().map(Cow::Owned)
183    }
184}
185impl RoomParam for Sid {
186    type IntoIter = std::iter::Once<Room>;
187    #[inline(always)]
188    fn into_room_iter(self) -> Self::IntoIter {
189        std::iter::once(Cow::Owned(self.to_string()))
190    }
191}
192
193/// A item yield by the ack stream.
194pub type AckStreamItem<E> = (Sid, Result<Value, E>);
195/// The [`SocketEmitter`] will be implemented by the socketioxide library.
196/// It is simply used as an abstraction to allow the adapter to communicate
197/// with the socket server without the need to depend on the socketioxide lib.
198pub trait SocketEmitter: Send + Sync + 'static {
199    /// An error that can occur when sending data an acknowledgment.
200    type AckError: StdError + Send + Serialize + DeserializeOwned + 'static;
201    /// A stream that emits the acknowledgments of multiple sockets.
202    type AckStream: Stream<Item = AckStreamItem<Self::AckError>> + FusedStream + Send + 'static;
203
204    /// Get all the socket ids in the namespace.
205    fn get_all_sids(&self, filter: impl Fn(&Sid) -> bool) -> Vec<Sid>;
206    /// Get the socket data that match the list of socket ids.
207    fn get_remote_sockets(&self, sids: BroadcastIter<'_>) -> Vec<RemoteSocketData>;
208    /// Send data to the list of socket ids.
209    fn send_many(&self, sids: BroadcastIter<'_>, data: Value) -> Result<(), Vec<SocketError>>;
210    /// Send data to the list of socket ids with volatile semantics.
211    /// Errors are silently discarded; packets may be dropped if the
212    /// transport is not ready.
213    fn send_many_volatile(&self, sids: BroadcastIter<'_>, data: Value);
214    /// Send data to the list of socket ids and get a stream of acks and the number of expected acks.
215    fn send_many_with_ack(
216        &self,
217        sids: BroadcastIter<'_>,
218        packet: Packet,
219        timeout: Option<Duration>,
220    ) -> (Self::AckStream, u32);
221    /// Disconnect all the sockets in the list.
222    /// TODO: take a [`BroadcastIter`]. Currently it is impossible because it may create deadlocks
223    /// with Adapter::del_all call.
224    fn disconnect_many(&self, sids: Vec<Sid>) -> Result<(), Vec<SocketError>>;
225    /// Get the path of the namespace.
226    fn path(&self) -> &Str;
227    /// Get the parser of the namespace.
228    fn parser(&self) -> impl Parse;
229    /// Get the unique server id.
230    fn server_id(&self) -> Uid;
231    /// Get the default configured ack timeout.
232    fn ack_timeout(&self) -> Duration;
233}
234
235/// For static namespaces, the init response will be managed by the user.
236/// However, for dynamic namespaces, the socket.io client will manage the response.
237/// As it does not know the type of the response, the spawnable trait is used to spawn the response.
238/// Without the client having to know the type of the response.
239pub trait Spawnable {
240    /// Spawn the response. Implementors should spawn the future with `tokio::spawn` if it is an async function.
241    /// They should also print a `tracing::error` log in case of an error.
242    fn spawn(self);
243}
244impl Spawnable for () {
245    fn spawn(self) {}
246}
247
248/// A trait to add a "defined" bound to adapter types.
249/// This allow the socket io library to implement function given a *defined* adapter
250/// and not a generic `A: Adapter`.
251///
252/// This is useful to force the user to handle potential init response type [`CoreAdapter::InitRes`].
253pub trait DefinedAdapter {}
254
255/// An adapter is responsible for managing the state of the namespace.
256/// This adapter can be implemented to share the state between multiple servers.
257///
258/// A [`CoreLocalAdapter`] instance will be given when constructing this type, it will allow
259/// you to manipulate local sockets (emitting, fetching data, broadcasting).
260pub trait CoreAdapter<E: SocketEmitter>: Sized + Send + Sync + 'static {
261    /// An error that can occur when using the adapter.
262    type Error: StdError + Into<AdapterError> + Send + 'static;
263    /// A shared state between all the namespace [`CoreAdapter`].
264    /// This can be used to share a connection for example.
265    type State: Send + Sync + 'static;
266    /// A stream that emits the acknowledgments of multiple sockets.
267    type AckStream: Stream<Item = AckStreamItem<E::AckError>> + FusedStream + Send + 'static;
268    /// A named result type for the initialization of the adapter.
269    type InitRes: Spawnable + Send;
270
271    /// Creates a new adapter with the given state and local adapter.
272    ///
273    /// The state is used to share a common state between all your adapters. E.G. a connection to a remote system.
274    /// The local adapter is used to manipulate the local sockets.
275    fn new(state: &Self::State, local: CoreLocalAdapter<E>) -> Self;
276
277    /// Initializes the adapter. The on_success callback should be called when the adapter ready.
278    fn init(self: Arc<Self>, on_success: impl FnOnce() + Send + 'static) -> Self::InitRes;
279
280    /// Closes the adapter.
281    fn close(&self) -> impl Future<Output = Result<(), Self::Error>> + Send {
282        future::ready(Ok(()))
283    }
284
285    /// Returns the number of servers.
286    fn server_count(&self) -> impl Future<Output = Result<u16, Self::Error>> + Send {
287        future::ready(Ok(1))
288    }
289
290    /// Broadcasts the packet to the sockets that match the [`BroadcastOptions`].
291    fn broadcast(
292        &self,
293        packet: Packet,
294        opts: BroadcastOptions,
295    ) -> impl Future<Output = Result<(), BroadcastError>> + Send {
296        future::ready(
297            self.get_local()
298                .broadcast(packet, opts)
299                .map_err(BroadcastError::from),
300        )
301    }
302
303    /// Broadcasts the packet to the sockets that match the [`BroadcastOptions`]
304    /// and return a stream of ack responses.
305    ///
306    /// This method does not have default implementation because GAT cannot have default impls.
307    /// <https://github.com/rust-lang/rust/issues/29661>
308    fn broadcast_with_ack(
309        &self,
310        packet: Packet,
311        opts: BroadcastOptions,
312        timeout: Option<Duration>,
313    ) -> impl Future<Output = Result<Self::AckStream, Self::Error>> + Send;
314
315    /// Adds the sockets that match the [`BroadcastOptions`] to the rooms.
316    fn add_sockets(
317        &self,
318        opts: BroadcastOptions,
319        rooms: impl RoomParam,
320    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
321        self.get_local().add_sockets(opts, rooms);
322        future::ready(Ok(()))
323    }
324
325    /// Removes the sockets that match the [`BroadcastOptions`] from the rooms.
326    fn del_sockets(
327        &self,
328        opts: BroadcastOptions,
329        rooms: impl RoomParam,
330    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
331        self.get_local().del_sockets(opts, rooms);
332        future::ready(Ok(()))
333    }
334
335    /// Disconnects the sockets that match the [`BroadcastOptions`].
336    fn disconnect_socket(
337        &self,
338        opts: BroadcastOptions,
339    ) -> impl Future<Output = Result<(), BroadcastError>> + Send {
340        future::ready(
341            self.get_local()
342                .disconnect_socket(opts)
343                .map_err(BroadcastError::Socket),
344        )
345    }
346
347    /// Fetches rooms that match the [`BroadcastOptions`]
348    fn rooms(
349        &self,
350        opts: BroadcastOptions,
351    ) -> impl Future<Output = Result<Vec<Room>, Self::Error>> + Send {
352        future::ready(Ok(self.get_local().rooms(opts).into_iter().collect()))
353    }
354
355    /// Fetches remote sockets that match the [`BroadcastOptions`].
356    fn fetch_sockets(
357        &self,
358        opts: BroadcastOptions,
359    ) -> impl Future<Output = Result<Vec<RemoteSocketData>, Self::Error>> + Send {
360        future::ready(Ok(self.get_local().fetch_sockets(opts)))
361    }
362
363    /// Returns the local adapter. Used to enable default behaviors.
364    fn get_local(&self) -> &CoreLocalAdapter<E>;
365
366    //TODO: implement
367    // fn server_side_emit(&self, packet: Packet, opts: BroadcastOptions) -> Result<u64, Error>;
368    // fn persist_session(&self, sid: i64);
369    // fn restore_session(&self, sid: i64) -> Session;
370}
371
372/// The default adapter. Store the state in memory.
373pub struct CoreLocalAdapter<E> {
374    rooms: RwLock<HashMap<Room, HashSet<Sid>>>,
375    sockets: RwLock<HashMap<Sid, HashSet<Room>>>,
376    emitter: E,
377}
378
379impl<E: SocketEmitter> CoreLocalAdapter<E> {
380    /// Create a new local adapter with the given sockets interface.
381    pub fn new(emitter: E) -> Self {
382        Self {
383            rooms: RwLock::new(HashMap::new()),
384            sockets: RwLock::new(HashMap::new()),
385            emitter,
386        }
387    }
388
389    /// Clears all the rooms and sockets.
390    pub fn close(&self) {
391        let mut rooms = self.rooms.write().unwrap();
392        rooms.clear();
393        rooms.shrink_to_fit();
394    }
395
396    /// Adds the socket to all the rooms.
397    pub fn add_all(&self, sid: Sid, rooms: impl RoomParam) {
398        let mut rooms_map = self.rooms.write().unwrap();
399        let mut socket_map = self.sockets.write().unwrap();
400        for room in rooms.into_room_iter() {
401            rooms_map.entry(room.clone()).or_default().insert(sid);
402            socket_map.entry(sid).or_default().insert(room);
403        }
404    }
405
406    /// Removes the socket from the rooms.
407    pub fn del(&self, sid: Sid, rooms: impl RoomParam) {
408        let mut rooms_map = self.rooms.write().unwrap();
409        let mut socket_map = self.sockets.write().unwrap();
410        for room in rooms.into_room_iter() {
411            remove_and_clean_entry(rooms_map.entry(room.clone()), &sid, || {
412                socket_map.entry(sid).and_modify(|r| {
413                    r.remove(&room);
414                });
415            });
416        }
417    }
418
419    /// Removes the socket from all the rooms.
420    pub fn del_all(&self, sid: Sid) {
421        let mut rooms_map = self.rooms.write().unwrap();
422        if let Some(rooms) = self.sockets.write().unwrap().remove(&sid) {
423            for room in rooms {
424                remove_and_clean_entry(rooms_map.entry(room.clone()), &sid, || ());
425            }
426        }
427    }
428
429    /// Broadcasts the packet to the sockets that match the [`BroadcastOptions`].
430    pub fn broadcast(
431        &self,
432        packet: Packet,
433        opts: BroadcastOptions,
434    ) -> Result<(), Vec<SocketError>> {
435        let room_map = self.rooms.read().unwrap();
436        let sids = self.apply_opts(&opts, &room_map);
437
438        if sids.is_empty() {
439            return Ok(());
440        }
441
442        let is_volatile = opts.has_flag(BroadcastFlags::Volatile);
443        let data = self.emitter.parser().encode(packet);
444        if is_volatile {
445            self.emitter.send_many_volatile(sids, data);
446            Ok(())
447        } else {
448            self.emitter.send_many(sids, data)
449        }
450    }
451
452    /// Broadcasts the packet to the sockets that match the [`BroadcastOptions`] and return a stream of ack responses.
453    /// Also returns the number of local expected aknowledgements to know when to stop waiting.
454    pub fn broadcast_with_ack(
455        &self,
456        packet: Packet,
457        opts: BroadcastOptions,
458        timeout: Option<Duration>,
459    ) -> (E::AckStream, u32) {
460        let room_map = self.rooms.read().unwrap();
461        let sids = self.apply_opts(&opts, &room_map);
462        // We cannot pre-serialize the packet because we need to change the ack id.
463        self.emitter.send_many_with_ack(sids, packet, timeout)
464    }
465
466    /// Returns the sockets ids that match the [`BroadcastOptions`].
467    pub fn sockets(&self, opts: BroadcastOptions) -> Vec<Sid> {
468        self.apply_opts(&opts, &self.rooms.read().unwrap())
469            .collect()
470    }
471
472    /// Returns the sockets ids that match the [`BroadcastOptions`].
473    pub fn fetch_sockets(&self, opts: BroadcastOptions) -> Vec<RemoteSocketData> {
474        let rooms = self.rooms.read().unwrap();
475        let sids = self.apply_opts(&opts, &rooms);
476        self.emitter.get_remote_sockets(sids)
477    }
478
479    /// Returns the rooms of the socket.
480    pub fn socket_rooms(&self, sid: Sid) -> HashSet<Room> {
481        self.sockets
482            .read()
483            .unwrap()
484            .get(&sid)
485            .cloned()
486            .unwrap_or_default()
487    }
488
489    /// Adds the sockets that match the [`BroadcastOptions`] to the rooms.
490    pub fn add_sockets(&self, opts: BroadcastOptions, rooms: impl RoomParam) {
491        let rooms: Vec<Room> = rooms.into_room_iter().collect();
492        let mut room_map = self.rooms.write().unwrap();
493        let mut socket_map = self.sockets.write().unwrap();
494        // Here we have to collect sids, because we are going to modify the rooms map.
495        let sids = self.apply_opts(&opts, &room_map).collect::<Vec<_>>();
496        for sid in &sids {
497            let entry = socket_map.entry(*sid).or_default();
498            for room in &rooms {
499                entry.insert(room.clone());
500            }
501        }
502        for room in rooms {
503            let entry = room_map.entry(room).or_default();
504            for sid in &sids {
505                entry.insert(*sid);
506            }
507        }
508    }
509
510    /// Removes the sockets that match the [`BroadcastOptions`] from the rooms.
511    pub fn del_sockets(&self, opts: BroadcastOptions, rooms: impl RoomParam) {
512        let rooms: Vec<Room> = rooms.into_room_iter().collect();
513        let mut rooms_map = self.rooms.write().unwrap();
514        let mut socket_map = self.sockets.write().unwrap();
515        let sids = self.apply_opts(&opts, &rooms_map).collect::<Vec<_>>();
516        for room in rooms {
517            for sid in &sids {
518                remove_and_clean_entry(socket_map.entry(*sid), &room, || ());
519                remove_and_clean_entry(rooms_map.entry(room.clone()), sid, || ());
520            }
521        }
522    }
523
524    /// Disconnects the sockets that match the [`BroadcastOptions`].
525    pub fn disconnect_socket(&self, opts: BroadcastOptions) -> Result<(), Vec<SocketError>> {
526        let sids = self
527            .apply_opts(&opts, &self.rooms.read().unwrap())
528            .collect();
529        self.emitter.disconnect_many(sids)
530    }
531
532    /// Returns all the matching rooms
533    pub fn rooms(&self, opts: BroadcastOptions) -> HashSet<Room> {
534        let rooms = self.rooms.read().unwrap();
535        let sockets = self.sockets.read().unwrap();
536        let sids = self.apply_opts(&opts, &rooms);
537        sids.filter_map(|id| sockets.get(&id))
538            .flatten()
539            .cloned()
540            .collect()
541    }
542
543    /// Get the namespace path.
544    pub fn path(&self) -> &Str {
545        self.emitter.path()
546    }
547
548    /// Get the parser of the namespace.
549    pub fn parser(&self) -> impl Parse + '_ {
550        self.emitter.parser()
551    }
552    /// Get the unique server identifier
553    pub fn server_id(&self) -> Uid {
554        self.emitter.server_id()
555    }
556    /// Get the default configured ack timeout.
557    pub fn ack_timeout(&self) -> Duration {
558        self.emitter.ack_timeout()
559    }
560}
561
562/// The default broadcast iterator.
563/// Extract, flatten and filter a list of sid from a room list
564struct BroadcastRooms<'a> {
565    rooms: slice::Iter<'a, Room>,
566    rooms_map: &'a HashMap<Room, HashSet<Sid>>,
567    except: HashSet<Sid>,
568    flatten_iter: Option<hash_set::Iter<'a, Sid>>,
569}
570impl<'a> BroadcastRooms<'a> {
571    fn new(
572        rooms: &'a [Room],
573        rooms_map: &'a HashMap<Room, HashSet<Sid>>,
574        except: HashSet<Sid>,
575    ) -> Self {
576        BroadcastRooms {
577            rooms: rooms.iter(),
578            rooms_map,
579            except,
580            flatten_iter: None,
581        }
582    }
583}
584impl Iterator for BroadcastRooms<'_> {
585    type Item = Sid;
586    fn next(&mut self) -> Option<Self::Item> {
587        loop {
588            match self.flatten_iter.as_mut().and_then(Iterator::next) {
589                Some(sid) if !self.except.contains(sid) => return Some(*sid),
590                Some(_) => continue,
591                None => self.flatten_iter = None,
592            }
593
594            let room = self.rooms.next()?;
595            self.flatten_iter = self.rooms_map.get(room).map(HashSet::iter);
596        }
597    }
598}
599
600impl<E: SocketEmitter> CoreLocalAdapter<E> {
601    /// Applies the given `opts` and return the sockets that match.
602    fn apply_opts<'a>(
603        &self,
604        opts: &'a BroadcastOptions,
605        rooms: &'a HashMap<Room, HashSet<Sid>>,
606    ) -> BroadcastIter<'a> {
607        let is_broadcast = opts.has_flag(BroadcastFlags::Broadcast);
608
609        let mut except = get_except_sids(&opts.except, rooms);
610        // In case of broadcast flag + if the sender is set,
611        // we should not broadcast to it.
612        if is_broadcast {
613            //FIXME(1.88): switch to if let chains when available
614            if let Some(sid) = opts.sid {
615                except.insert(sid);
616            }
617        }
618
619        if !opts.rooms.is_empty() {
620            let iter = BroadcastRooms::new(&opts.rooms, rooms, except);
621            InnerBroadcastIter::BroadcastRooms(iter).into()
622        } else if is_broadcast {
623            let sids = self.emitter.get_all_sids(|id| !except.contains(id));
624            InnerBroadcastIter::GlobalBroadcast(sids.into_iter()).into()
625        } else if let Some(id) = opts.sid {
626            InnerBroadcastIter::Single(id).into()
627        } else {
628            InnerBroadcastIter::None.into()
629        }
630    }
631}
632
633#[inline]
634fn get_except_sids(except: &[Room], rooms: &HashMap<Room, HashSet<Sid>>) -> HashSet<Sid> {
635    let mut except_sids = HashSet::new();
636    for room in except {
637        if let Some(sockets) = rooms.get(room) {
638            except_sids.extend(sockets);
639        }
640    }
641    except_sids
642}
643
644/// Remove a field from a HashSet value and remove it if empty.
645/// Call `cleanup` fn if the entry exists
646#[inline]
647fn remove_and_clean_entry<K, T: Hash + Eq>(
648    entry: hash_map::Entry<'_, K, HashSet<T>>,
649    el: &T,
650    cleanup: impl FnOnce(),
651) {
652    //TODO: use hashmap raw entry when stabilized to avoid entry clone.
653    // https://github.com/rust-lang/rust/issues/56167
654    match entry {
655        hash_map::Entry::Occupied(mut entry) => {
656            entry.get_mut().remove(el);
657            if entry.get().is_empty() {
658                entry.remove_entry();
659            }
660            cleanup();
661        }
662        hash_map::Entry::Vacant(_) => (),
663    }
664}
665
666/// An iterator that yields the socket ids that match the broadcast options.
667/// Used with the [`SocketEmitter`] interface.
668pub struct BroadcastIter<'a> {
669    inner: InnerBroadcastIter<'a>,
670}
671enum InnerBroadcastIter<'a> {
672    BroadcastRooms(BroadcastRooms<'a>),
673    GlobalBroadcast(<Vec<Sid> as IntoIterator>::IntoIter),
674    Single(Sid),
675    None,
676}
677impl BroadcastIter<'_> {
678    fn is_empty(&self) -> bool {
679        matches!(self.inner, InnerBroadcastIter::None)
680    }
681}
682impl<'a> From<InnerBroadcastIter<'a>> for BroadcastIter<'a> {
683    fn from(inner: InnerBroadcastIter<'a>) -> Self {
684        BroadcastIter { inner }
685    }
686}
687
688impl Iterator for BroadcastIter<'_> {
689    type Item = Sid;
690
691    #[inline(always)]
692    fn next(&mut self) -> Option<Self::Item> {
693        self.inner.next()
694    }
695}
696impl Iterator for InnerBroadcastIter<'_> {
697    type Item = Sid;
698
699    fn next(&mut self) -> Option<Self::Item> {
700        match self {
701            InnerBroadcastIter::BroadcastRooms(inner) => inner.next(),
702            InnerBroadcastIter::GlobalBroadcast(inner) => inner.next(),
703            InnerBroadcastIter::Single(sid) => {
704                let sid = *sid;
705                *self = InnerBroadcastIter::None;
706                Some(sid)
707            }
708            InnerBroadcastIter::None => None,
709        }
710    }
711}
712
713/// Represent the data of a remote socket.
714#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Default, Clone)]
715pub struct RemoteSocketData {
716    /// The id of the remote socket.
717    pub id: Sid,
718    /// The server id this socket is connected to.
719    pub server_id: Uid,
720    /// The namespace this socket is connected to.
721    pub ns: Str,
722}
723
724#[cfg(test)]
725mod test {
726
727    use smallvec::smallvec;
728    use std::{
729        array,
730        pin::Pin,
731        task::{Context, Poll},
732    };
733
734    use super::*;
735
736    struct StubSockets {
737        sockets: HashSet<Sid>,
738        path: Str,
739    }
740    impl StubSockets {
741        fn new(sockets: &[Sid]) -> Self {
742            let sockets = HashSet::from_iter(sockets.iter().copied());
743            Self {
744                sockets,
745                path: Str::from("/"),
746            }
747        }
748    }
749
750    struct StubAckStream;
751    impl Stream for StubAckStream {
752        type Item = (Sid, Result<Value, StubError>);
753        fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
754            Poll::Ready(None)
755        }
756    }
757    impl FusedStream for StubAckStream {
758        fn is_terminated(&self) -> bool {
759            true
760        }
761    }
762    #[derive(Debug, Serialize, Deserialize)]
763    struct StubError;
764    impl std::fmt::Display for StubError {
765        fn fmt(&self, _: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
766            Ok(())
767        }
768    }
769    impl std::error::Error for StubError {}
770
771    impl SocketEmitter for StubSockets {
772        type AckError = StubError;
773        type AckStream = StubAckStream;
774        fn get_all_sids(&self, filter: impl Fn(&Sid) -> bool) -> Vec<Sid> {
775            self.sockets
776                .iter()
777                .copied()
778                .filter(|id| filter(id))
779                .collect()
780        }
781
782        fn get_remote_sockets(&self, sids: BroadcastIter<'_>) -> Vec<RemoteSocketData> {
783            sids.map(|id| RemoteSocketData {
784                id,
785                server_id: Uid::ZERO,
786                ns: self.path.clone(),
787            })
788            .collect()
789        }
790
791        fn send_many(&self, _: BroadcastIter<'_>, _: Value) -> Result<(), Vec<SocketError>> {
792            Ok(())
793        }
794
795        fn send_many_volatile(&self, _: BroadcastIter<'_>, _: Value) {}
796
797        fn send_many_with_ack(
798            &self,
799            _: BroadcastIter<'_>,
800            _: Packet,
801            _: Option<Duration>,
802        ) -> (Self::AckStream, u32) {
803            (StubAckStream, 0)
804        }
805
806        fn disconnect_many(&self, _: Vec<Sid>) -> Result<(), Vec<SocketError>> {
807            Ok(())
808        }
809
810        fn path(&self) -> &Str {
811            &self.path
812        }
813        fn parser(&self) -> impl Parse {
814            crate::parser::test::StubParser
815        }
816        fn server_id(&self) -> Uid {
817            Uid::ZERO
818        }
819        fn ack_timeout(&self) -> Duration {
820            Duration::ZERO
821        }
822    }
823
824    fn create_adapter<const S: usize>(sockets: [Sid; S]) -> CoreLocalAdapter<StubSockets> {
825        CoreLocalAdapter::new(StubSockets::new(&sockets))
826    }
827
828    #[test]
829    fn add_all() {
830        let socket = Sid::new();
831        let adapter = create_adapter([socket]);
832        adapter.add_all(socket, ["room1", "room2"]);
833        let rooms_map = adapter.rooms.read().unwrap();
834        let socket_map = adapter.sockets.read().unwrap();
835        assert_eq!(rooms_map.len(), 2);
836        assert_eq!(socket_map.len(), 1);
837        assert_eq!(rooms_map.get("room1").unwrap().len(), 1);
838        assert_eq!(rooms_map.get("room2").unwrap().len(), 1);
839
840        let rooms = socket_map.get(&socket).unwrap();
841        assert!(rooms.contains("room1"));
842        assert!(rooms.contains("room2"));
843    }
844
845    #[test]
846    fn del() {
847        let socket = Sid::new();
848        let adapter = create_adapter([socket]);
849        adapter.add_all(socket, ["room1", "room2"]);
850        {
851            let rooms_map = adapter.rooms.read().unwrap();
852            assert_eq!(rooms_map.len(), 2);
853            assert_eq!(rooms_map.get("room1").unwrap().len(), 1);
854            assert_eq!(rooms_map.get("room2").unwrap().len(), 1);
855            let socket_map = adapter.sockets.read().unwrap();
856            let rooms = socket_map.get(&socket).unwrap();
857            assert!(rooms.contains("room1"));
858            assert!(rooms.contains("room2"));
859        }
860        adapter.del(socket, "room1");
861        let rooms_map = adapter.rooms.read().unwrap();
862        let socket_map = adapter.sockets.read().unwrap();
863        assert_eq!(rooms_map.len(), 1);
864        assert!(rooms_map.get("room1").is_none());
865        assert_eq!(rooms_map.get("room2").unwrap().len(), 1);
866        assert_eq!(socket_map.get(&socket).unwrap().len(), 1);
867    }
868    #[test]
869    fn del_all() {
870        let socket = Sid::new();
871        let adapter = create_adapter([socket]);
872        adapter.add_all(socket, ["room1", "room2"]);
873
874        {
875            let rooms_map = adapter.rooms.read().unwrap();
876            assert_eq!(rooms_map.len(), 2);
877            assert_eq!(rooms_map.get("room1").unwrap().len(), 1);
878            assert_eq!(rooms_map.get("room2").unwrap().len(), 1);
879
880            let socket_map = adapter.sockets.read().unwrap();
881            let rooms = socket_map.get(&socket).unwrap();
882            assert!(rooms.contains("room1"));
883            assert!(rooms.contains("room2"));
884        }
885
886        adapter.del_all(socket);
887
888        {
889            let rooms_map = adapter.rooms.read().unwrap();
890            assert_eq!(rooms_map.len(), 0);
891
892            let socket_map = adapter.sockets.read().unwrap();
893            assert!(socket_map.get(&socket).is_none());
894        }
895    }
896
897    #[test]
898    fn socket_room() {
899        let sid1 = Sid::new();
900        let sid2 = Sid::new();
901        let sid3 = Sid::new();
902        let adapter = create_adapter([sid1, sid2, sid3]);
903        adapter.add_all(sid1, ["room1", "room2"]);
904        adapter.add_all(sid2, ["room1"]);
905        adapter.add_all(sid3, ["room2"]);
906        assert!(adapter.socket_rooms(sid1).contains(&Cow::Borrowed("room1")));
907        assert!(adapter.socket_rooms(sid1).contains(&Cow::Borrowed("room2")));
908        assert_eq!(
909            adapter.socket_rooms(sid2).into_iter().collect::<Vec<_>>(),
910            ["room1"]
911        );
912        assert_eq!(
913            adapter.socket_rooms(sid3).into_iter().collect::<Vec<_>>(),
914            ["room2"]
915        );
916    }
917
918    #[test]
919    fn add_socket() {
920        let socket = Sid::new();
921        let adapter = create_adapter([socket]);
922        adapter.add_all(socket, ["room1"]);
923
924        let mut opts = BroadcastOptions::new(socket);
925        opts.rooms = smallvec!["room1".into()];
926        adapter.add_sockets(opts, "room2");
927        let rooms_map = adapter.rooms.read().unwrap();
928
929        assert_eq!(rooms_map.len(), 2);
930        assert!(rooms_map.get("room1").unwrap().contains(&socket));
931        assert!(rooms_map.get("room2").unwrap().contains(&socket));
932    }
933
934    #[test]
935    fn del_socket() {
936        let socket = Sid::new();
937        let adapter = create_adapter([socket]);
938        adapter.add_all(socket, ["room1"]);
939
940        let mut opts = BroadcastOptions::new(socket);
941        opts.rooms = smallvec!["room1".into()];
942        adapter.add_sockets(opts, "room2");
943
944        {
945            let rooms_map = adapter.rooms.read().unwrap();
946
947            assert_eq!(rooms_map.len(), 2);
948            assert!(rooms_map.get("room1").unwrap().contains(&socket));
949            assert!(rooms_map.get("room2").unwrap().contains(&socket));
950        }
951
952        let mut opts = BroadcastOptions::new(socket);
953        opts.rooms = smallvec!["room1".into()];
954        adapter.del_sockets(opts, "room2");
955
956        {
957            let rooms_map = adapter.rooms.read().unwrap();
958
959            assert_eq!(rooms_map.len(), 1);
960            assert!(rooms_map.get("room1").unwrap().contains(&socket));
961            assert!(rooms_map.get("room2").is_none());
962        }
963    }
964
965    #[test]
966    fn sockets() {
967        let socket0 = Sid::new();
968        let socket1 = Sid::new();
969        let socket2 = Sid::new();
970        let adapter = create_adapter([socket0, socket1, socket2]);
971        adapter.add_all(socket0, ["room1", "room2"]);
972        adapter.add_all(socket1, ["room1", "room3"]);
973        adapter.add_all(socket2, ["room2", "room3"]);
974
975        let mut opts = BroadcastOptions {
976            rooms: smallvec!["room1".into()],
977            ..Default::default()
978        };
979        let sockets = adapter.sockets(opts.clone());
980        assert_eq!(sockets.len(), 2);
981        assert!(sockets.contains(&socket0));
982        assert!(sockets.contains(&socket1));
983
984        opts.rooms = smallvec!["room2".into()];
985        let sockets = adapter.sockets(opts.clone());
986        assert_eq!(sockets.len(), 2);
987        assert!(sockets.contains(&socket0));
988        assert!(sockets.contains(&socket2));
989
990        opts.rooms = smallvec!["room3".into()];
991        let sockets = adapter.sockets(opts.clone());
992        assert_eq!(sockets.len(), 2);
993        assert!(sockets.contains(&socket1));
994        assert!(sockets.contains(&socket2));
995    }
996
997    #[test]
998    fn disconnect_socket() {
999        let socket0 = Sid::new();
1000        let socket1 = Sid::new();
1001        let socket2 = Sid::new();
1002        let adapter = create_adapter([socket0, socket1, socket2]);
1003        adapter.add_all(socket0, ["room1", "room2", "room4"]);
1004        adapter.add_all(socket1, ["room1", "room3", "room5"]);
1005        adapter.add_all(socket2, ["room2", "room3", "room6"]);
1006
1007        let mut opts = BroadcastOptions::new(socket0);
1008        opts.rooms = smallvec!["room5".into()];
1009        adapter.disconnect_socket(opts).unwrap();
1010
1011        let mut opts = BroadcastOptions::default();
1012        opts.rooms.push("room2".into());
1013        let sockets = adapter.sockets(opts.clone());
1014        assert_eq!(sockets.len(), 2);
1015        assert!(sockets.contains(&socket2));
1016        assert!(sockets.contains(&socket0));
1017    }
1018    #[test]
1019    fn disconnect_empty_opts() {
1020        let adapter = create_adapter([]);
1021        let opts = BroadcastOptions::default();
1022        adapter.disconnect_socket(opts).unwrap();
1023    }
1024    #[test]
1025    fn rooms() {
1026        let socket0 = Sid::new();
1027        let socket1 = Sid::new();
1028        let socket2 = Sid::new();
1029        let adapter = create_adapter([socket0, socket1, socket2]);
1030        adapter.add_all(socket0, ["room1", "room2", "room4"]);
1031        adapter.add_all(socket1, ["room1", "room3", "room5"]);
1032        adapter.add_all(socket2, ["room2", "room3", "room6"]);
1033
1034        let mut opts = BroadcastOptions::new(socket0);
1035        opts.rooms = smallvec!["room5".into()];
1036        opts.add_flag(BroadcastFlags::Broadcast);
1037        let rooms = adapter.rooms(opts);
1038        assert_eq!(rooms.len(), 3);
1039        assert!(rooms.contains(&Cow::Borrowed("room1")));
1040        assert!(rooms.contains(&Cow::Borrowed("room3")));
1041        assert!(rooms.contains(&Cow::Borrowed("room5")));
1042
1043        let mut opts = BroadcastOptions::default();
1044        opts.rooms.push("room2".into());
1045        let rooms = adapter.rooms(opts.clone());
1046        assert_eq!(rooms.len(), 5);
1047        assert!(rooms.contains(&Cow::Borrowed("room1")));
1048        assert!(rooms.contains(&Cow::Borrowed("room2")));
1049        assert!(rooms.contains(&Cow::Borrowed("room3")));
1050        assert!(rooms.contains(&Cow::Borrowed("room4")));
1051        assert!(rooms.contains(&Cow::Borrowed("room6")));
1052    }
1053
1054    #[test]
1055    fn apply_opts() {
1056        let mut sockets: [Sid; 3] = array::from_fn(|_| Sid::new());
1057        sockets.sort();
1058        let adapter = create_adapter(sockets);
1059
1060        adapter.add_all(sockets[0], ["room1", "room2"]);
1061        adapter.add_all(sockets[1], ["room1", "room3"]);
1062        adapter.add_all(sockets[2], ["room1", "room2", "room3"]);
1063
1064        // socket 2 is the sender
1065        let mut opts = BroadcastOptions::new(sockets[2]);
1066        opts.rooms = smallvec!["room1".into()];
1067        opts.except = smallvec!["room2".into()];
1068        let sids = adapter
1069            .apply_opts(&opts, &adapter.rooms.read().unwrap())
1070            .collect::<Vec<_>>();
1071        assert_eq!(sids, [sockets[1]]);
1072
1073        let mut opts = BroadcastOptions::new(sockets[2]);
1074        opts.add_flag(BroadcastFlags::Broadcast);
1075        let mut sids = adapter
1076            .apply_opts(&opts, &adapter.rooms.read().unwrap())
1077            .collect::<Vec<_>>();
1078        sids.sort();
1079        assert_eq!(sids, [sockets[0], sockets[1]]);
1080
1081        let mut opts = BroadcastOptions::new(sockets[2]);
1082        opts.add_flag(BroadcastFlags::Broadcast);
1083        opts.except = smallvec!["room2".into()];
1084        let sids = adapter
1085            .apply_opts(&opts, &adapter.rooms.read().unwrap())
1086            .collect::<Vec<_>>();
1087        assert_eq!(sids.len(), 1);
1088
1089        let opts = BroadcastOptions::new(sockets[2]);
1090        let sids = adapter
1091            .apply_opts(&opts, &adapter.rooms.read().unwrap())
1092            .collect::<Vec<_>>();
1093        assert_eq!(sids.len(), 1);
1094        assert_eq!(sids[0], sockets[2]);
1095
1096        let opts = BroadcastOptions::new(Sid::new());
1097        let sids = adapter
1098            .apply_opts(&opts, &adapter.rooms.read().unwrap())
1099            .collect::<Vec<_>>();
1100        assert_eq!(sids.len(), 1);
1101    }
1102
1103    #[test]
1104    fn test_is_local_opts() {
1105        let server_id = Uid::new();
1106        let remote = RemoteSocketData {
1107            id: Sid::new(),
1108            server_id,
1109            ns: "/".into(),
1110        };
1111        let opts = BroadcastOptions::new_remote(&remote);
1112        assert!(opts.is_local(server_id));
1113        assert!(!opts.is_local(Uid::new()));
1114        let opts = BroadcastOptions::new(Sid::new());
1115        assert!(!opts.is_local(Uid::new()));
1116    }
1117}