From e76983d912be53d796788623c64f419025505e91 Mon Sep 17 00:00:00 2001 From: oech3 <79379754+oech3@users.noreply.github.com> Date: Wed, 11 Mar 2026 20:29:24 +0900 Subject: [PATCH] yes: Add zero copy fast-path --- .../workspace.wordlist.txt | 1 + Cargo.lock | 31 ++++++ Cargo.toml | 2 + src/uu/yes/Cargo.toml | 3 + src/uu/yes/src/yes.rs | 101 ++++++++++++++---- 5 files changed, 119 insertions(+), 19 deletions(-) diff --git a/.vscode/cspell.dictionaries/workspace.wordlist.txt b/.vscode/cspell.dictionaries/workspace.wordlist.txt index 1d3929832a2..cc682aa8fbe 100644 --- a/.vscode/cspell.dictionaries/workspace.wordlist.txt +++ b/.vscode/cspell.dictionaries/workspace.wordlist.txt @@ -8,6 +8,7 @@ advapi32-sys aho-corasick backtrace blake2b_simd +rustix # * uutils project uutils diff --git a/Cargo.lock b/Cargo.lock index a34facd0566..2a6afc6fe5a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "aligned-vec" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b" +dependencies = [ + "equator", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -978,6 +987,26 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" +[[package]] +name = "equator" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" +dependencies = [ + "equator-macro", +] + +[[package]] +name = "equator-macro" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -4407,9 +4436,11 @@ dependencies = [ name = "uu_yes" version = "0.7.0" dependencies = [ + "aligned-vec", "clap", "fluent", "itertools 0.14.0", + "rustix", "uucore", ] diff --git a/Cargo.toml b/Cargo.toml index 92013b0af57..6a840755829 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -372,6 +372,7 @@ version = "0.7.0" [workspace.dependencies] ansi-width = "0.1.0" +aligned-vec = "0.6.4" bigdecimal = "0.4" binary-heap-plus = "0.5.0" bstr = "1.9.1" @@ -434,6 +435,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" diff --git a/src/uu/yes/Cargo.toml b/src/uu/yes/Cargo.toml index b8703bb55d1..93a2631c8ce 100644 --- a/src/uu/yes/Cargo.toml +++ b/src/uu/yes/Cargo.toml @@ -19,9 +19,12 @@ workspace = true path = "src/yes.rs" [dependencies] +aligned-vec = { workspace = true } clap = { workspace = true } itertools = { workspace = true } fluent = { workspace = true } +rustix = { workspace = true } + [target.'cfg(unix)'.dependencies] uucore = { workspace = true, features = ["pipes", "signals"] } diff --git a/src/uu/yes/src/yes.rs b/src/uu/yes/src/yes.rs index 88d2c7ed113..7c762a13939 100644 --- a/src/uu/yes/src/yes.rs +++ b/src/uu/yes/src/yes.rs @@ -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,6 +14,15 @@ 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; @@ -20,11 +30,32 @@ 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> = + 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::("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, +fn args_into_buffer<'a, A: Alignment>( + buf: &mut AVec, i: impl Iterator, ) -> Result<(), Box> { // On Unix (and wasi), OsStrs are just &[u8]'s underneath... @@ -91,23 +122,29 @@ 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) { - if buf.len() * 2 > BUF_SIZE { +fn prepare_buffer(buf: &mut AVec, 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); + #[allow( + clippy::unnecessary_to_owned, + reason = "needs useless copy without unsafe" + )] + 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 +154,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::*; @@ -142,8 +198,9 @@ mod tests { ]; for (line, final_len) in tests { - let mut v = std::iter::repeat_n(b'a', line).collect::>(); - prepare_buffer(&mut v); + let mut v: AVec> = + AVec::from_iter(ROOTLESS_MAX_PIPE_SIZE, std::iter::repeat_n(b'a', line)); + prepare_buffer(&mut v, BUF_SIZE); assert_eq!(v.len(), final_len); } } @@ -151,24 +208,30 @@ mod tests { #[test] fn test_args_into_buf() { { - let mut v = Vec::with_capacity(BUF_SIZE); + let mut v: AVec> = + AVec::with_capacity(ROOTLESS_MAX_PIPE_SIZE, BUF_SIZE); let default_args = ["y".into()]; args_into_buffer(&mut v, default_args.iter()).unwrap(); - assert_eq!(String::from_utf8(v).unwrap(), "y\n"); + assert_eq!(String::from_utf8(v.to_vec()).unwrap(), "y\n"); } { - let mut v = Vec::with_capacity(BUF_SIZE); + let mut v: AVec> = + AVec::with_capacity(ROOTLESS_MAX_PIPE_SIZE, BUF_SIZE); let args = ["foo".into()]; args_into_buffer(&mut v, args.iter()).unwrap(); - assert_eq!(String::from_utf8(v).unwrap(), "foo\n"); + assert_eq!(String::from_utf8(v.to_vec()).unwrap(), "foo\n"); } { - let mut v = Vec::with_capacity(BUF_SIZE); + let mut v: AVec> = + AVec::with_capacity(ROOTLESS_MAX_PIPE_SIZE, BUF_SIZE); let args = ["foo".into(), "bar baz".into(), "qux".into()]; args_into_buffer(&mut v, args.iter()).unwrap(); - assert_eq!(String::from_utf8(v).unwrap(), "foo bar baz qux\n"); + assert_eq!( + String::from_utf8(v.to_vec()).unwrap(), + "foo bar baz qux\n" + ); } } }