//! In-memory pub/sub dispatcher (feature `mem`), backed by //! [`tokio::sync::broadcast`]. //! //! Suitable for communication between components inside a single process //! (monolith). Each topic lazily gets its own broadcast channel on first use. use std::collections::HashMap; use std::sync::{Arc, Mutex}; use async_trait::async_trait; use tokio::sync::broadcast; use crate::traits::{EventBus, EventError, Subscription, SubscriptionInner}; const DEFAULT_CAPACITY: usize = 256; /// An in-process [`EventBus`] backed by a `tokio::sync::broadcast` channel per /// topic. #[derive(Clone)] pub struct InMemoryEventBus { channels: Arc>>>>, capacity: usize, } impl Default for InMemoryEventBus { fn default() -> Self { Self::new(DEFAULT_CAPACITY) } } impl InMemoryEventBus { /// Builds a bus whose per-topic channels buffer up to `capacity` messages /// for slow subscribers before older ones are dropped (lagged). pub fn new(capacity: usize) -> Self { Self { channels: Arc::new(Mutex::new(HashMap::new())), capacity, } } fn sender_for(&self, topic: &str) -> broadcast::Sender> { let mut channels = self.channels.lock().unwrap(); channels .entry(topic.to_string()) .or_insert_with(|| broadcast::channel(self.capacity).0) .clone() } /// The number of currently subscribed receivers for `topic` (0 if the /// topic has never been touched or has no active subscribers). pub fn subscriber_count(&self, topic: &str) -> usize { self.channels .lock() .unwrap() .get(topic) .map(|tx| tx.receiver_count()) .unwrap_or(0) } } #[async_trait] impl EventBus for InMemoryEventBus { async fn publish(&self, topic: &str, payload: Vec) -> Result<(), EventError> { let tx = self.sender_for(topic); // A `SendError` here just means there are currently no subscribers, // which is not an error condition for a pub/sub bus. let _ = tx.send(payload); Ok(()) } async fn subscribe(&self, topic: &str) -> Result { let rx = self.sender_for(topic).subscribe(); Ok(Subscription::new(Box::new(BroadcastSubscription(rx)))) } } struct BroadcastSubscription(broadcast::Receiver>); #[async_trait] impl SubscriptionInner for BroadcastSubscription { async fn recv(&mut self) -> Result, EventError> { loop { match self.0.recv().await { Ok(payload) => return Ok(payload), // A slow subscriber missed some messages; skip ahead rather // than erroring, matching typical broadcast semantics. Err(broadcast::error::RecvError::Lagged(_)) => continue, Err(broadcast::error::RecvError::Closed) => return Err(EventError::Closed), } } } } #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn publish_subscribe_roundtrip() { let bus = InMemoryEventBus::default(); let mut sub = bus.subscribe("orders").await.unwrap(); bus.publish("orders", b"order-created".to_vec()) .await .unwrap(); let msg = sub.recv().await.unwrap(); assert_eq!(msg, b"order-created"); } #[tokio::test] async fn multiple_subscribers_all_receive() { let bus = InMemoryEventBus::default(); let mut sub_a = bus.subscribe("t").await.unwrap(); let mut sub_b = bus.subscribe("t").await.unwrap(); bus.publish("t", b"hello".to_vec()).await.unwrap(); assert_eq!(sub_a.recv().await.unwrap(), b"hello"); assert_eq!(sub_b.recv().await.unwrap(), b"hello"); } #[tokio::test] async fn publish_with_no_subscribers_is_ok() { let bus = InMemoryEventBus::default(); assert!(bus.publish("nobody-listening", b"x".to_vec()).await.is_ok()); } #[tokio::test] async fn different_topics_are_isolated() { let bus = InMemoryEventBus::default(); let mut sub = bus.subscribe("a").await.unwrap(); bus.publish("b", b"for-b".to_vec()).await.unwrap(); // Nothing arrives on `a`'s subscription; use a short timeout to prove it. let result = tokio::time::timeout(std::time::Duration::from_millis(30), sub.recv()).await; assert!( result.is_err(), "subscription on `a` should not receive `b`'s event" ); } #[tokio::test] async fn subscriber_count_tracks_active_subscriptions() { let bus = InMemoryEventBus::default(); assert_eq!(bus.subscriber_count("t"), 0); let sub = bus.subscribe("t").await.unwrap(); assert_eq!(bus.subscriber_count("t"), 1); drop(sub); assert_eq!(bus.subscriber_count("t"), 0); } }