From 666257f6a84bb690db36d80a116a6a601503fa16 Mon Sep 17 00:00:00 2001
From: Timmy Xiao <34635512+tzx@users.noreply.github.com>
Date: Sat, 8 Apr 2023 16:58:51 -0400
Subject: [PATCH] Don't use bool in reserve and use stack allocated buffer

---
 tokio/src/io/util/read_to_end.rs          | 28 +++++++++++++----------
 tokio/src/io/util/vec_with_initialized.rs | 16 +++++++------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/tokio/src/io/util/read_to_end.rs b/tokio/src/io/util/read_to_end.rs
index c43a060175b..9f66cb0aa37 100644
--- a/tokio/src/io/util/read_to_end.rs
+++ b/tokio/src/io/util/read_to_end.rs
@@ -5,7 +5,7 @@ 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};
 
@@ -68,25 +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.
-    let try_small_read = buf.reserve(32);
+    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 poll_result;
     let filled_before = read_buf.filled().len();
-    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);
+    let poll_result;
+
+    let filled_after = if try_small_read {
+        let mut small_buf: [MaybeUninit<u8>; NUM_BYTES] =
+            unsafe { MaybeUninit::uninit().assume_init() };
+        let mut small_read_buf = ReadBuf::uninit(&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;
+        let filled = small_read_buf.filled();
+        read_buf.put_slice(filled);
+        filled_before + filled.len()
     } else {
+        buf.reserve(NUM_BYTES);
+        read_buf = buf.get_read_buf();
         poll_result = read.poll_read(cx, &mut read_buf);
-        filled_after = read_buf.filled().len();
+        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);
diff --git a/tokio/src/io/util/vec_with_initialized.rs b/tokio/src/io/util/vec_with_initialized.rs
index 952915c4c7c..f932db738ee 100644
--- a/tokio/src/io/util/vec_with_initialized.rs
+++ b/tokio/src/io/util/vec_with_initialized.rs
@@ -55,21 +55,16 @@ where
 
     // 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 {
+    pub(crate) fn reserve(&mut self, num_bytes: usize) {
         let vec = self.vec.as_mut();
         if vec.capacity() - vec.len() >= num_bytes {
-            return false;
-        }
-
-        if self.starting_capacity == vec.capacity() && self.starting_capacity >= num_bytes {
-            return true;
+            return;
         }
 
         // 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")]
@@ -121,6 +116,13 @@ where
             vec.set_len(parts.len);
         }
     }
+
+    pub(crate) fn try_small_read_first(&self, num_bytes: usize) -> bool {
+        let vec = self.vec.as_ref();
+        vec.capacity() - vec.len() < num_bytes
+            && self.starting_capacity == vec.capacity()
+            && self.starting_capacity >= num_bytes
+    }
 }
 
 pub(crate) struct ReadBufParts {