io: add read_exact_arc to safely read a new uninitialized Arc (#7165)

Author: joshtriplett

tokio-util/Cargo.toml

@@ -35,7 +35,7 @@ __docs_rs = ["futures-util"]
 
 [dependencies]
 tokio = { version = "1.28.0", path = "../tokio", features = ["sync"] }
-bytes = "1.2.1"
+bytes = "1.5.0"
 futures-core = "0.3.0"
 futures-sink = "0.3.0"
 futures-io = { version = "0.3.0", optional = true }

tokio-util/src/io/mod.rs

@@ -18,6 +18,9 @@ mod sink_writer;
 mod stream_reader;
 
 cfg_io_util! {
+    mod read_arc;
+    pub use self::read_arc::read_exact_arc;
+
     mod sync_bridge;
     pub use self::sync_bridge::SyncIoBridge;
 }

tokio-util/src/io/read_arc.rs

@@ -0,0 +1,44 @@
+use std::io;
+use std::mem::MaybeUninit;
+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<R: AsyncRead>(read: R, len: usize) -> io::Result<Arc<[u8]>> {
+    tokio::pin!(read);
+    // TODO(MSRV 1.82): When bumping MSRV, switch to `Arc::new_uninit_slice(len)`. The following is
+    // equivalent, and generates the same assembly, but works without requiring MSRV 1.82.
+    let arc: Arc<[MaybeUninit<u8>]> = (0..len).map(|_| MaybeUninit::uninit()).collect();
+    // TODO(MSRV future): Use `Arc::get_mut_unchecked` once it's stabilized.
+    // SAFETY: We're the only owner of the `Arc`, and we keep the `Arc` valid throughout this loop
+    // as we write through this reference.
+    let mut buf = unsafe { &mut *(Arc::as_ptr(&arc) as *mut [MaybeUninit<u8>]) };
+    while !buf.is_empty() {
+        if read.read_buf(&mut buf).await? == 0 {
+            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "early eof"));
+        }
+    }
+    // TODO(MSRV 1.82): When bumping MSRV, switch to `arc.assume_init()`. The following is
+    // equivalent, and generates the same assembly, but works without requiring MSRV 1.82.
+    // SAFETY: This changes `[MaybeUninit<u8>]` to `[u8]`, and we've initialized all the bytes in
+    // the loop above.
+    Ok(unsafe { Arc::from_raw(Arc::into_raw(arc) as *const [u8]) })
+}