1 files changed, 42 insertions, 11 deletions

diff --git a/vendor/tokio/src/io/util/read_to_end.rs b/vendor/tokio/src/io/util/read_to_end.rs
index f4a564d7d..8edba2a17 100644
--- a/vendor/tokio/src/io/util/read_to_end.rs
+++ b/vendor/tokio/src/io/util/read_to_end.rs
@@ -1,11 +1,11 @@
 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;
 use std::io;
 use std::marker::PhantomPinned;
-use std::mem;
+use std::mem::{self, MaybeUninit};
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
@@ -67,16 +67,47 @@ fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>(
     // has 4 bytes while still making large reads if the reader does have a ton
     // 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);
+    // amount of data to return. When the vector is full with its starting
+    // capacity, we first try to read into a small buffer to see if we reached
+    // an EOF. This only happens when the starting capacity is >= NUM_BYTES, since
+    // we allocate at least NUM_BYTES each time. This avoids the unnecessary
+    // allocation that we attempt before reading into the vector.
+
+    const NUM_BYTES: usize = 32;
+    let try_small_read = buf.try_small_read_first(NUM_BYTES);
 
     // Get a ReadBuf into the vector.
-    let mut read_buf = buf.get_read_buf();
+    let mut read_buf;
+    let poll_result;
+
+    let n = if try_small_read {
+        // Read some bytes using a small read.
+        let mut small_buf: [MaybeUninit<u8>; NUM_BYTES] = [MaybeUninit::uninit(); NUM_BYTES];
+        let mut small_read_buf = ReadBuf::uninit(&mut small_buf);
+        poll_result = read.poll_read(cx, &mut small_read_buf);
+        let to_write = small_read_buf.filled();
+
+        // Ensure we have enough space to fill our vector with what we read.
+        read_buf = buf.get_read_buf();
+        if to_write.len() > read_buf.remaining() {
+            buf.reserve(NUM_BYTES);
+            read_buf = buf.get_read_buf();
+        }
+        read_buf.put_slice(to_write);
+
+        to_write.len()
+    } else {
+        // Ensure we have enough space for reading.
+        buf.reserve(NUM_BYTES);
+        read_buf = buf.get_read_buf();
+
+        // Read data directly into vector.
+        let filled_before = read_buf.filled().len();
+        poll_result = read.poll_read(cx, &mut read_buf);
 
-    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;
+        // Compute the number of bytes read.
+        read_buf.filled().len() - filled_before
+    };
 
     // Update the length of the vector using the result of poll_read.
     let read_buf_parts = into_read_buf_parts(read_buf);
@@ -87,11 +118,11 @@ fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>(
             // In this case, nothing should have been read. However we still
             // update the vector in case the poll_read call initialized parts of
             // the vector's unused capacity.
-            debug_assert_eq!(filled_before, filled_after);
+            debug_assert_eq!(n, 0);
             Poll::Pending
         }
         Poll::Ready(Err(err)) => {
-            debug_assert_eq!(filled_before, filled_after);
+            debug_assert_eq!(n, 0);
             Poll::Ready(Err(err))
         }
         Poll::Ready(Ok(())) => Poll::Ready(Ok(n)),