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

6 files changed, 147 insertions, 8 deletions

diff --git a/src/libstd/sys/unix/ext/net.rs b/src/libstd/sys/unix/ext/net.rs
index acc064acfcd..4b60ea654c1 100644
--- a/src/libstd/sys/unix/ext/net.rs
+++ b/src/libstd/sys/unix/ext/net.rs
@@ -18,7 +18,7 @@ mod libc {
 use ascii;
 use ffi::OsStr;
 use fmt;
-use io::{self, Initializer};
+use io::{self, Initializer, IoVec, IoVecMut};
 use mem;
 use net::{self, Shutdown};
 use os::unix::ffi::OsStrExt;
@@ -551,6 +551,10 @@ impl io::Read for UnixStream {
         io::Read::read(&mut &*self, buf)
     }
 
+    fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+        io::Read::read_vectored(&mut &*self, bufs)
+    }
+
     #[inline]
     unsafe fn initializer(&self) -> Initializer {
         Initializer::nop()
@@ -563,6 +567,10 @@ impl<'a> io::Read for &'a UnixStream {
         self.0.read(buf)
     }
 
+    fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+        self.0.read_vectored(bufs)
+    }
+
     #[inline]
     unsafe fn initializer(&self) -> Initializer {
         Initializer::nop()
@@ -575,6 +583,10 @@ impl io::Write for UnixStream {
         io::Write::write(&mut &*self, buf)
     }
 
+    fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+        io::Write::write_vectored(&mut &*self, bufs)
+    }
+
     fn flush(&mut self) -> io::Result<()> {
         io::Write::flush(&mut &*self)
     }
@@ -586,6 +598,10 @@ impl<'a> io::Write for &'a UnixStream {
         self.0.write(buf)
     }
 
+    fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+        self.0.write_vectored(bufs)
+    }
+
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
     }
@@ -1511,6 +1527,25 @@ mod test {
     }
 
     #[test]
+    fn vectored() {
+        let (mut s1, mut s2) = or_panic!(UnixStream::pair());
+
+        let len = or_panic!(s1.write_vectored(
+            &[IoVec::new(b"hello"), IoVec::new(b" "), IoVec::new(b"world!")],
+        ));
+        assert_eq!(len, 12);
+
+        let mut buf1 = [0; 6];
+        let mut buf2 = [0; 7];
+        let len = or_panic!(s2.read_vectored(
+            &mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
+        ));
+        assert_eq!(len, 12);
+        assert_eq!(&buf1, b"hello ");
+        assert_eq!(&buf2, b"world!\0");
+    }
+
+    #[test]
     fn pair() {
         let msg1 = b"hello";
         let msg2 = b"world!";
diff --git a/src/libstd/sys/unix/fd.rs b/src/libstd/sys/unix/fd.rs
index 2cbd9536f4d..6946b7b5dfa 100644
--- a/src/libstd/sys/unix/fd.rs
+++ b/src/libstd/sys/unix/fd.rs
@@ -1,7 +1,7 @@
 #![unstable(reason = "not public", issue = "0", feature = "fd")]
 
 use cmp;
-use io::{self, Read, Initializer};
+use io::{self, Read, Initializer, IoVec, IoVecMut};
 use libc::{self, c_int, c_void, ssize_t};
 use mem;
 use sync::atomic::{AtomicBool, Ordering};
@@ -52,6 +52,15 @@ impl FileDesc {
         Ok(ret as usize)
     }
 
+    pub fn read_vectored(&self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+        let ret = cvt(unsafe {
+            libc::readv(self.fd,
+                        bufs.as_ptr() as *const libc::iovec,
+                        cmp::min(bufs.len(), c_int::max_value() as usize) as c_int)
+        })?;
+        Ok(ret as usize)
+    }
+
     pub fn read_to_end(&self, buf: &mut Vec<u8>) -> io::Result<usize> {
         let mut me = self;
         (&mut me).read_to_end(buf)
@@ -105,6 +114,15 @@ impl FileDesc {
         Ok(ret as usize)
     }
 
+    pub fn write_vectored(&self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+        let ret = cvt(unsafe {
+            libc::writev(self.fd,
+                         bufs.as_ptr() as *const libc::iovec,
+                         cmp::min(bufs.len(), c_int::max_value() as usize) as c_int)
+        })?;
+        Ok(ret as usize)
+    }
+
     pub fn write_at(&self, buf: &[u8], offset: u64) -> io::Result<usize> {
         #[cfg(target_os = "android")]
         use super::android::cvt_pwrite64;
diff --git a/src/libstd/sys/unix/io.rs b/src/libstd/sys/unix/io.rs
new file mode 100644
index 00000000000..65e4c6e0577
--- /dev/null
+++ b/src/libstd/sys/unix/io.rs
@@ -0,0 +1,61 @@
+use marker::PhantomData;
+use libc::{iovec, c_void};
+use slice;
+
+#[repr(transparent)]
+pub struct IoVec<'a> {
+    vec: iovec,
+    _p: PhantomData<&'a [u8]>,
+}
+
+impl<'a> IoVec<'a> {
+    #[inline]
+    pub fn new(buf: &'a [u8]) -> IoVec<'a> {
+        IoVec {
+            vec: iovec {
+                iov_base: buf.as_ptr() as *mut u8 as *mut c_void,
+                iov_len: buf.len()
+            },
+            _p: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn as_slice(&self) -> &[u8] {
+        unsafe {
+            slice::from_raw_parts(self.vec.iov_base as *mut u8, self.vec.iov_len)
+        }
+    }
+}
+
+pub struct IoVecMut<'a> {
+    vec: iovec,
+    _p: PhantomData<&'a mut [u8]>,
+}
+
+impl<'a> IoVecMut<'a> {
+    #[inline]
+    pub fn new(buf: &'a mut [u8]) -> IoVecMut<'a> {
+        IoVecMut {
+            vec: iovec {
+                iov_base: buf.as_mut_ptr() as *mut c_void,
+                iov_len: buf.len()
+            },
+            _p: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn as_slice(&self) -> &[u8] {
+        unsafe {
+            slice::from_raw_parts(self.vec.iov_base as *mut u8, self.vec.iov_len)
+        }
+    }
+
+    #[inline]
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        unsafe {
+            slice::from_raw_parts_mut(self.vec.iov_base as *mut u8, self.vec.iov_len)
+        }
+    }
+}
diff --git a/src/libstd/sys/unix/l4re.rs b/src/libstd/sys/unix/l4re.rs
index 48037310c8d..4775e29fb57 100644
--- a/src/libstd/sys/unix/l4re.rs
+++ b/src/libstd/sys/unix/l4re.rs
@@ -5,7 +5,7 @@ macro_rules! unimpl {
 pub mod net {
     #![allow(warnings)]
     use fmt;
-    use io;
+    use io::{self, IoVec, IoVecMut};
     use libc;
     use net::{SocketAddr, Shutdown, Ipv4Addr, Ipv6Addr};
     use sys_common::{AsInner, FromInner, IntoInner};
@@ -46,6 +46,10 @@ pub mod net {
             unimpl!();
         }
 
+        pub fn read_vectored(&self, _: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+            unimpl!();
+        }
+
         pub fn peek(&self, _: &mut [u8]) -> io::Result<usize> {
             unimpl!();
         }
@@ -62,6 +66,10 @@ pub mod net {
             unimpl!();
         }
 
+        pub fn write_vectored(&self, _: &[IoVec<'_>]) -> io::Result<usize> {
+            unimpl!();
+        }
+
         pub fn set_timeout(&self, _: Option<Duration>, _: libc::c_int) -> io::Result<()> {
             unimpl!();
         }
@@ -144,10 +152,18 @@ pub mod net {
             unimpl!();
         }
 
+        pub fn read_vectored(&self, _: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+            unimpl!();
+        }
+
         pub fn write(&self, _: &[u8]) -> io::Result<usize> {
             unimpl!();
         }
 
+        pub fn write_vectored(&self, _: &[IoVec<'_>]) -> io::Result<usize> {
+            unimpl!();
+        }
+
         pub fn peer_addr(&self) -> io::Result<SocketAddr> {
             unimpl!();
         }
diff --git a/src/libstd/sys/unix/mod.rs b/src/libstd/sys/unix/mod.rs
index b36c117fd09..0de1a223fbd 100644
--- a/src/libstd/sys/unix/mod.rs
+++ b/src/libstd/sys/unix/mod.rs
@@ -1,6 +1,6 @@
 #![allow(missing_docs, nonstandard_style)]
 
-use io::{self, ErrorKind};
+use io::ErrorKind;
 use libc;
 
 #[cfg(any(rustdoc, target_os = "linux"))] pub use os::linux as platform;
@@ -39,6 +39,7 @@ pub mod fast_thread_local;
 pub mod fd;
 pub mod fs;
 pub mod memchr;
+pub mod io;
 pub mod mutex;
 #[cfg(not(target_os = "l4re"))]
 pub mod net;
@@ -126,15 +127,15 @@ macro_rules! impl_is_minus_one {
 
 impl_is_minus_one! { i8 i16 i32 i64 isize }
 
-pub fn cvt<T: IsMinusOne>(t: T) -> io::Result<T> {
+pub fn cvt<T: IsMinusOne>(t: T) -> ::io::Result<T> {
     if t.is_minus_one() {
-        Err(io::Error::last_os_error())
+        Err(::io::Error::last_os_error())
     } else {
         Ok(t)
     }
 }
 
-pub fn cvt_r<T, F>(mut f: F) -> io::Result<T>
+pub fn cvt_r<T, F>(mut f: F) -> ::io::Result<T>
     where T: IsMinusOne,
           F: FnMut() -> T
 {
diff --git a/src/libstd/sys/unix/net.rs b/src/libstd/sys/unix/net.rs
index d780d71c376..521d9b42517 100644
--- a/src/libstd/sys/unix/net.rs
+++ b/src/libstd/sys/unix/net.rs
@@ -1,5 +1,5 @@
 use ffi::CStr;
-use io;
+use io::{self, IoVec, IoVecMut};
 use libc::{self, c_int, c_void, size_t, sockaddr, socklen_t, EAI_SYSTEM, MSG_PEEK};
 use mem;
 use net::{SocketAddr, Shutdown};
@@ -241,6 +241,10 @@ impl Socket {
         self.recv_with_flags(buf, MSG_PEEK)
     }
 
+    pub fn read_vectored(&self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+        self.0.read_vectored(bufs)
+    }
+
     fn recv_from_with_flags(&self, buf: &mut [u8], flags: c_int)
                             -> io::Result<(usize, SocketAddr)> {
         let mut storage: libc::sockaddr_storage = unsafe { mem::zeroed() };
@@ -269,6 +273,10 @@ impl Socket {
         self.0.write(buf)
     }
 
+    pub fn write_vectored(&self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+        self.0.write_vectored(bufs)
+    }
+
     pub fn set_timeout(&self, dur: Option<Duration>, kind: libc::c_int) -> io::Result<()> {
         let timeout = match dur {
             Some(dur) => {