diff --git a/src/uu/tee/src/tee.rs b/src/uu/tee/src/tee.rs index 29409fff26f..4e8a855e7e1 100644 --- a/src/uu/tee/src/tee.rs +++ b/src/uu/tee/src/tee.rs @@ -160,59 +160,91 @@ impl MultiWriter { 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::() - && let Ok((mut pipe2_read, mut pipe2_write)) = uucore::pipes::pipe::() + 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 + // use pre-allocated memfd to avoid raising errors about file size e.g. ENOSPC during copy + } else if let Ok(memfd) = uucore::pipes::memfd_sized() + && let Ok((mut pipe_reader, mut pipe_writer)) = uucore::pipes::pipe::() { - use uucore::pipes::MAX_ROOTLESS_PIPE_SIZE; - // improve throughput when stdout is pipe, but no needed to be same size with broker. + use uucore::pipes::{MAX_ROOTLESS_PIPE_SIZE, splice}; + // 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) => { - if let Err(e) = uucore::pipes::drain_pipe(&$pipe_read, &$writer, $len) { + while let Ok(s_in) = splice(&input, &pipe_writer, MAX_ROOTLESS_PIPE_SIZE) { + use rustix::pipe::{SpliceFlags, splice as splice_off}; + if s_in == 0 { + return Ok(()); + } + // Own page caches. This is not a zero-copy, but it still copies in kernel. + let mut remaining = s_in; + let mut off = 0; + while remaining > 0 { + remaining -= splice_off( + &pipe_reader, + None, + &memfd, + Some(&mut off), + remaining, + SpliceFlags::empty(), + ) + .expect("copy between private fds failed"); + } + for writer in &mut self.writers { + let mut remaining = s_in; + let mut off = 0; + while remaining > 0 { + remaining -= splice_off( + &memfd, + Some(&mut off), + &pipe_writer, + None, + remaining, + SpliceFlags::empty(), + ) + .expect("splice between private fds failed"); + } + if let Err(e) = uucore::pipes::drain_pipe(&pipe_reader, &writer, s_in) { self.aborted |= process_error( self.output_error_mode, e, - &$writer, + writer, &mut self.ignored_errors, - ).is_err(); - $writer.name.clear(); //mark as exited - // the failed write can leave bytes in the pipe: replace it with an empty one. + ) + .is_err(); + writer.name.clear(); //mark as exited + // pipes can be dirty. Regenerate them. // Free it first, so that the new one does not need more file descriptors. - drop($pipe_read); - drop($pipe_write); - match uucore::pipes::pipe::() { - Ok(pipe) => ($pipe_read, $pipe_write) = pipe, + drop(pipe_reader); + drop(pipe_writer); + match uucore::pipes::pipe::() { + Ok(pipe) => (pipe_reader, pipe_writer) = pipe, Err(e) => { show_error!("{}", strip_errno(&e)); return Err(()); } } } - }; - } - 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, pipe2_write, *other, s); } - // last one consumes input - splice_or_detach!(pipe_read, pipe_write, *last, s); self.writers.retain(|w| !w.name.is_empty()); if self.aborted { return Err(()); } + if self.writers.is_empty() { + return Ok(()); + } } } diff --git a/src/uucore/src/lib/features/pipes.rs b/src/uucore/src/lib/features/pipes.rs index ef3b3f7446b..8477fb9f4ce 100644 --- a/src/uucore/src/lib/features/pipes.rs +++ b/src/uucore/src/lib/features/pipes.rs @@ -238,3 +238,15 @@ pub fn dev_null() -> Option { pub fn tee(source: &impl AsFd, target: &impl AsFd, len: usize) -> rustix::io::Result { rustix::pipe::tee(source, target, len, SpliceFlags::empty()) } + +/// return 1MiB memfd. +pub fn memfd_sized() -> rustix::io::Result { + 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) +} diff --git a/tests/by-util/test_tee.rs b/tests/by-util/test_tee.rs index 8f7b7e455e4..7fab1c38caf 100644 --- a/tests/by-util/test_tee.rs +++ b/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) { use rlimit::Resource; @@ -756,7 +757,7 @@ 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, pipes 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() {