diff --git a/Cargo.lock b/Cargo.lock index 3b6c0bfaf62..e4c1cc744f2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4487,6 +4487,7 @@ dependencies = [ "clap", "fluent", "itertools 0.14.0", + "rustix", "uucore", ] diff --git a/src/uu/yes/Cargo.toml b/src/uu/yes/Cargo.toml index 0f1737f225a..314c2bd067f 100644 --- a/src/uu/yes/Cargo.toml +++ b/src/uu/yes/Cargo.toml @@ -22,7 +22,8 @@ doctest = false clap = { workspace = true } itertools = { workspace = true } fluent = { workspace = true } -uucore = { workspace = true } +rustix = { workspace = true, features = ["pipe"] } +uucore = { workspace = true, features = ["pipes"] } [[bin]] name = "yes" diff --git a/src/uu/yes/src/yes.rs b/src/uu/yes/src/yes.rs index 06e53f65c4b..3200d997b88 100644 --- a/src/uu/yes/src/yes.rs +++ b/src/uu/yes/src/yes.rs @@ -12,7 +12,14 @@ use uucore::error::{UResult, USimpleError, strip_errno}; use uucore::format_usage; use uucore::translate; +#[cfg(any(target_os = "linux", target_os = "android"))] +const PAGE_SIZE: usize = 4096; +#[cfg(any(target_os = "linux", target_os = "android"))] +use uucore::pipes::MAX_ROOTLESS_PIPE_SIZE; +#[cfg(any(target_os = "linux", target_os = "android"))] +const BUF_SIZE: usize = MAX_ROOTLESS_PIPE_SIZE; // it's possible that using a smaller or larger buffer might provide better performance +#[cfg(not(any(target_os = "linux", target_os = "android")))] const BUF_SIZE: usize = 16 * 1024; #[uucore::main] @@ -21,9 +28,13 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> { #[allow(clippy::unwrap_used, reason = "clap provides 'y' by default")] let mut buffer = args_into_buffer(matches.get_many::("STRING").unwrap())?; - prepare_buffer(&mut buffer); + #[cfg(any(target_os = "linux", target_os = "android"))] + let aligned = PAGE_SIZE.is_multiple_of(buffer.len()); + #[cfg(not(any(target_os = "linux", target_os = "android")))] + let aligned = false; - match exec(&buffer) { + prepare_buffer(&mut buffer); + match exec(&buffer, aligned) { Ok(()) => Ok(()), // On Windows, silently handle broken pipe since there's no SIGPIPE #[cfg(windows)] @@ -97,8 +108,40 @@ fn prepare_buffer(buf: &mut Vec) { } } -pub fn exec(bytes: &[u8]) -> io::Result<()> { +pub fn exec(bytes: &[u8], aligned: bool) -> io::Result<()> { + #[cfg(not(any(target_os = "linux", target_os = "android")))] + let _ = aligned; let stdout = io::stdout(); + #[cfg(any(target_os = "linux", target_os = "android"))] + let mut stdout = stdout; + #[cfg(any(target_os = "linux", target_os = "android"))] + { + use uucore::pipes::{pipe, splice, tee}; + // don't show any error from fast-path and fallback to write for proper message + if let Ok((p_read, mut p_write)) = pipe() + // todo: zero-copy with default size when fcntl failed + && rustix::pipe::fcntl_setpipe_size(&stdout, MAX_ROOTLESS_PIPE_SIZE).is_ok() + && p_write.write_all(bytes).is_ok() + { + if aligned && tee(&p_read, &stdout, PAGE_SIZE).is_ok() { + while let Ok(1..) = tee(&p_read, &stdout, usize::MAX) {} + } else if let Ok((broker_read, broker_write)) = pipe() { + // tee() cannot control offset and write to non-pipe + 'hybrid: while let Ok(mut remain) = tee(&p_read, &broker_write, usize::MAX) { + debug_assert!(remain == bytes.len(), "splice() should cleanup pipe"); + while remain > 0 { + if let Ok(s) = splice(&broker_read, &stdout, remain) { + remain -= s; + } else { + // avoid output breakage with reduced remain even if it would not happen + stdout.write_all(&bytes[bytes.len() - remain..])?; + break 'hybrid; + } + } + } + } + } + } let mut stdout = stdout.lock(); loop { @@ -111,6 +154,7 @@ mod tests { use super::*; #[test] + #[cfg(not(any(target_os = "linux", target_os = "android")))] // Linux uses different buffer size fn test_prepare_buffer() { let tests = [ (150, 16350), diff --git a/src/uucore/src/lib/features/pipes.rs b/src/uucore/src/lib/features/pipes.rs index c91fa4a0dd0..bd5e422a1b6 100644 --- a/src/uucore/src/lib/features/pipes.rs +++ b/src/uucore/src/lib/features/pipes.rs @@ -88,3 +88,10 @@ pub fn dev_null() -> Option { None } } + +// Less noisy wrapper around [`rustix::pipe::tee`] +#[inline] +#[cfg(any(target_os = "linux", target_os = "android"))] +pub fn tee(source: &impl AsFd, target: &impl AsFd, len: usize) -> rustix::io::Result { + rustix::pipe::tee(source, target, len, SpliceFlags::empty()) +} diff --git a/tests/by-util/test_yes.rs b/tests/by-util/test_yes.rs index db460c998fc..c638915be94 100644 --- a/tests/by-util/test_yes.rs +++ b/tests/by-util/test_yes.rs @@ -70,7 +70,7 @@ fn test_long_input() { #[cfg(windows)] const TIMES: usize = 500; let arg = "abcdef".repeat(TIMES) + "\n"; - let expected_out = arg.repeat(30); + let expected_out = arg.repeat(5); run(&[&arg[..arg.len() - 1]], expected_out.as_bytes()); }