1use std::any::{Any, TypeId};
2use std::collections::HashMap;
3use std::fmt;
4use std::sync::{Arc, RwLock};
5use tokio::sync::broadcast;
6
7const CHANNEL_CAPACITY: usize = 256;
9
10pub use tokio::sync::broadcast::error::{RecvError, TryRecvError};
11
12pub trait Event: fmt::Debug + Clone + Send + Sync + 'static {
17 fn name(&self) -> &'static str;
21}
22
23#[derive(Debug, thiserror::Error)]
24pub enum PublishError {
25 #[error("event bus lock poisoned")]
26 Poisoned,
27}
28
29#[derive(Clone, Default)]
31pub struct EventBus {
32 channels: Arc<RwLock<HashMap<TypeId, Box<dyn Any + Send + Sync>>>>,
33}
34
35impl fmt::Debug for EventBus {
36 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37 let channels = self.channels.read().map(|c| c.len()).unwrap_or(0);
38 f.debug_struct("EventBus")
39 .field("event_types", &channels)
40 .finish()
41 }
42}
43
44impl EventBus {
45 pub fn new() -> Self {
46 Self::default()
47 }
48
49 pub fn publish<E: Event>(&self, event: E) -> Result<usize, PublishError> {
54 let name = event.name();
55 let sender = self.sender::<E>()?;
56 let delivered = sender.send(event).unwrap_or(0);
57 tracing::debug!(event = name, subscribers = delivered, "event published");
58 Ok(delivered)
59 }
60
61 pub fn subscribe<E: Event>(&self) -> Result<EventStream<E>, PublishError> {
65 Ok(EventStream {
66 receiver: self.sender::<E>()?.subscribe(),
67 })
68 }
69
70 fn sender<E: Event>(&self) -> Result<broadcast::Sender<E>, PublishError> {
71 let type_id = TypeId::of::<E>();
72
73 {
74 let channels = self.channels.read().map_err(|_| PublishError::Poisoned)?;
75 if let Some(existing) = channels.get(&type_id) {
76 return Ok(Self::downcast::<E>(existing));
77 }
78 }
79
80 let mut channels = self.channels.write().map_err(|_| PublishError::Poisoned)?;
81 let entry = channels.entry(type_id).or_insert_with(|| {
82 let (sender, _) = broadcast::channel::<E>(CHANNEL_CAPACITY);
83 Box::new(sender)
84 });
85 Ok(Self::downcast::<E>(entry))
86 }
87
88 fn downcast<E: Event>(entry: &Box<dyn Any + Send + Sync>) -> broadcast::Sender<E> {
89 entry
90 .downcast_ref::<broadcast::Sender<E>>()
91 .expect("event channel registered under a mismatched type id")
94 .clone()
95 }
96}
97
98#[derive(Debug)]
100pub struct EventStream<E: Event> {
101 receiver: broadcast::Receiver<E>,
102}
103
104impl<E: Event> EventStream<E> {
105 pub async fn recv(&mut self) -> Result<E, RecvError> {
110 self.receiver.recv().await
111 }
112
113 pub fn try_recv(&mut self) -> Result<E, TryRecvError> {
118 self.receiver.try_recv()
119 }
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125
126 #[derive(Debug, Clone, PartialEq)]
127 struct Ping(u32);
128 impl Event for Ping {
129 fn name(&self) -> &'static str {
130 "test.ping"
131 }
132 }
133
134 #[derive(Debug, Clone, PartialEq)]
135 struct Pong;
136 impl Event for Pong {
137 fn name(&self) -> &'static str {
138 "test.pong"
139 }
140 }
141
142 #[tokio::test]
143 async fn subscribers_receive_events_of_their_own_type() {
144 let bus = EventBus::new();
145 let mut pings = bus.subscribe::<Ping>().unwrap();
146
147 assert_eq!(bus.publish(Ping(7)).unwrap(), 1);
148
149 assert_eq!(pings.recv().await.unwrap(), Ping(7));
150 }
151
152 #[tokio::test]
153 async fn events_of_a_different_type_are_not_delivered() {
154 let bus = EventBus::new();
155 let mut pings = bus.subscribe::<Ping>().unwrap();
156
157 bus.publish(Pong).unwrap();
158 bus.publish(Ping(1)).unwrap();
159
160 assert_eq!(pings.recv().await.unwrap(), Ping(1));
162 }
163
164 #[tokio::test]
165 async fn publishing_without_subscribers_is_not_an_error() {
166 let bus = EventBus::new();
167 assert_eq!(bus.publish(Ping(1)).unwrap(), 0);
168 }
169
170 #[tokio::test]
171 async fn every_subscriber_receives_a_copy() {
172 let bus = EventBus::new();
173 let mut first = bus.subscribe::<Ping>().unwrap();
174 let mut second = bus.subscribe::<Ping>().unwrap();
175
176 assert_eq!(bus.publish(Ping(42)).unwrap(), 2);
177
178 assert_eq!(first.recv().await.unwrap(), Ping(42));
179 assert_eq!(second.recv().await.unwrap(), Ping(42));
180 }
181}