use futures::Stream;
use std::any::Any;
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::task::JoinHandle;
use tracing::warn;
pub const DEFAULT_BLOCKING_CHUNK_SIZE: usize = 1024;
pub fn stream_blocking_iter<Init, Iter, Item, Error>(
chunk_size: usize,
init: Init,
) -> impl Stream<Item = Result<Item, Error>> + Send + Unpin + 'static
where
Init: FnOnce() -> Result<Iter, Error> + Send + 'static,
Iter: Iterator<Item = Result<Item, Error>> + Send + 'static,
Item: Send + 'static,
Error: Send + 'static,
{
assert!(
chunk_size > 0,
"stream_blocking_iter requires a non-zero chunk size"
);
BlockingIterStream {
chunk_size,
init: Some(init),
decoding: None,
ready: Vec::new().into_iter(),
panic: None,
}
}
enum Source<Init, Iter> {
Init(Init),
Iter(Iter),
}
struct Chunk<Iter, Item, Error> {
items: Vec<Result<Item, Error>>,
next: Option<Iter>,
panic: Option<Box<dyn Any + Send>>,
}
fn decode_chunk<Init, Iter, Item, Error>(
source: Source<Init, Iter>,
chunk_size: usize,
) -> Chunk<Iter, Item, Error>
where
Init: FnOnce() -> Result<Iter, Error>,
Iter: Iterator<Item = Result<Item, Error>>,
{
let mut items = Vec::with_capacity(chunk_size.min(DEFAULT_BLOCKING_CHUNK_SIZE));
let decoded = std::panic::catch_unwind(AssertUnwindSafe(|| {
let mut iter = match source {
Source::Iter(iter) => iter,
Source::Init(init) => match init() {
Ok(iter) => iter,
Err(error) => {
items.push(Err(error));
return None;
}
},
};
while items.len() < chunk_size {
match iter.next() {
Some(item) => items.push(item),
None => return None,
}
}
Some(iter)
}));
match decoded {
Ok(next) => Chunk {
items,
next,
panic: None,
},
Err(payload) => Chunk {
items,
next: None,
panic: Some(payload),
},
}
}
struct BlockingIterStream<Init, Iter, Item, Error> {
chunk_size: usize,
init: Option<Init>,
decoding: Option<JoinHandle<Chunk<Iter, Item, Error>>>,
ready: std::vec::IntoIter<Result<Item, Error>>,
panic: Option<Box<dyn Any + Send>>,
}
impl<Init, Iter, Item, Error> Unpin for BlockingIterStream<Init, Iter, Item, Error> {}
impl<Init, Iter, Item, Error> BlockingIterStream<Init, Iter, Item, Error>
where
Init: FnOnce() -> Result<Iter, Error> + Send + 'static,
Iter: Iterator<Item = Result<Item, Error>> + Send + 'static,
Item: Send + 'static,
Error: Send + 'static,
{
fn start_chunk(&mut self, source: Source<Init, Iter>) {
let chunk_size = self.chunk_size;
self.decoding = Some(tokio::task::spawn_blocking(move || {
decode_chunk(source, chunk_size)
}));
}
}
impl<Init, Iter, Item, Error> Stream for BlockingIterStream<Init, Iter, Item, Error>
where
Init: FnOnce() -> Result<Iter, Error> + Send + 'static,
Iter: Iterator<Item = Result<Item, Error>> + Send + 'static,
Item: Send + 'static,
Error: Send + 'static,
{
type Item = Result<Item, Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
if let Some(init) = this.init.take() {
this.start_chunk(Source::Init(init));
}
loop {
if let Some(item) = this.ready.next() {
return Poll::Ready(Some(item));
}
if let Some(payload) = this.panic.take() {
std::panic::resume_unwind(payload);
}
let Some(handle) = this.decoding.as_mut() else {
return Poll::Ready(None);
};
let result = std::task::ready!(Pin::new(handle).poll(cx));
this.decoding = None;
let chunk = match result {
Ok(chunk) => chunk,
Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()),
Err(error) => {
warn!(
%error,
"stream_blocking_iter: the runtime cancelled a decode chunk - ending the \
stream early"
);
return Poll::Ready(None);
}
};
if let Some(iter) = chunk.next {
this.start_chunk(Source::Iter(iter));
}
this.ready = chunk.items.into_iter();
this.panic = chunk.panic;
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)] mod tests {
use super::*;
use futures::StreamExt;
use std::{
io::{Error, ErrorKind},
sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
fn ok_iter(count: usize) -> Result<std::vec::IntoIter<Result<usize, Error>>, Error> {
Ok((0..count).map(Ok).collect::<Vec<_>>().into_iter())
}
#[tokio::test]
async fn forwards_every_item_in_order() {
let stream = stream_blocking_iter(4, || ok_iter(50));
let items = stream.map(|item| item.unwrap()).collect::<Vec<_>>().await;
assert_eq!(items, (0..50).collect::<Vec<_>>());
}
#[tokio::test]
async fn an_init_failure_becomes_a_single_err_item() {
let stream = stream_blocking_iter(4, || {
Err::<std::vec::IntoIter<Result<usize, Error>>, _>(Error::new(
ErrorKind::NotFound,
"no such artifact",
))
});
let items = stream.collect::<Vec<_>>().await;
assert_eq!(items.len(), 1);
assert_eq!(items[0].as_ref().unwrap_err().kind(), ErrorKind::NotFound);
}
#[tokio::test]
async fn a_mid_stream_failure_is_surfaced_not_skipped() {
let stream = stream_blocking_iter(4, || {
Ok(vec![
Ok(1_usize),
Err(Error::new(ErrorKind::InvalidData, "bad row")),
Ok(3),
]
.into_iter())
});
let items = stream.collect::<Vec<_>>().await;
assert_eq!(items.len(), 3);
assert!(items[1].is_err());
}
#[tokio::test]
#[should_panic(expected = "corrupt page header")]
async fn a_panic_in_the_iterator_reaches_the_consumer() {
let stream = stream_blocking_iter(4, || {
Ok((0..10).map(|index: usize| {
assert!(index < 3, "corrupt page header");
Ok::<usize, Error>(index)
}))
});
let _items = stream.collect::<Vec<_>>().await;
}
#[tokio::test]
#[should_panic(expected = "unreadable footer")]
async fn a_panic_in_init_reaches_the_consumer() {
let stream = stream_blocking_iter(4, || {
panic!("unreadable footer");
#[allow(unreachable_code)] ok_iter(0)
});
let _items = stream.collect::<Vec<_>>().await;
}
#[tokio::test]
async fn items_before_a_panic_are_still_yielded() {
let delivered = Arc::new(Mutex::new(Vec::new()));
let sink = Arc::clone(&delivered);
let collect = std::panic::AssertUnwindSafe(async move {
let mut stream = Box::pin(stream_blocking_iter(4, || {
Ok((0..10).map(|index: usize| {
assert!(index < 3, "boom");
Ok::<usize, Error>(index)
}))
}));
while let Some(item) = stream.next().await {
sink.lock().unwrap().push(item.unwrap());
}
});
let outcome = futures::FutureExt::catch_unwind(collect).await;
assert!(outcome.is_err(), "the panic must not be swallowed");
assert_eq!(
*delivered.lock().unwrap(),
vec![0, 1, 2],
"every item decoded before the panic is real data and must still reach the consumer"
);
}
#[tokio::test]
async fn a_normal_end_of_stream_does_not_panic() {
let mut stream = Box::pin(stream_blocking_iter(4, || ok_iter(3)));
for expected in 0..3 {
assert_eq!(stream.next().await.unwrap().unwrap(), expected);
}
assert!(stream.next().await.is_none());
assert!(stream.next().await.is_none());
}
#[tokio::test]
async fn the_producer_is_bounded_by_two_chunks() {
const CHUNK: usize = 2;
let produced = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&produced);
let mut stream = Box::pin(stream_blocking_iter(CHUNK, move || {
Ok((0..10_000).map(move |index| {
counter.fetch_add(1, Ordering::SeqCst);
Ok::<usize, Error>(index)
}))
}));
assert_eq!(stream.next().await.unwrap().unwrap(), 0);
let mut waited = 0;
while produced.load(Ordering::SeqCst) < 2 * CHUNK {
assert!(
waited < 500,
"the next chunk was not started ahead of the consumer"
);
tokio::time::sleep(Duration::from_millis(10)).await;
waited += 1;
}
tokio::time::sleep(Duration::from_millis(50)).await;
let produced = produced.load(Ordering::SeqCst);
assert_eq!(
produced,
2 * CHUNK,
"producer ran {produced} items ahead with chunks of {CHUNK}"
);
}
#[test]
#[should_panic(expected = "non-zero chunk size")]
fn a_zero_chunk_size_is_rejected() {
let _stream = stream_blocking_iter(0, || ok_iter(1));
}
#[tokio::test]
async fn every_chunk_size_forwards_every_item_in_order() {
for chunk_size in [1, 2, 3, 5, 10, usize::MAX] {
let items = stream_blocking_iter(chunk_size, || ok_iter(10))
.map(|item| item.unwrap())
.collect::<Vec<_>>()
.await;
assert_eq!(
items,
(0..10).collect::<Vec<_>>(),
"chunk size {chunk_size}"
);
}
let empty = stream_blocking_iter(4, || ok_iter(0))
.collect::<Vec<_>>()
.await;
assert!(empty.is_empty());
}
#[tokio::test]
async fn a_panic_in_a_later_chunk_follows_every_earlier_item() {
let delivered = Arc::new(Mutex::new(Vec::new()));
let sink = Arc::clone(&delivered);
let mut stream = Box::pin(stream_blocking_iter(2, || {
Ok((0..10).map(|index: usize| {
assert!(index < 5, "boom");
Ok::<usize, Error>(index)
}))
}));
let collect = std::panic::AssertUnwindSafe(async {
while let Some(item) = stream.next().await {
sink.lock().unwrap().push(item.unwrap());
}
});
let outcome = futures::FutureExt::catch_unwind(collect).await;
assert!(outcome.is_err(), "the panic must not be swallowed");
assert_eq!(*delivered.lock().unwrap(), vec![0, 1, 2, 3, 4]);
assert!(
stream.next().await.is_none(),
"polled again after the panic, the stream reads as ended"
);
}
#[tokio::test]
async fn nothing_is_decoded_before_the_first_poll() {
let opened = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&opened);
let mut stream = Box::pin(stream_blocking_iter(4, move || {
counter.fetch_add(1, Ordering::SeqCst);
ok_iter(3)
}));
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(opened.load(Ordering::SeqCst), 0, "init ran before any poll");
assert_eq!(stream.next().await.unwrap().unwrap(), 0);
assert_eq!(opened.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn items_before_a_panic_in_the_same_chunk_are_still_yielded() {
let delivered = Arc::new(Mutex::new(Vec::new()));
let sink = Arc::clone(&delivered);
let collect = std::panic::AssertUnwindSafe(async move {
let mut stream = Box::pin(stream_blocking_iter(8, || {
Ok((0..10).map(|index: usize| {
assert!(index < 3, "boom");
Ok::<usize, Error>(index)
}))
}));
while let Some(item) = stream.next().await {
sink.lock().unwrap().push(item.unwrap());
}
});
let outcome = futures::FutureExt::catch_unwind(collect).await;
assert!(outcome.is_err(), "the panic must not be swallowed");
assert_eq!(*delivered.lock().unwrap(), vec![0, 1, 2]);
}
#[tokio::test]
async fn dropping_the_stream_stops_the_producer() {
const CHUNK: usize = 2;
const ITEMS: usize = 10_000;
let produced = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&produced);
let mut stream = Box::pin(stream_blocking_iter(CHUNK, move || {
Ok((0..ITEMS).map(move |index| {
counter.fetch_add(1, Ordering::SeqCst);
Ok::<usize, Error>(index)
}))
}));
assert!(stream.next().await.is_some());
drop(stream);
let mut settled = produced.load(Ordering::SeqCst);
let mut stopped = false;
for _ in 0..50 {
tokio::time::sleep(Duration::from_millis(10)).await;
let current = produced.load(Ordering::SeqCst);
if current == settled {
stopped = true;
break;
}
settled = current;
}
assert!(
stopped,
"the producer was still running 500ms after the stream was dropped ({settled} items)"
);
assert!(
settled < ITEMS,
"the producer ran to completion ({settled} items) despite the stream being dropped"
);
assert!(
settled <= 2 * CHUNK,
"{settled} items were decoded, more than the two chunks started before the drop"
);
}
}