Fix perf regression in Read::read_to_end on short reads due to not checking if the cursor has initialized bytes

library/alloc/src/io/read.rs

@@ -1,4 +1,3 @@
-use core::cmp;
 use core::mem::{DropGuard, MaybeUninit};
 
 use crate::io::{
@@ -822,9 +821,11 @@ where
 /// - avoid allocating unless necessary
 /// - avoid overallocating if we know the exact size (#89165)
 /// - avoid passing large buffers to readers that always initialize the free capacity if they perform short reads (#23815, #23820)
+/// - avoid re-initializing unfilled bytes into the spare buffer if we initialized >PROBE_SIZE unfilled bytes in a previous loop (#158008)
 /// - pass large buffers to readers that do not initialize the spare capacity. this can amortize per-call overheads
-/// - and finally pass not-too-small and not-too-large buffers to Windows read APIs because they manage to suffer from both problems
+/// - pass not-too-small and not-too-large buffers to Windows read APIs because they manage to suffer from both problems
 ///   at the same time, i.e. small reads suffer from syscall overhead, all reads incur costs proportional to buffer size (#110650)
+/// - also avoid <4 byte reads as this may split UTF-8 code points, which can be a problem for Windows console reads (#142847)
 #[doc(hidden)]
 #[unstable(feature = "core_io_internals", reason = "exposed only for libstd", issue = "none")]
 pub fn default_read_to_end<R: Read + ?Sized>(
@@ -840,6 +841,9 @@ pub fn default_read_to_end<R: Read + ?Sized>(
         .and_then(|s| s.checked_add(1024)?.checked_next_multiple_of(DEFAULT_BUF_SIZE))
         .unwrap_or(DEFAULT_BUF_SIZE);
 
+    // Tracks how many bytes are initialized in the buffer
+    let mut init_until = buf.len();
+
     const PROBE_SIZE: usize = 32;
 
     fn small_probe_read<R: Read + ?Sized>(r: &mut R, buf: &mut Vec<u8>) -> Result<usize> {
@@ -887,7 +891,7 @@ pub fn default_read_to_end<R: Read + ?Sized>(
     }
 
     loop {
-        if buf.len() == buf.capacity() && buf.capacity() == start_cap {
+        if buf.spare_capacity_mut().len() < PROBE_SIZE && buf.capacity() == start_cap {
             // The buffer might be an exact fit. Let's read into a probe buffer
             // and see if it returns `Ok(0)`. If so, we've avoided an
             // unnecessary doubling of the capacity. But if not, append the
@@ -897,20 +901,42 @@ pub fn default_read_to_end<R: Read + ?Sized>(
             if read == 0 {
                 return Ok(buf.len() - start_len);
             }
+
+            init_until = buf.len();
+            // In the case of very short reads, continue to use the stack buffer
+            // until either we reach the end or we need to reallocate.
+            continue;
         }
 
-        if buf.len() == buf.capacity() {
-            // buf is full, need more space
+        // Avoid unnecessarily short reads by ensuring there's at least PROBE_SIZE space available.
+        // And assert that PROBE_SIZE is always at least large enough to fit any UTF-8 encoded code point.
+        const { assert!(PROBE_SIZE >= char::MAX_LEN_UTF8) }
+        if buf.spare_capacity_mut().len() < PROBE_SIZE {
             buf.try_reserve(PROBE_SIZE)?;
+            // When reallocation occurs, we have to update init_until accordingly
+            // to re-calibrate how many bytes are actually initialized in the buffer
+            init_until = buf.len();
         }
 
+        // We set a threshold of >PROBE_SIZE initialized yet unfilled bytes left in the
+        // spare buffer before determining that we need to initialize more bytes into
+        // the spare buffer
+        let buf_len = if init_until > buf.len() + PROBE_SIZE {
+            init_until - buf.len()
+        } else {
+            usize::min(max_read_size, buf.capacity() - buf.len())
+        };
+        let was_init = init_until >= buf.len() + buf_len;
+
         let mut spare = buf.spare_capacity_mut();
-        let buf_len = cmp::min(spare.len(), max_read_size);
         spare = &mut spare[..buf_len];
         let mut read_buf: BorrowedBuf<'_, u8> = spare.into();
 
-        // Note that we don't track already initialized bytes here, but this is fine
-        // because we explicitly limit the read size
+        if was_init {
+            // SAFETY: These bytes were initialized but not filled in the previous loop
+            unsafe { read_buf.set_init() };
+        }
+
         let mut cursor = read_buf.unfilled();
         let result = loop {
             match r.read_buf(cursor.reborrow()) {
@@ -924,6 +950,10 @@ pub fn default_read_to_end<R: Read + ?Sized>(
         let bytes_read = cursor.written();
         let is_init = read_buf.is_init();
 
+        if is_init {
+            init_until = buf.len() + buf_len;
+        }
+
         // SAFETY: BorrowedBuf's invariants mean this much memory is initialized.
         unsafe {
             let new_len = bytes_read + buf.len();
@@ -948,9 +978,11 @@ pub fn default_read_to_end<R: Read + ?Sized>(
             if !is_init {
                 max_read_size = usize::MAX;
             }
-            // we have passed a larger buffer than previously and the
-            // reader still hasn't returned a short read
-            else if buf_len >= max_read_size && bytes_read == buf_len {
+            // the spare buffer has initialized and read in `max_read_size` bytes.
+            // it's possible that we have more than `max_read_size` bytes to read
+            // left, so a larger buffer may be necessary to minimize the number of
+            // iterations of reading in bytes to the buffer
+            else if bytes_read == max_read_size {
                 max_read_size = max_read_size.saturating_mul(2);
             }
         }