diff --git a/tokio-stream/src/stream_ext/chunks_timeout.rs b/tokio-stream/src/stream_ext/chunks_timeout.rs index 89b8d49a6..0278387c2 100644 --- a/tokio-stream/src/stream_ext/chunks_timeout.rs +++ b/tokio-stream/src/stream_ext/chunks_timeout.rs @@ -33,6 +33,12 @@ impl ChunksTimeout { cap: max_size, } } + + /// Consumes the [`ChunksTimeout`] and then returns all buffered items. + pub fn into_remainder(mut self: Pin<&mut Self>) -> Vec { + let me = self.as_mut().project(); + std::mem::take(me.items) + } } impl Stream for ChunksTimeout { diff --git a/tokio-stream/tests/stream_chunks_timeout.rs b/tokio-stream/tests/stream_chunks_timeout.rs new file mode 100644 index 000000000..33da658b8 --- /dev/null +++ b/tokio-stream/tests/stream_chunks_timeout.rs @@ -0,0 +1,30 @@ +#![warn(rust_2018_idioms)] + +use futures::FutureExt; +use std::error::Error; +use tokio::time; +use tokio::time::Duration; +use tokio_stream::{self as stream, StreamExt}; +use tokio_test::assert_pending; +use tokio_test::task; + +#[tokio::test(start_paused = true)] +async fn stream_chunks_remainder() -> Result<(), Box> { + let stream1 = + stream::iter([5]).then(move |n| time::sleep(Duration::from_secs(1)).map(move |_| n)); + + let inner = stream::iter([1, 2, 3, 4]).chain(stream1); + tokio::pin!(inner); + + let chunked = (&mut inner).chunks_timeout(10, Duration::from_millis(20)); + + let mut chunked = task::spawn(chunked); + assert_pending!(chunked.poll_next()); + + let remainder = chunked.enter(|_, stream| stream.into_remainder()); + + assert_eq!(remainder, vec![1, 2, 3, 4]); + time::advance(Duration::from_secs(2)).await; + assert_eq!(inner.next().await, Some(5)); + Ok(()) +}