use std::sync::Arc;
use dashmap::DashMap;
use tokio::sync::broadcast;
use crate::codec::Frame;
use crate::config::WsConfig;
use crate::error::WsError;
use crate::stats::StatsRecorder;
pub struct Room {
pub id: String,
tx: broadcast::Sender<Frame>,
participants: DashMap<u32, String>,
stats: Option<StatsRecorder>,
}
impl std::fmt::Debug for Room {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Room")
.field("id", &self.id)
.field("participants", &self.participant_count())
.field("receiver_count", &self.receiver_count())
.finish()
}
}
impl Room {
pub fn new(id: impl Into<String>, capacity: usize) -> Self {
Self::with_stats(id, capacity, None)
}
pub fn with_stats(
id: impl Into<String>,
capacity: usize,
stats: Option<StatsRecorder>,
) -> Self {
let (tx, _) = broadcast::channel(capacity);
Self {
id: id.into(),
tx,
participants: DashMap::new(),
stats,
}
}
pub fn subscribe(&self) -> broadcast::Receiver<Frame> {
self.tx.subscribe()
}
pub fn broadcast(&self, msg: impl Into<Frame>) -> Result<usize, WsError> {
let frame = msg.into();
let n = self
.tx
.send(frame.clone())
.map_err(|_| self.record_drop())?;
if let Some(stats) = &self.stats {
stats.record_room_message_out(&self.id, n as u64, frame.len() as u64 * n as u64);
}
Ok(n)
}
pub fn join(&self, participant_id: u32, name: String) {
self.participants.insert(participant_id, name);
}
pub fn leave(&self, participant_id: u32) -> bool {
self.participants.remove(&participant_id).is_some()
}
pub fn contains(&self, participant_id: u32) -> bool {
self.participants.contains_key(&participant_id)
}
pub fn participants(&self) -> Vec<(u32, String)> {
self.participants
.iter()
.map(|r| (*r.key(), r.value().clone()))
.collect()
}
pub fn participant_names(&self) -> Vec<String> {
self.participants
.iter()
.map(|r| r.value().clone())
.collect()
}
pub fn participant_count(&self) -> usize {
self.participants.len()
}
pub fn receiver_count(&self) -> usize {
self.tx.receiver_count()
}
pub fn is_empty(&self) -> bool {
self.participants.is_empty() && self.tx.receiver_count() == 0
}
pub fn sender(&self) -> &broadcast::Sender<Frame> {
&self.tx
}
fn record_drop(&self) -> WsError {
if let Some(stats) = &self.stats {
stats.record_room_drop(&self.id);
}
WsError::BroadcastFull
}
}
pub trait RoomRegistry: Send + Sync {
fn get_or_create(&self, room_id: &str) -> Arc<Room>;
fn get(&self, room_id: &str) -> Option<Arc<Room>>;
fn remove(&self, room_id: &str) -> Option<Arc<Room>>;
fn room_count(&self) -> usize;
fn room_ids(&self) -> Vec<String>;
fn total_participants(&self) -> usize;
fn cleanup(&self) -> usize;
fn cleanup_empty(&self) -> usize;
}
#[derive(Debug)]
pub struct RoomManager {
rooms: DashMap<String, Arc<Room>>,
default_capacity: usize,
stats: Option<StatsRecorder>,
}
impl Default for RoomManager {
fn default() -> Self {
Self::new()
}
}
impl RoomManager {
pub fn new() -> Self {
Self {
rooms: DashMap::new(),
default_capacity: 1024,
stats: None,
}
}
pub fn with_config(config: &WsConfig) -> Self {
Self {
rooms: DashMap::new(),
default_capacity: config.broadcast_capacity,
stats: None,
}
}
pub fn with_stats(config: &WsConfig, stats: StatsRecorder) -> Self {
Self {
rooms: DashMap::new(),
default_capacity: config.broadcast_capacity,
stats: Some(stats),
}
}
pub fn stats(&self) -> Option<&StatsRecorder> {
self.stats.as_ref()
}
}
impl RoomRegistry for RoomManager {
fn get_or_create(&self, room_id: &str) -> Arc<Room> {
RoomManager::get_or_create(self, room_id)
}
fn get(&self, room_id: &str) -> Option<Arc<Room>> {
RoomManager::get(self, room_id)
}
fn remove(&self, room_id: &str) -> Option<Arc<Room>> {
RoomManager::remove(self, room_id)
}
fn room_count(&self) -> usize {
RoomManager::room_count(self)
}
fn room_ids(&self) -> Vec<String> {
RoomManager::room_ids(self)
}
fn total_participants(&self) -> usize {
RoomManager::total_participants(self)
}
fn cleanup(&self) -> usize {
RoomManager::cleanup(self)
}
fn cleanup_empty(&self) -> usize {
RoomManager::cleanup_empty(self)
}
}
impl RoomManager {
pub fn get_or_create(&self, room_id: &str) -> Arc<Room> {
self.get_or_create_with_capacity(room_id, self.default_capacity)
}
pub fn get(&self, room_id: &str) -> Option<Arc<Room>> {
self.rooms.get(room_id).map(|r| Arc::clone(&r))
}
pub fn remove(&self, room_id: &str) -> Option<Arc<Room>> {
let removed = self.rooms.remove(room_id).map(|(_, v)| v);
if removed.is_some() {
if let Some(stats) = &self.stats {
stats.room_closed();
}
}
removed
}
pub fn room_count(&self) -> usize {
self.rooms.len()
}
pub fn total_participants(&self) -> usize {
self.rooms
.iter()
.map(|e| e.value().participant_count())
.sum()
}
pub fn cleanup(&self) -> usize {
let mut to_remove = Vec::new();
for entry in self.rooms.iter() {
if entry.value().is_empty() {
to_remove.push(entry.key().clone());
}
}
let n = to_remove.len();
for k in to_remove {
self.rooms.remove(&k);
}
if n > 0 {
if let Some(stats) = &self.stats {
for _ in 0..n {
stats.room_closed();
}
}
}
n
}
pub fn cleanup_empty(&self) -> usize {
let to_remove: Vec<String> = self
.rooms
.iter()
.filter(|e| e.value().participant_count() == 0)
.map(|e| e.key().clone())
.collect();
let n = to_remove.len();
for k in to_remove {
self.rooms.remove(&k);
}
if n > 0 {
if let Some(stats) = &self.stats {
for _ in 0..n {
stats.room_closed();
}
}
}
n
}
pub fn room_ids(&self) -> Vec<String> {
self.rooms.iter().map(|e| e.key().clone()).collect()
}
pub fn get_or_create_with_capacity(&self, room_id: &str, capacity: usize) -> Arc<Room> {
if let Some(r) = self.rooms.get(room_id) {
return Arc::clone(&r);
}
let room = Arc::new(Room::with_stats(room_id, capacity, self.stats.clone()));
if let Some(stats) = &self.stats {
stats.room_opened();
}
let entry = self
.rooms
.entry(room_id.to_string())
.or_insert_with(|| Arc::clone(&room));
Arc::clone(entry.value())
}
}
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic
)]
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
fn name(s: &str) -> String {
s.to_string()
}
#[test]
fn room_join_leave() {
let room = Room::new("lobby", 16);
assert_eq!(room.participant_count(), 0);
room.join(1, name("alice"));
room.join(2, name("bob"));
assert_eq!(room.participant_count(), 2);
assert!(room.contains(1));
assert!(room.leave(1));
assert!(!room.contains(1));
assert_eq!(room.participant_count(), 1);
}
#[test]
fn room_rejoin_updates_name() {
let room = Room::new("lobby", 16);
room.join(1, name("alice"));
room.join(1, name("alice2"));
assert_eq!(room.participant_count(), 1);
assert_eq!(room.participant_names(), vec!["alice2".to_string()]);
}
#[test]
fn room_participant_snapshots() {
let room = Room::new("lobby", 16);
room.join(7, name("alice"));
room.join(3, name("bob"));
let mut pairs = room.participants();
pairs.sort();
assert_eq!(
pairs,
vec![(3, "bob".to_string()), (7, "alice".to_string())]
);
let mut names = room.participant_names();
names.sort();
assert_eq!(names, vec!["alice".to_string(), "bob".to_string()]);
assert_eq!(room.participant_count(), 2);
}
#[tokio::test]
async fn room_broadcast() {
let room = Room::new("chat", 16);
let mut rx = room.subscribe();
let n = room.broadcast("hello".to_string()).unwrap();
assert_eq!(n, 1);
assert_eq!(rx.recv().await.unwrap(), Frame::from("hello"));
}
#[tokio::test]
async fn room_broadcast_binary() {
let room = Room::new("binroom", 16);
let mut rx = room.subscribe();
let payload = vec![0xde_u8, 0xad, 0xbe, 0xef];
let n = room.broadcast(Frame::from(payload.clone())).unwrap();
assert_eq!(n, 1);
let got = rx.recv().await.unwrap();
assert!(got.is_binary());
assert_eq!(got.as_bytes(), &payload[..]);
}
#[tokio::test]
async fn room_broadcast_bytes_vec_convenience() {
let room = Room::new("binroom", 16);
let mut rx = room.subscribe();
room.broadcast(vec![1u8, 2, 3]).unwrap();
let got = rx.recv().await.unwrap();
assert!(got.is_binary());
assert_eq!(got.as_bytes(), &[1, 2, 3]);
}
#[tokio::test]
async fn room_text_and_binary_share_one_channel_without_cross_talk() {
let room = Room::new("mixed", 16);
let mut rx = room.subscribe();
room.broadcast("text".to_string()).unwrap();
room.broadcast(Bytes::from_static(b"\x00\x01")).unwrap();
assert!(rx.recv().await.unwrap().is_text());
assert!(rx.recv().await.unwrap().is_binary());
}
#[tokio::test]
async fn room_broadcast_no_receiver() {
let room = Room::new("empty", 16);
let err = room.broadcast("hi".to_string()).unwrap_err();
assert_eq!(err, WsError::BroadcastFull);
}
#[tokio::test]
async fn room_stats_records_outbound_and_drops() {
let stats = StatsRecorder::new();
let room = Room::with_stats("metrics", 16, Some(stats.clone()));
let mut rx = room.subscribe();
room.broadcast("12345".to_string()).unwrap(); rx.recv().await.unwrap();
let s = stats.snapshot();
let room_stats = s.per_room.get("metrics").unwrap();
assert_eq!(room_stats.messages_out, 1);
assert_eq!(room_stats.bytes_out, 5);
assert_eq!(s.messages_out, 1);
drop(rx);
room.broadcast("nobody".to_string()).unwrap_err();
assert_eq!(stats.snapshot().per_room.get("metrics").unwrap().drops, 1);
assert_eq!(stats.snapshot().drops, 1);
}
#[test]
fn manager_get_or_create_idempotent() {
let m = RoomManager::new();
let r1 = m.get_or_create("room1");
let r2 = m.get_or_create("room1");
assert!(Arc::ptr_eq(&r1, &r2));
assert_eq!(m.room_count(), 1);
}
#[test]
fn manager_with_config_capacity() {
let cfg = WsConfig::builder().broadcast_capacity(2048).build();
let m = RoomManager::with_config(&cfg);
let r = m.get_or_create("room1");
let r2 = m.get_or_create("room1");
assert!(Arc::ptr_eq(&r, &r2));
assert_eq!(m.room_count(), 1);
}
#[test]
fn manager_cleanup_removes_empty() {
let m = RoomManager::new();
let _r = m.get_or_create("a");
assert_eq!(m.cleanup(), 1);
assert_eq!(m.room_count(), 0);
let r = m.get_or_create("b");
r.join(1, name("user1"));
assert_eq!(m.cleanup(), 0);
assert_eq!(m.room_count(), 1);
}
#[test]
fn manager_cleanup_empty_removes_zero_participant_rooms() {
let m = RoomManager::new();
let _a = m.get_or_create("a");
assert_eq!(m.cleanup_empty(), 1);
assert_eq!(m.room_count(), 0);
let a = m.get_or_create("a");
a.join(1, name("alice"));
assert_eq!(m.cleanup_empty(), 0);
assert_eq!(m.room_count(), 1);
let kept = m.get_or_create("a");
assert!(Arc::ptr_eq(&a, &kept));
a.leave(1);
assert_eq!(m.cleanup_empty(), 1);
assert_eq!(m.room_count(), 0);
}
#[test]
fn manager_total_participants() {
let m = RoomManager::new();
let r1 = m.get_or_create("r1");
let r2 = m.get_or_create("r2");
assert_eq!(m.total_participants(), 0);
r1.join(1, name("alice"));
r1.join(2, name("bob"));
r2.join(3, name("carol"));
assert_eq!(m.total_participants(), 3);
r2.leave(3);
assert_eq!(m.total_participants(), 2);
}
#[test]
fn manager_get_and_remove() {
let m = RoomManager::new();
assert!(m.get("nonexistent").is_none());
let r = m.get_or_create("x");
assert!(m.get("x").is_some());
let removed = m.remove("x").unwrap();
assert!(Arc::ptr_eq(&r, &removed));
assert!(m.get("x").is_none());
}
#[test]
fn manager_stats_room_lifecycle() {
let stats = StatsRecorder::new();
let cfg = WsConfig::default();
let m = RoomManager::with_stats(&cfg, stats.clone());
let _r = m.get_or_create("a");
assert_eq!(stats.snapshot().rooms, 1);
m.get_or_create("a"); assert_eq!(stats.snapshot().rooms, 1);
m.remove("a");
assert_eq!(stats.snapshot().rooms, 0);
m.get_or_create("b");
m.get_or_create("c");
m.cleanup();
assert_eq!(stats.snapshot().rooms, 0);
}
#[test]
fn trait_object_room_registry() {
let manager = RoomManager::new();
let registry: &dyn RoomRegistry = &manager;
let r = registry.get_or_create("dyn_room");
r.join(1, name("alice"));
assert_eq!(registry.room_count(), 1);
assert_eq!(registry.total_participants(), 1);
assert_eq!(registry.room_ids(), vec!["dyn_room".to_string()]);
assert!(registry.get("dyn_room").is_some());
assert!(registry.remove("dyn_room").is_some());
assert_eq!(registry.cleanup(), 0);
assert_eq!(registry.cleanup_empty(), 0);
}
#[tokio::test]
async fn manager_broadcast_isolation() {
let m = RoomManager::new();
let r1 = m.get_or_create("r1");
let r2 = m.get_or_create("r2");
let mut rx1 = r1.subscribe();
let mut rx2 = r2.subscribe();
r1.broadcast("for r1".to_string()).unwrap();
assert_eq!(rx1.recv().await.unwrap(), Frame::from("for r1"));
use tokio::time::{timeout, Duration};
let res = timeout(Duration::from_millis(10), rx2.recv()).await;
assert!(res.is_err());
}
}