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)>; type Tx = mpsc::Sender<Msg>;
14
15pub struct LocalEventStream {
16 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 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 let Some(senders) = self.subs.get(&subject) else {
48 return Ok(()); };
50
51 for tx in senders.iter() {
55 if tx.send(msg.clone()).await.is_err() {
56 }
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 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 while let Some(msg) = rx.recv().await {
87 let (subj, payload) = &*msg; 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 if let Ok(mut tasks) = self._tasks.try_lock() {
101 tasks.abort_all();
102 }
103 }
104}