#![deny(missing_docs)]
use std::{
collections::{BTreeMap, VecDeque},
sync::mpsc,
sync::{Arc, Mutex, MutexGuard, Weak},
};
#[derive(Clone, Debug)]
struct SubscriberBuffer<T> {
events: SubscriberEvents<T>,
}
#[derive(Clone, Debug)]
enum SubscriberEvents<T> {
Unbounded(Vec<T>),
Bounded {
events: VecDeque<T>,
capacity: usize,
},
}
impl<T> Default for SubscriberBuffer<T> {
fn default() -> Self {
Self {
events: SubscriberEvents::Unbounded(Vec::new()),
}
}
}
impl<T> SubscriberBuffer<T> {
fn push(&mut self, event: T) {
match &mut self.events {
SubscriberEvents::Unbounded(events) => events.push(event),
SubscriberEvents::Bounded { events, capacity } => {
if *capacity == 0 {
return;
}
if events.len() == *capacity {
events.pop_front();
}
events.push_back(event);
}
}
}
fn drain(&mut self) -> Vec<T> {
match &mut self.events {
SubscriberEvents::Unbounded(events) => std::mem::take(events),
SubscriberEvents::Bounded { events, .. } => std::mem::take(events).into(),
}
}
}
type SubscriberQueue<T> = Arc<Mutex<SubscriberBuffer<T>>>;
type SubscriberHandle<T> = Weak<Mutex<SubscriberBuffer<T>>>;
type SubscriberList<T> = Arc<Mutex<Vec<SubscriberHandle<T>>>>;
struct LiveSubscribers<T> {
first: Option<SubscriberQueue<T>>,
additional: Vec<SubscriberQueue<T>>,
}
impl<T> LiveSubscribers<T> {
fn new() -> Self {
Self {
first: None,
additional: Vec::new(),
}
}
fn push(&mut self, subscriber: SubscriberQueue<T>) {
if self.first.is_none() {
self.first = Some(subscriber);
} else {
self.additional.push(subscriber);
}
}
}
#[derive(Clone, Debug)]
pub struct EventBus<T> {
subscribers: SubscriberList<T>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ObservedEventContext {
operation_id: String,
event_name: String,
attributes: Arc<BTreeMap<String, String>>,
}
impl ObservedEventContext {
pub fn new(operation_id: impl Into<String>, event_name: impl Into<String>) -> Self {
Self {
operation_id: operation_id.into(),
event_name: event_name.into(),
attributes: Arc::new(BTreeMap::new()),
}
}
pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
Arc::make_mut(&mut self.attributes).insert(key.into(), value.into());
self
}
pub fn operation_id(&self) -> &str {
&self.operation_id
}
pub fn event_name(&self) -> &str {
&self.event_name
}
pub fn attributes(&self) -> &BTreeMap<String, String> {
&self.attributes
}
}
pub trait EventObserver<T>: Clone + Send + Sync + 'static
where
T: Clone + Send + Sync + 'static,
{
fn on_event_published(&self, context: &ObservedEventContext);
}
impl<T> EventObserver<T> for ()
where
T: Clone + Send + Sync + 'static,
{
fn on_event_published(&self, _context: &ObservedEventContext) {}
}
#[derive(Clone)]
pub struct EventObserverChannel {
sender: mpsc::Sender<ObservedEventContext>,
}
impl EventObserverChannel {
pub fn new(sender: mpsc::Sender<ObservedEventContext>) -> Self {
Self { sender }
}
}
impl<T> EventObserver<T> for EventObserverChannel
where
T: Clone + Send + Sync + 'static,
{
fn on_event_published(&self, context: &ObservedEventContext) {
let _ = self.sender.send(context.clone());
}
}
pub fn event_observer_channel() -> (EventObserverChannel, mpsc::Receiver<ObservedEventContext>) {
let (sender, receiver) = mpsc::channel();
(EventObserverChannel::new(sender), receiver)
}
#[derive(Clone)]
pub struct ObservedEventBus<T, O = ()>
where
T: Clone + Send + Sync + 'static,
O: EventObserver<T>,
{
bus: EventBus<T>,
observer: O,
attributes: Arc<BTreeMap<String, String>>,
operation_id_generator: Arc<dyn Fn() -> String + Send + Sync>,
}
impl<T, O> ObservedEventBus<T, O>
where
T: Clone + Send + Sync + 'static,
O: EventObserver<T>,
{
pub fn new(bus: EventBus<T>, observer: O) -> Self {
Self {
bus,
observer,
attributes: Arc::new(BTreeMap::new()),
operation_id_generator: Arc::new(|| uuid::Uuid::new_v4().to_string()),
}
}
pub fn context(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
Arc::make_mut(&mut self.attributes).insert(key.into(), value.into());
self
}
pub fn operation_id_generator(
mut self,
generator: impl Fn() -> String + Send + Sync + 'static,
) -> Self {
self.operation_id_generator = Arc::new(generator);
self
}
pub fn publish_named(&self, event_name: impl Into<String>, event: T) {
let event_name = event_name.into();
let context = self.context_for(event_name);
let span = tracing::info_span!(
"event.publish",
event.name = %context.event_name(),
event.operation_id = %context.operation_id()
);
let _entered = span.enter();
self.bus.publish(event);
self.observer.on_event_published(&context);
}
pub fn bus(&self) -> &EventBus<T> {
&self.bus
}
fn context_for(&self, event_name: String) -> ObservedEventContext {
ObservedEventContext {
operation_id: (self.operation_id_generator)(),
event_name,
attributes: Arc::clone(&self.attributes),
}
}
}
impl<T> EventBus<T>
where
T: Clone,
{
pub fn new() -> Self {
Self {
subscribers: Arc::new(Mutex::new(Vec::new())),
}
}
pub fn subscribe(&self) -> EventSubscriber<T> {
self.subscribe_with_buffer(SubscriberBuffer::default())
}
pub fn subscribe_with_capacity(&self, capacity: usize) -> EventSubscriber<T> {
self.subscribe_with_buffer(SubscriberBuffer {
events: SubscriberEvents::Bounded {
events: VecDeque::new(),
capacity,
},
})
}
fn subscribe_with_buffer(&self, buffer: SubscriberBuffer<T>) -> EventSubscriber<T> {
let queue = Arc::new(Mutex::new(buffer));
lock_unpoisoned(&self.subscribers).push(Arc::downgrade(&queue));
EventSubscriber { queue }
}
pub fn publish(&self, event: T) {
let LiveSubscribers {
first,
mut additional,
} = self.live_subscribers();
let Some(first) = first else {
return;
};
let Some(last) = additional.pop() else {
lock_unpoisoned(&first).push(event);
return;
};
lock_unpoisoned(&first).push(event.clone());
for subscriber in additional {
lock_unpoisoned(&subscriber).push(event.clone());
}
lock_unpoisoned(&last).push(event);
}
pub fn observed<O>(self, observer: O) -> ObservedEventBus<T, O>
where
T: Send + Sync + 'static,
O: EventObserver<T>,
{
ObservedEventBus::new(self, observer)
}
pub fn subscriber_count(&self) -> usize {
let mut subscribers = lock_unpoisoned(&self.subscribers);
subscribers.retain(|subscriber| subscriber.upgrade().is_some());
subscribers.len()
}
fn live_subscribers(&self) -> LiveSubscribers<T> {
let mut subscribers = lock_unpoisoned(&self.subscribers);
let mut live = LiveSubscribers::new();
subscribers.retain(|subscriber| {
if let Some(queue) = subscriber.upgrade() {
live.push(queue);
true
} else {
false
}
});
live
}
}
impl<T> Default for EventBus<T>
where
T: Clone,
{
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Debug)]
pub struct EventSubscriber<T> {
queue: SubscriberQueue<T>,
}
impl<T> EventSubscriber<T> {
pub fn drain(&self) -> Vec<T> {
lock_unpoisoned(&self.queue).drain()
}
}
fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(|poisoned| {
tracing::warn!("event bus mutex poisoned; recovering inner state");
poisoned.into_inner()
})
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, thread};
use super::*;
use tracing::Level;
use tracing_subscriber::{Layer, fmt::MakeWriter, layer::SubscriberExt};
#[derive(Clone, Default)]
struct SharedLogWriter {
output: Arc<Mutex<Vec<u8>>>,
}
impl SharedLogWriter {
fn contents(&self) -> String {
String::from_utf8(self.output.lock().unwrap().clone()).unwrap()
}
fn clear(&self) {
self.output.lock().unwrap().clear();
}
}
impl<'writer> MakeWriter<'writer> for SharedLogWriter {
type Writer = SharedLogGuard;
fn make_writer(&'writer self) -> Self::Writer {
SharedLogGuard {
output: Arc::clone(&self.output),
}
}
}
struct SharedLogGuard {
output: Arc<Mutex<Vec<u8>>>,
}
impl std::io::Write for SharedLogGuard {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.output.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct UserCreated(u64);
#[test]
fn event_bus_recovers_from_poisoned_subscriber_list() {
let bus = EventBus::<UserCreated>::new();
let subscribers = Arc::clone(&bus.subscribers);
let panic = thread::spawn(move || {
let _subscribers = subscribers.lock().unwrap();
panic!("poison subscriber list");
});
assert!(panic.join().is_err());
let subscriber = bus.subscribe();
bus.publish(UserCreated(42));
assert_eq!(subscriber.drain(), vec![UserCreated(42)]);
}
#[test]
fn event_bus_warns_when_recovering_from_poisoned_subscriber_list() {
let bus = EventBus::<UserCreated>::new();
let subscribers = Arc::clone(&bus.subscribers);
let panic = thread::spawn(move || {
let _subscribers = subscribers.lock().unwrap();
panic!("poison subscriber list");
});
assert!(panic.join().is_err());
let writer = SharedLogWriter::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.with_writer(writer.clone())
.with_ansi(false)
.with_target(false)
.with_filter(tracing_subscriber::filter::LevelFilter::from_level(
Level::WARN,
)),
);
tracing::subscriber::with_default(subscriber, || {
for _ in 0..16 {
writer.clear();
tracing_core::callsite::rebuild_interest_cache();
let _subscriber = bus.subscribe();
let logs = writer.contents();
if logs.contains("event bus mutex poisoned") {
return;
}
std::thread::yield_now();
}
});
let logs = writer.contents();
assert!(logs.contains("event bus mutex poisoned"), "{logs}");
}
#[test]
fn event_bus_recovers_from_poisoned_subscriber_queue() {
let bus = EventBus::<UserCreated>::new();
let subscriber = bus.subscribe();
let queue = Arc::clone(&subscriber.queue);
let panic = thread::spawn(move || {
let _queue = queue.lock().unwrap();
panic!("poison subscriber queue");
});
assert!(panic.join().is_err());
bus.publish(UserCreated(42));
assert_eq!(subscriber.drain(), vec![UserCreated(42)]);
}
#[test]
fn bounded_subscriber_drops_oldest_events_beyond_capacity() {
let bus = EventBus::<UserCreated>::new();
let bounded = bus.subscribe_with_capacity(2);
bus.publish(UserCreated(1));
bus.publish(UserCreated(2));
bus.publish(UserCreated(3));
assert_eq!(bounded.drain(), vec![UserCreated(2), UserCreated(3)]);
bus.publish(UserCreated(4));
bus.publish(UserCreated(5));
bus.publish(UserCreated(6));
assert_eq!(bounded.drain(), vec![UserCreated(5), UserCreated(6)]);
}
#[test]
fn zero_capacity_subscriber_never_retains_events() {
let bus = EventBus::<UserCreated>::new();
let subscriber = bus.subscribe_with_capacity(0);
bus.publish(UserCreated(1));
bus.publish(UserCreated(2));
assert!(subscriber.drain().is_empty());
}
#[test]
fn unbounded_subscriber_keeps_all_events_by_default() {
let bus = EventBus::<UserCreated>::new();
let subscriber = bus.subscribe();
for id in 1..=50u64 {
bus.publish(UserCreated(id));
}
let drained: Vec<u64> = subscriber
.drain()
.into_iter()
.map(|event| event.0)
.collect();
assert_eq!(drained, (1..=50).collect::<Vec<_>>());
}
#[test]
fn observed_context_shares_configured_attributes_until_enriched() {
let observed = EventBus::<UserCreated>::new()
.observed(())
.operation_id_generator(|| "event-run".to_owned())
.context("service", "users-api");
let enriched_observed = observed.clone().context("region", "sa-east-1");
assert!(!Arc::ptr_eq(
&enriched_observed.attributes,
&observed.attributes
));
assert!(!observed.attributes.contains_key("region"));
assert_eq!(
enriched_observed.attributes.get("region").unwrap(),
"sa-east-1"
);
let context = observed.context_for("user.created".to_owned());
assert!(Arc::ptr_eq(&context.attributes, &observed.attributes));
let enriched = context.clone().with_attribute("request_id", "request-42");
assert!(!Arc::ptr_eq(&enriched.attributes, &context.attributes));
assert_eq!(context.attributes().get("service").unwrap(), "users-api");
assert!(!context.attributes().contains_key("request_id"));
assert_eq!(
enriched.attributes().get("request_id").unwrap(),
"request-42"
);
}
}