use serde::Serialize;
use tokio::sync::broadcast;
use crate::counter::ConnectionCounter;
use crate::error::WsError;
pub struct BroadcastHub<T>
where
T: Clone + Serialize + Send + Sync + 'static,
{
tx: broadcast::Sender<T>,
connection_count: ConnectionCounter,
}
impl<T> Clone for BroadcastHub<T>
where
T: Clone + Serialize + Send + Sync + 'static,
{
fn clone(&self) -> Self {
Self {
tx: self.tx.clone(),
connection_count: self.connection_count.clone(),
}
}
}
impl<T> std::fmt::Debug for BroadcastHub<T>
where
T: Clone + Serialize + Send + Sync + 'static,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BroadcastHub")
.field("capacity", &self.capacity())
.field("connection_count", &self.connection_count())
.field("receiver_count", &self.receiver_count())
.finish()
}
}
impl<T> BroadcastHub<T>
where
T: Clone + Serialize + Send + Sync + 'static,
{
pub fn new(capacity: usize) -> Self {
let (tx, _) = broadcast::channel(capacity);
Self {
tx,
connection_count: ConnectionCounter::new(),
}
}
pub fn from_config(config: &crate::config::WsConfig) -> Self {
Self::new(config.broadcast_capacity)
}
pub fn subscribe(&self) -> broadcast::Receiver<T> {
self.tx.subscribe()
}
pub fn broadcast(&self, msg: T) -> Result<usize, WsError> {
self.tx.send(msg).map_err(|_| WsError::BroadcastFull)
}
pub fn try_broadcast(&self, msg: T) -> usize {
self.tx.send(msg).unwrap_or(0)
}
pub fn connection_count(&self) -> u64 {
self.connection_count.get()
}
pub fn increment_connections(&self, max: Option<usize>) -> Result<u64, WsError> {
self.connection_count.increment(max)
}
pub fn decrement_connections(&self) -> u64 {
self.connection_count.decrement()
}
pub fn receiver_count(&self) -> usize {
self.tx.receiver_count()
}
pub fn capacity(&self) -> usize {
0
}
pub fn sender(&self) -> &broadcast::Sender<T> {
&self.tx
}
}
impl<T> BroadcastHub<T>
where
T: Clone + Serialize + Send + Sync + 'static,
{
pub fn with_capacity(capacity: usize) -> Self {
Self::new(capacity)
}
}
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic
)]
#[cfg(test)]
mod tests {
use super::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
struct TestMsg {
id: u32,
body: String,
}
#[tokio::test]
async fn broadcast_and_subscribe() {
let hub = BroadcastHub::<TestMsg>::new(16);
let mut rx1 = hub.subscribe();
let mut rx2 = hub.subscribe();
let msg = TestMsg {
id: 1,
body: "hi".into(),
};
let n = hub.broadcast(msg.clone()).unwrap();
assert_eq!(n, 2);
assert_eq!(rx1.recv().await.unwrap(), msg);
assert_eq!(rx2.recv().await.unwrap(), msg);
}
#[tokio::test]
async fn broadcast_no_receivers_error() {
let hub = BroadcastHub::<String>::new(16);
let err = hub.broadcast("hello".to_string()).unwrap_err();
assert_eq!(err, WsError::BroadcastFull);
assert_eq!(hub.try_broadcast("hello".to_string()), 0);
}
#[test]
fn connection_count_atomic() {
let hub = BroadcastHub::<String>::new(8);
assert_eq!(hub.connection_count(), 0);
hub.increment_connections(None).unwrap();
assert_eq!(hub.connection_count(), 1);
hub.increment_connections(None).unwrap();
assert_eq!(hub.connection_count(), 2);
hub.decrement_connections();
assert_eq!(hub.connection_count(), 1);
hub.decrement_connections();
hub.decrement_connections(); assert_eq!(hub.connection_count(), 0);
}
#[test]
fn max_connections_enforced() {
let hub = BroadcastHub::<String>::new(8);
hub.increment_connections(Some(2)).unwrap();
hub.increment_connections(Some(2)).unwrap();
let err = hub.increment_connections(Some(2)).unwrap_err();
assert_eq!(err, WsError::TooManyConnections);
}
#[test]
fn clone_shares_counter() {
let hub = BroadcastHub::<String>::new(8);
let hub2 = hub.clone();
hub.increment_connections(None).unwrap();
assert_eq!(hub2.connection_count(), 1);
}
}