uucore, yes: add zero-copy on supported platforms (only Linux currently)

Cargo.lock

@@ -879,6 +879,18 @@ dependencies = [
  "void 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
+[[package]]
+name = "nix"
+version = "0.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "bitflags 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)",
+ "cc 1.0.36 (registry+https://github.com/rust-lang/crates.io-index)",
+ "cfg-if 0.1.7 (registry+https://github.com/rust-lang/crates.io-index)",
+ "libc 0.2.53 (registry+https://github.com/rust-lang/crates.io-index)",
+ "void 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "nl"
 version = "0.0.1"
@@ -1876,7 +1888,10 @@ dependencies = [
  "failure 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)",
  "failure_derive 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)",
  "getopts 0.2.19 (registry+https://github.com/rust-lang/crates.io-index)",
+ "lazy_static 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "libc 0.2.53 (registry+https://github.com/rust-lang/crates.io-index)",
+ "nix 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "platform-info 0.0.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "termion 1.5.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "time 0.1.42 (registry+https://github.com/rust-lang/crates.io-index)",
  "wild 2.0.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -2163,6 +2178,7 @@ dependencies = [
 "checksum md5 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)" = "79c56d6a0b07f9e19282511c83fc5b086364cbae4ba8c7d5f190c3d9b0425a48"
 "checksum memchr 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)" = "148fab2e51b4f1cfc66da2a7c32981d1d3c083a803978268bb11fe4b86925e7a"
 "checksum memchr 2.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "2efc7bc57c883d4a4d6e3246905283d8dae951bb3bd32f49d6ef297f546e1c39"
+"checksum nix 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)" = "46f0f3210768d796e8fa79ec70ee6af172dacbe7147f5e69be5240a47778302b"
 "checksum nix 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "47e49f6982987135c5e9620ab317623e723bd06738fd85377e8d55f57c8b6487"
 "checksum nodrop 0.1.13 (registry+https://github.com/rust-lang/crates.io-index)" = "2f9667ddcc6cc8a43afc9b7917599d7216aa09c463919ea32c59ed6cac8bc945"
 "checksum num-integer 0.1.39 (registry+https://github.com/rust-lang/crates.io-index)" = "e83d528d2677f0518c570baf2b7abdcf0cd2d248860b68507bdcb3e91d4c0cea"

src/uucore/Cargo.toml

@@ -11,6 +11,9 @@ time = { version = "0.1.40", optional = true }
 data-encoding = { version = "^2.1", optional = true }
 libc = { version = "0.2.42", optional = true }
 wild = "2.0.1"
+nix = { version = "0.13", optional = true }
+lazy_static = { version = "1.3", optional = true }
+platform-info = { version = "0.0.1", optional = true }
 
 [target.'cfg(target_os = "redox")'.dependencies]
 termion = "1.5"
@@ -25,6 +28,7 @@ utmpx = ["time", "libc"]
 process = ["libc"]
 signals = []
 entries = ["libc"]
+zero-copy = ["nix", "libc", "lazy_static", "platform-info"]
 wide = []
 default = []
 

src/uucore/lib.rs

@@ -13,6 +13,13 @@ extern crate failure;
 #[cfg(feature = "failure_derive")]
 #[macro_use]
 extern crate failure_derive;
+#[cfg(feature = "nix")]
+extern crate nix;
+#[cfg(all(feature = "lazy_static", target_os = "linux"))]
+#[macro_use]
+extern crate lazy_static;
+#[cfg(feature = "platform-info")]
+extern crate platform_info;
 
 #[macro_use]
 mod macros;
@@ -40,5 +47,8 @@ pub mod process;
 #[cfg(all(unix, not(target_os = "fuchsia"), feature = "signals"))]
 pub mod signals;
 
+#[cfg(feature = "zero-copy")]
+pub mod zero_copy;
+
 #[cfg(all(windows, feature = "wide"))]
 pub mod wide;

src/uucore/zero_copy/mod.rs

@@ -0,0 +1,139 @@
+use self::platform::*;
+
+use std::io::{self, Write};
+
+mod platform;
+
+pub trait AsRawObject {
+   fn as_raw_object(&self) -> RawObject;
+}
+
+pub trait FromRawObject : Sized {
+   unsafe fn from_raw_object(obj: RawObject) -> Option<Self>;
+}
+
+// TODO: also make a SpliceWriter that takes an input fd and and output fd and uses splice() to
+//       transfer data
+// TODO: make a TeeWriter or something that takes an input fd and two output fds and uses tee() to
+//       transfer to both output fds
+
+enum InnerZeroCopyWriter<T: Write + Sized> {
+   Platform(PlatformZeroCopyWriter),
+   Standard(T),
+}
+
+impl<T: Write + Sized> Write for InnerZeroCopyWriter<T> {
+   fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+      match self {
+         InnerZeroCopyWriter::Platform(ref mut writer) => writer.write(buf),
+         InnerZeroCopyWriter::Standard(ref mut writer) => writer.write(buf),
+      }
+   }
+
+   fn flush(&mut self) -> io::Result<()> {
+      match self {
+         InnerZeroCopyWriter::Platform(ref mut writer) => writer.flush(),
+         InnerZeroCopyWriter::Standard(ref mut writer) => writer.flush(),
+      }
+   }
+}
+
+pub struct ZeroCopyWriter<T: Write + AsRawObject + Sized> {
+   /// This field is never used, but we need it to drop file descriptors
+   #[allow(dead_code)]
+   raw_obj_owner: Option<T>,
+
+   inner: InnerZeroCopyWriter<T>,
+}
+
+struct TransformContainer<'a, A: Write + AsRawObject + Sized, B: Write + Sized> {
+   /// This field is never used and probably could be converted into PhantomData, but might be
+   /// useful for restructuring later (at the moment it's basically left over from an earlier
+   /// design)
+   #[allow(dead_code)]
+   original: Option<&'a mut A>,
+
+   transformed: Option<B>,
+}
+
+impl<'a, A: Write + AsRawObject + Sized, B: Write + Sized> Write for TransformContainer<'a, A, B> {
+   fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
+      self.transformed.as_mut().unwrap().write(bytes)
+   }
+
+   fn flush(&mut self) -> io::Result<()> {
+      self.transformed.as_mut().unwrap().flush()
+   }
+}
+
+impl<'a, A: Write + AsRawObject + Sized, B: Write + Sized> AsRawObject for TransformContainer<'a, A, B> {
+   fn as_raw_object(&self) -> RawObject {
+      panic!("Test should never be used")
+   }
+}
+
+impl<T: Write + AsRawObject + Sized> ZeroCopyWriter<T> {
+   pub fn new(writer: T) -> Self {
+      let raw_obj = writer.as_raw_object();
+      match unsafe { PlatformZeroCopyWriter::new(raw_obj) } {
+         Ok(inner) => {
+            ZeroCopyWriter {
+               raw_obj_owner: Some(writer),
+               inner: InnerZeroCopyWriter::Platform(inner),
+            }
+         }
+         _ => {
+            // creating the splice writer failed for whatever reason, so just make a default
+            // writer
+            ZeroCopyWriter {
+               raw_obj_owner: None,
+               inner: InnerZeroCopyWriter::Standard(writer),
+            }
+         }
+      }
+   }
+
+   pub fn with_default<'a: 'b, 'b, F, W>(writer: &'a mut T, func: F) -> ZeroCopyWriter<impl Write + AsRawObject + Sized + 'b>
+   where
+      F: Fn(&'a mut T) -> W,
+      W: Write + Sized + 'b,
+   {
+      let raw_obj = writer.as_raw_object();
+      match unsafe { PlatformZeroCopyWriter::new(raw_obj) } {
+         Ok(inner) => {
+            ZeroCopyWriter {
+               raw_obj_owner: Some(TransformContainer { original: Some(writer), transformed: None, }),
+               inner: InnerZeroCopyWriter::Platform(inner),
+            }
+         }
+         _ => {
+            // XXX: should func actually consume writer and leave it up to the user to save the value?
+            //      maybe provide a default stdin method then?  in some cases it would make more sense for the
+            //      value to be consumed
+            let real_writer = func(writer);
+            ZeroCopyWriter {
+               raw_obj_owner: None,
+               inner: InnerZeroCopyWriter::Standard(TransformContainer { original: None, transformed: Some(real_writer) }),
+            }
+         }
+      }
+   }
+
+   // XXX: unsure how to get something like this working without allocating, so not providing it
+   /*pub fn stdout() -> ZeroCopyWriter<impl Write + AsRawObject + Sized> {
+      let mut stdout = io::stdout();
+      ZeroCopyWriter::with_default(&mut stdout, |stdout| {
+         stdout.lock()
+      })
+   }*/
+}
+
+impl<T: Write + AsRawObject + Sized> Write for ZeroCopyWriter<T> {
+   fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+      self.inner.write(buf)
+   }
+
+   fn flush(&mut self) -> io::Result<()> {
+      self.inner.flush()
+   }
+}

src/uucore/zero_copy/platform/default.rs

@@ -0,0 +1,25 @@
+use crate::zero_copy::RawObject;
+
+use std::io::{self, Write};
+
+/// A "zero-copy" writer used on platforms for which we have no actual zero-copy implementation (or
+/// which use standard read/write operations for zero-copy I/O).  This writer just delegates to the
+/// inner writer used to create it.  Using this struct avoids going through the machinery used to
+/// handle the case where a given writer does not support zero-copy on a platform.
+pub struct PlatformZeroCopyWriter;
+
+impl PlatformZeroCopyWriter {
+   pub unsafe fn new(_obj: RawObject) -> Result<Self, ()> {
+      Err(())
+   }
+}
+
+impl Write for PlatformZeroCopyWriter {
+   fn write(&mut self, _bytes: &[u8]) -> io::Result<usize> {
+      panic!("should never occur")
+   }
+
+   fn flush(&mut self) -> io::Result<()> {
+      panic!("should never occur")
+   }
+}

src/uucore/zero_copy/platform/linux.rs

@@ -0,0 +1,105 @@
+use std::io::{self, Write};
+use std::os::unix::io::RawFd;
+
+use libc::{O_APPEND, S_IFIFO, S_IFREG};
+use nix::errno::Errno;
+use nix::fcntl::{fcntl, splice, vmsplice, FcntlArg, SpliceFFlags};
+use nix::sys::uio::IoVec;
+use nix::sys::stat::{fstat, FileStat};
+use nix::unistd::pipe;
+use platform_info::{Uname, PlatformInfo};
+
+use crate::zero_copy::{FromRawObject, RawObject};
+
+lazy_static! {
+   static ref IN_WSL: bool = {
+      let info = PlatformInfo::new().unwrap();
+      info.release().contains("Microsoft")
+   };
+}
+
+pub struct PlatformZeroCopyWriter {
+   raw_obj: RawObject,
+   read_pipe: RawFd,
+   write_pipe: RawFd,
+   write_fn: fn(&mut PlatformZeroCopyWriter, &[IoVec<&[u8]>], usize) -> io::Result<usize>,
+}
+
+impl PlatformZeroCopyWriter {
+   pub unsafe fn new(raw_obj: RawObject) -> nix::Result<Self> {
+      if *IN_WSL {
+         // apparently WSL hasn't implemented vmsplice(), causing writes to fail
+         // thus, we will just say zero-copy doesn't work there rather than working
+         // around it
+         return Err(nix::Error::from(Errno::EOPNOTSUPP))
+      }
+
+      let stat_info: FileStat = fstat(raw_obj)?;
+      let access_mode: libc::c_int = fcntl(raw_obj, FcntlArg::F_GETFL)?;
+
+      let is_regular = (stat_info.st_mode & S_IFREG) != 0;
+      let is_append = (access_mode & O_APPEND) != 0;
+      let is_fifo = (stat_info.st_mode & S_IFIFO) != 0;
+
+      if is_regular && !is_append {
+         let (read_pipe, write_pipe) = pipe()?;
+
+         Ok(PlatformZeroCopyWriter {
+            raw_obj,
+            read_pipe,
+            write_pipe,
+            write_fn: write_regular,
+         })
+      } else if is_fifo {
+         Ok(PlatformZeroCopyWriter {
+            raw_obj,
+            read_pipe: Default::default(),
+            write_pipe: Default::default(),
+            write_fn: write_fifo,
+         })
+      } else {
+         // FIXME: how to error?
+         Err(nix::Error::from(Errno::UnknownErrno))
+      }
+   }
+}
+
+impl FromRawObject for PlatformZeroCopyWriter {
+   unsafe fn from_raw_object(obj: RawObject) -> Option<Self> {
+      PlatformZeroCopyWriter::new(obj).ok()
+   }
+}
+
+impl Write for PlatformZeroCopyWriter {
+   fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+      let iovec = &[IoVec::from_slice(buf)];
+
+      let func = self.write_fn;
+      func(self, iovec, buf.len())
+   }
+
+   fn flush(&mut self) -> io::Result<()> {
+      // XXX: not sure if we need anything else
+      Ok(())
+   }
+}
+
+fn write_regular(writer: &mut PlatformZeroCopyWriter, iovec: &[IoVec<&[u8]>], len: usize) -> io::Result<usize> {
+   vmsplice(writer.write_pipe, iovec, SpliceFFlags::empty())
+      .and_then(|_| 
+         splice(
+            writer.read_pipe,
+            None,
+            writer.raw_obj,
+            None,
+            len,
+            SpliceFFlags::empty()
+         )
+      )
+      .map_err(|_| io::Error::last_os_error())
+}
+
+fn write_fifo(writer: &mut PlatformZeroCopyWriter, iovec: &[IoVec<&[u8]>], _len: usize) -> io::Result<usize> {
+   vmsplice(writer.raw_obj, iovec, SpliceFFlags::empty())
+      .map_err(|_| io::Error::last_os_error())
+}

src/uucore/zero_copy/platform/mod.rs

@@ -0,0 +1,21 @@
+#[cfg(unix)]
+pub use self::unix::*;
+#[cfg(target_os = "linux")]
+pub use self::linux::*;
+#[cfg(windows)]
+pub use self::windows::*;
+
+// Add any operating systems we support here
+#[cfg(not(any(target_os = "linux")))]
+pub use self::default::*;
+
+#[cfg(unix)]
+mod unix;
+#[cfg(target_os = "linux")]
+mod linux;
+#[cfg(windows)]
+mod windows;
+
+// Add any operating systems we support here
+#[cfg(not(any(target_os = "linux")))]
+mod default;

src/uucore/zero_copy/platform/unix.rs

@@ -0,0 +1,18 @@
+use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
+
+use crate::zero_copy::{AsRawObject, FromRawObject};
+
+pub type RawObject = RawFd;
+
+impl<T: AsRawFd> AsRawObject for T {
+   fn as_raw_object(&self) -> RawObject {
+      self.as_raw_fd()
+   }
+}
+
+// FIXME: check if this works right
+impl<T: FromRawFd> FromRawObject for T {
+   unsafe fn from_raw_object(obj: RawObject) -> Option<Self> {
+      Some(T::from_raw_fd(obj))
+   }
+}

src/uucore/zero_copy/platform/windows.rs

@@ -0,0 +1,19 @@
+use std::os::windows::io::{AsRawHandle, FromRawHandle, RawHandle};
+
+use crate::zero_copy::{AsRawObject, FromRawObject};
+
+pub type RawObject = RawHandle;
+
+impl<T: AsRawHandle> AsRawObject for T {
+   fn as_raw_object(&self) -> RawObject {
+      self.as_raw_handle()
+   }
+}
+
+impl<T: FromRawHandle> FromRawObject for T {
+   unsafe fn from_raw_object(obj: RawObject) -> Option<Self> {
+      Some(T::from_raw_handle(obj))
+   }
+}
+
+// TODO: see if there's some zero-copy stuff in Windows

src/yes/Cargo.toml

@@ -11,7 +11,7 @@ path = "yes.rs"
 
 [dependencies]
 clap = "2.32"
-uucore = { path = "../uucore" }
+uucore = { path = "../uucore", features = ["zero-copy"] }
 
 [features]
 latency = []

src/yes/yes.rs

@@ -17,6 +17,7 @@ extern crate clap;
 extern crate uucore;
 
 use clap::Arg;
+use uucore::zero_copy::ZeroCopyWriter;
 use std::borrow::Cow;
 use std::io::{self, Write};
 
@@ -83,9 +84,10 @@ fn prepare_buffer<'a>(input: &'a str, _buffer: &'a mut [u8; BUF_SIZE]) -> &'a [u
 }
 
 pub fn exec(bytes: &[u8]) {
-    let stdout_raw = io::stdout();
-    let mut stdout = stdout_raw.lock();
+    let mut stdin_raw = io::stdout();
+    let mut writer = ZeroCopyWriter::with_default(&mut stdin_raw, |stdin| stdin.lock());
     loop {
-        stdout.write_all(bytes).unwrap();
+        // TODO: needs to check if pipe fails
+        writer.write_all(bytes).unwrap();
     }
 }