yes: Add zero copy fast-path

Author: oech3

.vscode/cspell.dictionaries/workspace.wordlist.txt

@@ -8,6 +8,7 @@ advapi32-sys
 aho-corasick
 backtrace
 blake2b_simd
+rustix
 
 # * uutils project
 uutils

Cargo.lock

@@ -434,6 +434,7 @@ rlimit = "0.11.0"
 rstest = "0.26.0"
 rustc-hash = "2.1.1"
 rust-ini = "0.21.0"
+rustix = { version = "1.1.4", features = ["fs", "param", "pipe"] }
 same-file = "1.0.6"
 self_cell = "1.0.4"
 selinux = "=0.6.0"

Cargo.toml

@@ -19,9 +19,12 @@ workspace = true
 path = "src/yes.rs"
 
 [dependencies]
+aligned-vec = "0.6.4"
 clap = { workspace = true }
 itertools = { workspace = true }
 fluent = { workspace = true }
+rustix = { workspace = true }
+
 
 [target.'cfg(unix)'.dependencies]
 uucore = { workspace = true, features = ["pipes", "signals"] }

src/uu/yes/Cargo.toml

@@ -3,8 +3,9 @@
 // For the full copyright and license information, please view the LICENSE
 // file that was distributed with this source code.
 
-// cSpell:ignore strs
+// cSpell:ignore strs setpipe
 
+use aligned_vec::{AVec, Alignment, ConstAlign};
 use clap::{Arg, ArgAction, Command, builder::ValueParser};
 use std::error::Error;
 use std::ffi::OsString;
@@ -13,18 +14,48 @@ use uucore::error::{UResult, USimpleError};
 use uucore::format_usage;
 use uucore::translate;
 
+#[cfg(target_os = "linux")]
+use rustix::{
+    fd::{AsRawFd, BorrowedFd},
+    param::page_size,
+    pipe::{IoSliceRaw, SpliceFlags, fcntl_setpipe_size, vmsplice},
+};
+
+// Should be multiple of page size
+const ROOTLESS_MAX_PIPE_SIZE: usize = 1024 * 1024;
 // it's possible that using a smaller or larger buffer might provide better performance on some
 // systems, but honestly this is good enough
 const BUF_SIZE: usize = 16 * 1024;
 
 #[uucore::main]
 pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     let matches = uucore::clap_localization::handle_clap_result(uu_app(), args)?;
-
-    let mut buffer = Vec::with_capacity(BUF_SIZE);
+    // use larger pipe if zero-copy is possible
+    // todo: deduplicate logic
+    #[cfg(target_os = "linux")]
+    let buf_size = {
+        use std::os::unix::fs::FileTypeExt;
+        // todo: detect pipe under masked /dev. This is really bad detection.
+        if let Ok(m) = std::fs::metadata("/dev/stdout")
+            && m.file_type().is_fifo()
+        {
+            let fd_raw = io::stdout().as_raw_fd();
+            fcntl_setpipe_size(
+                unsafe { BorrowedFd::borrow_raw(fd_raw) },
+                ROOTLESS_MAX_PIPE_SIZE,
+            )
+            .unwrap_or(BUF_SIZE)
+        } else {
+            BUF_SIZE
+        }
+    };
+    #[cfg(not(target_os = "linux"))]
+    let buf_size = BUF_SIZE;
+    let mut buffer: AVec<u8, ConstAlign<ROOTLESS_MAX_PIPE_SIZE>> =
+        AVec::with_capacity(ROOTLESS_MAX_PIPE_SIZE, buf_size);
     #[allow(clippy::unwrap_used, reason = "clap provides 'y' by default")]
     let _ = args_into_buffer(&mut buffer, matches.get_many::<OsString>("STRING").unwrap());
-    prepare_buffer(&mut buffer);
+    prepare_buffer(&mut buffer, buf_size);
 
     match exec(&buffer) {
         Ok(()) => Ok(()),
@@ -55,8 +86,8 @@ pub fn uu_app() -> Command {
 
 /// Copies words from `i` into `buf`, separated by spaces.
 #[allow(clippy::unnecessary_wraps, reason = "needed on some platforms")]
-fn args_into_buffer<'a>(
-    buf: &mut Vec<u8>,
+fn args_into_buffer<'a, A: Alignment>(
+    buf: &mut AVec<u8, A>,
     i: impl Iterator<Item = &'a OsString>,
 ) -> Result<(), Box<dyn Error>> {
     // On Unix (and wasi), OsStrs are just &[u8]'s underneath...
@@ -91,23 +122,25 @@ fn args_into_buffer<'a>(
 
 /// Assumes buf holds a single output line forged from the command line arguments, copies it
 /// repeatedly until the buffer holds as many copies as it can under [`BUF_SIZE`].
-fn prepare_buffer(buf: &mut Vec<u8>) {
-    if buf.len() * 2 > BUF_SIZE {
+fn prepare_buffer<A: Alignment>(buf: &mut AVec<u8, A>, buf_size: usize) {
+    if buf.len() * 2 > buf_size {
         return;
     }
 
     assert!(!buf.is_empty());
 
     let line_len = buf.len();
-    let target_size = line_len * (BUF_SIZE / line_len);
+    let target_size = line_len * (buf_size / line_len);
 
     while buf.len() < target_size {
-        let to_copy = std::cmp::min(target_size - buf.len(), buf.len());
+        let current_len = buf.len();
+        let to_copy = std::cmp::min(target_size - current_len, current_len);
         debug_assert_eq!(to_copy % line_len, 0);
-        buf.extend_from_within(..to_copy);
+        buf.extend_from_slice(&buf[..to_copy].to_vec());
     }
 }
 
+#[cfg(not(target_os = "linux"))]
 pub fn exec(bytes: &[u8]) -> io::Result<()> {
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
@@ -117,6 +150,25 @@ pub fn exec(bytes: &[u8]) -> io::Result<()> {
     }
 }
 
+#[cfg(target_os = "linux")]
+pub fn exec(bytes: &[u8]) -> io::Result<()> {
+    let stdout = io::stdout();
+    //zero copy fast-path
+    //needed for  large args
+    //todo: align instead of giving up fast-path
+    let aligned = bytes.len().is_multiple_of(page_size());
+    if aligned {
+        let fd_raw = stdout.as_raw_fd();
+        let fd = unsafe { BorrowedFd::borrow_raw(fd_raw) };
+        let iovec = [IoSliceRaw::from_slice(bytes)];
+        while unsafe { vmsplice(fd, &iovec, SpliceFlags::empty()) }.is_ok() {}
+    }
+    let mut stdout = stdout.lock();
+    loop {
+        stdout.write_all(bytes)?;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -143,7 +195,7 @@ mod tests {
 
         for (line, final_len) in tests {
             let mut v = std::iter::repeat_n(b'a', line).collect::<Vec<_>>();
-            prepare_buffer(&mut v);
+            prepare_buffer(&mut v, BUF_SIZE);
             assert_eq!(v.len(), final_len);
         }
     }