From 87c514d1d14960dd04e74cf74870f64426537424 Mon Sep 17 00:00:00 2001 From: oech3 <79379754+oech3@users.noreply.github.com> Date: Sat, 4 Apr 2026 00:11:28 +0900 Subject: [PATCH] head: splice fast-path for -c --- Cargo.lock | 1 + src/uu/head/Cargo.toml | 4 ++ src/uu/head/src/head.rs | 77 ++++++++++++++++++++++++++++ src/uucore/src/lib/features/pipes.rs | 12 +++++ 4 files changed, 94 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 45107199eec..fde0ac97b92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3638,6 +3638,7 @@ dependencies = [ "clap", "fluent", "memchr", + "rustix", "thiserror 2.0.18", "uucore", ] diff --git a/src/uu/head/Cargo.toml b/src/uu/head/Cargo.toml index 353404a775a..77bfd32c76e 100644 --- a/src/uu/head/Cargo.toml +++ b/src/uu/head/Cargo.toml @@ -30,6 +30,10 @@ uucore = { workspace = true, features = [ ] } fluent = { workspace = true } +[target.'cfg(any(target_os = "linux", target_os = "android"))'.dependencies] +rustix = { workspace = true, features = ["fs"] } +uucore = { workspace = true, features = ["pipes"] } + [[bin]] name = "head" path = "src/main.rs" diff --git a/src/uu/head/src/head.rs b/src/uu/head/src/head.rs index 235ea52587b..26c2539d63e 100644 --- a/src/uu/head/src/head.rs +++ b/src/uu/head/src/head.rs @@ -166,6 +166,7 @@ fn wrap_in_stdout_error(err: io::Error) -> io::Error { ) } +#[cfg(not(any(target_os = "linux", target_os = "android")))] fn read_n_bytes(input: impl Read, n: u64) -> io::Result { // Read the first `n` bytes from the `input` reader. let mut reader = input.take(n); @@ -184,6 +185,81 @@ fn read_n_bytes(input: impl Read, n: u64) -> io::Result { Ok(bytes_written) } +// read_n_bytes with zero-copy fast-path +#[cfg(any(target_os = "linux", target_os = "android"))] +fn read_n_bytes(input: impl Read + AsFd, n: u64) -> io::Result { + use uucore::pipes::splice; + let pipe_size = uucore::pipes::MAX_ROOTLESS_PIPE_SIZE.min(n as usize); + let mut stdout = io::stdout(); + let mut n = n; + let mut bytes_written: u64 = 0; + // io::copy's internal splice() breaks FUSE + // we cannot always fallback to write as it needs 2 Ctrl+D on tty + let mut needs_fallback = uucore::pipes::might_fuse(&input); + if let Ok(b) = splice(&input, &stdout, n as usize) { + bytes_written = b as u64; + n -= bytes_written; + if n == 0 { + // avoid fcntl overhead for small input + return Ok(bytes_written); + } + // improves throughput (expected that input is already extended if it is coming from splice) + let _ = rustix::pipe::fcntl_setpipe_size(&stdout, pipe_size); + loop { + match splice(&input, &stdout, n as usize) { + Ok(0) => break, + Ok(s @ 1..) => { + n -= s as u64; + bytes_written += s as u64; + } + _ => { + needs_fallback = true; + break; + } + } + } + } else if let Ok((broker_r, broker_w)) = uucore::pipes::pipe_with_size(pipe_size) { + // both of in/output are not pipe. needs broker to use splice() with additional cost + loop { + match splice(&input, &broker_w, n as usize) { + Ok(0) => break, + Ok(s @ 1..) => { + if uucore::pipes::splice_exact(&broker_r, &stdout, s).is_ok() { + n -= s as u64; + bytes_written += s as u64; + } else { + let mut drain = Vec::with_capacity(s); // bounded by pipe size + broker_r.take(s as u64).read_to_end(&mut drain)?; + stdout.write_all(&drain).map_err(wrap_in_stdout_error)?; + needs_fallback = true; + break; + } + } + _ => { + needs_fallback = true; + break; + } + } + } + } + + if !needs_fallback { + return Ok(bytes_written); + } + let mut reader = input.take(n); + let mut stdout = stdout.lock(); + let mut buf = vec![0u8; (32 * 1024).min(n as usize)]; //use heap to avoid early allocation + loop { + match reader.read(&mut buf).map_err(wrap_in_stdout_error)? { + 0 => return Ok(bytes_written), + n => { + stdout.write_all(&buf[..n]).map_err(wrap_in_stdout_error)?; + bytes_written += n as u64; + } + } + } +} + fn read_n_lines(input: &mut impl io::BufRead, n: u64, separator: u8) -> io::Result { // Read the first `n` lines from the `input` reader. let mut reader = take_lines(input, n, separator); @@ -608,6 +684,7 @@ mod tests { } #[test] + #[cfg(not(any(target_os = "linux", target_os = "android")))] // missing trait for rustix::fd::AsFd fn read_early_exit() { let mut empty = io::BufReader::new(Cursor::new(Vec::new())); assert!(read_n_bytes(&mut empty, 0).is_ok()); diff --git a/src/uucore/src/lib/features/pipes.rs b/src/uucore/src/lib/features/pipes.rs index 24890efca94..c952966cc6a 100644 --- a/src/uucore/src/lib/features/pipes.rs +++ b/src/uucore/src/lib/features/pipes.rs @@ -30,6 +30,18 @@ pub fn pipe() -> std::io::Result<(File, File)> { Ok((File::from(read), File::from(write))) } +#[inline] +#[cfg(any(target_os = "linux", target_os = "android"))] +pub fn pipe_with_size(s: usize) -> std::io::Result<(File, File)> { + let (read, write) = rustix::pipe::pipe()?; + const DEFAULT_SIZE: usize = 64 * 1024; + if s > DEFAULT_SIZE { + let _ = fcntl_setpipe_size(&read, s); + } + + Ok((File::from(read), File::from(write))) +} + /// Less noisy wrapper around [`rustix::pipe::splice`]. /// /// Up to `len` bytes are moved from `source` to `target`. Returns the number