1use std::pin::Pin;
10
11use tokio::sync::broadcast;
12use tokio_stream::Stream;
13
14use crate::event::{FleetEvent, stream_key_for};
15
16#[derive(Debug, Clone, PartialEq)]
18pub enum BusEvent {
19 Event(FleetEvent),
21 Lagged(u64),
26}
27
28pub type BusSubscription = Pin<Box<dyn Stream<Item = BusEvent> + Send>>;
31
32pub trait EventBus: Send + Sync {
34 fn publish(&self, event: FleetEvent);
37
38 fn subscribe(&self, stream_key: &str) -> BusSubscription;
42
43 fn subscribe_all(&self) -> BusSubscription;
48}
49
50pub struct InProcessBus {
52 tx: broadcast::Sender<FleetEvent>,
53}
54
55impl InProcessBus {
56 pub fn new(capacity: usize) -> Self {
58 let (tx, _rx) = broadcast::channel(capacity);
59 Self { tx }
60 }
61}
62
63impl EventBus for InProcessBus {
64 fn publish(&self, event: FleetEvent) {
65 let _ = self.tx.send(event);
67 }
68
69 fn subscribe(&self, stream_key: &str) -> BusSubscription {
70 let mut rx = self.tx.subscribe();
74 let key = stream_key.to_string();
75 Box::pin(async_stream::stream! {
76 loop {
77 match rx.recv().await {
78 Ok(ev) if stream_key_for(&ev.repo, &ev.agent_id) == key => {
79 yield BusEvent::Event(ev);
80 }
81 Ok(_) => continue, Err(broadcast::error::RecvError::Lagged(n)) => {
83 yield BusEvent::Lagged(n);
84 }
85 Err(broadcast::error::RecvError::Closed) => break,
86 }
87 }
88 })
89 }
90
91 fn subscribe_all(&self) -> BusSubscription {
92 let mut rx = self.tx.subscribe();
95 Box::pin(async_stream::stream! {
96 loop {
97 match rx.recv().await {
98 Ok(ev) => yield BusEvent::Event(ev),
99 Err(broadcast::error::RecvError::Lagged(n)) => {
100 yield BusEvent::Lagged(n);
101 }
102 Err(broadcast::error::RecvError::Closed) => break,
103 }
104 }
105 })
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112 use crate::event::EventKind;
113 use tokio_stream::StreamExt;
114
115 fn ev_for(seq: u64, agent: &str) -> FleetEvent {
116 FleetEvent {
117 seq,
118 ts: "t".into(),
119 repo: "r".into(),
120 agent_id: agent.into(),
121 kind: EventKind::AgentSpawned,
122 }
123 }
124
125 #[tokio::test]
126 async fn publish_reaches_a_subscriber_on_its_stream() {
127 let bus = InProcessBus::new(8);
128 let mut sub = bus.subscribe("a");
129 bus.publish(ev_for(1, "a"));
130 assert_eq!(sub.next().await, Some(BusEvent::Event(ev_for(1, "a"))));
131 }
132
133 #[tokio::test]
134 async fn subscriber_only_sees_its_own_stream() {
135 let bus = InProcessBus::new(8);
136 let mut sub = bus.subscribe("a");
137 bus.publish(ev_for(1, "b")); bus.publish(ev_for(2, "a")); assert_eq!(sub.next().await, Some(BusEvent::Event(ev_for(2, "a"))));
140 }
141
142 #[tokio::test]
143 async fn publish_with_no_subscriber_is_a_noop() {
144 let bus = InProcessBus::new(8);
145 bus.publish(ev_for(1, "a")); }
147
148 #[tokio::test]
149 async fn subscribe_all_sees_every_stream_unfiltered() {
150 let bus = InProcessBus::new(8);
151 let mut sub = bus.subscribe_all();
152 bus.publish(ev_for(1, "a"));
153 bus.publish(ev_for(2, "b"));
154 assert_eq!(sub.next().await, Some(BusEvent::Event(ev_for(1, "a"))));
155 assert_eq!(sub.next().await, Some(BusEvent::Event(ev_for(2, "b"))));
156 }
157}