Skip to main content

rustrade_data/streams/
blocking.rs

1use futures::Stream;
2use std::any::Any;
3use std::future::Future;
4use std::panic::AssertUnwindSafe;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7use tokio::task::JoinHandle;
8use tracing::warn;
9
10/// Chunk size [`stream_blocking_iter`] is documented for, and a sensible default for a decoder
11/// feeding a backtest.
12///
13/// Large enough that a blocking task is started once per thousand items rather than per item, small
14/// enough that the buffered events are an accounting rounding error next to the dataset. At ~160
15/// bytes per market event, the two chunks one stream can hold are ~320 KiB. That is per stream: a
16/// merge of N streams holds up to N times as much, see [`stream_blocking_iter`].
17pub const DEFAULT_BLOCKING_CHUNK_SIZE: usize = 1024;
18
19/// Bridge a **blocking**, fallible iterator into a bounded [`Stream`], decoding on Tokio's blocking
20/// pool.
21///
22/// Historical data usually arrives as something that blocks: a Parquet artifact, a compressed
23/// archive, a CSV on a slow disk. Two things then go wrong if such an iterator is simply wrapped in
24/// [`futures::stream::iter`]:
25///
26/// 1. **It stalls a runtime worker for the whole decode.** Nothing else scheduled on that worker —
27///    the engine, another instrument's decode, a heartbeat — makes progress until the file ends.
28/// 2. **It never yields `Pending`, so nothing downstream can pace it.** A consumer that forwards
29///    into an unbounded channel (the engine feed is one) therefore takes the *entire* artifact into
30///    memory before processing the first event, which is the opposite of the streaming the caller
31///    asked for. Laziness alone does not bound memory; the producer has to be made to wait.
32///
33/// This helper fixes both. The decode runs in chunks of `chunk_size` items, each on a
34/// [`tokio::task::spawn_blocking`] task that decodes the chunk, returns it, and gives its thread
35/// back. When a chunk arrives the stream starts the next one, so decoding overlaps consumption, and
36/// it starts no further chunk until the consumer has drained the one it holds. That is the
37/// back-pressure: the decoder runs at most two chunks, `2 × chunk_size` items, ahead of whatever
38/// polls this stream, whatever the dataset's size. `init` runs in the first chunk's task, so opening
39/// the source is off the runtime's workers too, as it usually needs to be.
40///
41/// **That bound is local to this hand-off.** It constrains nothing further downstream: a caller that
42/// forwards this stream into another *unbounded* queue — a backtest engine feed is exactly that shape
43/// — reintroduces unbounded growth there whenever the far end is the slower side. What this helper
44/// guarantees is "the decoder will not get more than two chunks ahead of its own consumer", not an
45/// end-to-end memory bound for whatever pipeline it is embedded in.
46///
47/// The bound is also **per stream**. A [`merge_time_sorted`](super::merge::merge_time_sorted) of N
48/// of these polls every input until each has delivered a chunk, so it holds up to
49/// `2 × N × chunk_size` decoded items. For a merge of thousands of inputs, lower `chunk_size`.
50///
51/// # No thread is held while waiting, so any number of these can be merged
52/// A blocking thread is held only while a chunk is being decoded, never while the stream waits for
53/// its consumer. Building a stream starts nothing: the first chunk is started on first poll. Any
54/// number of these can therefore be driven together, for example by
55/// [`merge_time_sorted`](super::merge::merge_time_sorted), which cannot emit until every input has
56/// buffered an event. Past Tokio's `max_blocking_threads` (512 by default), chunks queue for a thread
57/// and every one of them still finishes, so the merge is slower but always progresses.
58///
59/// # Failure model
60/// The item type is `Result`, because a source that reads incrementally can fail *after* it opened
61/// successfully. An `init` that fails yields exactly one `Err` and ends the stream, so a caller
62/// handles open and mid-stream failures on one path.
63///
64/// A **panic** is not an `Err` and is not reported as one: `Error` is an unconstrained generic, so
65/// there is no value this helper could construct to represent one. It is re-raised on the consumer
66/// instead — see `# Panics`. What it must never do is end the stream quietly, which would leave a
67/// truncated decode indistinguishable from a source that finished, and drive a backtest to a
68/// normal-looking summary over partial data.
69///
70/// A [`tokio::task::JoinError`] that is *not* a panic ends the stream with `None`, after a `warn!`.
71/// This one's handle is never exposed, so nothing aborts the task, and the only way to reach that is
72/// runtime shutdown cancelling a chunk still queued for a thread. The consumer is normally being torn
73/// down for the same reason; the `warn!` says the decode was cut short in case it is not.
74///
75/// # Cancellation
76/// Dropping the returned stream starts no further chunk. The chunk already started is not
77/// cancelled, whether it is being decoded or still waiting for a thread: it runs to its end, at most
78/// `chunk_size` items, and its result is discarded with the iterator. Letting it run keeps the
79/// iterator, and whatever its `Drop` does, on the blocking pool.
80///
81/// # Panics
82/// Panics if `chunk_size` is zero (a chunk must deliver at least one item), and, when first polled,
83/// outside a Tokio runtime, like any [`tokio::task::spawn_blocking`] caller. `chunk_size` bounds the
84/// items per chunk and is not allocated up front, so a large value such as `usize::MAX` is safe.
85///
86/// **A panic in `init` or in the iterator is re-raised on the task that polls this stream**, with
87/// the original payload, at the point the stream would otherwise have ended. This mirrors
88/// [`futures::stream::iter`] over the same iterator, where the panic surfaces on the poller
89/// directly: moving the decode to another thread changes where the work happens, not whether the
90/// caller hears about it. Items decoded before the panic, including those earlier in the same
91/// chunk, are still yielded first, so the panic arrives after the last successfully decoded item
92/// rather than discarding the batch.
93///
94/// # Polling past the end
95/// `None` is terminal and repeatable: polling again yields `None` rather than panicking, so a
96/// consumer that does not track termination itself needs no `.fuse()`. The returned stream is also
97/// [`Unpin`], so awaiting `next()` needs no pinning.
98///
99/// Both are part of the contract, not accidents of the current body — a rewrite that broke either
100/// (an `async_stream` generator breaks both) would break callers. [`Unpin`] is therefore in the
101/// return type, so such a rewrite fails to compile here rather than at some downstream call site;
102/// the terminal `None` cannot be spelled in a signature and is pinned by test instead. `FusedStream`
103/// is deliberately *not* implemented: the terminal `None` above already gives a consumer what
104/// fusing would, the intended consumer
105/// [`merge_time_sorted`](super::merge::merge_time_sorted) fuses its inputs itself, and widening the
106/// return type later is an additive change. That is a design judgement rather than an observation —
107/// this helper has no in-tree production callers, so there is no consumer set to appeal to.
108///
109/// # The iterator must be `Send + 'static`
110/// The iterator moves from one chunk's task to the next, so it has to cross threads. A decoder with
111/// `!Send` internals is not accepted. Holding one thread for the whole decode would allow it, and
112/// is exactly what made a merge of more inputs than the pool has threads deadlock.
113///
114/// # Examples
115/// ```no_run
116/// # use rustrade_data::streams::blocking::{stream_blocking_iter, DEFAULT_BLOCKING_CHUNK_SIZE};
117/// # fn decode() -> Result<std::vec::IntoIter<Result<u64, std::io::Error>>, std::io::Error> {
118/// #     unimplemented!()
119/// # }
120/// let _stream = stream_blocking_iter(DEFAULT_BLOCKING_CHUNK_SIZE, decode);
121/// ```
122pub fn stream_blocking_iter<Init, Iter, Item, Error>(
123    chunk_size: usize,
124    init: Init,
125) -> impl Stream<Item = Result<Item, Error>> + Send + Unpin + 'static
126where
127    Init: FnOnce() -> Result<Iter, Error> + Send + 'static,
128    Iter: Iterator<Item = Result<Item, Error>> + Send + 'static,
129    Item: Send + 'static,
130    Error: Send + 'static,
131{
132    assert!(
133        chunk_size > 0,
134        "stream_blocking_iter requires a non-zero chunk size"
135    );
136
137    BlockingIterStream {
138        chunk_size,
139        init: Some(init),
140        decoding: None,
141        ready: Vec::new().into_iter(),
142        panic: None,
143    }
144}
145
146/// Where a chunk's task picks the decode up: the source not yet opened, or the iterator a previous
147/// chunk returned.
148enum Source<Init, Iter> {
149    Init(Init),
150    Iter(Iter),
151}
152
153/// What one chunk's task hands back.
154struct Chunk<Iter, Item, Error> {
155    /// The items decoded, in order.
156    items: Vec<Result<Item, Error>>,
157    /// The iterator, if it may have more items. `None` once it is exhausted, `init` failed, or the
158    /// decode panicked.
159    next: Option<Iter>,
160    /// A panic caught in `init` or the iterator, re-raised once `items` have been yielded.
161    panic: Option<Box<dyn Any + Send>>,
162}
163
164/// Decodes up to `chunk_size` items from `source`, catching a panic so that the items decoded
165/// before it are still returned.
166fn decode_chunk<Init, Iter, Item, Error>(
167    source: Source<Init, Iter>,
168    chunk_size: usize,
169) -> Chunk<Iter, Item, Error>
170where
171    Init: FnOnce() -> Result<Iter, Error>,
172    Iter: Iterator<Item = Result<Item, Error>>,
173{
174    // Capped: `chunk_size` bounds the chunk, but a source may yield far fewer items, and a huge
175    // value would otherwise fail the allocation before decoding anything.
176    let mut items = Vec::with_capacity(chunk_size.min(DEFAULT_BLOCKING_CHUNK_SIZE));
177
178    // `AssertUnwindSafe`: after a panic, the only state used again is `items`, which holds whole
179    // items pushed before it. The iterator, the one value a panic could leave half-updated, is
180    // never touched again.
181    let decoded = std::panic::catch_unwind(AssertUnwindSafe(|| {
182        let mut iter = match source {
183            Source::Iter(iter) => iter,
184            Source::Init(init) => match init() {
185                Ok(iter) => iter,
186                Err(error) => {
187                    items.push(Err(error));
188                    return None;
189                }
190            },
191        };
192
193        while items.len() < chunk_size {
194            match iter.next() {
195                Some(item) => items.push(item),
196                None => return None,
197            }
198        }
199        Some(iter)
200    }));
201
202    match decoded {
203        Ok(next) => Chunk {
204            items,
205            next,
206            panic: None,
207        },
208        Err(payload) => Chunk {
209            items,
210            next: None,
211            panic: Some(payload),
212        },
213    }
214}
215
216/// The [`Stream`] returned by [`stream_blocking_iter`].
217///
218/// Private: it exists to sequence the chunk tasks and re-raise a decode panic, and nothing about
219/// that needs to be nameable by a caller.
220struct BlockingIterStream<Init, Iter, Item, Error> {
221    chunk_size: usize,
222    /// The source, until the first poll starts the first chunk with it.
223    init: Option<Init>,
224    /// The chunk being decoded, if any. At most one is in flight.
225    decoding: Option<JoinHandle<Chunk<Iter, Item, Error>>>,
226    /// Items of the last chunk received, not yet yielded.
227    ready: std::vec::IntoIter<Result<Item, Error>>,
228    /// A decode panic, re-raised once `ready` is drained.
229    panic: Option<Box<dyn Any + Send>>,
230}
231
232// No field is ever pinned: `poll_next` reaches them through `get_mut` and moves `init` out by value.
233// Without this, `Unpin` would depend on `Init`, `Item` and `Error`, and a closure capturing a
234// `!Unpin` value would make the stream `!Unpin`, breaking the `Unpin` the return type promises.
235impl<Init, Iter, Item, Error> Unpin for BlockingIterStream<Init, Iter, Item, Error> {}
236
237impl<Init, Iter, Item, Error> BlockingIterStream<Init, Iter, Item, Error>
238where
239    Init: FnOnce() -> Result<Iter, Error> + Send + 'static,
240    Iter: Iterator<Item = Result<Item, Error>> + Send + 'static,
241    Item: Send + 'static,
242    Error: Send + 'static,
243{
244    fn start_chunk(&mut self, source: Source<Init, Iter>) {
245        let chunk_size = self.chunk_size;
246        self.decoding = Some(tokio::task::spawn_blocking(move || {
247            decode_chunk(source, chunk_size)
248        }));
249    }
250}
251
252impl<Init, Iter, Item, Error> Stream for BlockingIterStream<Init, Iter, Item, Error>
253where
254    Init: FnOnce() -> Result<Iter, Error> + Send + 'static,
255    Iter: Iterator<Item = Result<Item, Error>> + Send + 'static,
256    Item: Send + 'static,
257    Error: Send + 'static,
258{
259    type Item = Result<Item, Error>;
260
261    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
262        let this = self.get_mut();
263
264        if let Some(init) = this.init.take() {
265            this.start_chunk(Source::Init(init));
266        }
267
268        loop {
269            if let Some(item) = this.ready.next() {
270                return Poll::Ready(Some(item));
271            }
272
273            // Every item decoded before a panic has now been yielded, so it is re-raised here,
274            // with the original payload, on the task polling this stream. Losing it would be the
275            // silent truncation the `# Failure model` rules out.
276            if let Some(payload) = this.panic.take() {
277                std::panic::resume_unwind(payload);
278            }
279
280            let Some(handle) = this.decoding.as_mut() else {
281                return Poll::Ready(None);
282            };
283            let result = std::task::ready!(Pin::new(handle).poll(cx));
284            this.decoding = None;
285
286            let chunk = match result {
287                Ok(chunk) => chunk,
288                // `decode_chunk` catches panics in the decode itself, so this is a panic outside
289                // it, still re-raised rather than lost.
290                Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()),
291                // Nothing aborts the task, so this is runtime shutdown cancelling a chunk still
292                // queued for a thread. The consumer is normally going away too.
293                Err(error) => {
294                    warn!(
295                        %error,
296                        "stream_blocking_iter: the runtime cancelled a decode chunk - ending the \
297                         stream early"
298                    );
299                    return Poll::Ready(None);
300                }
301            };
302
303            // Start the next chunk before yielding this one, so decoding overlaps consumption. No
304            // chunk after that starts until this one is drained: the back-pressure.
305            if let Some(iter) = chunk.next {
306                this.start_chunk(Source::Iter(iter));
307            }
308            this.ready = chunk.items.into_iter();
309            this.panic = chunk.panic;
310        }
311    }
312}
313
314#[cfg(test)]
315#[allow(clippy::unwrap_used)] // Test code: panics on bad input are acceptable
316mod tests {
317    use super::*;
318    use futures::StreamExt;
319    use std::{
320        io::{Error, ErrorKind},
321        sync::{
322            Arc, Mutex,
323            atomic::{AtomicUsize, Ordering},
324        },
325        time::Duration,
326    };
327
328    fn ok_iter(count: usize) -> Result<std::vec::IntoIter<Result<usize, Error>>, Error> {
329        Ok((0..count).map(Ok).collect::<Vec<_>>().into_iter())
330    }
331
332    #[tokio::test]
333    async fn forwards_every_item_in_order() {
334        let stream = stream_blocking_iter(4, || ok_iter(50));
335
336        let items = stream.map(|item| item.unwrap()).collect::<Vec<_>>().await;
337
338        assert_eq!(items, (0..50).collect::<Vec<_>>());
339    }
340
341    #[tokio::test]
342    async fn an_init_failure_becomes_a_single_err_item() {
343        let stream = stream_blocking_iter(4, || {
344            Err::<std::vec::IntoIter<Result<usize, Error>>, _>(Error::new(
345                ErrorKind::NotFound,
346                "no such artifact",
347            ))
348        });
349
350        let items = stream.collect::<Vec<_>>().await;
351        assert_eq!(items.len(), 1);
352        assert_eq!(items[0].as_ref().unwrap_err().kind(), ErrorKind::NotFound);
353    }
354
355    #[tokio::test]
356    async fn a_mid_stream_failure_is_surfaced_not_skipped() {
357        let stream = stream_blocking_iter(4, || {
358            Ok(vec![
359                Ok(1_usize),
360                Err(Error::new(ErrorKind::InvalidData, "bad row")),
361                Ok(3),
362            ]
363            .into_iter())
364        });
365
366        let items = stream.collect::<Vec<_>>().await;
367        assert_eq!(items.len(), 3);
368        assert!(items[1].is_err());
369    }
370
371    /// A panicking decoder must not look like a decoder that reached the end of its data. Silently
372    /// ending here would let a truncated artifact drive a backtest to a normal-looking summary.
373    #[tokio::test]
374    #[should_panic(expected = "corrupt page header")]
375    async fn a_panic_in_the_iterator_reaches_the_consumer() {
376        let stream = stream_blocking_iter(4, || {
377            Ok((0..10).map(|index: usize| {
378                assert!(index < 3, "corrupt page header");
379                Ok::<usize, Error>(index)
380            }))
381        });
382
383        let _items = stream.collect::<Vec<_>>().await;
384    }
385
386    #[tokio::test]
387    #[should_panic(expected = "unreadable footer")]
388    async fn a_panic_in_init_reaches_the_consumer() {
389        let stream = stream_blocking_iter(4, || {
390            panic!("unreadable footer");
391            #[allow(unreachable_code)] // Pins the closure's return type for inference.
392            ok_iter(0)
393        });
394
395        let _items = stream.collect::<Vec<_>>().await;
396    }
397
398    /// The items decoded before the panic are real data and are still delivered; the panic arrives
399    /// after them, not instead of them.
400    #[tokio::test]
401    async fn items_before_a_panic_are_still_yielded() {
402        // Collected through a handle that OUTLIVES the unwind. A `Vec` local to the caught future
403        // is destroyed by the panic along with the rest of that frame, so nothing can be asserted
404        // about it afterwards -- which left the "items before the panic are still yielded" half of
405        // this contract untested, and only the "the panic escapes" half pinned.
406        let delivered = Arc::new(Mutex::new(Vec::new()));
407
408        let sink = Arc::clone(&delivered);
409        let collect = std::panic::AssertUnwindSafe(async move {
410            let mut stream = Box::pin(stream_blocking_iter(4, || {
411                Ok((0..10).map(|index: usize| {
412                    assert!(index < 3, "boom");
413                    Ok::<usize, Error>(index)
414                }))
415            }));
416
417            while let Some(item) = stream.next().await {
418                // The lock is never held across the await, so the panic cannot poison it.
419                sink.lock().unwrap().push(item.unwrap());
420            }
421        });
422
423        let outcome = futures::FutureExt::catch_unwind(collect).await;
424
425        assert!(outcome.is_err(), "the panic must not be swallowed");
426        assert_eq!(
427            *delivered.lock().unwrap(),
428            vec![0, 1, 2],
429            "every item decoded before the panic is real data and must still reach the consumer"
430        );
431    }
432
433    /// The clean path must stay clean: a normal end is still a normal end, not a panic.
434    #[tokio::test]
435    async fn a_normal_end_of_stream_does_not_panic() {
436        let mut stream = Box::pin(stream_blocking_iter(4, || ok_iter(3)));
437
438        for expected in 0..3 {
439            assert_eq!(stream.next().await.unwrap().unwrap(), expected);
440        }
441
442        // Polled past the end twice: the join handle is consumed on the first, and the second must
443        // not poll it again.
444        assert!(stream.next().await.is_none());
445        assert!(stream.next().await.is_none());
446    }
447
448    /// The property the type exists for: the producer must not run ahead of the consumer without
449    /// bound. With chunks of 2 and one item drained, at most two chunks can have been decoded --
450    /// emphatically not all 10,000, however long the decoder is left to run.
451    #[tokio::test]
452    async fn the_producer_is_bounded_by_two_chunks() {
453        const CHUNK: usize = 2;
454        let produced = Arc::new(AtomicUsize::new(0));
455
456        let counter = Arc::clone(&produced);
457        let mut stream = Box::pin(stream_blocking_iter(CHUNK, move || {
458            Ok((0..10_000).map(move |index| {
459                counter.fetch_add(1, Ordering::SeqCst);
460                Ok::<usize, Error>(index)
461            }))
462        }));
463
464        // Take one item, then wait for the next chunk, which starts as soon as the first arrives
465        // so that decoding overlaps consumption.
466        assert_eq!(stream.next().await.unwrap().unwrap(), 0);
467        let mut waited = 0;
468        while produced.load(Ordering::SeqCst) < 2 * CHUNK {
469            assert!(
470                waited < 500,
471                "the next chunk was not started ahead of the consumer"
472            );
473            tokio::time::sleep(Duration::from_millis(10)).await;
474            waited += 1;
475        }
476
477        // Then give a third chunk every chance to appear: it must not.
478        tokio::time::sleep(Duration::from_millis(50)).await;
479        let produced = produced.load(Ordering::SeqCst);
480        assert_eq!(
481            produced,
482            2 * CHUNK,
483            "producer ran {produced} items ahead with chunks of {CHUNK}"
484        );
485    }
486
487    #[test]
488    #[should_panic(expected = "non-zero chunk size")]
489    fn a_zero_chunk_size_is_rejected() {
490        let _stream = stream_blocking_iter(0, || ok_iter(1));
491    }
492
493    /// Every chunk size delivers every item once, in order, including sizes that divide the item
494    /// count exactly and a size far larger than the source.
495    #[tokio::test]
496    async fn every_chunk_size_forwards_every_item_in_order() {
497        for chunk_size in [1, 2, 3, 5, 10, usize::MAX] {
498            let items = stream_blocking_iter(chunk_size, || ok_iter(10))
499                .map(|item| item.unwrap())
500                .collect::<Vec<_>>()
501                .await;
502            assert_eq!(
503                items,
504                (0..10).collect::<Vec<_>>(),
505                "chunk size {chunk_size}"
506            );
507        }
508
509        let empty = stream_blocking_iter(4, || ok_iter(0))
510            .collect::<Vec<_>>()
511            .await;
512        assert!(empty.is_empty());
513    }
514
515    /// A panic after earlier chunks were handed over arrives after all of their items, and the
516    /// stream then reads as ended.
517    #[tokio::test]
518    async fn a_panic_in_a_later_chunk_follows_every_earlier_item() {
519        let delivered = Arc::new(Mutex::new(Vec::new()));
520
521        let sink = Arc::clone(&delivered);
522        let mut stream = Box::pin(stream_blocking_iter(2, || {
523            Ok((0..10).map(|index: usize| {
524                assert!(index < 5, "boom");
525                Ok::<usize, Error>(index)
526            }))
527        }));
528        let collect = std::panic::AssertUnwindSafe(async {
529            while let Some(item) = stream.next().await {
530                sink.lock().unwrap().push(item.unwrap());
531            }
532        });
533
534        let outcome = futures::FutureExt::catch_unwind(collect).await;
535
536        assert!(outcome.is_err(), "the panic must not be swallowed");
537        assert_eq!(*delivered.lock().unwrap(), vec![0, 1, 2, 3, 4]);
538        assert!(
539            stream.next().await.is_none(),
540            "polled again after the panic, the stream reads as ended"
541        );
542    }
543
544    /// Building a stream starts nothing: the first chunk, `init` included, starts on first poll.
545    #[tokio::test]
546    async fn nothing_is_decoded_before_the_first_poll() {
547        let opened = Arc::new(AtomicUsize::new(0));
548
549        let counter = Arc::clone(&opened);
550        let mut stream = Box::pin(stream_blocking_iter(4, move || {
551            counter.fetch_add(1, Ordering::SeqCst);
552            ok_iter(3)
553        }));
554
555        tokio::time::sleep(Duration::from_millis(50)).await;
556        assert_eq!(opened.load(Ordering::SeqCst), 0, "init ran before any poll");
557
558        assert_eq!(stream.next().await.unwrap().unwrap(), 0);
559        assert_eq!(opened.load(Ordering::SeqCst), 1);
560    }
561
562    /// The items decoded earlier in the same chunk as a panic are real data too.
563    #[tokio::test]
564    async fn items_before_a_panic_in_the_same_chunk_are_still_yielded() {
565        let delivered = Arc::new(Mutex::new(Vec::new()));
566
567        let sink = Arc::clone(&delivered);
568        let collect = std::panic::AssertUnwindSafe(async move {
569            // One chunk of 8 covers items 0..=2 and the panic at 3.
570            let mut stream = Box::pin(stream_blocking_iter(8, || {
571                Ok((0..10).map(|index: usize| {
572                    assert!(index < 3, "boom");
573                    Ok::<usize, Error>(index)
574                }))
575            }));
576
577            while let Some(item) = stream.next().await {
578                sink.lock().unwrap().push(item.unwrap());
579            }
580        });
581
582        let outcome = futures::FutureExt::catch_unwind(collect).await;
583
584        assert!(outcome.is_err(), "the panic must not be swallowed");
585        assert_eq!(*delivered.lock().unwrap(), vec![0, 1, 2]);
586    }
587
588    #[tokio::test]
589    async fn dropping_the_stream_stops_the_producer() {
590        const CHUNK: usize = 2;
591        const ITEMS: usize = 10_000;
592        let produced = Arc::new(AtomicUsize::new(0));
593
594        let counter = Arc::clone(&produced);
595        let mut stream = Box::pin(stream_blocking_iter(CHUNK, move || {
596            Ok((0..ITEMS).map(move |index| {
597                counter.fetch_add(1, Ordering::SeqCst);
598                Ok::<usize, Error>(index)
599            }))
600        }));
601
602        assert!(stream.next().await.is_some());
603        drop(stream);
604
605        // Wait for the count to STOP climbing, rather than assuming it already has. A chunk already
606        // being decoded runs to its end on a blocking thread, so a `yield_now` on this side proves
607        // nothing about where that thread has got to -- sampling twice around one yield can catch it
608        // mid-chunk and fail spuriously. Polling until two samples separated by a real
609        // delay agree is the same assertion made soundly.
610        let mut settled = produced.load(Ordering::SeqCst);
611        let mut stopped = false;
612        for _ in 0..50 {
613            tokio::time::sleep(Duration::from_millis(10)).await;
614            let current = produced.load(Ordering::SeqCst);
615            if current == settled {
616                stopped = true;
617                break;
618            }
619            settled = current;
620        }
621
622        assert!(
623            stopped,
624            "the producer was still running 500ms after the stream was dropped ({settled} items)"
625        );
626        // The property that matters, and the one a "count stopped changing" assertion alone does
627        // not establish: it stopped *early*, rather than decoding the whole iterator for a consumer
628        // that was gone. Checked on its own first so that a drop never observed at all fails with
629        // that message, rather than as a bound exceeded.
630        assert!(
631            settled < ITEMS,
632            "the producer ran to completion ({settled} items) despite the stream being dropped"
633        );
634        // Tighter: the first chunk and the one started ahead of it, and nothing after the drop.
635        assert!(
636            settled <= 2 * CHUNK,
637            "{settled} items were decoded, more than the two chunks started before the drop"
638        );
639    }
640}