diff --git a/library/std/src/io/buffered/linewritershim.rs b/library/std/src/io/buffered/linewritershim.rs index 967e24812b9ff..2135eb6d2e578 100644 --- a/library/std/src/io/buffered/linewritershim.rs +++ b/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 { + // 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 { diff --git a/library/std/src/io/buffered/tests.rs b/library/std/src/io/buffered/tests.rs index 742e753f9b951..8888ffa0c4cb7 100644 --- a/library/std/src/io/buffered/tests.rs +++ b/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, + } + + impl Write for OneSliceSink { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.buffer.extend_from_slice(buf); + Ok(buf.len()) + } + fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result { + 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 = (0..N).map(|i| if i == NEWLINE_AT { b'\n' } else { b'a' }).collect(); + let mut io_slices: Vec> = 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, + } + + impl Write for FullSink { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.buffer.extend_from_slice(buf); + Ok(buf.len()) + } + fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result { + 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 = (0..N) + .map(|i| if i == FIRST_NEWLINE || i == LAST_NEWLINE { b'\n' } else { b'a' }) + .collect(); + let mut io_slices: Vec> = 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