Skip to content

Commit 8212fc5

Browse files
blyxxyzblyxxyz
committed
Fix quadratic behavior of repeated vectored writes
Some implementations of `Write::write_vectored` in the standard library (`BufWriter`, `LineWriter`, `Stdout`, `Stderr`) check all buffers to calculate the total length. This is O(n) over the number of buffers. It's common that only a limited number of buffers is written at a time (e.g. 1024 for `writev(2)`). `write_vectored_all` will then call `write_vectored` repeatedly, leading to a runtime of O(n²) over the number of buffers. The fix is to only calculate as much as needed if it's needed.
1 parent 3793e5b commit 8212fc5

File tree

3 files changed

+50
-31
lines changed

3 files changed

+50
-31
lines changed

‎library/std/src/io/buffered/bufwriter.rs

+27-24
Original file line numberDiff line numberDiff line change
@@ -554,35 +554,38 @@ impl<W: ?Sized + Write> Write for BufWriter<W> {
554554
// same underlying buffer, as otherwise the buffers wouldn't fit in memory). If the
555555
// computation overflows, then surely the input cannot fit in our buffer, so we forward
556556
// to the inner writer's `write_vectored` method to let it handle it appropriately.
557-
let saturated_total_len =
558-
bufs.iter().fold(0usize, |acc, b| acc.saturating_add(b.len()));
557+
let mut saturated_total_len: usize = 0;
559558

560-
if saturated_total_len > self.spare_capacity() {
561-
// Flush if the total length of the input exceeds our buffer's spare capacity.
562-
// If we would have overflowed, this condition also holds, and we need to flush.
563-
self.flush_buf()?;
559+
for buf in bufs {
560+
saturated_total_len = saturated_total_len.saturating_add(buf.len());
561+
562+
if saturated_total_len > self.spare_capacity() && !self.buf.is_empty() {
563+
// Flush if the total length of the input exceeds our buffer's spare capacity.
564+
// If we would have overflowed, this condition also holds, and we need to flush.
565+
self.flush_buf()?;
566+
}
567+
568+
if saturated_total_len >= self.buf.capacity() {
569+
// Forward to our inner writer if the total length of the input is greater than or
570+
// equal to our buffer capacity. If we would have overflowed, this condition also
571+
// holds, and we punt to the inner writer.
572+
self.panicked = true;
573+
let r = self.get_mut().write_vectored(bufs);
574+
self.panicked = false;
575+
return r;
576+
}
564577
}
565578

566-
if saturated_total_len >= self.buf.capacity() {
567-
// Forward to our inner writer if the total length of the input is greater than or
568-
// equal to our buffer capacity. If we would have overflowed, this condition also
569-
// holds, and we punt to the inner writer.
570-
self.panicked = true;
571-
let r = self.get_mut().write_vectored(bufs);
572-
self.panicked = false;
573-
r
574-
} else {
575-
// `saturated_total_len < self.buf.capacity()` implies that we did not saturate.
579+
// `saturated_total_len < self.buf.capacity()` implies that we did not saturate.
576580

577-
// SAFETY: We checked whether or not the spare capacity was large enough above. If
578-
// it was, then we're safe already. If it wasn't, we flushed, making sufficient
579-
// room for any input <= the buffer size, which includes this input.
580-
unsafe {
581-
bufs.iter().for_each(|b| self.write_to_buffer_unchecked(b));
582-
};
581+
// SAFETY: We checked whether or not the spare capacity was large enough above. If
582+
// it was, then we're safe already. If it wasn't, we flushed, making sufficient
583+
// room for any input <= the buffer size, which includes this input.
584+
unsafe {
585+
bufs.iter().for_each(|b| self.write_to_buffer_unchecked(b));
586+
};
583587

584-
Ok(saturated_total_len)
585-
}
588+
Ok(saturated_total_len)
586589
} else {
587590
let mut iter = bufs.iter();
588591
let mut total_written = if let Some(buf) = iter.by_ref().find(|&buf| !buf.is_empty()) {

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

+12-3
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,10 @@ impl<'a, W: ?Sized + Write> Write for LineWriterShim<'a, W> {
175175
}
176176

177177
// Find the buffer containing the last newline
178+
// FIXME: This is overly slow if there are very many bufs and none contain
179+
// newlines. e.g. writev() on Linux only writes up to 1024 slices, so
180+
// scanning the rest is wasted effort. This makes write_all_vectored()
181+
// quadratic.
178182
let last_newline_buf_idx = bufs
179183
.iter()
180184
.enumerate()
@@ -215,9 +219,14 @@ impl<'a, W: ?Sized + Write> Write for LineWriterShim<'a, W> {
215219

216220
// Don't try to reconstruct the exact amount written; just bail
217221
// in the event of a partial write
218-
let lines_len = lines.iter().map(|buf| buf.len()).sum();
219-
if flushed < lines_len {
220-
return Ok(flushed);
222+
let mut lines_len: usize = 0;
223+
for buf in lines {
224+
// With overlapping/duplicate slices the total length may in theory
225+
// exceed usize::MAX
226+
lines_len = lines_len.saturating_add(buf.len());
227+
if flushed < lines_len {
228+
return Ok(flushed);
229+
}
221230
}
222231

223232
// Now that the write has succeeded, buffer the rest (or as much of the

‎library/std/src/io/stdio.rs

+11-4
Original file line numberDiff line numberDiff line change
@@ -128,8 +128,8 @@ impl Write for StdoutRaw {
128128
}
129129

130130
fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
131-
let total = bufs.iter().map(|b| b.len()).sum();
132-
handle_ebadf(self.0.write_vectored(bufs), total)
131+
let total = || bufs.iter().map(|b| b.len()).sum();
132+
handle_ebadf_lazy(self.0.write_vectored(bufs), total)
133133
}
134134

135135
#[inline]
@@ -160,8 +160,8 @@ impl Write for StderrRaw {
160160
}
161161

162162
fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
163-
let total = bufs.iter().map(|b| b.len()).sum();
164-
handle_ebadf(self.0.write_vectored(bufs), total)
163+
let total = || bufs.iter().map(|b| b.len()).sum();
164+
handle_ebadf_lazy(self.0.write_vectored(bufs), total)
165165
}
166166

167167
#[inline]
@@ -193,6 +193,13 @@ fn handle_ebadf<T>(r: io::Result<T>, default: T) -> io::Result<T> {
193193
}
194194
}
195195

196+
fn handle_ebadf_lazy<T>(r: io::Result<T>, default: impl FnOnce() -> T) -> io::Result<T> {
197+
match r {
198+
Err(ref e) if stdio::is_ebadf(e) => Ok(default()),
199+
r => r,
200+
}
201+
}
202+
196203
/// A handle to the standard input stream of a process.
197204
///
198205
/// Each handle is a shared reference to a global buffer of input data to this

0 commit comments

Comments
 (0)