diff --git a/tokio-util/src/io/mod.rs b/tokio-util/src/io/mod.rs index 6c40d739014..f0bc676845a 100644 --- a/tokio-util/src/io/mod.rs +++ b/tokio-util/src/io/mod.rs @@ -12,6 +12,7 @@ mod copy_to_bytes; mod inspect; +mod read_arc; mod read_buf; mod reader_stream; mod sink_writer; @@ -24,6 +25,7 @@ cfg_io_util! { pub use self::copy_to_bytes::CopyToBytes; pub use self::inspect::{InspectReader, InspectWriter}; +pub use self::read_arc::read_exact_arc; pub use self::read_buf::read_buf; pub use self::reader_stream::ReaderStream; pub use self::sink_writer::SinkWriter; diff --git a/tokio-util/src/io/read_arc.rs b/tokio-util/src/io/read_arc.rs new file mode 100644 index 00000000000..a81b173d8ee --- /dev/null +++ b/tokio-util/src/io/read_arc.rs @@ -0,0 +1,39 @@ +use std::io; +use std::sync::Arc; +use tokio::io::{AsyncRead, AsyncReadExt}; + +/// Read data from an `AsyncRead` into an `Arc`. +/// +/// This uses `Arc::new_uninit_slice` and reads into the resulting uninitialized `Arc`. +/// +/// # Example +/// +/// ``` +/// # #[tokio::main] +/// # async fn main() -> std::io::Result<()> { +/// use tokio_util::io::read_exact_arc; +/// +/// let read = tokio::io::repeat(42); +/// +/// let arc = read_exact_arc(read, 4).await?; +/// +/// assert_eq!(&arc[..], &[42; 4]); +/// # Ok(()) +/// # } +/// ``` +pub async fn read_exact_arc(mut read: R, len: usize) -> io::Result> +where + R: AsyncRead + Unpin, +{ + let mut arc = Arc::new_uninit_slice(len); + let mut buf = unsafe { Arc::get_mut(&mut arc).unwrap_unchecked() }; + let mut bytes_remaining = len; + while bytes_remaining != 0 { + let bytes_read = read.read_buf(&mut buf).await?; + if bytes_read == 0 { + return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "early eof")); + } + bytes_remaining -= bytes_read; + } + Ok(unsafe { arc.assume_init() }) +}