use std::sync::Arc;
use dashmap::{DashMap, DashSet};
use serde::Serialize;
use tokio::sync::mpsc::Sender;
use tokio_tungstenite::tungstenite::Message;
use uuid::Uuid;
use crate::error::{Result, RustWaveError};
use crate::utils::send_data::Response;
pub(crate) struct ServerStates {
clients: Arc<DashMap<Uuid, Sender<Message>>>,
rooms: Arc<DashMap<String, DashSet<Uuid>>>,
}
impl ServerStates {
pub(crate) fn new() -> Self {
Self {
clients: Arc::new(DashMap::new()),
rooms: Arc::new(DashMap::new()),
}
}
pub(crate) fn insert_client(&self, id: Uuid, sender: Sender<Message>) {
self.clients.insert(id, sender);
}
pub(crate) fn remove_client(&self, id: &Uuid) {
self.clients.remove(id);
}
pub(crate) fn insert_room(&self, room: &str, id: Uuid) {
self.rooms
.entry(room.to_string())
.or_insert_with(DashSet::new)
.insert(id);
}
pub(crate) fn remove_room(&self, room: &str, id: &Uuid) {
let remove = if let Some(room) = self.rooms.get_mut(room) {
room.remove(id);
room.is_empty()
} else {
false
};
if remove {
self.rooms.remove(room);
}
}
pub(crate) async fn room_broadcast<T: Serialize>(
&self,
room: &str,
event: &str,
data: T,
) -> Result<()> {
let json = serde_json::to_string(&Response {
event: event.to_string(),
data,
})?;
let mess = Message::Text(json.into());
let socket_ids: Vec<Uuid> = match self.rooms.get(room) {
Some(con) => con.value().iter().map(|c| *c.key()).collect(),
None => return Ok(()),
};
for id in socket_ids {
let tx = self.clients.get(&id).map(|c| c.value().clone());
if let Some(tx) = tx {
let _ = tx.send(mess.clone()).await;
}
}
Ok(())
}
pub(crate) async fn room_broadcast_except<T: Serialize>(
&self,
room: &str,
socket_id: &Uuid,
event: &str,
data: T,
) -> Result<()> {
let json = serde_json::to_string(&Response {
event: event.to_string(),
data,
})?;
let mess = Message::Text(json.into());
let socket_ids: Vec<Uuid> = match self.rooms.get(room) {
Some(con) => con.value().iter().map(|c| *c.key()).collect(),
None => return Ok(()),
};
for id in socket_ids {
if socket_id == &id {
continue;
}
let tx = self.clients.get(&id).map(|c| c.value().clone());
if let Some(tx) = tx {
let _ = tx.send(mess.clone()).await;
}
}
Ok(())
}
pub(crate) async fn broadcast<T: Serialize>(&self, event: &str, data: T) -> Result<()> {
let json = serde_json::to_string(&Response {
event: event.to_string(),
data,
})?;
let mess = Message::Text(json.into());
let senders = self
.clients
.iter()
.map(|v| v.value().clone())
.collect::<Vec<_>>();
for sender in senders {
let _ = sender.send(mess.clone()).await;
}
Ok(())
}
pub(crate) async fn emit_to<T: Serialize>(
&self,
socket_id: &Uuid,
event: &str,
message: T,
) -> Result<()> {
let client = match self.clients.get(socket_id) {
Some(client) => client.value().clone(),
None => return Err(RustWaveError::ClientNotFound(socket_id.to_string())),
};
let json = serde_json::to_string(&Response {
event: event.to_string(),
data: message,
})?;
match client.send(Message::Text(json.into())).await {
Ok(_) => Ok(()),
Err(_) => Err(RustWaveError::ConnectionClosed),
}
}
pub(crate) fn is_connected(&self, socket_id: &Uuid) -> bool {
self.clients.contains_key(socket_id)
}
pub(crate) fn get_client_count(&self) -> usize {
self.clients.len()
}
pub(crate) async fn disconnect(&self, socket_id: &Uuid) -> Result<()> {
let sender = match self.clients.get(socket_id) {
Some(s) => s.value().clone(),
None => return Err(RustWaveError::ClientNotFound(socket_id.to_string())),
};
let _ = sender.send(Message::Close(None)).await;
Ok(())
}
pub(crate) async fn disconnect_all(&self) -> Result<()> {
let senders = self
.clients
.iter()
.map(|v| v.value().clone())
.collect::<Vec<Sender<Message>>>();
for sender in senders {
let _ = sender.send(Message::Close(None)).await;
}
Ok(())
}
}