Skip to main content

visual_cortex/
stream.rs

1use std::collections::HashSet;
2
3use tokio::sync::broadcast;
4
5use crate::event::Event;
6
7/// An async stream of events for one or more watchers.
8/// All session events flow through one broadcast channel; each `EventStream`
9/// filters to its subscribed watcher names, so `merge` is just a name union.
10#[derive(Debug)]
11pub struct EventStream {
12    rx: broadcast::Receiver<Event>,
13    names: HashSet<String>,
14}
15
16impl EventStream {
17    pub(crate) fn new(rx: broadcast::Receiver<Event>, names: HashSet<String>) -> Self {
18        Self { rx, names }
19    }
20
21    /// Next event, or `None` when the session has shut down.
22    /// If this consumer lags behind the broadcast buffer, missed events are
23    /// skipped silently and the stream continues with the newest.
24    pub async fn next(&mut self) -> Option<Event> {
25        loop {
26            match self.rx.recv().await {
27                Ok(ev) if ev.kind.is_session_level() || self.names.contains(&*ev.watcher) => {
28                    return Some(ev)
29                }
30                Ok(_) => continue,
31                Err(broadcast::error::RecvError::Lagged(skipped)) => {
32                    tracing::warn!(skipped, "EventStream lagged; dropping missed events");
33                    continue;
34                }
35                Err(broadcast::error::RecvError::Closed) => return None,
36            }
37        }
38    }
39
40    /// Combine two streams into one that yields events from both.
41    ///
42    /// Both streams must originate from the same `Session`: the merged stream
43    /// keeps only `self`'s receiver and unions the name filters, which is
44    /// correct because all of a session's events flow through one broadcast
45    /// channel. Merging streams from different sessions would silently drop
46    /// the other session's events.
47    pub fn merge(mut self, other: EventStream) -> EventStream {
48        self.names.extend(other.names);
49        self
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use crate::event::EventKind;
57    use std::sync::Arc;
58    use std::time::SystemTime;
59    use visual_cortex_vision::DetectorOutput;
60
61    fn ev(name: &str) -> Event {
62        Event {
63            watcher: Arc::from(name),
64            timestamp: SystemTime::now(),
65            kind: EventKind::Matched {
66                output: DetectorOutput::Bool(true),
67            },
68        }
69    }
70
71    #[tokio::test]
72    async fn filters_to_subscribed_watcher_names() {
73        let (tx, rx) = broadcast::channel(16);
74        let mut stream = EventStream::new(rx, HashSet::from(["a".to_string()]));
75        tx.send(ev("b")).unwrap();
76        tx.send(ev("a")).unwrap();
77        let got = stream.next().await.unwrap();
78        assert_eq!(&*got.watcher, "a");
79    }
80
81    #[tokio::test]
82    async fn merge_unions_names() {
83        let (tx, rx) = broadcast::channel(16);
84        let a = EventStream::new(rx, HashSet::from(["a".to_string()]));
85        let b = EventStream::new(tx.subscribe(), HashSet::from(["b".to_string()]));
86        let mut merged = a.merge(b);
87        tx.send(ev("b")).unwrap();
88        tx.send(ev("a")).unwrap();
89        assert_eq!(&*merged.next().await.unwrap().watcher, "b");
90        assert_eq!(&*merged.next().await.unwrap().watcher, "a");
91    }
92
93    #[tokio::test]
94    async fn session_level_events_bypass_name_filter() {
95        let (tx, rx) = broadcast::channel(16);
96        // Stream only subscribes to "a", but a session-level event for the
97        // reserved "$session" name must still come through.
98        let mut stream = EventStream::new(rx, HashSet::from(["a".to_string()]));
99        tx.send(Event {
100            watcher: Arc::from("$session"),
101            timestamp: SystemTime::now(),
102            kind: EventKind::TargetLost,
103        })
104        .unwrap();
105        let got = stream.next().await.unwrap();
106        assert!(matches!(got.kind, EventKind::TargetLost));
107    }
108
109    #[tokio::test]
110    async fn returns_none_when_sender_dropped() {
111        let (tx, rx) = broadcast::channel(16);
112        let mut stream = EventStream::new(rx, HashSet::from(["a".to_string()]));
113        drop(tx);
114        assert!(stream.next().await.is_none());
115    }
116}