Skip to content

Commit fb162e6

Browse files
borsbors
committed
Auto merge of #58357 - sfackler:vectored-io, r=alexcrichton
Add vectored read and write support This functionality has lived for a while in the tokio ecosystem, where it can improve performance by minimizing copies. r? @alexcrichton
2 parents 55c173c + 4785c74 commit fb162e6

File tree

31 files changed

+1052
-95
lines changed

31 files changed

+1052
-95
lines changed

‎src/libstd/io/buffered.rs

+31-2
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use io::prelude::*;
55
use cmp;
66
use error;
77
use fmt;
8-
use io::{self, Initializer, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom};
8+
use io::{self, Initializer, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom, IoVec, IoVecMut};
99
use memchr;
1010

1111
/// The `BufReader` struct adds buffering to any reader.
@@ -235,6 +235,19 @@ impl<R: Read> Read for BufReader<R> {
235235
Ok(nread)
236236
}
237237

238+
fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
239+
let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
240+
if self.pos == self.cap && total_len >= self.buf.len() {
241+
return self.inner.read_vectored(bufs);
242+
}
243+
let nread = {
244+
let mut rem = self.fill_buf()?;
245+
rem.read_vectored(bufs)?
246+
};
247+
self.consume(nread);
248+
Ok(nread)
249+
}
250+
238251
// we can't skip unconditionally because of the large buffer case in read.
239252
unsafe fn initializer(&self) -> Initializer {
240253
self.inner.initializer()
@@ -577,9 +590,25 @@ impl<W: Write> Write for BufWriter<W> {
577590
self.panicked = false;
578591
r
579592
} else {
580-
Write::write(&mut self.buf, buf)
593+
self.buf.write(buf)
594+
}
595+
}
596+
597+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
598+
let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
599+
if self.buf.len() + total_len > self.buf.capacity() {
600+
self.flush_buf()?;
601+
}
602+
if total_len >= self.buf.capacity() {
603+
self.panicked = true;
604+
let r = self.inner.as_mut().unwrap().write_vectored(bufs);
605+
self.panicked = false;
606+
r
607+
} else {
608+
self.buf.write_vectored(bufs)
581609
}
582610
}
611+
583612
fn flush(&mut self) -> io::Result<()> {
584613
self.flush_buf().and_then(|()| self.get_mut().flush())
585614
}

‎src/libstd/io/cursor.rs

+210-5
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use io::prelude::*;
22

33
use core::convert::TryInto;
44
use cmp;
5-
use io::{self, Initializer, SeekFrom, Error, ErrorKind};
5+
use io::{self, Initializer, SeekFrom, Error, ErrorKind, IoVec, IoVecMut};
66

77
/// A `Cursor` wraps an in-memory buffer and provides it with a
88
/// [`Seek`] implementation.
@@ -221,6 +221,18 @@ impl<T> Read for Cursor<T> where T: AsRef<[u8]> {
221221
Ok(n)
222222
}
223223

224+
fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
225+
let mut nread = 0;
226+
for buf in bufs {
227+
let n = self.read(buf)?;
228+
nread += n;
229+
if n < buf.len() {
230+
break;
231+
}
232+
}
233+
Ok(nread)
234+
}
235+
224236
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
225237
let n = buf.len();
226238
Read::read_exact(&mut self.fill_buf()?, buf)?;
@@ -251,6 +263,23 @@ fn slice_write(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<us
251263
Ok(amt)
252264
}
253265

266+
fn slice_write_vectored(
267+
pos_mut: &mut u64,
268+
slice: &mut [u8],
269+
bufs: &[IoVec<'_>],
270+
) -> io::Result<usize>
271+
{
272+
let mut nwritten = 0;
273+
for buf in bufs {
274+
let n = slice_write(pos_mut, slice, buf)?;
275+
nwritten += n;
276+
if n < buf.len() {
277+
break;
278+
}
279+
}
280+
Ok(nwritten)
281+
}
282+
254283
// Resizing write implementation
255284
fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usize> {
256285
let pos: usize = (*pos_mut).try_into().map_err(|_| {
@@ -278,12 +307,31 @@ fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usi
278307
Ok(buf.len())
279308
}
280309

310+
fn vec_write_vectored(
311+
pos_mut: &mut u64,
312+
vec: &mut Vec<u8>,
313+
bufs: &[IoVec<'_>],
314+
) -> io::Result<usize>
315+
{
316+
let mut nwritten = 0;
317+
for buf in bufs {
318+
nwritten += vec_write(pos_mut, vec, buf)?;
319+
}
320+
Ok(nwritten)
321+
}
322+
281323
#[stable(feature = "rust1", since = "1.0.0")]
282324
impl Write for Cursor<&mut [u8]> {
283325
#[inline]
284326
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
285327
slice_write(&mut self.pos, self.inner, buf)
286328
}
329+
330+
#[inline]
331+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
332+
slice_write_vectored(&mut self.pos, self.inner, bufs)
333+
}
334+
287335
fn flush(&mut self) -> io::Result<()> { Ok(()) }
288336
}
289337

@@ -292,6 +340,11 @@ impl Write for Cursor<&mut Vec<u8>> {
292340
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
293341
vec_write(&mut self.pos, self.inner, buf)
294342
}
343+
344+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
345+
vec_write_vectored(&mut self.pos, self.inner, bufs)
346+
}
347+
295348
fn flush(&mut self) -> io::Result<()> { Ok(()) }
296349
}
297350

@@ -300,6 +353,11 @@ impl Write for Cursor<Vec<u8>> {
300353
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
301354
vec_write(&mut self.pos, &mut self.inner, buf)
302355
}
356+
357+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
358+
vec_write_vectored(&mut self.pos, &mut self.inner, bufs)
359+
}
360+
303361
fn flush(&mut self) -> io::Result<()> { Ok(()) }
304362
}
305363

@@ -309,21 +367,30 @@ impl Write for Cursor<Box<[u8]>> {
309367
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
310368
slice_write(&mut self.pos, &mut self.inner, buf)
311369
}
370+
371+
#[inline]
372+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
373+
slice_write_vectored(&mut self.pos, &mut self.inner, bufs)
374+
}
375+
312376
fn flush(&mut self) -> io::Result<()> { Ok(()) }
313377
}
314378

315379
#[cfg(test)]
316380
mod tests {
317381
use io::prelude::*;
318-
use io::{Cursor, SeekFrom};
382+
use io::{Cursor, SeekFrom, IoVec, IoVecMut};
319383

320384
#[test]
321385
fn test_vec_writer() {
322386
let mut writer = Vec::new();
323387
assert_eq!(writer.write(&[0]).unwrap(), 1);
324388
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
325389
assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
326-
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
390+
assert_eq!(writer.write_vectored(
391+
&[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
392+
).unwrap(), 3);
393+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
327394
assert_eq!(writer, b);
328395
}
329396

@@ -333,7 +400,10 @@ mod tests {
333400
assert_eq!(writer.write(&[0]).unwrap(), 1);
334401
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
335402
assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
336-
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
403+
assert_eq!(writer.write_vectored(
404+
&[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
405+
).unwrap(), 3);
406+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
337407
assert_eq!(&writer.get_ref()[..], b);
338408
}
339409

@@ -344,7 +414,10 @@ mod tests {
344414
assert_eq!(writer.write(&[0]).unwrap(), 1);
345415
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
346416
assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
347-
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
417+
assert_eq!(writer.write_vectored(
418+
&[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
419+
).unwrap(), 3);
420+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
348421
assert_eq!(&writer.get_ref()[..], b);
349422
}
350423

@@ -366,6 +439,26 @@ mod tests {
366439
assert_eq!(&**writer.get_ref(), b);
367440
}
368441

442+
#[test]
443+
fn test_box_slice_writer_vectored() {
444+
let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
445+
assert_eq!(writer.position(), 0);
446+
assert_eq!(writer.write_vectored(&[IoVec::new(&[0])]).unwrap(), 1);
447+
assert_eq!(writer.position(), 1);
448+
assert_eq!(
449+
writer.write_vectored(&[IoVec::new(&[1, 2, 3]), IoVec::new(&[4, 5, 6, 7])]).unwrap(),
450+
7,
451+
);
452+
assert_eq!(writer.position(), 8);
453+
assert_eq!(writer.write_vectored(&[]).unwrap(), 0);
454+
assert_eq!(writer.position(), 8);
455+
456+
assert_eq!(writer.write_vectored(&[IoVec::new(&[8, 9])]).unwrap(), 1);
457+
assert_eq!(writer.write_vectored(&[IoVec::new(&[10])]).unwrap(), 0);
458+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
459+
assert_eq!(&**writer.get_ref(), b);
460+
}
461+
369462
#[test]
370463
fn test_buf_writer() {
371464
let mut buf = [0 as u8; 9];
@@ -387,6 +480,31 @@ mod tests {
387480
assert_eq!(buf, b);
388481
}
389482

483+
#[test]
484+
fn test_buf_writer_vectored() {
485+
let mut buf = [0 as u8; 9];
486+
{
487+
let mut writer = Cursor::new(&mut buf[..]);
488+
assert_eq!(writer.position(), 0);
489+
assert_eq!(writer.write_vectored(&[IoVec::new(&[0])]).unwrap(), 1);
490+
assert_eq!(writer.position(), 1);
491+
assert_eq!(
492+
writer.write_vectored(
493+
&[IoVec::new(&[1, 2, 3]), IoVec::new(&[4, 5, 6, 7])],
494+
).unwrap(),
495+
7,
496+
);
497+
assert_eq!(writer.position(), 8);
498+
assert_eq!(writer.write_vectored(&[]).unwrap(), 0);
499+
assert_eq!(writer.position(), 8);
500+
501+
assert_eq!(writer.write_vectored(&[IoVec::new(&[8, 9])]).unwrap(), 1);
502+
assert_eq!(writer.write_vectored(&[IoVec::new(&[10])]).unwrap(), 0);
503+
}
504+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
505+
assert_eq!(buf, b);
506+
}
507+
390508
#[test]
391509
fn test_buf_writer_seek() {
392510
let mut buf = [0 as u8; 8];
@@ -447,6 +565,35 @@ mod tests {
447565
assert_eq!(reader.read(&mut buf).unwrap(), 0);
448566
}
449567

568+
#[test]
569+
fn test_mem_reader_vectored() {
570+
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7]);
571+
let mut buf = [];
572+
assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
573+
assert_eq!(reader.position(), 0);
574+
let mut buf = [0];
575+
assert_eq!(
576+
reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
577+
1,
578+
);
579+
assert_eq!(reader.position(), 1);
580+
let b: &[_] = &[0];
581+
assert_eq!(buf, b);
582+
let mut buf1 = [0; 4];
583+
let mut buf2 = [0; 4];
584+
assert_eq!(
585+
reader.read_vectored(
586+
&mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
587+
).unwrap(),
588+
7,
589+
);
590+
let b1: &[_] = &[1, 2, 3, 4];
591+
let b2: &[_] = &[5, 6, 7];
592+
assert_eq!(buf1, b1);
593+
assert_eq!(&buf2[..3], b2);
594+
assert_eq!(reader.read(&mut buf).unwrap(), 0);
595+
}
596+
450597
#[test]
451598
fn test_boxed_slice_reader() {
452599
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice());
@@ -469,6 +616,35 @@ mod tests {
469616
assert_eq!(reader.read(&mut buf).unwrap(), 0);
470617
}
471618

619+
#[test]
620+
fn test_boxed_slice_reader_vectored() {
621+
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice());
622+
let mut buf = [];
623+
assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
624+
assert_eq!(reader.position(), 0);
625+
let mut buf = [0];
626+
assert_eq!(
627+
reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
628+
1,
629+
);
630+
assert_eq!(reader.position(), 1);
631+
let b: &[_] = &[0];
632+
assert_eq!(buf, b);
633+
let mut buf1 = [0; 4];
634+
let mut buf2 = [0; 4];
635+
assert_eq!(
636+
reader.read_vectored(
637+
&mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
638+
).unwrap(),
639+
7,
640+
);
641+
let b1: &[_] = &[1, 2, 3, 4];
642+
let b2: &[_] = &[5, 6, 7];
643+
assert_eq!(buf1, b1);
644+
assert_eq!(&buf2[..3], b2);
645+
assert_eq!(reader.read(&mut buf).unwrap(), 0);
646+
}
647+
472648
#[test]
473649
fn read_to_end() {
474650
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7]);
@@ -499,6 +675,35 @@ mod tests {
499675
assert_eq!(reader.read(&mut buf).unwrap(), 0);
500676
}
501677

678+
#[test]
679+
fn test_slice_reader_vectored() {
680+
let in_buf = vec![0, 1, 2, 3, 4, 5, 6, 7];
681+
let reader = &mut &in_buf[..];
682+
let mut buf = [];
683+
assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
684+
let mut buf = [0];
685+
assert_eq!(
686+
reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
687+
1,
688+
);
689+
assert_eq!(reader.len(), 7);
690+
let b: &[_] = &[0];
691+
assert_eq!(buf, b);
692+
let mut buf1 = [0; 4];
693+
let mut buf2 = [0; 4];
694+
assert_eq!(
695+
reader.read_vectored(
696+
&mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
697+
).unwrap(),
698+
7,
699+
);
700+
let b1: &[_] = &[1, 2, 3, 4];
701+
let b2: &[_] = &[5, 6, 7];
702+
assert_eq!(buf1, b1);
703+
assert_eq!(&buf2[..3], b2);
704+
assert_eq!(reader.read(&mut buf).unwrap(), 0);
705+
}
706+
502707
#[test]
503708
fn test_read_exact() {
504709
let in_buf = vec![0, 1, 2, 3, 4, 5, 6, 7];

0 commit comments

Comments
 (0)