Proposal: add a poll_read_exact method to AsyncRead

tokio-util/src/either.rs

@@ -104,6 +104,14 @@ where
     ) -> Poll<Result<()>> {
         delegate_call!(self.poll_read(cx, buf))
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<std::io::Result<()>> {
+        delegate_call!(self.poll_read_exact(cx, buf))
+    }
 }
 
 impl<L, R> AsyncBufRead for Either<L, R>

tokio-util/src/io/inspect.rs

@@ -176,4 +176,12 @@ impl<W: AsyncRead, F> AsyncRead for InspectWriter<W, F> {
     ) -> Poll<std::io::Result<()>> {
         self.project().writer.poll_read(cx, buf)
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<std::io::Result<()>> {
+        self.project().writer.poll_read_exact(cx, buf)
+    }
 }

tokio-util/src/io/sink_writer.rs

@@ -131,4 +131,12 @@ impl<S: AsyncRead> AsyncRead for SinkWriter<S> {
     ) -> Poll<io::Result<()>> {
         self.project().inner.poll_read(cx, buf)
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.project().inner.poll_read_exact(cx, buf)
+    }
 }

tokio/src/io/async_read.rs

@@ -2,7 +2,31 @@ use super::ReadBuf;
 use std::io;
 use std::ops::DerefMut;
 use std::pin::Pin;
-use std::task::{Context, Poll};
+use std::task::{ready, Context, Poll};
+
+#[cold]
+pub(crate) fn eof() -> io::Error {
+    io::Error::new(io::ErrorKind::UnexpectedEof, "early eof")
+}
+
+pub(crate) fn default_poll_read_exact<R: AsyncRead + ?Sized>(
+    mut reader: Pin<&mut R>,
+    cx: &mut Context<'_>,
+    buf: &mut ReadBuf<'_>,
+) -> Poll<io::Result<()>> {
+    loop {
+        // if our buffer is empty, then we need to read some data to continue.
+        let rem = buf.remaining();
+        if rem != 0 {
+            ready!(reader.as_mut().poll_read(cx, buf))?;
+            if buf.remaining() == rem {
+                return Poll::Ready(Err(eof()));
+            }
+        } else {
+            return Poll::Ready(Ok(()));
+        }
+    }
+}
 
 /// Reads bytes from a source.
 ///
@@ -56,6 +80,25 @@ pub trait AsyncRead {
         cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<io::Result<()>>;
+
+    /// Attempts to read the exact number of bytes required to fill `buf`.
+    ///
+    /// On success, returns `Poll::Ready(Ok(()))` and places data in the
+    /// unfilled portion of `buf`.
+    ///
+    /// If EOF was reached, this function returns an error of kind
+    /// `ErrorKind::UnexpectedEof`.
+    ///
+    /// If an error occurs or `Poll::Pending` is returned, `buf` contains
+    /// all bytes read so far. This means that calling `poll_read_exact`
+    /// with the same `ReadBuf` until it succeeds will work as intended.
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        default_poll_read_exact(self, cx, buf)
+    }
 }
 
 macro_rules! deref_async_read {
@@ -67,6 +110,14 @@ macro_rules! deref_async_read {
         ) -> Poll<io::Result<()>> {
             Pin::new(&mut **self).poll_read(cx, buf)
         }
+
+        fn poll_read_exact(
+            mut self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            buf: &mut ReadBuf<'_>,
+        ) -> Poll<io::Result<()>> {
+            Pin::new(&mut **self).poll_read_exact(cx, buf)
+        }
     };
 }
 
@@ -90,6 +141,14 @@ where
     ) -> Poll<io::Result<()>> {
         crate::util::pin_as_deref_mut(self).poll_read(cx, buf)
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        crate::util::pin_as_deref_mut(self).poll_read_exact(cx, buf)
+    }
 }
 
 impl AsyncRead for &[u8] {
@@ -104,6 +163,21 @@ impl AsyncRead for &[u8] {
         *self = b;
         Poll::Ready(Ok(()))
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let _ = self.poll_read(cx, buf);
+
+        // If there was enough data, `poll_read` has filled the buffer
+        if buf.remaining() == 0 {
+            Poll::Ready(Ok(()))
+        } else {
+            Poll::Ready(Err(eof()))
+        }
+    }
 }
 
 impl<T: AsRef<[u8]> + Unpin> AsyncRead for io::Cursor<T> {
@@ -129,4 +203,19 @@ impl<T: AsRef<[u8]> + Unpin> AsyncRead for io::Cursor<T> {
 
         Poll::Ready(Ok(()))
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let _ = self.poll_read(cx, buf);
+
+        // If there was enough data, `poll_read` has filled the buffer
+        if buf.remaining() == 0 {
+            Poll::Ready(Ok(()))
+        } else {
+            Poll::Ready(Err(eof()))
+        }
+    }
 }

tokio/src/io/join.rs

@@ -81,6 +81,14 @@ where
     ) -> Poll<Result<(), io::Error>> {
         self.project().reader.poll_read(cx, buf)
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.project().reader.poll_read_exact(cx, buf)
+    }
 }
 
 impl<R, W> AsyncWrite for Join<R, W>

tokio/src/io/split.rs

@@ -115,6 +115,15 @@ impl<T: AsyncRead> AsyncRead for ReadHalf<T> {
     ) -> Poll<io::Result<()>> {
         self.inner.with_lock(|stream| stream.poll_read(cx, buf))
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.inner
+            .with_lock(|stream| stream.poll_read_exact(cx, buf))
+    }
 }
 
 impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> {

tokio/src/io/stdin.rs

@@ -95,4 +95,12 @@ impl AsyncRead for Stdin {
     ) -> Poll<io::Result<()>> {
         Pin::new(&mut self.std).poll_read(cx, buf)
     }
+
+    fn poll_read_exact(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.std).poll_read_exact(cx, buf)
+    }
 }

tokio/src/io/util/buf_reader.rs

@@ -95,6 +95,15 @@ impl<R: AsyncRead> BufReader<R> {
         *me.pos = 0;
         *me.cap = 0;
     }
+
+    #[cold]
+    fn poll_read_exact_slow(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        crate::io::async_read::default_poll_read_exact(self, cx, buf)
+    }
 }
 
 impl<R: AsyncRead> AsyncRead for BufReader<R> {
@@ -117,6 +126,22 @@ impl<R: AsyncRead> AsyncRead for BufReader<R> {
         self.consume(amt);
         Poll::Ready(Ok(()))
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        // Fast path: all the data we need is already in the buffer
+        if let Some(data) = self.buffer().get(..buf.remaining()) {
+            buf.put_slice(data);
+            *self.project().pos += data.len();
+            debug_assert!(buf.remaining() == 0);
+            return Poll::Ready(Ok(()));
+        }
+
+        self.poll_read_exact_slow(cx, buf)
+    }
 }
 
 impl<R: AsyncRead> AsyncBufRead for BufReader<R> {

tokio/src/io/util/buf_stream.rs

@@ -156,6 +156,14 @@ impl<RW: AsyncRead + AsyncWrite> AsyncRead for BufStream<RW> {
     ) -> Poll<io::Result<()>> {
         self.project().inner.poll_read(cx, buf)
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.project().inner.poll_read_exact(cx, buf)
+    }
 }
 
 /// Seek to an offset, in bytes, in the underlying stream.

tokio/src/io/util/buf_writer.rs

@@ -274,6 +274,14 @@ impl<W: AsyncWrite + AsyncRead> AsyncRead for BufWriter<W> {
     ) -> Poll<io::Result<()>> {
         self.get_pin_mut().poll_read(cx, buf)
     }
+
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.get_pin_mut().poll_read_exact(cx, buf)
+    }
 }
 
 impl<W: AsyncWrite + AsyncBufRead> AsyncBufRead for BufWriter<W> {

tokio/src/io/util/empty.rs

@@ -75,6 +75,22 @@ impl AsyncRead for Empty {
         ready!(poll_proceed_and_make_progress(cx));
         Poll::Ready(Ok(()))
     }
+
+    #[inline]
+    fn poll_read_exact(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        ready!(crate::trace::trace_leaf(cx));
+        ready!(poll_proceed_and_make_progress(cx));
+
+        if buf.remaining() == 0 {
+            Poll::Ready(Ok(()))
+        } else {
+            Poll::Ready(Err(crate::io::async_read::eof()))
+        }
+    }
 }
 
 impl AsyncBufRead for Empty {

tokio/src/io/util/mem.rs

@@ -131,6 +131,15 @@ impl AsyncRead for DuplexStream {
     ) -> Poll<std::io::Result<()>> {
         Pin::new(&mut *self.read.lock()).poll_read(cx, buf)
     }
+
+    #[allow(unused_mut)]
+    fn poll_read_exact(
+        mut self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut *self.read.lock()).poll_read_exact(cx, buf)
+    }
 }
 
 impl AsyncWrite for DuplexStream {

tokio/src/io/util/read_exact.rs

@@ -40,10 +40,6 @@ pin_project! {
     }
 }
 
-fn eof() -> io::Error {
-    io::Error::new(io::ErrorKind::UnexpectedEof, "early eof")
-}
-
 impl<A> Future for ReadExact<'_, A>
 where
     A: AsyncRead + Unpin + ?Sized,
@@ -53,17 +49,8 @@ where
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
         let me = self.project();
 
-        loop {
-            // if our buffer is empty, then we need to read some data to continue.
-            let rem = me.buf.remaining();
-            if rem != 0 {
-                ready!(Pin::new(&mut *me.reader).poll_read(cx, me.buf))?;
-                if me.buf.remaining() == rem {
-                    return Err(eof()).into();
-                }
-            } else {
-                return Poll::Ready(Ok(me.buf.capacity()));
-            }
-        }
+        ready!(Pin::new(&mut *me.reader).poll_read_exact(cx, me.buf))?;
+
+        Poll::Ready(Ok(me.buf.capacity()))
     }
 }