diff --git a/tokio/src/io/mod.rs b/tokio/src/io/mod.rs index 9e0e06319..37da942ff 100644 --- a/tokio/src/io/mod.rs +++ b/tokio/src/io/mod.rs @@ -236,6 +236,7 @@ cfg_io_util! { cfg_stream! { pub use util::{stream_reader, StreamReader}; + pub use util::{reader_stream, ReaderStream}; } } diff --git a/tokio/src/io/util/mod.rs b/tokio/src/io/util/mod.rs index 609ff2386..782a02a8f 100644 --- a/tokio/src/io/util/mod.rs +++ b/tokio/src/io/util/mod.rs @@ -66,6 +66,9 @@ cfg_io_util! { cfg_stream! { mod stream_reader; pub use stream_reader::{stream_reader, StreamReader}; + + mod reader_stream; + pub use reader_stream::{reader_stream, ReaderStream}; } mod take; diff --git a/tokio/src/io/util/reader_stream.rs b/tokio/src/io/util/reader_stream.rs new file mode 100644 index 000000000..51651cede --- /dev/null +++ b/tokio/src/io/util/reader_stream.rs @@ -0,0 +1,105 @@ +use crate::io::AsyncRead; +use crate::stream::Stream; +use bytes::{Bytes, BytesMut}; +use pin_project_lite::pin_project; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// Convert an [`AsyncRead`] implementor into a + /// [`Stream`] of Result<[`Bytes`], std::io::Error>. + /// After first error it will stop. + /// Additionally, this stream is fused: after it returns None at some + /// moment, it is guaranteed that further `next()`, `poll_next()` and + /// similar functions will instantly return None. + /// + /// This type can be created using the [`reader_stream`] function + /// + /// [`AsyncRead`]: crate::io::AsyncRead + /// [`Stream`]: crate::stream::Stream + /// [`Bytes`]: bytes::Bytes + /// [`reader_stream`]: crate::io::reader_stream + #[derive(Debug)] + #[cfg_attr(docsrs, doc(cfg(feature = "stream")))] + #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] + pub struct ReaderStream { + // Reader itself. + // None if we had error reading from the `reader` in the past. + #[pin] + reader: Option, + // Working buffer, used to optimize allocations. + // # Capacity behavior + // Initially `buf` is empty. Also it's getting smaller and smaller + // during polls (because its chunks are returned to stream user). + // But when it's capacity reaches 0, it is growed. + buf: BytesMut, + } +} + +/// Convert an [`AsyncRead`] implementor into a +/// [`Stream`] of Result<[`Bytes`], std::io::Error>. +/// +/// # Example +/// +/// ``` +/// # #[tokio::main] +/// # async fn main() -> std::io::Result<()> { +/// use tokio::stream::StreamExt; +/// +/// let data: &[u8] = b"hello, world!"; +/// let mut stream = tokio::io::reader_stream(data); +/// let mut stream_contents = Vec::new(); +/// while let Some(chunk) = stream.next().await { +/// stream_contents.extend_from_slice(chunk?.as_ref()); +/// } +/// assert_eq!(stream_contents, data); +/// # Ok(()) +/// # } +/// ``` +/// +/// [`AsyncRead`]: crate::io::AsyncRead +/// [`Stream`]: crate::stream::Stream +/// [`Bytes`]: bytes::Bytes +pub fn reader_stream(reader: R) -> ReaderStream +where + R: AsyncRead, +{ + ReaderStream { + reader: Some(reader), + buf: BytesMut::new(), + } +} + +const CAPACITY: usize = 4096; + +impl Stream for ReaderStream +where + R: AsyncRead, +{ + type Item = std::io::Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.as_mut().project(); + let reader = match this.reader.as_pin_mut() { + Some(r) => r, + None => return Poll::Ready(None), + }; + if this.buf.capacity() == 0 { + this.buf.reserve(CAPACITY); + } + match reader.poll_read_buf(cx, &mut this.buf) { + Poll::Pending => Poll::Pending, + Poll::Ready(Err(err)) => { + self.project().reader.set(None); + Poll::Ready(Some(Err(err))) + } + Poll::Ready(Ok(0)) => { + self.project().reader.set(None); + Poll::Ready(None) + } + Poll::Ready(Ok(_)) => { + let chunk = this.buf.split(); + Poll::Ready(Some(Ok(chunk.freeze()))) + } + } + } +} diff --git a/tokio/tests/io_reader_stream.rs b/tokio/tests/io_reader_stream.rs new file mode 100644 index 000000000..6546a0ef4 --- /dev/null +++ b/tokio/tests/io_reader_stream.rs @@ -0,0 +1,64 @@ +#![warn(rust_2018_idioms)] +#![cfg(feature = "full")] + +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::AsyncRead; +use tokio::stream::StreamExt; + +/// produces at most `remaining` zeros, that returns error. +/// each time it reads at most 31 byte. +struct Reader { + remaining: usize, +} + +impl AsyncRead for Reader { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll> { + let this = Pin::into_inner(self); + assert_ne!(buf.len(), 0); + if this.remaining > 0 { + let n = std::cmp::min(this.remaining, buf.len()); + let n = std::cmp::min(n, 31); + for x in &mut buf[..n] { + *x = 0; + } + this.remaining -= n; + Poll::Ready(Ok(n)) + } else { + Poll::Ready(Err(std::io::Error::from_raw_os_error(22))) + } + } +} + +#[tokio::test] +async fn correct_behavior_on_errors() { + let reader = Reader { remaining: 8000 }; + let mut stream = tokio::io::reader_stream(reader); + let mut zeros_received = 0; + let mut had_error = false; + loop { + let item = stream.next().await.unwrap(); + match item { + Ok(bytes) => { + let bytes = &*bytes; + for byte in bytes { + assert_eq!(*byte, 0); + zeros_received += 1; + } + } + Err(_) => { + assert!(!had_error); + had_error = true; + break; + } + } + } + + assert!(had_error); + assert_eq!(zeros_received, 8000); + assert!(stream.next().await.is_none()); +}