visual-cortex 0.2.0

use std::collections::HashSet;

use tokio::sync::broadcast;

use crate::event::Event;

/// An async stream of events for one or more watchers.
/// All session events flow through one broadcast channel; each `EventStream`
/// filters to its subscribed watcher names, so `merge` is just a name union.
#[derive(Debug)]
pub struct EventStream {
    rx: broadcast::Receiver<Event>,
    names: HashSet<String>,
}

impl EventStream {
    pub(crate) fn new(rx: broadcast::Receiver<Event>, names: HashSet<String>) -> Self {
        Self { rx, names }
    }

    /// Next event, or `None` when the session has shut down.
    /// If this consumer lags behind the broadcast buffer, missed events are
    /// skipped silently and the stream continues with the newest.
    pub async fn next(&mut self) -> Option<Event> {
        loop {
            match self.rx.recv().await {
                Ok(ev) if ev.kind.is_session_level() || self.names.contains(&*ev.watcher) => {
                    return Some(ev)
                }
                Ok(_) => continue,
                Err(broadcast::error::RecvError::Lagged(skipped)) => {
                    tracing::warn!(skipped, "EventStream lagged; dropping missed events");
                    continue;
                }
                Err(broadcast::error::RecvError::Closed) => return None,
            }
        }
    }

    /// Combine two streams into one that yields events from both.
    ///
    /// Both streams must originate from the same `Session`: the merged stream
    /// keeps only `self`'s receiver and unions the name filters, which is
    /// correct because all of a session's events flow through one broadcast
    /// channel. Merging streams from different sessions would silently drop
    /// the other session's events.
    pub fn merge(mut self, other: EventStream) -> EventStream {
        self.names.extend(other.names);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::event::EventKind;
    use std::sync::Arc;
    use std::time::SystemTime;
    use visual_cortex_vision::DetectorOutput;

    fn ev(name: &str) -> Event {
        Event {
            watcher: Arc::from(name),
            timestamp: SystemTime::now(),
            kind: EventKind::Matched {
                output: DetectorOutput::Bool(true),
            },
        }
    }

    #[tokio::test]
    async fn filters_to_subscribed_watcher_names() {
        let (tx, rx) = broadcast::channel(16);
        let mut stream = EventStream::new(rx, HashSet::from(["a".to_string()]));
        tx.send(ev("b")).unwrap();
        tx.send(ev("a")).unwrap();
        let got = stream.next().await.unwrap();
        assert_eq!(&*got.watcher, "a");
    }

    #[tokio::test]
    async fn merge_unions_names() {
        let (tx, rx) = broadcast::channel(16);
        let a = EventStream::new(rx, HashSet::from(["a".to_string()]));
        let b = EventStream::new(tx.subscribe(), HashSet::from(["b".to_string()]));
        let mut merged = a.merge(b);
        tx.send(ev("b")).unwrap();
        tx.send(ev("a")).unwrap();
        assert_eq!(&*merged.next().await.unwrap().watcher, "b");
        assert_eq!(&*merged.next().await.unwrap().watcher, "a");
    }

    #[tokio::test]
    async fn session_level_events_bypass_name_filter() {
        let (tx, rx) = broadcast::channel(16);
        // Stream only subscribes to "a", but a session-level event for the
        // reserved "$session" name must still come through.
        let mut stream = EventStream::new(rx, HashSet::from(["a".to_string()]));
        tx.send(Event {
            watcher: Arc::from("$session"),
            timestamp: SystemTime::now(),
            kind: EventKind::TargetLost,
        })
        .unwrap();
        let got = stream.next().await.unwrap();
        assert!(matches!(got.kind, EventKind::TargetLost));
    }

    #[tokio::test]
    async fn returns_none_when_sender_dropped() {
        let (tx, rx) = broadcast::channel(16);
        let mut stream = EventStream::new(rx, HashSet::from(["a".to_string()]));
        drop(tx);
        assert!(stream.next().await.is_none());
    }
}