tee: add zero-copy fast-path

Author: oech3

src/uu/tee/Cargo.toml

@@ -30,6 +30,10 @@ divan = { workspace = true }
 tempfile = { workspace = true }
 uucore = { workspace = true, features = ["benchmark"] }
 
+[target.'cfg(any(target_os = "linux", target_os = "android"))'.dependencies]
+rustix = { workspace = true, features = ["pipe"] }
+uucore = { workspace = true, features = ["pipes"] }
+
 [[bin]]
 name = "tee"
 path = "src/main.rs"

src/uu/tee/src/tee.rs

@@ -143,22 +143,67 @@ impl MultiWriter {
     /// Copies all bytes from the input buffer to the output buffer
     /// without buffering which is POSIX requirement.
     pub fn copy_unbuffered(&mut self) -> Result<(), ()> {
-        // todo: support splice() and tee() fast-path at here
         #[cfg(not(any(unix, target_os = "wasi")))]
         use io::Read as _;
         const BUF_SIZE: usize = 32 * 1024;
+        #[cfg(any(unix, target_os = "wasi"))]
+        let input = rustix::stdio::stdin();
+        #[cfg(not(any(unix, target_os = "wasi")))]
+        let mut input = io::stdin();
+        #[cfg(any(target_os = "linux", target_os = "android"))]
+        macro_rules! splice_or_detach {
+            ($pipe:expr, $writer:expr, $len:expr) => {
+                if let Err(e) = uucore::pipes::drain_pipe($pipe, $writer, $len) {
+                    self.aborted |=
+                        process_error(self.output_error_mode, e, $writer, &mut self.ignored_errors)
+                            .is_err();
+                    $writer.name.clear(); //mark as exited
+                }
+            };
+        }
+        // needs 2 pipes to duplicate input multiple times
+        #[cfg(any(target_os = "linux", target_os = "android"))]
+        if let Ok((pipe_read, pipe_write)) = io::pipe()
+            && let Ok((pipe2_read, pipe2_write)) = io::pipe()
+        {
+            use rustix::pipe::fcntl_setpipe_size;
+            use uucore::pipes::MAX_ROOTLESS_PIPE_SIZE;
+            // improve throughput. 2nd pipe should be larger than 1st one for proper tee() length.
+            if fcntl_setpipe_size(&pipe2_read, MAX_ROOTLESS_PIPE_SIZE).is_ok() {
+                let _ = fcntl_setpipe_size(&pipe_read, MAX_ROOTLESS_PIPE_SIZE);
+                let _ = fcntl_setpipe_size(&self.writers[0], MAX_ROOTLESS_PIPE_SIZE); // stdout
+            }
+            while let Ok(s) = uucore::pipes::splice(&input, &pipe_write, MAX_ROOTLESS_PIPE_SIZE) {
+                if s == 0 {
+                    return Ok(());
+                }
+                let Some((last, others)) = self.writers.split_last_mut() else {
+                    // all writers exited
+                    return Ok(());
+                };
+                for other in others {
+                    // do not consume input
+                    let tee_res = uucore::pipes::tee(&pipe_read, &pipe2_write, s);
+                    assert_eq!(tee_res, Ok(s), "2nd pipe should have enough spare");
+                    splice_or_detach!(&pipe2_read, other, s);
+                }
+                // last one consumes input
+                splice_or_detach!(&pipe_read, last, s);
+                self.writers.retain(|w| !w.name.is_empty());
+                if self.aborted {
+                    return Err(());
+                }
+            }
+        }
+
         #[cfg(any(unix, target_os = "wasi"))]
         let mut buf = [std::mem::MaybeUninit::<u8>::uninit(); BUF_SIZE];
         // todo: avoid cost by 0-fill keeping throughput
         #[cfg(not(any(unix, target_os = "wasi")))]
         let mut buf = [0u8; BUF_SIZE];
-
-        let input = io::stdin();
-        #[cfg(not(any(unix, target_os = "wasi")))]
-        let mut input = input;
         loop {
             #[cfg(any(unix, target_os = "wasi"))]
-            let res = rustix::io::read(&input, &mut buf)
+            let res = rustix::io::read(input, &mut buf)
                 .map(|f| f.0)
                 .map_err(Error::from);
             #[cfg(not(any(unix, target_os = "wasi")))]
@@ -267,3 +312,13 @@ struct NamedWriter {
     inner: Writer,
     pub name: OsString,
 }
+
+#[cfg(any(target_os = "linux", target_os = "android"))]
+impl rustix::fd::AsFd for NamedWriter {
+    fn as_fd(&self) -> rustix::fd::BorrowedFd<'_> {
+        match &self.inner {
+            Writer::File(f) => f.as_fd(),
+            Writer::Stdout(s) => s.0,
+        }
+    }
+}