tee: add splice() and tee() 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

@@ -153,22 +153,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 std::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 = std::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) {
+                    if let Err(e) =
+                        process_error(self.output_error_mode, e, $writer, &mut self.ignored_errors)
+                    {
+                        self.aborted.get_or_insert(e);
+                    }
+                    $writer.name.clear(); //mark as exited
+                }
+            };
+        }
+        #[cfg(any(target_os = "linux", target_os = "android"))]
+        {
+            use rustix::pipe::fcntl_setpipe_size;
+            use uucore::pipes::MAX_ROOTLESS_PIPE_SIZE;
+            let (pipe_read, pipe_write) = std::io::pipe()?; // needed to duplicate input
+            let (pipe2_read, pipe2_write) = std::io::pipe()?; // force-tee() even output is not pipe
+            // 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);
+            }
+            while let Ok(s) = uucore::pipes::splice(&input, &pipe_write, MAX_ROOTLESS_PIPE_SIZE) {
+                if s == 0 {
+                    return Ok(());
+                }
+                let (last, others) = self
+                    .writers
+                    .split_last_mut()
+                    .ok_or_else(|| Error::from(ErrorKind::Other))?; // all writers exited
+                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());
+                self.aborted.take().map_or(Ok(()), 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 = std::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")))]
@@ -276,3 +321,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,
+        }
+    }
+}