leviath_runtime/embed/
stream.rs

1//! The in-process event stream an embedder consumes.
2
3use tokio_stream::wrappers::BroadcastStream;
4use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
5
6use crate::host::WorldEvent;
7
8/// An async stream of [`WorldEvent`]s from an embedded world.
9///
10/// Wraps the host's broadcast channel: a slow consumer that falls more than
11/// the channel capacity behind skips the missed events (with a warning)
12/// rather than erroring, and the stream ends (`None`) when the world shuts
13/// down. Implements [`futures_core::Stream`], and offers an inherent
14/// [`next`](Self::next) so the common loop needs no extra imports:
15///
16/// ```ignore
17/// while let Some(event) = events.next().await {
18///     // ...
19/// }
20/// ```
21pub struct EventStream {
22    inner: BroadcastStream<WorldEvent>,
23}
24
25impl EventStream {
26    pub(crate) fn new(rx: tokio::sync::broadcast::Receiver<WorldEvent>) -> Self {
27        Self {
28            inner: BroadcastStream::new(rx),
29        }
30    }
31
32    /// The next event, or `None` once the world has shut down.
33    pub async fn next(&mut self) -> Option<WorldEvent> {
34        use futures_core::Stream;
35        std::future::poll_fn(|cx| std::pin::Pin::new(&mut *self).poll_next(cx)).await
36    }
37}
38
39impl futures_core::Stream for EventStream {
40    type Item = WorldEvent;
41
42    fn poll_next(
43        mut self: std::pin::Pin<&mut Self>,
44        cx: &mut std::task::Context<'_>,
45    ) -> std::task::Poll<Option<Self::Item>> {
46        use std::task::Poll;
47        loop {
48            match std::pin::Pin::new(&mut self.inner).poll_next(cx) {
49                Poll::Ready(Some(Ok(event))) => return Poll::Ready(Some(event)),
50                Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n)))) => {
51                    tracing::warn!("event stream lagged, skipped {n} events");
52                    continue; // resubscribed at the live edge; keep polling
53                }
54                Poll::Ready(None) => return Poll::Ready(None),
55                Poll::Pending => return Poll::Pending,
56            }
57        }
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    fn log(line: &str) -> WorldEvent {
66        WorldEvent::Log {
67            run_id: "r".to_string(),
68            agent_id: "a".to_string(),
69            line: line.to_string(),
70        }
71    }
72
73    #[tokio::test]
74    async fn yields_events_then_ends_when_the_sender_drops() {
75        let (tx, rx) = tokio::sync::broadcast::channel(16);
76        let mut stream = EventStream::new(rx);
77        tx.send(log("one")).unwrap();
78        tx.send(log("two")).unwrap();
79        assert_eq!(stream.next().await, Some(log("one")));
80        assert_eq!(stream.next().await, Some(log("two")));
81        drop(tx);
82        assert_eq!(stream.next().await, None);
83    }
84
85    #[tokio::test]
86    async fn skips_over_a_lag_instead_of_erroring() {
87        // Capacity 1: sending twice before reading overwrites the first event,
88        // which surfaces as a Lagged error the stream must swallow.
89        let (tx, rx) = tokio::sync::broadcast::channel(1);
90        let mut stream = EventStream::new(rx);
91        tx.send(log("dropped")).unwrap();
92        tx.send(log("kept")).unwrap();
93        assert_eq!(stream.next().await, Some(log("kept")));
94        drop(tx);
95        assert_eq!(stream.next().await, None);
96    }
97}
leviath_runtime/embed/stream.rs

leviath_runtime/embed/
stream.rs
