tee: own page caches using memfd for 2+ outputs

src/uu/tee/src/tee.rs

@@ -149,7 +149,7 @@
         const BUF_SIZE: usize = 32 * 1024;
         #[cfg(all(unix, not(target_os = "fuchsia")))]
         if uucore::signals::stdin_was_closed() {
             let e = Error::from(rustix::io::Errno::BADF);
             show_error!(
                 "{}",
                 translate!("tee-error-stdin", "error" => strip_errno(&e))
@@ -160,14 +160,30 @@
         let input = rustix::stdio::stdin();
         #[cfg(not(any(unix, target_os = "wasi")))]
         let mut input = io::stdin();
-        // needs 2 pipes to duplicate input multiple times
-        // freeze by max size since system allocates pipe smaller than default size when user has pages more than fs.pipe-user-pages-soft
         #[cfg(any(target_os = "linux", target_os = "android"))]
-        if let Ok((mut pipe_read, mut pipe_write)) = uucore::pipes::pipe::<true>()
+        // use trivial splice for one writer.
+        if let [writer] = &mut self.writers[..] {
+            match uucore::pipes::splice_unbounded_auto(&input, writer) {
+                Ok(Ok(())) => return Ok(()),
+                Ok(Err(())) => {} // no splice support
+                Err(e) => {
+                    return process_error(
+                        self.output_error_mode,
+                        e.into(),
+                        writer,
+                        &mut self.ignored_errors,
+                    )
+                    .map_err(|_| ());
+                }
+            }
+        // For 2+ outputs, we need to own page caches of input for the case input was changed after writing to 1st writer
+        } else if let Ok(memfd) = uucore::pipes::memfd_sized()
+            // freeze pipe size since system allocates different sizes when fs.pipe-user-pages-soft was changed between 2 pipe generations.
+            && let Ok((mut pipe_read, mut pipe_write)) = uucore::pipes::pipe::<true>()
             && let Ok((mut pipe2_read, mut pipe2_write)) = uucore::pipes::pipe::<true>()
         {
             use uucore::pipes::MAX_ROOTLESS_PIPE_SIZE;
-            // improve throughput when stdout is pipe, but no needed to be same size with broker.
+            // improve throughput when stdout is pipe, but no needed to guarantee same size with broker.
             let _ = rustix::pipe::fcntl_setpipe_size(&self.writers[0], MAX_ROOTLESS_PIPE_SIZE);
             macro_rules! splice_or_detach {
                 ($pipe_read:ident, $pipe_write:ident, $writer:expr, $len:expr) => {
@@ -201,6 +217,15 @@
                     // all writers exited
                     return Ok(());
                 };
+                // Own page caches using the fact that splicing to file is not a zero-copy.
+                // This is still faster than read/write since it is done in kernel space.
+                // memfd is pre-allocated to avoid errors about file size e.g. ENOSPC during copy
+                uucore::pipes::splice_to_top(&pipe_read, &memfd, s)
+                    .expect("operation between private fds failed");
+                // repeating memfd -> pipe -> writer saves 1 pipe fd pair,
+                // but it is slower than memfd -> pipe and repeating pipe -> pipe -> writer
+                uucore::pipes::splice_from_top(&memfd, &pipe_write, s)
+                    .expect("operation between private fds failed");
                 for other in others {
                     // do not consume input
                     let tee_res = uucore::pipes::tee(&pipe_read, &pipe2_write, s);

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

@@ -238,3 +238,57 @@ pub fn dev_null() -> Option<std::fs::File> {
 pub fn tee(source: &impl AsFd, target: &impl AsFd, len: usize) -> rustix::io::Result<usize> {
     rustix::pipe::tee(source, target, len, SpliceFlags::empty())
 }
+
+/// return 1MiB memfd.
+pub fn memfd_sized() -> rustix::io::Result<rustix::fd::OwnedFd> {
+    let memfd = rustix::fs::memfd_create("", rustix::fs::MemfdFlags::empty())?;
+    rustix::fs::fallocate(
+        &memfd,
+        rustix::fs::FallocateFlags::empty(),
+        0,
+        MAX_ROOTLESS_PIPE_SIZE as u64,
+    )?;
+    Ok(memfd)
+}
+
+/// override 1st `len` bytes of target by splice from source
+pub fn splice_to_top(
+    source: &PipeReader,
+    target: &impl AsFd,
+    len: usize,
+) -> rustix::io::Result<()> {
+    let mut remaining = len;
+    let mut off = 0;
+    while remaining > 0 {
+        remaining -= rustix::pipe::splice(
+            source,
+            None,
+            target,
+            Some(&mut off),
+            remaining,
+            SpliceFlags::empty(),
+        )?;
+    }
+    Ok(())
+}
+
+/// splice `len` bytes to target from top of source
+pub fn splice_from_top(
+    source: &impl AsFd,
+    target: &PipeWriter,
+    len: usize,
+) -> rustix::io::Result<()> {
+    let mut remaining = len;
+    let mut off = 0;
+    while remaining > 0 {
+        remaining -= rustix::pipe::splice(
+            source,
+            Some(&mut off),
+            target,
+            None,
+            remaining,
+            SpliceFlags::empty(),
+        )?;
+    }
+    Ok(())
+}

tests/by-util/test_tee.rs

@@ -716,6 +716,7 @@ mod linux_only {
 
     /// Runs tee with 2 outputs that fail part-way, and checks that the bytes they did
     /// not take do not end up in the other outputs.
+    /// Also check that memfd does not reach to size limit.
     fn check_failed_outputs_do_not_spill(open_file_limit: Option<u64>) {
         use rlimit::Resource;
 
@@ -756,11 +757,11 @@ mod linux_only {
         check_failed_outputs_do_not_spill(None);
     }
 
-    // Just enough file descriptors for stdio, the 3 outputs and tee's 2 pipes: a
+    // Just enough file descriptors for stdio, the 3 outputs, 2 pipe pairs and memfd: a
     // failed output's pipe has to be freed before it is replaced.
     #[test]
     fn test_tee_failed_output_at_open_file_limit() {
-        check_failed_outputs_do_not_spill(Some(10));
+        check_failed_outputs_do_not_spill(Some(11));
     }
 
     #[test]