Auto merge of #78768 - mzabaluev:optimize-buf-writer, r=cramertj

Use is_write_vectored to optimize the write_vectored implementation for BufWriter

In case when the underlying writer does not have an efficient implementation `write_vectored`, the present implementation of
`write_vectored` for `BufWriter` may still forward vectored writes directly to the writer depending on the total length of the data. This misses the advantage of buffering, as the actually written slice may be small.

Provide an alternative code path for the non-vectored case, where the slices passed to `BufWriter` are coalesced in the buffer before being flushed to the underlying writer with plain `write` calls. The buffer is only bypassed if an individual slice's length is at least as large as the buffer.

Remove a FIXME comment referring to #72919 as the issue has been closed with an explanation provided.

2 files changed, 53 insertions, 14 deletions

diff --git a/library/std/src/io/buffered/bufwriter.rs b/library/std/src/io/buffered/bufwriter.rs
index 067ed6ba7ff..3b3399860ba 100644
--- a/library/std/src/io/buffered/bufwriter.rs
+++ b/library/std/src/io/buffered/bufwriter.rs
@@ -328,24 +328,57 @@ impl<W: Write> Write for BufWriter<W> {
     }
 
     fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
-        let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
-        if self.buf.len() + total_len > self.buf.capacity() {
-            self.flush_buf()?;
-        }
-        // FIXME: Why no len > capacity? Why not buffer len == capacity? #72919
-        if total_len >= self.buf.capacity() {
-            self.panicked = true;
-            let r = self.get_mut().write_vectored(bufs);
-            self.panicked = false;
-            r
+        if self.get_ref().is_write_vectored() {
+            let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
+            if self.buf.len() + total_len > self.buf.capacity() {
+                self.flush_buf()?;
+            }
+            if total_len >= self.buf.capacity() {
+                self.panicked = true;
+                let r = self.get_mut().write_vectored(bufs);
+                self.panicked = false;
+                r
+            } else {
+                bufs.iter().for_each(|b| self.buf.extend_from_slice(b));
+                Ok(total_len)
+            }
         } else {
-            bufs.iter().for_each(|b| self.buf.extend_from_slice(b));
-            Ok(total_len)
+            let mut iter = bufs.iter();
+            let mut total_written = if let Some(buf) = iter.by_ref().find(|&buf| !buf.is_empty()) {
+                // This is the first non-empty slice to write, so if it does
+                // not fit in the buffer, we still get to flush and proceed.
+                if self.buf.len() + buf.len() > self.buf.capacity() {
+                    self.flush_buf()?;
+                }
+                if buf.len() >= self.buf.capacity() {
+                    // The slice is at least as large as the buffering capacity,
+                    // so it's better to write it directly, bypassing the buffer.
+                    self.panicked = true;
+                    let r = self.get_mut().write(buf);
+                    self.panicked = false;
+                    return r;
+                } else {
+                    self.buf.extend_from_slice(buf);
+                    buf.len()
+                }
+            } else {
+                return Ok(0);
+            };
+            debug_assert!(total_written != 0);
+            for buf in iter {
+                if self.buf.len() + buf.len() > self.buf.capacity() {
+                    break;
+                } else {
+                    self.buf.extend_from_slice(buf);
+                    total_written += buf.len();
+                }
+            }
+            Ok(total_written)
         }
     }
 
     fn is_write_vectored(&self) -> bool {
-        self.get_ref().is_write_vectored()
+        true
     }
 
     fn flush(&mut self) -> io::Result<()> {
diff --git a/library/std/src/io/buffered/linewritershim.rs b/library/std/src/io/buffered/linewritershim.rs
index a80d08db869..d0c859d2e0c 100644
--- a/library/std/src/io/buffered/linewritershim.rs
+++ b/library/std/src/io/buffered/linewritershim.rs
@@ -20,6 +20,12 @@ impl<'a, W: Write> LineWriterShim<'a, W> {
         Self { buffer }
     }
 
+    /// Get a reference to the inner writer (that is, the writer
+    /// wrapped by the BufWriter).
+    fn inner(&self) -> &W {
+        self.buffer.get_ref()
+    }
+
     /// Get a mutable reference to the inner writer (that is, the writer
     /// wrapped by the BufWriter). Be careful with this writer, as writes to
     /// it will bypass the buffer.
@@ -227,7 +233,7 @@ impl<'a, W: Write> Write for LineWriterShim<'a, W> {
     }
 
     fn is_write_vectored(&self) -> bool {
-        self.buffer.is_write_vectored()
+        self.inner().is_write_vectored()
     }
 
     /// Write some data into this BufReader with line buffering. This means