1use std::collections::HashSet;
2
3use tokio::sync::broadcast;
4
5use crate::event::Event;
6
7#[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 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 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 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}