Auto merge of rust-lang#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

Author: bors

src/libstd/io/buffered.rs

@@ -5,7 +5,7 @@ use io::prelude::*;
 use cmp;
 use error;
 use fmt;
-use io::{self, Initializer, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom};
+use io::{self, Initializer, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom, IoVec, IoVecMut};
 use memchr;
 
 /// The `BufReader` struct adds buffering to any reader.
@@ -235,6 +235,19 @@ impl<R: Read> Read for BufReader<R> {
  Ok(nread)
  }
 
+ fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+ let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
+ if self.pos == self.cap && total_len >= self.buf.len() {
+ return self.inner.read_vectored(bufs);
+ }
+ let nread = {
+ let mut rem = self.fill_buf()?;
+ rem.read_vectored(bufs)?
+ };
+ self.consume(nread);
+ Ok(nread)
+ }
+
  // we can't skip unconditionally because of the large buffer case in read.
  unsafe fn initializer(&self) -> Initializer {
  self.inner.initializer()
@@ -577,9 +590,25 @@ impl<W: Write> Write for BufWriter<W> {
  self.panicked = false;
  r
  } else {
- Write::write(&mut self.buf, buf)
+ self.buf.write(buf)
+ }
+ }
+
+ fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> 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()?;
+ }
+ if total_len >= self.buf.capacity() {
+ self.panicked = true;
+ let r = self.inner.as_mut().unwrap().write_vectored(bufs);
+ self.panicked = false;
+ r
+ } else {
+ self.buf.write_vectored(bufs)
  }
  }
+
  fn flush(&mut self) -> io::Result<()> {
  self.flush_buf().and_then(|()| self.get_mut().flush())
  }

src/libstd/io/cursor.rs

@@ -2,7 +2,7 @@ use io::prelude::*;
 
 use core::convert::TryInto;
 use cmp;
-use io::{self, Initializer, SeekFrom, Error, ErrorKind};
+use io::{self, Initializer, SeekFrom, Error, ErrorKind, IoVec, IoVecMut};
 
 /// A `Cursor` wraps an in-memory buffer and provides it with a
 /// [`Seek`] implementation.
@@ -221,6 +221,18 @@ impl<T> Read for Cursor<T> where T: AsRef<[u8]> {
  Ok(n)
  }
 
+ fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+ let mut nread = 0;
+ for buf in bufs {
+ let n = self.read(buf)?;
+ nread += n;
+ if n < buf.len() {
+ break;
+ }
+ }
+ Ok(nread)
+ }
+
  fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
  let n = buf.len();
  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
  Ok(amt)
 }
 
+fn slice_write_vectored(
+ pos_mut: &mut u64,
+ slice: &mut [u8],
+ bufs: &[IoVec<'_>],
+) -> io::Result<usize>
+{
+ let mut nwritten = 0;
+ for buf in bufs {
+ let n = slice_write(pos_mut, slice, buf)?;
+ nwritten += n;
+ if n < buf.len() {
+ break;
+ }
+ }
+ Ok(nwritten)
+}
+
 // Resizing write implementation
 fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usize> {
  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
  Ok(buf.len())
 }
 
+fn vec_write_vectored(
+ pos_mut: &mut u64,
+ vec: &mut Vec<u8>,
+ bufs: &[IoVec<'_>],
+) -> io::Result<usize>
+{
+ let mut nwritten = 0;
+ for buf in bufs {
+ nwritten += vec_write(pos_mut, vec, buf)?;
+ }
+ Ok(nwritten)
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl Write for Cursor<&mut [u8]> {
  #[inline]
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
  slice_write(&mut self.pos, self.inner, buf)
  }
+
+ #[inline]
+ fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+ slice_write_vectored(&mut self.pos, self.inner, bufs)
+ }
+
  fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }
 
@@ -292,6 +340,11 @@ impl Write for Cursor<&mut Vec<u8>> {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
  vec_write(&mut self.pos, self.inner, buf)
  }
+
+ fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+ vec_write_vectored(&mut self.pos, self.inner, bufs)
+ }
+
  fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }
 
@@ -300,6 +353,11 @@ impl Write for Cursor<Vec<u8>> {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
  vec_write(&mut self.pos, &mut self.inner, buf)
  }
+
+ fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+ vec_write_vectored(&mut self.pos, &mut self.inner, bufs)
+ }
+
  fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }
 
@@ -309,21 +367,30 @@ impl Write for Cursor<Box<[u8]>> {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
  slice_write(&mut self.pos, &mut self.inner, buf)
  }
+
+ #[inline]
+ fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+ slice_write_vectored(&mut self.pos, &mut self.inner, bufs)
+ }
+
  fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }
 
 #[cfg(test)]
 mod tests {
  use io::prelude::*;
- use io::{Cursor, SeekFrom};
+ use io::{Cursor, SeekFrom, IoVec, IoVecMut};
 
  #[test]
  fn test_vec_writer() {
  let mut writer = Vec::new();
  assert_eq!(writer.write(&[0]).unwrap(), 1);
  assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
  assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
- let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
+ assert_eq!(writer.write_vectored(
+ &[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
+ ).unwrap(), 3);
+ let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
  assert_eq!(writer, b);
  }
 
@@ -333,7 +400,10 @@ mod tests {
  assert_eq!(writer.write(&[0]).unwrap(), 1);
  assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
  assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
- let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
+ assert_eq!(writer.write_vectored(
+ &[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
+ ).unwrap(), 3);
+ let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
  assert_eq!(&writer.get_ref()[..], b);
  }
 
@@ -344,7 +414,10 @@ mod tests {
  assert_eq!(writer.write(&[0]).unwrap(), 1);
  assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
  assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
- let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
+ assert_eq!(writer.write_vectored(
+ &[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
+ ).unwrap(), 3);
+ let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
  assert_eq!(&writer.get_ref()[..], b);
  }
 
@@ -366,6 +439,26 @@ mod tests {
  assert_eq!(&**writer.get_ref(), b);
  }
 
+ #[test]
+ fn test_box_slice_writer_vectored() {
+ let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
+ assert_eq!(writer.position(), 0);
+ assert_eq!(writer.write_vectored(&[IoVec::new(&[0])]).unwrap(), 1);
+ assert_eq!(writer.position(), 1);
+ assert_eq!(
+ writer.write_vectored(&[IoVec::new(&[1, 2, 3]), IoVec::new(&[4, 5, 6, 7])]).unwrap(),
+ 7,
+ );
+ assert_eq!(writer.position(), 8);
+ assert_eq!(writer.write_vectored(&[]).unwrap(), 0);
+ assert_eq!(writer.position(), 8);
+
+ assert_eq!(writer.write_vectored(&[IoVec::new(&[8, 9])]).unwrap(), 1);
+ assert_eq!(writer.write_vectored(&[IoVec::new(&[10])]).unwrap(), 0);
+ let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
+ assert_eq!(&**writer.get_ref(), b);
+ }
+
  #[test]
  fn test_buf_writer() {
  let mut buf = [0 as u8; 9];
@@ -387,6 +480,31 @@ mod tests {
  assert_eq!(buf, b);
  }
 
+ #[test]
+ fn test_buf_writer_vectored() {
+ let mut buf = [0 as u8; 9];
+ {
+ let mut writer = Cursor::new(&mut buf[..]);
+ assert_eq!(writer.position(), 0);
+ assert_eq!(writer.write_vectored(&[IoVec::new(&[0])]).unwrap(), 1);
+ assert_eq!(writer.position(), 1);
+ assert_eq!(
+ writer.write_vectored(
+ &[IoVec::new(&[1, 2, 3]), IoVec::new(&[4, 5, 6, 7])],
+ ).unwrap(),
+ 7,
+ );
+ assert_eq!(writer.position(), 8);
+ assert_eq!(writer.write_vectored(&[]).unwrap(), 0);
+ assert_eq!(writer.position(), 8);
+
+ assert_eq!(writer.write_vectored(&[IoVec::new(&[8, 9])]).unwrap(), 1);
+ assert_eq!(writer.write_vectored(&[IoVec::new(&[10])]).unwrap(), 0);
+ }
+ let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
+ assert_eq!(buf, b);
+ }
+
  #[test]
  fn test_buf_writer_seek() {
  let mut buf = [0 as u8; 8];
@@ -447,6 +565,35 @@ mod tests {
  assert_eq!(reader.read(&mut buf).unwrap(), 0);
  }
 
+ #[test]
+ fn test_mem_reader_vectored() {
+ let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7]);
+ let mut buf = [];
+ assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
+ assert_eq!(reader.position(), 0);
+ let mut buf = [0];
+ assert_eq!(
+ reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
+ 1,
+ );
+ assert_eq!(reader.position(), 1);
+ let b: &[_] = &[0];
+ assert_eq!(buf, b);
+ let mut buf1 = [0; 4];
+ let mut buf2 = [0; 4];
+ assert_eq!(
+ reader.read_vectored(
+ &mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
+ ).unwrap(),
+ 7,
+ );
+ let b1: &[_] = &[1, 2, 3, 4];
+ let b2: &[_] = &[5, 6, 7];
+ assert_eq!(buf1, b1);
+ assert_eq!(&buf2[..3], b2);
+ assert_eq!(reader.read(&mut buf).unwrap(), 0);
+ }
+
  #[test]
  fn test_boxed_slice_reader() {
  let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice());
@@ -469,6 +616,35 @@ mod tests {
  assert_eq!(reader.read(&mut buf).unwrap(), 0);
  }
 
+ #[test]
+ fn test_boxed_slice_reader_vectored() {
+ let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice());
+ let mut buf = [];
+ assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
+ assert_eq!(reader.position(), 0);
+ let mut buf = [0];
+ assert_eq!(
+ reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
+ 1,
+ );
+ assert_eq!(reader.position(), 1);
+ let b: &[_] = &[0];
+ assert_eq!(buf, b);
+ let mut buf1 = [0; 4];
+ let mut buf2 = [0; 4];
+ assert_eq!(
+ reader.read_vectored(
+ &mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
+ ).unwrap(),
+ 7,
+ );
+ let b1: &[_] = &[1, 2, 3, 4];
+ let b2: &[_] = &[5, 6, 7];
+ assert_eq!(buf1, b1);
+ assert_eq!(&buf2[..3], b2);
+ assert_eq!(reader.read(&mut buf).unwrap(), 0);
+ }
+
  #[test]
  fn read_to_end() {
  let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7]);
@@ -499,6 +675,35 @@ mod tests {
  assert_eq!(reader.read(&mut buf).unwrap(), 0);
  }
 
+ #[test]
+ fn test_slice_reader_vectored() {
+ let in_buf = vec![0, 1, 2, 3, 4, 5, 6, 7];
+ let reader = &mut &in_buf[..];
+ let mut buf = [];
+ assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
+ let mut buf = [0];
+ assert_eq!(
+ reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
+ 1,
+ );
+ assert_eq!(reader.len(), 7);
+ let b: &[_] = &[0];
+ assert_eq!(buf, b);
+ let mut buf1 = [0; 4];
+ let mut buf2 = [0; 4];
+ assert_eq!(
+ reader.read_vectored(
+ &mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
+ ).unwrap(),
+ 7,
+ );
+ let b1: &[_] = &[1, 2, 3, 4];
+ let b2: &[_] = &[5, 6, 7];
+ assert_eq!(buf1, b1);
+ assert_eq!(&buf2[..3], b2);
+ assert_eq!(reader.read(&mut buf).unwrap(), 0);
+ }
+
  #[test]
  fn test_read_exact() {
  let in_buf = vec![0, 1, 2, 3, 4, 5, 6, 7];