io: make read_to_end not grow unnecessarily

Fixes: #5594

tokio/src/io/util/read_to_end.rs

@@ -1,5 +1,5 @@
 use crate::io::util::vec_with_initialized::{into_read_buf_parts, VecU8, VecWithInitialized};
-use crate::io::AsyncRead;
+use crate::io::{AsyncRead, ReadBuf};
 
 use pin_project_lite::pin_project;
 use std::future::Future;
@@ -68,19 +68,29 @@ fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>(
     // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every
     // time is 4,500 times (!) slower than this if the reader has a very small
     // amount of data to return.
-    buf.reserve(32);
+    let try_small_read = buf.reserve(32);
 
     // Get a ReadBuf into the vector.
     let mut read_buf = buf.get_read_buf();
 
+    let poll_result;
     let filled_before = read_buf.filled().len();
-    let poll_result = read.poll_read(cx, &mut read_buf);
-    let filled_after = read_buf.filled().len();
-    let n = filled_after - filled_before;
-
+    let filled_after;
+    if try_small_read {
+        let mut small_buf = Vec::with_capacity(32);
+        let mut small_read_buf = ReadBuf::new(&mut small_buf);
+        poll_result = read.poll_read(cx, &mut small_read_buf);
+        let filled = small_read_buf.filled().len();
+        read_buf.put_slice(&small_buf[..filled]);
+        filled_after = filled_before + filled;
+    } else {
+        poll_result = read.poll_read(cx, &mut read_buf);
+        filled_after = read_buf.filled().len();
+    };
     // Update the length of the vector using the result of poll_read.
     let read_buf_parts = into_read_buf_parts(read_buf);
     buf.apply_read_buf(read_buf_parts);
+    let n = filled_after - filled_before;
 
     match poll_result {
         Poll::Pending => {

tokio/src/io/util/vec_with_initialized.rs

@@ -28,6 +28,7 @@ pub(crate) struct VecWithInitialized<V> {
     // The number of initialized bytes in the vector.
     // Always between `vec.len()` and `vec.capacity()`.
     num_initialized: usize,
+    starting_capacity: usize,
 }
 
 impl VecWithInitialized<Vec<u8>> {
@@ -47,19 +48,28 @@ where
         // to its length are initialized.
         Self {
             num_initialized: vec.as_mut().len(),
+            starting_capacity: vec.as_ref().capacity(),
             vec,
         }
     }
 
-    pub(crate) fn reserve(&mut self, num_bytes: usize) {
+    // Returns a boolean telling the caller to try reading into a small local buffer first if true.
+    // Doing so would avoid overallocating when vec is filled to capacity and we reached EOF.
+    pub(crate) fn reserve(&mut self, num_bytes: usize) -> bool {
         let vec = self.vec.as_mut();
         if vec.capacity() - vec.len() >= num_bytes {
-            return;
+            return false;
         }
+
+        if self.starting_capacity == vec.capacity() && self.starting_capacity >= num_bytes {
+            return true;
+        }
+
         // SAFETY: Setting num_initialized to `vec.len()` is correct as
         // `reserve` does not change the length of the vector.
         self.num_initialized = vec.len();
         vec.reserve(num_bytes);
+        false
     }
 
     #[cfg(feature = "io-util")]

tokio/tests/io_read_to_end.rs

@@ -5,6 +5,7 @@ use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
 use tokio_test::assert_ok;
+use tokio_test::io::Builder;
 
 #[tokio::test]
 async fn read_to_end() {
@@ -76,3 +77,14 @@ async fn read_to_end_uninit() {
     test.read_to_end(&mut buf).await.unwrap();
     assert_eq!(buf.len(), 33);
 }
+
+#[tokio::test]
+async fn read_to_end_doesnt_grow_with_capacity() {
+    let bytes = b"imlargerthan32bytessoIcanhelpwiththetest";
+    let mut mock = Builder::new().read(bytes).build();
+    let mut buf = Vec::with_capacity(bytes.len());
+    AsyncReadExt::read_to_end(&mut mock, &mut buf)
+        .await
+        .unwrap();
+    assert_eq!(bytes.len(), buf.capacity());
+}