Skip to main content

prospero_core/
bus.rs

1//! Live event distribution behind a trait.
2//!
3//! Standalone uses [`InProcessBus`] (a tokio broadcast channel); the clustered
4//! [`crate::DistributedBus`] (Postgres `LISTEN/NOTIFY`) drops in behind the same
5//! trait — see the topology design spec §3.2. Both expose the live tail as a
6//! per-stream [`BusSubscription`]; consumers dedup the history/live overlap on
7//! `seq`.
8
9use std::pin::Pin;
10
11use tokio::sync::broadcast;
12use tokio_stream::Stream;
13
14use crate::event::{FleetEvent, stream_key_for};
15
16/// One item from a per-stream live subscription (transport-agnostic).
17#[derive(Debug, Clone, PartialEq)]
18pub enum BusEvent {
19    /// A live event on the subscribed stream.
20    Event(FleetEvent),
21    /// `skipped` events were dropped for a slow local subscriber; the consumer
22    /// must self-heal by replaying from the durable store. Only [`InProcessBus`]
23    /// emits this — the clustered bus reads the store on every doorbell, so it
24    /// cannot lag.
25    Lagged(u64),
26}
27
28/// A live, per-stream subscription: an ordered stream of [`BusEvent`]s for one
29/// stream key. Ends (`None`) when the bus is gone.
30pub type BusSubscription = Pin<Box<dyn Stream<Item = BusEvent> + Send>>;
31
32/// Publishes events to live subscribers.
33pub trait EventBus: Send + Sync {
34    /// Fan an event out to current subscribers. Never blocks on slow/absent
35    /// receivers (delivery is best-effort relative to the durable store).
36    fn publish(&self, event: FleetEvent);
37
38    /// A live subscription to one stream's events. The returned stream yields
39    /// only events whose stream key equals `stream_key`; consumers dedup the
40    /// initial-history/live overlap on `seq`.
41    fn subscribe(&self, stream_key: &str) -> BusSubscription;
42
43    /// A live subscription to EVERY stream's events, unfiltered. Needed by
44    /// fleet-wide watchers (e.g. `FleetManager::watch_changes`) that can't name
45    /// a stream key in advance — a brand-new agent's own id keys its
46    /// `AgentDiscovered` event, so no one can pre-subscribe to it by key.
47    fn subscribe_all(&self) -> BusSubscription;
48}
49
50/// In-process broadcast bus — the standalone implementation.
51pub struct InProcessBus {
52    tx: broadcast::Sender<FleetEvent>,
53}
54
55impl InProcessBus {
56    /// A bus buffering up to `capacity` events for slow subscribers.
57    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        // No subscribers is fine; ignore the send error.
66        let _ = self.tx.send(event);
67    }
68
69    fn subscribe(&self, stream_key: &str) -> BusSubscription {
70        // Register the broadcast receiver EAGERLY (synchronously, here) so it
71        // captures events from this point — before the caller reads initial
72        // history — even though the stream body below is polled lazily.
73        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, // an event on a different stream
82                    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        // Same eager-registration discipline as `subscribe`, just without the
93        // per-key filter.
94        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")); // other stream — filtered out
138        bus.publish(ev_for(2, "a")); // our stream — delivered
139        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")); // must not panic
146    }
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}