use std::{
any::{
Any,
TypeId,
},
collections::hash_map::Entry,
fmt::{
Debug as FmtDebug,
Formatter,
Result as FmtResult,
},
marker::PhantomData,
pin::Pin,
sync::{
Arc,
LazyLock,
atomic::{
AtomicU32,
Ordering,
},
},
};
use anyhow::Result;
use kikiutils::types::fx_collections::FxHashMap;
use parking_lot::RwLock;
use serde::de::DeserializeOwned;
use crate::{
packet::codecs::WsIoPacketCodec,
traits::task::spawner::TaskSpawner,
};
type DataDecoder = fn(&[u8], WsIoPacketCodec) -> Result<Arc<dyn Any + Send + Sync>>;
type Handler<C> = Arc<
dyn Fn(Arc<C>, Arc<dyn Any + Send + Sync>) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>
+ Send
+ Sync
+ 'static,
>;
static EMPTY_EVENT_DATA_ANY_ARC: LazyLock<Arc<dyn Any + Send + Sync>> = LazyLock::new(|| Arc::new(()));
struct EventEntry<C> {
data_decoder: DataDecoder,
data_type_id: TypeId,
handlers: RwLock<FxHashMap<u32, Handler<C>>>,
}
impl<C> FmtDebug for EventEntry<C> {
fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
let mut debug = f.debug_struct("EventEntry");
debug
.field("data_decoder", &self.data_decoder)
.field("data_type_id", &self.data_type_id);
match self.handlers.try_read() {
Some(handlers) => {
debug.field("handlers_len", &handlers.len());
},
None => {
debug.field("handlers", &"<locked>");
},
}
debug.finish()
}
}
#[derive(Debug)]
pub struct WsIoEventRegistry<C: Send + Sync + 'static, S: TaskSpawner> {
_task_spawner: PhantomData<S>,
event_entries: RwLock<FxHashMap<String, Arc<EventEntry<C>>>>,
next_handler_id: AtomicU32,
}
impl<C: Send + Sync + 'static, S: TaskSpawner> Default for WsIoEventRegistry<C, S> {
fn default() -> Self {
Self::new()
}
}
impl<C: Send + Sync + 'static, S: TaskSpawner> WsIoEventRegistry<C, S> {
#[inline]
pub fn new() -> Self {
Self {
_task_spawner: PhantomData,
event_entries: RwLock::new(FxHashMap::default()),
next_handler_id: AtomicU32::new(0),
}
}
#[inline]
pub fn dispatch_event_packet(
&self,
ctx: Arc<C>,
event: impl AsRef<str>,
packet_codec: &WsIoPacketCodec,
packet_data: Option<Vec<u8>>,
task_spawner: &Arc<S>,
) {
let event = event.as_ref();
let Some(event_entry) = self.event_entries.read().get(event).cloned() else {
#[cfg(feature = "tracing")]
tracing::trace!(event, "dropping event packet without registered handlers");
return;
};
#[cfg(feature = "tracing")]
tracing::trace!(event, has_data = packet_data.is_some(), "dispatching event packet");
let packet_codec = *packet_codec;
let task_spawner_clone = task_spawner.clone();
#[cfg(feature = "tracing")]
let event_name = event.to_owned();
task_spawner.spawn_task(async move {
let data = match packet_data {
Some(bytes) => match (event_entry.data_decoder)(&bytes, packet_codec) {
Ok(data) => data,
Err(_err) => {
#[cfg(feature = "tracing")]
tracing::debug!(event = %event_name, error = %_err, "failed to decode event packet data");
return Ok(());
},
},
None => EMPTY_EVENT_DATA_ANY_ARC.clone(),
};
let handlers = event_entry.handlers.read().values().cloned().collect::<Vec<_>>();
#[cfg(feature = "tracing")]
tracing::trace!(
event = %event_name,
handler_count = handlers.len(),
"spawning event handlers"
);
for handler in handlers {
let ctx = ctx.clone();
let data = data.clone();
task_spawner_clone.spawn_task(handler(ctx, data));
}
Ok(())
});
}
#[inline]
pub fn off(&self, event: impl AsRef<str>) {
let event = event.as_ref();
let _removed = self.event_entries.write().remove(event).is_some();
#[cfg(feature = "tracing")]
tracing::trace!(event, removed = _removed, "removed event handlers");
}
#[inline]
pub fn off_by_handler_id(&self, event: impl AsRef<str>, handler_id: u32) {
let event = event.as_ref();
if let Some(event_entry) = self.event_entries.read().get(event) {
let _removed = event_entry.handlers.write().remove(&handler_id).is_some();
#[cfg(feature = "tracing")]
tracing::trace!(event, handler_id, removed = _removed, "removed event handler by id");
if !event_entry.handlers.read().is_empty() {
return;
}
}
if let Entry::Occupied(entry) = self.event_entries.write().entry(event.to_owned())
&& entry.get().handlers.read().is_empty()
{
entry.remove();
}
}
#[inline]
pub fn on<H, Fut, D>(&self, event: impl AsRef<str>, handler: H) -> u32
where
H: Fn(Arc<C>, Arc<D>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
D: DeserializeOwned + Send + Sync + 'static,
{
let event = event.as_ref();
let data_type_id = TypeId::of::<D>();
let mut event_entries = self.event_entries.write();
let event_entry = match event_entries.entry(event.to_owned()) {
Entry::Occupied(occupied) => {
let event_entry = occupied.into_mut();
assert_eq!(
event_entry.data_type_id, data_type_id,
"Event '{}' already registered with a different data type — each event name must correspond to exactly one payload type.",
event
);
event_entry
},
Entry::Vacant(vacant) => vacant.insert(Arc::new(EventEntry {
data_decoder: decode_data_as_any_arc::<D>,
data_type_id,
handlers: RwLock::new(FxHashMap::default()),
})),
};
let handler_id = self.next_handler_id.fetch_add(1, Ordering::Relaxed);
#[cfg(feature = "tracing")]
tracing::trace!(event, handler_id, "registered event handler");
event_entry.handlers.write().insert(
handler_id,
Arc::new(move |connection, data| {
if (*data).type_id() != data_type_id {
return Box::pin(async { Ok(()) });
}
Box::pin(handler(
connection,
#[allow(clippy::expect_used)]
data.downcast()
.expect("data type id matched handler registration but Arc::downcast failed"),
))
}),
);
handler_id
}
}
#[inline]
fn decode_data_as_any_arc<D: DeserializeOwned + Send + Sync + 'static>(
bytes: &[u8],
packet_codec: WsIoPacketCodec,
) -> Result<Arc<dyn Any + Send + Sync>> {
Ok(Arc::new(packet_codec.decode_data::<D>(bytes)?))
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use tokio::{
spawn,
sync::mpsc::unbounded_channel,
time::timeout,
};
use tokio_util::sync::CancellationToken;
use super::*;
struct DummyConnection;
struct DummySpawner {
cancel_token: Arc<CancellationToken>,
}
impl TaskSpawner for DummySpawner {
fn cancel_token(&self) -> Arc<CancellationToken> {
self.cancel_token.clone()
}
fn spawn_task<F: Future<Output = Result<()>> + Send + 'static>(&self, future: F) {
spawn(future);
}
}
#[tokio::test]
async fn test_registry_dispatch() {
let registry = WsIoEventRegistry::<DummyConnection, DummySpawner>::new();
let spawner = Arc::new(DummySpawner {
cancel_token: Arc::new(CancellationToken::new()),
});
let ctx = Arc::new(DummyConnection);
let (handled_tx, mut handled_rx) = unbounded_channel();
let first_handler_tx = handled_tx.clone();
registry.on("ping", move |_ctx, payload: Arc<String>| {
assert_eq!(*payload, "hello");
first_handler_tx.send("first").unwrap();
async move { Ok(()) }
});
registry.on("ping", move |_ctx, payload: Arc<String>| {
assert_eq!(*payload, "hello");
handled_tx.send("second").unwrap();
async move { Ok(()) }
});
let packet_codec = WsIoPacketCodec::SerdeJson;
let packet_data = packet_codec.encode_data(&"hello").unwrap();
registry.dispatch_event_packet(ctx, "ping", &packet_codec, Some(packet_data), &spawner);
let mut handlers = Vec::with_capacity(2);
for _ in 0..2 {
let handler = timeout(Duration::from_secs(1), handled_rx.recv())
.await
.expect("handler should run before timeout")
.expect("handler channel should remain open");
handlers.push(handler);
}
handlers.sort_unstable();
assert_eq!(handlers, ["first", "second"]);
}
#[test]
fn test_registry_on_off() {
let registry = WsIoEventRegistry::<DummyConnection, DummySpawner>::new();
let handler_id = registry.on("test_event", |_ctx, _data: Arc<String>| async { Ok(()) });
assert_eq!(handler_id, 0);
assert!(registry.event_entries.read().contains_key("test_event"));
assert_eq!(
registry
.event_entries
.read()
.get("test_event")
.unwrap()
.handlers
.read()
.len(),
1
);
registry.off_by_handler_id("test_event", handler_id);
assert!(!registry.event_entries.read().contains_key("test_event"));
registry.on("multi_event", |_ctx, _data: Arc<String>| async { Ok(()) });
registry.on("multi_event", |_ctx, _data: Arc<String>| async { Ok(()) });
assert_eq!(
registry
.event_entries
.read()
.get("multi_event")
.unwrap()
.handlers
.read()
.len(),
2
);
registry.off("multi_event");
assert!(!registry.event_entries.read().contains_key("multi_event"));
}
}