use std::{
collections::HashSet,
sync::Arc,
};
use anyhow::Result;
use futures_util::{
StreamExt,
future::ready,
stream::iter,
};
use roaring::RoaringTreemap;
use serde::Serialize;
use super::super::{
NamespaceStatus,
WsIoServerNamespace,
};
use crate::{
connection::WsIoServerConnection,
core::packet::WsIoPacket,
};
#[derive(Clone, Debug)]
pub struct WsIoServerNamespaceBroadcastOperator {
exclude_connection_ids: HashSet<u64>,
exclude_rooms: HashSet<String>,
include_rooms: HashSet<String>,
namespace: Arc<WsIoServerNamespace>,
}
impl WsIoServerNamespaceBroadcastOperator {
#[inline]
pub(in super::super) fn new(namespace: Arc<WsIoServerNamespace>) -> Self {
Self {
exclude_connection_ids: HashSet::new(),
exclude_rooms: HashSet::new(),
include_rooms: HashSet::new(),
namespace,
}
}
async fn for_each_target_connections<F, Fut>(&self, f: F)
where
F: Fn(Arc<WsIoServerConnection>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let target_connection_ids = self.target_connection_ids();
if target_connection_ids.is_empty() {
return;
}
iter(target_connection_ids)
.filter_map(|target_connection_id| {
ready(
self.namespace
.connections
.get(&target_connection_id)
.map(|entry| entry.value().clone()),
)
})
.for_each_concurrent(self.namespace.config.broadcast_concurrency_limit, |connection| async {
let _ = f(connection).await;
})
.await;
}
fn target_connection_ids(&self) -> RoaringTreemap {
let mut target_connection_ids = if self.include_rooms.is_empty() {
(**self.namespace.connection_ids.load()).clone()
} else {
let mut connection_ids = RoaringTreemap::new();
for room_name in &self.include_rooms {
if let Some(room) = self.namespace.rooms.get(room_name) {
connection_ids |= room.value();
}
}
connection_ids
};
for room_name in &self.exclude_rooms {
if let Some(room) = self.namespace.rooms.get(room_name) {
target_connection_ids -= room.value();
if target_connection_ids.is_empty() {
break;
}
}
}
for exclude_connection_id in &self.exclude_connection_ids {
target_connection_ids.remove(*exclude_connection_id);
}
#[cfg(feature = "tracing")]
tracing::trace!(
namespace = self.namespace.path(),
target_count = target_connection_ids.len(),
include_room_count = self.include_rooms.len(),
exclude_room_count = self.exclude_rooms.len(),
exclude_connection_count = self.exclude_connection_ids.len(),
"resolved broadcast targets"
);
target_connection_ids
}
pub async fn close(self) {
#[cfg(feature = "tracing")]
tracing::debug!(
namespace = self.namespace.path(),
"broadcasting close to namespace targets"
);
self.for_each_target_connections(|connection| async move {
connection.close();
Ok(())
})
.await;
}
pub async fn disconnect(self) -> Result<()> {
#[cfg(feature = "tracing")]
tracing::debug!(
namespace = self.namespace.path(),
"broadcasting disconnect to namespace targets"
);
let message = self.namespace.encode_packet_to_message(&WsIoPacket::new_disconnect())?;
self.for_each_target_connections(move |connection| {
let message = message.clone();
async move { connection.send_message(message).await }
})
.await;
Ok(())
}
pub async fn emit<D: Serialize>(self, event: impl AsRef<str>, data: Option<&D>) -> Result<()> {
self.namespace.status.ensure(NamespaceStatus::Running, |status| {
format!("Cannot emit in invalid status: {status:?}")
})?;
let event = event.as_ref();
#[cfg(feature = "tracing")]
tracing::trace!(
namespace = self.namespace.path(),
event,
has_data = data.is_some(),
"broadcasting event to namespace targets"
);
let message = self.namespace.encode_packet_to_message(&WsIoPacket::new_event(
event,
data.map(|data| self.namespace.config.packet_codec.encode_data(data))
.transpose()?,
))?;
self.for_each_target_connections(move |connection| {
let message = message.clone();
async move { connection.emit_event_message(message).await }
})
.await;
Ok(())
}
#[inline]
pub fn except(mut self, room_names: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.exclude_rooms.extend(room_names.into_iter().map(Into::into));
self
}
pub fn except_connection_ids(mut self, connection_ids: impl IntoIterator<Item = u64>) -> Self {
self.exclude_connection_ids.extend(connection_ids);
self
}
#[inline]
pub fn to(mut self, room_names: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.include_rooms.extend(room_names.into_iter().map(Into::into));
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::WsIoServer;
fn namespace_with_rooms() -> Arc<WsIoServerNamespace> {
let namespace = WsIoServer::builder()
.build()
.new_namespace_builder("/test")
.register()
.unwrap();
namespace
.connection_ids
.store(Arc::new(RoaringTreemap::from_iter([1, 2, 3])));
namespace.add_connection_id_to_room("included", 1);
namespace.add_connection_id_to_room("included", 2);
namespace.add_connection_id_to_room("excluded", 2);
namespace.add_connection_id_to_room("excluded", 3);
namespace
}
#[test]
fn resolves_existing_include_room_members() {
let targets = namespace_with_rooms().to(["included"]).target_connection_ids();
assert_eq!(targets, RoaringTreemap::from_iter([1, 2]));
}
#[test]
fn resolves_included_and_excluded_room_intersection() {
let targets = namespace_with_rooms()
.to(["included"])
.except(["excluded"])
.target_connection_ids();
assert_eq!(targets, RoaringTreemap::from_iter([1]));
}
#[test]
fn resolves_empty_include_room_to_no_targets() {
let targets = namespace_with_rooms().to(["missing"]).target_connection_ids();
assert!(targets.is_empty());
}
}