Fix read_to_end to not grow an exact size buffer

If you know how much data to expect and use `Vec::with_capacity` to
pre-allocate a buffer of that capacity, `Read::read_to_end` will still
double its capacity. It needs some space to perform a read, even though
that read ends up returning `0`.

It's a bummer to carefully pre-allocate 1GB to read a 1GB file into
memory and end up using 2GB.

This fixes that behavior by special casing a full buffer and reading
into a small "probe" buffer instead. If that read returns `0` then it's
confirmed that the buffer was the perfect size. If it doesn't, the probe
buffer is appended to the normal buffer and the read loop continues.

Fixing this allows several workarounds in the standard library to be
removed:

- `Take` no longer needs to override `Read::read_to_end`.
- The `reservation_size` callback that allowed `Take` to inhibit the
  previous over-allocation behavior isn't needed.
- `fs::read` doesn't need to reserve an extra byte in
  `initial_buffer_size`.

Curiously, there was a unit test that specifically checked that
`Read::read_to_end` *does* over-allocate. I removed that test, too.

diff --git a/library/std/src/fs.rs b/library/std/src/fs.rs
index bdb172907ffed39b770d802093c2ea416d485365..bcc1f9a8b94dda77e4ea406d3d1d662b1b1549b1 100644 (file)
--- a/library/std/src/fs.rs
+++ b/library/std/src/fs.rs
@@ -200,10 +200,9 @@ pub struct DirBuilder {
 
 /// Indicates how large a buffer to pre-allocate before reading the entire file.
 fn initial_buffer_size(file: &File) -> usize {
-    // Allocate one extra byte so the buffer doesn't need to grow before the
-    // final `read` call at the end of the file.  Don't worry about `usize`
-    // overflow because reading will fail regardless in that case.
-    file.metadata().map(|m| m.len() as usize + 1).unwrap_or(0)
+    // Don't worry about `usize` overflow because reading will fail regardless
+    // in that case.
+    file.metadata().map(|m| m.len() as usize).unwrap_or(0)
 }
 
 /// Read the entire contents of a file into a bytes vector.

diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs
index 4a35d36a9def72f284b48244219b347a9d140e80..52b04680745c9a73d7b2fa98db2b1719e0ac1c98 100644 (file)
--- a/library/std/src/io/mod.rs
+++ b/library/std/src/io/mod.rs
@@ -362,22 +362,18 @@ fn append_to_string<F>(buf: &mut String, f: F) -> Result<usize>
 // Because we're extending the buffer with uninitialized data for trusted
 // readers, we need to make sure to truncate that if any of this panics.
 fn read_to_end<R: Read + ?Sized>(r: &mut R, buf: &mut Vec<u8>) -> Result<usize> {
-    read_to_end_with_reservation(r, buf, |_| 32)
-}
-
-fn read_to_end_with_reservation<R, F>(
-    r: &mut R,
-    buf: &mut Vec<u8>,
-    mut reservation_size: F,
-) -> Result<usize>
-where
-    R: Read + ?Sized,
-    F: FnMut(&R) -> usize,
-{
     let start_len = buf.len();
+    let start_cap = buf.capacity();
     let mut g = Guard { len: buf.len(), buf };
     loop {
-        if g.len == g.buf.len() {
+        // If we've read all the way up to the capacity, reserve more space.
+        if g.len == g.buf.capacity() {
+            g.buf.reserve(32);
+        }
+
+        // Initialize any excess capacity and adjust the length so we can write
+        // to it.
+        if g.buf.len() < g.buf.capacity() {
             unsafe {
                 // FIXME(danielhenrymantilla): #42788
                 //
@@ -387,7 +383,6 @@ fn read_to_end_with_reservation<R, F>(
                 //   - Only the standard library gets to soundly "ignore" this,
                 //     based on its privileged knowledge of unstable rustc
                 //     internals;
-                g.buf.reserve(reservation_size(r));
                 let capacity = g.buf.capacity();
                 g.buf.set_len(capacity);
                 r.initializer().initialize(&mut g.buf[g.len..]);
@@ -404,9 +399,30 @@ fn read_to_end_with_reservation<R, F>(
                 assert!(n <= buf.len());
                 g.len += n;
             }
-            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
+            Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
             Err(e) => return Err(e),
         }
+
+        if g.len == g.buf.capacity() && g.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
+            // probe buffer to the primary buffer and let its capacity grow.
+            let mut probe = [0u8; 32];
+
+            loop {
+                match r.read(&mut probe) {
+                    Ok(0) => return Ok(g.len - start_len),
+                    Ok(n) => {
+                        g.buf.extend_from_slice(&probe[..n]);
+                        g.len += n;
+                        break;
+                    }
+                    Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
+                    Err(e) => return Err(e),
+                }
+            }
+        }
     }
 }
 
@@ -2583,13 +2599,6 @@ fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
     unsafe fn initializer(&self) -> Initializer {
         self.inner.initializer()
     }
-
-    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
-        // Pass in a reservation_size closure that respects the current value
-        // of limit for each read. If we hit the read limit, this prevents the
-        // final zero-byte read from allocating again.
-        read_to_end_with_reservation(self, buf, |self_| cmp::min(self_.limit, 32) as usize)
-    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

diff --git a/library/std/src/io/tests.rs b/library/std/src/io/tests.rs
index 1beb72a9a5072344a696c13eca99646ea5d0e65d..58b2d1d14c9279cc9f8aabe1e1bb26c0269c1388 100644 (file)
--- a/library/std/src/io/tests.rs
+++ b/library/std/src/io/tests.rs
@@ -362,24 +362,12 @@ fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
 fn test_read_to_end_capacity() -> io::Result<()> {
     let input = &b"foo"[..];
 
-    // read_to_end() generally needs to over-allocate, both for efficiency
-    // and so that it can distinguish EOF. Assert that this is the case
-    // with this simple ExampleSliceReader struct, which uses the default
-    // implementation of read_to_end. Even though vec1 is allocated with
-    // exactly enough capacity for the read, read_to_end will allocate more
-    // space here.
+    // read_to_end() takes care not to over-allocate when a buffer is the
+    // exact size needed.
     let mut vec1 = Vec::with_capacity(input.len());
     ExampleSliceReader { slice: input }.read_to_end(&mut vec1)?;
     assert_eq!(vec1.len(), input.len());
-    assert!(vec1.capacity() > input.len(), "allocated more");
-
-    // However, std::io::Take includes an implementation of read_to_end
-    // that will not allocate when the limit has already been reached. In
-    // this case, vec2 never grows.
-    let mut vec2 = Vec::with_capacity(input.len());
-    ExampleSliceReader { slice: input }.take(input.len() as u64).read_to_end(&mut vec2)?;
-    assert_eq!(vec2.len(), input.len());
-    assert_eq!(vec2.capacity(), input.len(), "did not allocate more");
+    assert_eq!(vec1.capacity(), input.len(), "did not allocate more");
 
     Ok(())
 }