LineWriter: cap write_vectored newline scan to avoid quadratic write_all_vectored

library/std/src/io/buffered/linewritershim.rs

@@ -49,6 +49,78 @@ impl<'a, W: ?Sized + Write> LineWriterShim<'a, W> {
             _ => Ok(()),
         }
     }
+
+    /// Vectored line-buffered write over an already-capped list of buffers.
+    ///
+    /// The caller is responsible for trimming `bufs` to the prefix it is
+    /// willing to scan (see `MAX_BUFS_TO_SCAN`). This method only ever writes
+    /// or buffers bytes from `bufs`, so any newline it might bury in the
+    /// `BufWriter` is one it has itself scanned for -- buffers the caller
+    /// dropped past the cap can never end up stuck in the buffer. Bytes not
+    /// accounted for in the return value are left for the next call.
+    fn write_vectored_scanned(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        // Find the buffer containing the last newline.
+        let last_newline_buf_idx = bufs
+            .iter()
+            .enumerate()
+            .rev()
+            .find_map(|(i, buf)| memchr::memchr(b'\n', buf).map(|_| i));
+
+        // If there are no new newlines (that is, if this write is less than
+        // one line), just do a regular buffered write.
+        let last_newline_buf_idx = match last_newline_buf_idx {
+            None => {
+                self.flush_if_completed_line()?;
+                return self.buffer.write_vectored(bufs);
+            }
+            Some(i) => i,
+        };
+
+        // Flush existing content to prepare for our write.
+        self.buffer.flush_buf()?;
+
+        // This is what we're going to try to write directly to the inner
+        // writer. The rest will be buffered, if nothing goes wrong.
+        let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1);
+
+        // Write `lines` directly to the inner writer. In keeping with the
+        // `write` convention, make at most one attempt to add new (unbuffered)
+        // data. Because this write doesn't touch the BufWriter state directly,
+        // and the buffer is known to be empty, we don't need to worry about
+        // self.panicked here.
+        let flushed = self.inner_mut().write_vectored(lines)?;
+
+        // If inner returns Ok(0), propagate that to the caller without
+        // doing additional buffering; otherwise we're just guaranteeing
+        // an "ErrorKind::WriteZero" later.
+        if flushed == 0 {
+            return Ok(0);
+        }
+
+        // Don't try to reconstruct the exact amount written; just bail
+        // in the event of a partial write.
+        let mut lines_len: usize = 0;
+        for buf in lines {
+            // With overlapping/duplicate slices the total length may in theory
+            // exceed usize::MAX
+            lines_len = lines_len.saturating_add(buf.len());
+            if flushed < lines_len {
+                return Ok(flushed);
+            }
+        }
+
+        // Now that the write has succeeded, buffer the rest (or as much of the
+        // rest as possible). `tail` is the part of the scanned prefix after the
+        // last newline, so it cannot contain a newline of its own.
+        let buffered: usize = tail
+            .iter()
+            .filter(|buf| !buf.is_empty())
+            .map(|buf| self.buffer.write_to_buf(buf))
+            .take_while(|&n| n > 0)
+            .sum();
+
+        Ok(flushed + buffered)
+    }
 }
 
 impl<'a, W: ?Sized + Write> Write for LineWriterShim<'a, W> {
@@ -185,71 +257,26 @@ impl<'a, W: ?Sized + Write> Write for LineWriterShim<'a, W> {
             };
         }
 
-        // Find the buffer containing the last newline
-        // FIXME: This is overly slow if there are very many bufs and none contain
-        // newlines. e.g. writev() on Linux only writes up to 1024 slices, so
-        // scanning the rest is wasted effort. This makes write_all_vectored()
-        // quadratic.
-        let last_newline_buf_idx = bufs
-            .iter()
-            .enumerate()
-            .rev()
-            .find_map(|(i, buf)| memchr::memchr(b'\n', buf).map(|_| i));
-
-        // If there are no new newlines (that is, if this write is less than
-        // one line), just do a regular buffered write
-        let last_newline_buf_idx = match last_newline_buf_idx {
-            // No newlines; just do a normal buffered write
-            None => {
-                self.flush_if_completed_line()?;
-                return self.buffer.write_vectored(bufs);
-            }
-            Some(i) => i,
-        };
-
-        // Flush existing content to prepare for our write
-        self.buffer.flush_buf()?;
-
-        // This is what we're going to try to write directly to the inner
-        // writer. The rest will be buffered, if nothing goes wrong.
-        let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1);
-
-        // Write `lines` directly to the inner writer. In keeping with the
-        // `write` convention, make at most one attempt to add new (unbuffered)
-        // data. Because this write doesn't touch the BufWriter state directly,
-        // and the buffer is known to be empty, we don't need to worry about
-        // self.panicked here.
-        let flushed = self.inner_mut().write_vectored(lines)?;
-
-        // If inner returns Ok(0), propagate that to the caller without
-        // doing additional buffering; otherwise we're just guaranteeing
-        // an "ErrorKind::WriteZero" later.
-        if flushed == 0 {
-            return Ok(0);
-        }
-
-        // Don't try to reconstruct the exact amount written; just bail
-        // in the event of a partial write
-        let mut lines_len: usize = 0;
-        for buf in lines {
-            // With overlapping/duplicate slices the total length may in theory
-            // exceed usize::MAX
-            lines_len = lines_len.saturating_add(buf.len());
-            if flushed < lines_len {
-                return Ok(flushed);
-            }
-        }
-
-        // Now that the write has succeeded, buffer the rest (or as much of the
-        // rest as possible)
-        let buffered: usize = tail
-            .iter()
-            .filter(|buf| !buf.is_empty())
-            .map(|buf| self.buffer.write_to_buf(buf))
-            .take_while(|&n| n > 0)
-            .sum();
-
-        Ok(flushed + buffered)
+        // Only scan (and operate on) the first MAX_BUFS_TO_SCAN slices. The cap
+        // is what keeps write_all_vectored() from going quadratic when callers
+        // pass many newline-free slices -- without it, every iteration of the
+        // outer loop rescans every remaining buffer. 1024 is a portable,
+        // generous upper bound: it is the value of UIO_MAXIOV / IOV_MAX on
+        // Linux and the BSDs (and the hardcoded cap in
+        // sys::net::connection::socket::solid), so on those platforms it also
+        // lines up with the most a single writev() can retire. On platforms
+        // whose syscall cap is smaller (POSIX requires only 16) or that have no
+        // cap at all (Windows), the constant still serves its primary purpose
+        // of bounding scan work.
+        //
+        // Everything past the cap is left untouched for the next call; the
+        // outer loop in write_all_vectored() makes forward progress via the
+        // short return value, and correctness is preserved everywhere. We hand
+        // the capped prefix to a helper so the rest of the logic can only ever
+        // see -- and therefore only ever write or buffer -- buffers we have
+        // actually scanned for newlines.
+        const MAX_BUFS_TO_SCAN: usize = 1024;
+        self.write_vectored_scanned(&bufs[..bufs.len().min(MAX_BUFS_TO_SCAN)])
     }
 
     fn is_write_vectored(&self) -> bool {

library/std/src/io/buffered/tests.rs

@@ -752,6 +752,133 @@ fn line_vectored_partial_and_errors() {
     }
 }
 
+/// Regression test for the quadratic scan in `LineWriterShim::write_vectored`.
+///
+/// When given a long list of newline-free buffers and a sink that only
+/// retires a few slices per call, the previous implementation rescanned all
+/// remaining buffers on every iteration of `write_all_vectored`, producing
+/// O(N^2) work. The fix caps the scan to a constant prefix; this test
+/// verifies that bytes still come out in order when the only newline lives
+/// past that cap.
+#[test]
+fn line_vectored_long_input_past_scan_cap() {
+    /// Vectored sink that retires at most one non-empty slice per call,
+    /// mimicking a writev() bounded by IOV_MAX.
+    #[derive(Default)]
+    struct OneSliceSink {
+        buffer: Vec<u8>,
+    }
+
+    impl Write for OneSliceSink {
+        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+            self.buffer.extend_from_slice(buf);
+            Ok(buf.len())
+        }
+        fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+            for b in bufs {
+                if !b.is_empty() {
+                    self.buffer.extend_from_slice(b);
+                    return Ok(b.len());
+                }
+            }
+            Ok(0)
+        }
+        fn flush(&mut self) -> io::Result<()> {
+            Ok(())
+        }
+        fn is_write_vectored(&self) -> bool {
+            true
+        }
+    }
+
+    // Place the only newline well past the 1024-slice scan cap.
+    const N: usize = 1500;
+    const NEWLINE_AT: usize = 1400;
+    let bytes: Vec<u8> = (0..N).map(|i| if i == NEWLINE_AT { b'\n' } else { b'a' }).collect();
+    let mut io_slices: Vec<IoSlice<'_>> = bytes.chunks(1).map(IoSlice::new).collect();
+
+    let mut writer = LineWriter::new(OneSliceSink::default());
+    writer.write_all_vectored(&mut io_slices).unwrap();
+
+    // The newline past the scan cap must still trigger a flush through
+    // the inner writer; only the tail after the newline may remain buffered.
+    assert_eq!(writer.get_ref().buffer, bytes[..=NEWLINE_AT]);
+
+    writer.flush().unwrap();
+    assert_eq!(writer.get_ref().buffer, bytes);
+}
+
+/// Regression test for newlines buried past the `write_vectored` scan cap.
+///
+/// `LineWriterShim::write_vectored` only scans a bounded prefix of the slice
+/// list for newlines. An earlier version computed the head/tail split against
+/// the *full* list, so when more than the cap's worth of slices were passed in
+/// a single call and one of the unscanned slices held a newline, that newline
+/// was silently copied into the inner `BufWriter` instead of being flushed --
+/// leaving a completed line stuck in the buffer.
+///
+/// The invariant this checks is the core line-buffering guarantee: once all
+/// input has been written, everything up to and including the *last* newline
+/// must have reached the inner writer, and only the trailing partial line may
+/// remain buffered. It holds regardless of the exact scan-cap value, and it is
+/// driven through the public `write_all_vectored` entry point so the cap
+/// boundary is crossed naturally.
+#[test]
+fn line_vectored_flushes_newline_past_scan_cap() {
+    /// Vectored sink that accepts every slice in full, so any data that fails
+    /// to reach it must have been (incorrectly) left in the LineWriter buffer.
+    #[derive(Default)]
+    struct FullSink {
+        buffer: Vec<u8>,
+    }
+
+    impl Write for FullSink {
+        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+            self.buffer.extend_from_slice(buf);
+            Ok(buf.len())
+        }
+        fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+            let mut written = 0;
+            for b in bufs {
+                self.buffer.extend_from_slice(b);
+                written += b.len();
+            }
+            Ok(written)
+        }
+        fn flush(&mut self) -> io::Result<()> {
+            Ok(())
+        }
+        fn is_write_vectored(&self) -> bool {
+            true
+        }
+    }
+
+    // Two newlines: one comfortably inside the scan cap and one well past it,
+    // so a single `write_vectored` call sees the first but not the second.
+    const N: usize = 1100;
+    const FIRST_NEWLINE: usize = 5;
+    const LAST_NEWLINE: usize = 1050;
+    let bytes: Vec<u8> = (0..N)
+        .map(|i| if i == FIRST_NEWLINE || i == LAST_NEWLINE { b'\n' } else { b'a' })
+        .collect();
+    let mut io_slices: Vec<IoSlice<'_>> = bytes.chunks(1).map(IoSlice::new).collect();
+
+    // The buffer has to be large enough to hold the whole tail past the first
+    // newline; otherwise it fills up before the buggy split would even reach
+    // the second newline, masking the bug.
+    let mut writer = LineWriter::with_capacity(4096, FullSink::default());
+    writer.write_all_vectored(&mut io_slices).unwrap();
+
+    // Everything up to and including the last newline must have been flushed;
+    // only the trailing partial line is allowed to remain buffered. The buggy
+    // version left the last newline buried in the buffer, so the sink only saw
+    // up to the first newline.
+    assert_eq!(writer.get_ref().buffer, bytes[..=LAST_NEWLINE]);
+
+    writer.flush().unwrap();
+    assert_eq!(writer.get_ref().buffer, bytes);
+}
+
 /// Test that, in cases where vectored writing is not enabled, the
 /// LineWriter uses the normal `write` call, which more-correctly handles
 /// partial lines