yes: use tee syscall as fast-path

This commit is contained in:
oech3
2026-03-23 04:22:21 +09:00
committed by Dorian Péron
parent a07879b8ab
commit 4620c2f3b6
5 changed files with 63 additions and 6 deletions
Generated
+1
View File
@@ -4491,6 +4491,7 @@ dependencies = [
"clap",
"fluent",
"itertools 0.14.0",
"rustix",
"uucore",
]
+2 -1
View File
@@ -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"
+52 -4
View File
@@ -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,15 @@ 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::<OsString>("STRING").unwrap())?;
prepare_buffer(&mut buffer);
#[cfg(any(target_os = "linux", target_os = "android"))]
let aligned_before_growing = PAGE_SIZE.is_multiple_of(buffer.len());
match exec(&buffer) {
prepare_buffer(&mut buffer);
#[cfg(any(target_os = "linux", target_os = "android"))]
let res = exec(&buffer, aligned_before_growing);
#[cfg(not(any(target_os = "linux", target_os = "android")))]
let res = exec(&buffer);
match res {
Ok(()) => Ok(()),
// On Windows, silently handle broken pipe since there's no SIGPIPE
#[cfg(windows)]
@@ -97,10 +110,44 @@ fn prepare_buffer(buf: &mut Vec<u8>) {
}
}
#[cfg(not(any(target_os = "linux", target_os = "android")))]
pub fn exec(bytes: &[u8]) -> io::Result<()> {
let stdout = io::stdout();
let mut stdout = stdout.lock();
let mut stdout = io::stdout().lock();
loop {
stdout.write_all(bytes)?;
}
}
#[cfg(any(target_os = "linux", target_os = "android"))]
pub fn exec(bytes: &[u8], aligned: bool) -> io::Result<()> {
use uucore::pipes::{pipe, splice, tee};
let mut stdout = io::stdout(); // no need to lock with zero-copy
// 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;
}
}
}
}
}
// fallback
let mut stdout = stdout.lock();
loop {
stdout.write_all(bytes)?;
}
@@ -111,6 +158,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),
+7
View File
@@ -258,3 +258,10 @@ pub fn dev_null() -> Option<File> {
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<usize> {
rustix::pipe::tee(source, target, len, SpliceFlags::empty())
}
+1 -1
View File
@@ -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());
}