uucore: splice fast-path for head -c & preliminary for dd, tail, split

Author: oech3

src/uu/head/Cargo.toml

@@ -24,6 +24,7 @@ memchr = { workspace = true }
 thiserror = { workspace = true }
 uucore = { workspace = true, features = [
   "parser-size",
+  "pipes",
   "ringbuffer",
   "lines",
   "fs",

src/uu/head/src/head.rs

@@ -166,6 +166,14 @@ fn wrap_in_stdout_error(err: io::Error) -> io::Error {
     )
 }
 
+// zero-copy fast-path
+#[cfg(any(target_os = "linux", target_os = "android"))]
+fn read_n_bytes(input: impl Read + AsFd, n: u64) -> io::Result<u64> {
+    let out = io::stdout();
+    uucore::pipes::send_n_bytes(input, out, n).map_err(wrap_in_stdout_error)
+}
+
+#[cfg(not(any(target_os = "linux", target_os = "android")))]
 fn read_n_bytes(input: impl Read, n: u64) -> io::Result<u64> {
     // Read the first `n` bytes from the `input` reader.
     let mut reader = input.take(n);
@@ -608,6 +616,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg(not(any(target_os = "linux", target_os = "android")))] // missing trait for AsFd
     fn read_early_exit() {
         let mut empty = io::BufReader::new(Cursor::new(Vec::new()));
         assert!(read_n_bytes(&mut empty, 0).is_ok());

src/uucore/src/lib/features/pipes.rs

@@ -7,10 +7,10 @@
 
 #[cfg(any(target_os = "linux", target_os = "android"))]
 use rustix::pipe::{SpliceFlags, fcntl_setpipe_size};
-#[cfg(any(target_os = "linux", target_os = "android", test))]
-use std::fs::File;
 #[cfg(any(target_os = "linux", target_os = "android"))]
 use std::os::fd::AsFd;
+#[cfg(any(target_os = "linux", target_os = "android", test))]
+use std::{fs::File, io::Read};
 #[cfg(any(target_os = "linux", target_os = "android"))]
 pub const MAX_ROOTLESS_PIPE_SIZE: usize = 1024 * 1024;
 
@@ -30,6 +30,21 @@ pub fn pipe() -> std::io::Result<(File, File)> {
     Ok((File::from(read), File::from(write)))
 }
 
+/// return pipe with given size or kernel's default size
+///
+/// useful to save RAM usage
+#[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
@@ -72,6 +87,85 @@ pub fn might_fuse(source: &impl AsFd) -> bool {
         .unwrap_or(true)
 }
 
+/// splice `n` bytes with safe read/write fallback
+/// return actually sent bytes
+#[inline]
+#[cfg(any(target_os = "linux", target_os = "android"))]
+pub fn send_n_bytes(
+    input: impl Read + AsFd,
+    mut target: impl std::io::Write + AsFd,
+    n: u64,
+) -> std::io::Result<u64> {
+    let pipe_size = MAX_ROOTLESS_PIPE_SIZE.min(n as usize);
+    let mut n = n;
+    let mut bytes_written: u64 = 0;
+    // do not always fallback to write as it needs 2 Ctrl+D on tty
+    let mut needs_fallback = might_fuse(&input);
+    if let Ok(b) = splice(&input, &target, n as usize) {
+        bytes_written = b as u64;
+        n -= bytes_written;
+        if n > 0 {
+            // for throughput. expected that input is already extended if it is coming from splice
+            let _ = fcntl_setpipe_size(&target, pipe_size);
+        } else {
+            // avoid unnecessary syscalls
+            return Ok(bytes_written);
+        }
+
+        loop {
+            match splice(&input, &target, 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)) = pipe_with_size(pipe_size) {
+        // todo: allow reusing broker pipe. needed for using splice to uutils/split in the future
+        loop {
+            match splice(&input, &broker_w, n as usize) {
+                Ok(0) => break,
+                Ok(s @ 1..) => {
+                    if splice_exact(&broker_r, &target, 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)?;
+                        target.write_all(&drain)?;
+                        needs_fallback = true;
+                        break;
+                    }
+                }
+                _ => {
+                    needs_fallback = true;
+                    break;
+                }
+            }
+        }
+    }
+
+    if !needs_fallback {
+        return Ok(bytes_written);
+    }
+    let mut reader = input.take(n);
+    let mut buf = vec![0u8; (32 * 1024).min(n as usize)]; //use heap to avoid early allocation
+    loop {
+        match reader.read(&mut buf)? {
+            0 => return Ok(bytes_written),
+            n => {
+                target.write_all(&buf[..n])?;
+                bytes_written += n as u64;
+            }
+        }
+    }
+}
+
 /// Return verified /dev/null
 ///
 /// `splice` to /dev/null is faster than `read` when we skip or count the non-seekable input