Skip to main content

typed_eventbus/
local.rs

1use bytes::Bytes;
2use dashmap::DashMap;
3use std::pin::Pin;
4use std::sync::Arc;
5use tokio::sync::mpsc;
6use tokio::task::JoinSet;
7
8use crate::{EventError, EventStream, Handler};
9
10pub type BoxFuture<'a, T> = Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
11
12type Msg = Arc<(String, Bytes)>; // Arc so clones are cheap
13type Tx = mpsc::Sender<Msg>;
14
15pub struct LocalEventStream {
16    // subject -> list of senders. DashMap shard locks only on insert/remove.
17    subs: DashMap<String, Vec<Tx>>,
18    _tasks: tokio::sync::Mutex<JoinSet<()>>,
19    channel_capacity: usize,
20}
21
22impl LocalEventStream {
23    pub fn new(channel_capacity: usize) -> Self {
24        Self {
25            subs: DashMap::new(),
26            _tasks: tokio::sync::Mutex::new(JoinSet::new()),
27            channel_capacity,
28        }
29    }
30
31    // 8192 = ~8MB per subscriber if 1KB avg msg. Tune based on RAM.
32    pub fn reliable() -> Self {
33        Self::new(8192)
34    }
35}
36
37impl EventStream for LocalEventStream {
38    fn publish<'a>(
39        &'a self,
40        subject: String,
41        payload: Vec<u8>,
42    ) -> BoxFuture<'a, Result<(), EventError>> {
43        Box::pin(async move {
44            let msg = Arc::new((subject.clone(), Bytes::from(payload.clone())));
45
46            // Fast path: no global lock. Only shard lock for this subject.
47            let Some(senders) = self.subs.get(&subject) else {
48                return Ok(()); // No subscribers = drop. Not an error.
49            };
50
51            // Fan-out. If any queue is full, await = backpressure.
52            // This is the reliability guarantee: publish won't return until all
53            // subscribers have space in their queue.
54            for tx in senders.iter() {
55                if tx.send(msg.clone()).await.is_err() {
56                    // Receiver dropped. Remove it lazily on next publish or subscribe.
57                    // For now just ignore to avoid write lock during publish.
58                }
59            }
60            tracing::info!(
61                "published {subject} : message = {}",
62                String::from_utf8(payload).unwrap_or("invalid utf-8".to_string())
63            );
64            Ok(())
65        })
66    }
67
68    fn subscribe<'a>(
69        &'a self,
70        subject: String,
71        handler: Arc<dyn Handler>,
72    ) -> BoxFuture<'a, Result<(), EventError>> {
73        Box::pin(async move {
74            let (tx, mut rx) = mpsc::channel::<Msg>(self.channel_capacity);
75
76            // Insert sender. Need write lock on shard, but only briefly.
77            self.subs
78                .entry(subject.clone())
79                .or_insert_with(Vec::new)
80                .push(tx);
81
82            let mut tasks = self._tasks.lock().await;
83            tasks.spawn(async move {
84                // Ordered, lossless consumption. If handler is slow, rx will fill
85                // and publishers will await on send(). That's backpressure.
86                while let Some(msg) = rx.recv().await {
87                    let (subj, payload) = &*msg; // Arc deref
88                    handler.handle(subj.clone(), payload.to_vec()).await;
89                }
90            });
91
92            Ok(())
93        })
94    }
95}
96
97impl Drop for LocalEventStream {
98    fn drop(&mut self) {
99        // Abort all handler tasks when bus drops
100        if let Ok(mut tasks) = self._tasks.try_lock() {
101            tasks.abort_all();
102        }
103    }
104}