finance_query/streaming/
batch.rs

1//! Batched delivery for any streaming subscription.
2//!
3//! A consumer watching N symbols normally receives N separate messages per
4//! tick. [`Batched`] coalesces whatever arrived inside a time window into one
5//! `Vec`, cutting per-message overhead for wide watchlists. It is a plain
6//! `Stream` adapter, so it composes with every handle in this module instead
7//! of being a per-stream delivery mode.
8
9use std::pin::Pin;
10use std::task::{Context, Poll};
11use std::time::Duration;
12
13use futures::stream::Stream;
14use tokio_stream::adapters::ChunksTimeout;
15
16/// Default cap on items per batch when none is given.
17const DEFAULT_MAX_BATCH: usize = 512;
18
19/// A stream that yields `Vec<T>` batches collected over a time window.
20///
21/// A batch is emitted when the window since the batch's first item elapses, or
22/// as soon as `max_items` is reached — whichever comes first. Empty batches are
23/// never emitted, and any partial batch is flushed when the source ends.
24///
25/// Boxed rather than a plain newtype so the handle stays `Unpin` — callers
26/// `.next()` it directly, without pinning it first.
27pub struct Batched<S>
28where
29    S: Stream,
30{
31    inner: Pin<Box<ChunksTimeout<S>>>,
32}
33
34impl<S> Batched<S>
35where
36    S: Stream,
37{
38    /// Batch `inner` over `window`, emitting early once `max_items` accumulate.
39    pub fn new(inner: S, window: Duration, max_items: usize) -> Self {
40        // Scoped: a module-level import would collide with `futures::StreamExt`.
41        use tokio_stream::StreamExt as _;
42
43        Self {
44            // `chunks_timeout` panics on a zero cap; a zero-sized batch is
45            // meaningless, so clamp rather than propagate the panic.
46            inner: Box::pin(inner.chunks_timeout(max_items.max(1), window)),
47        }
48    }
49}
50
51impl<S> Stream for Batched<S>
52where
53    S: Stream,
54{
55    type Item = Vec<S::Item>;
56
57    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
58        self.inner.as_mut().poll_next(cx)
59    }
60}
61
62/// Adds [`batched`](StreamBatchExt::batched) to every `Stream`.
63///
64/// # Example
65///
66/// ```no_run
67/// use finance_query::streaming::{PriceStream, StreamBatchExt};
68/// use futures::StreamExt;
69/// use std::time::Duration;
70///
71/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
72/// let mut batches = PriceStream::subscribe(["AAPL", "NVDA", "TSLA"])
73///     .await?
74///     .batched(Duration::from_millis(250));
75///
76/// while let Some(batch) = batches.next().await {
77///     println!("{} updates in this window", batch.len());
78/// }
79/// # Ok(())
80/// # }
81/// ```
82pub trait StreamBatchExt: Stream + Sized + Unpin {
83    /// Coalesce items arriving within `window` into one `Vec` (max 512 items).
84    fn batched(self, window: Duration) -> Batched<Self> {
85        Batched::new(self, window, DEFAULT_MAX_BATCH)
86    }
87
88    /// Same as [`batched`](Self::batched) with an explicit per-batch cap.
89    fn batched_with_capacity(self, window: Duration, max_items: usize) -> Batched<Self> {
90        Batched::new(self, window, max_items)
91    }
92}
93
94impl<S> StreamBatchExt for S where S: Stream + Sized + Unpin {}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use futures::StreamExt;
100    use tokio::sync::mpsc;
101    use tokio_stream::wrappers::ReceiverStream;
102
103    #[tokio::test]
104    async fn window_coalesces_items_into_one_batch() {
105        let (tx, rx) = mpsc::channel(16);
106        let mut batched = ReceiverStream::new(rx).batched(Duration::from_millis(60));
107
108        for i in 0..5 {
109            tx.send(i).await.unwrap();
110        }
111
112        let batch = tokio::time::timeout(Duration::from_secs(2), batched.next())
113            .await
114            .expect("timed out")
115            .expect("stream ended");
116        assert_eq!(batch, vec![0, 1, 2, 3, 4]);
117        drop(tx);
118    }
119
120    #[tokio::test]
121    async fn max_items_flushes_before_the_window_elapses() {
122        let (tx, rx) = mpsc::channel(16);
123        let mut batched = ReceiverStream::new(rx).batched_with_capacity(Duration::from_secs(30), 2);
124
125        for i in 0..4 {
126            tx.send(i).await.unwrap();
127        }
128
129        // Would block for 30s if the size cap were not honored.
130        let batch = tokio::time::timeout(Duration::from_secs(2), batched.next())
131            .await
132            .expect("timed out")
133            .expect("stream ended");
134        assert_eq!(batch, vec![0, 1]);
135        drop(tx);
136    }
137
138    #[tokio::test]
139    async fn partial_batch_is_flushed_when_the_source_ends() {
140        let (tx, rx) = mpsc::channel(16);
141        let mut batched = ReceiverStream::new(rx).batched(Duration::from_secs(30));
142        tx.send(7).await.unwrap();
143        drop(tx);
144
145        let batch = tokio::time::timeout(Duration::from_secs(2), batched.next())
146            .await
147            .expect("timed out")
148            .expect("stream ended");
149        assert_eq!(batch, vec![7]);
150
151        let end = tokio::time::timeout(Duration::from_secs(2), batched.next())
152            .await
153            .expect("timed out");
154        assert!(end.is_none());
155    }
156
157    #[tokio::test]
158    async fn empty_batches_are_never_emitted() {
159        let (tx, rx) = mpsc::channel::<u8>(1);
160        let mut batched = ReceiverStream::new(rx).batched(Duration::from_millis(20));
161
162        let idle = tokio::time::timeout(Duration::from_millis(150), batched.next()).await;
163        assert!(idle.is_err(), "idle source must not emit empty batches");
164        drop(tx);
165    }
166}
finance_query/streaming/batch.rs

finance_query/streaming/
batch.rs
