main
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use serde_json::Value;
use tokio::sync::{Mutex, Notify, mpsc, oneshot, watch};
use tokio::task::JoinSet;
use crate::backend::{ConnectionPhase, SessionHandle, SocketBackend};
use super::super::session::Input;
use super::super::{EmitPacket, SocketEvent, SocketOptions, SocketPayload};
const EVENT_CAPACITY: usize = 1024;
const SHARDS: usize = 64;
#[derive(Clone)]
pub struct SocketHub {
pub(super) inner: Arc<Inner>,
}
pub(super) struct Inner {
pub(super) options: SocketOptions,
pub(super) backend: Arc<dyn SocketBackend>,
pub(super) events: mpsc::Sender<SocketEvent>,
pub(super) sessions: Vec<Mutex<HashMap<String, mpsc::Sender<Input>>>>,
pub(super) sockets: Vec<Mutex<HashMap<i64, UserSockets>>>,
pub(crate) shutdown: watch::Sender<bool>,
pub(super) transports: Vec<Mutex<JoinSet<()>>>,
pub(super) transport_next: AtomicUsize,
}
#[derive(Default)]
pub(super) struct UserSockets {
pub(super) namespaces: HashMap<String, HashMap<String, mpsc::Sender<Input>>>,
pub(super) changed: Arc<Notify>,
}
impl Inner {
pub(super) fn sessions_shard(&self, sid: &str) -> &Mutex<HashMap<String, mpsc::Sender<Input>>> {
&self.sessions[sid_shard(sid)]
}
pub(super) fn sockets_shard(&self, user_id: i64) -> &Mutex<HashMap<i64, UserSockets>> {
&self.sockets[(user_id as u64 as usize) & (SHARDS - 1)]
}
pub(super) fn transport_shard(&self) -> &Mutex<JoinSet<()>> {
let index = self.transport_next.fetch_add(1, Ordering::Relaxed) & (SHARDS - 1);
&self.transports[index]
}
}
impl SocketHub {
pub fn new(
options: SocketOptions,
backend: Arc<dyn SocketBackend>,
) -> Result<(Self, mpsc::Receiver<SocketEvent>), String> {
options.validate()?;
let (events, receiver) = mpsc::channel(EVENT_CAPACITY);
let hub = Self {
inner: Arc::new(Inner {
options,
backend,
events,
sessions: shards(),
sockets: shards(),
shutdown: watch::channel(false).0,
transports: shards(),
transport_next: AtomicUsize::new(0),
}),
};
Ok((hub, receiver))
}
pub fn shutdown(&self) {
self.inner.shutdown.send_replace(true);
}
pub fn shutdown_receiver(&self) -> watch::Receiver<bool> {
self.inner.shutdown.subscribe()
}
pub async fn emit(
&self,
user_id: i64,
namespace: &str,
event: &str,
args: Vec<Value>,
) -> usize {
self.emit_binary(user_id, namespace, event, args, Vec::new())
.await
}
pub async fn emit_binary(
&self,
user_id: i64,
namespace: &str,
event: &str,
args: Vec<Value>,
binaries: Vec<Vec<u8>>,
) -> usize {
let packet = EmitPacket {
namespace: namespace_key(namespace),
event: event.to_string(),
args,
binaries,
};
let namespace = packet.namespace.clone();
self.deliver(user_id, &namespace, packet).await
}
pub async fn emit_with_ack(
&self,
user_id: i64,
namespace: &str,
event: &str,
args: Vec<Value>,
timeout: Duration,
) -> Result<SocketPayload, String> {
let deadline = tokio::time::Instant::now() + timeout;
let packet = EmitPacket {
namespace: namespace_key(namespace),
event: event.to_string(),
args,
binaries: Vec::new(),
};
let inputs = self.inputs(user_id, &packet.namespace).await;
if inputs.is_empty() {
return Err("no connected socket".to_string());
}
let mut tasks = tokio::task::JoinSet::new();
for input in inputs {
let packet = packet.clone();
tasks.spawn(async move {
tokio::time::timeout_at(deadline, async {
let (ack_tx, ack_rx) = oneshot::channel();
input
.send(Input::Emit {
packet,
ack: Some((ack_tx, deadline)),
accepted: None,
})
.await
.map_err(|_| "socket closed")?;
ack_rx.await.map_err(|_| "socket closed")
})
.await
});
}
while let Some(result) = tasks.join_next().await {
if let Ok(Ok(Ok(payload))) = result {
tasks.abort_all();
return Ok(payload);
}
}
Err("ack timeout or socket closed".to_string())
}
pub async fn wake_user(&self, user_id: i64) {
let sockets = self.inner.sockets_shard(user_id).lock().await;
if let Some(user) = sockets.get(&user_id) {
user.changed.notify_waiters();
}
}
pub async fn notification(&self, user_id: i64) -> Arc<Notify> {
self.inner
.sockets_shard(user_id)
.lock()
.await
.entry(user_id)
.or_default()
.changed
.clone()
}
/// 聊天连接生命周期上报（Open / Heartbeat / Close）。
///
/// 失败只记日志：在线态是旁路信息，不能因为平台不可用就把 socket 踢掉。
pub async fn report_connection(&self, phase: ConnectionPhase, user_id: i64, session: &str) {
let _ = tokio::time::timeout(
self.inner.options.auth_timeout,
self.inner
.backend
.report_connection(phase, user_id, session),
)
.await;
}
/// 关闭某用户的全部在线会话（踢下线 / 权限撤销）。
pub async fn close_user(&self, user_id: i64) -> usize {
let sockets = self.inner.sockets_shard(user_id).lock().await;
let mut closed = 0;
if let Some(user) = sockets.get(&user_id) {
for entries in user.namespaces.values() {
for input in entries.values() {
if input.try_send(Input::Close).is_ok() {
closed += 1;
}
}
}
}
closed
}
pub(crate) fn backend(&self) -> &Arc<dyn SocketBackend> {
&self.inner.backend
}
pub(crate) fn session_handle(&self, input: mpsc::Sender<Input>) -> SessionHandle {
SessionHandle::new(input)
}
pub(crate) fn options(&self) -> &SocketOptions {
&self.inner.options
}
pub(crate) fn events(&self) -> &mpsc::Sender<SocketEvent> {
&self.inner.events
}
}
pub(super) fn namespace_key(namespace: &str) -> String {
if namespace.is_empty() {
"/".to_string()
} else if namespace.starts_with('/') {
namespace.to_string()
} else {
format!("/{namespace}")
}
}
fn shards<T: Default>() -> Vec<Mutex<T>> {
(0..SHARDS).map(|_| Mutex::new(T::default())).collect()
}
fn sid_shard(sid: &str) -> usize {
let mut hasher = std::collections::hash_map::DefaultHasher::new();
sid.hash(&mut hasher);
(hasher.finish() as usize) & (SHARDS - 1)
}