pipes.rs: deduplicate some code

Author: oech3

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

@@ -106,11 +106,7 @@ pub fn splice_unbounded(source: &impl AsFd, dest: &mut impl AsFd) -> std::io::Re
 /// This should not be used if one of them are pipe to save resources
 #[inline]
 #[cfg(any(target_os = "linux", target_os = "android"))]
-pub fn splice_unbounded_broker<R, S>(source: &R, dest: &mut S) -> std::io::Result<bool>
-where
-    R: Read + AsFd,
-    S: AsFd,
-{
+pub fn splice_unbounded_broker(source: &impl AsFd, dest: &mut impl AsFd) -> std::io::Result<bool> {
     static PIPE_CACHE: OnceLock<Option<(PipeReader, PipeWriter)>> = OnceLock::new();
     let Some((pipe_rd, pipe_wr)) = PIPE_CACHE
         .get_or_init(|| pipe::<false>(MAX_ROOTLESS_PIPE_SIZE).ok())
@@ -134,10 +130,7 @@ where
                     // we can recover by copying the data that we have from the
                     // intermediate pipe to stdout using unbuffered read/write. Then
                     // we tell the caller to fall back.
-                    debug_assert!(n <= MAX_ROOTLESS_PIPE_SIZE, "unexpected RAM usage");
-                    let mut drain = Vec::with_capacity(n);
-                    pipe_rd.take(n as u64).read_to_end(&mut drain)?;
-                    crate::io::RawWriter(&dest).write_all(&drain)?;
+                    send_n_rw(pipe_rd, dest, n as u64)?;
                     return Ok(true);
                 }
             }
@@ -150,11 +143,7 @@ where
 /// 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 Write + AsFd,
-    n: u64,
-) -> std::io::Result<u64> {
+pub fn send_n_bytes(input: impl AsFd, target: impl AsFd, n: u64) -> std::io::Result<u64> {
     static PIPE_CACHE: OnceLock<Option<(PipeReader, PipeWriter)>> = OnceLock::new();
     let pipe_size = MAX_ROOTLESS_PIPE_SIZE.min(n as usize);
     let mut n = n;
@@ -203,10 +192,8 @@ pub fn send_n_bytes(
                         }
                     } else {
                         debug_assert!(s <= MAX_ROOTLESS_PIPE_SIZE, "unexpected RAM usage");
-                        // drain pipe before fallback to raw write
-                        let mut drain = Vec::with_capacity(s);
-                        broker_r.take(s as u64).read_to_end(&mut drain)?;
-                        crate::io::RawWriter(&target).write_all(&drain)?;
+                        // cleanup before fallback to raw write
+                        send_n_rw(broker_r, &target, s as u64)?;
                         break true;
                     }
                 }
@@ -220,17 +207,29 @@ pub fn send_n_bytes(
     if !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 {
+    Ok(bytes_written + send_n_rw(&input, &target, n)?)
+}
+
+/// read n bytes from src and write to dest
+/// return sent bytes
+///
+// we cannot include this to splice_exact at a fallback since we need to catch I/O error
+#[cfg(any(target_os = "linux", target_os = "android"))]
+pub fn send_n_rw(src: &impl AsFd, dest: &impl AsFd, n: u64) -> std::io::Result<u64> {
+    let mut sent = 0;
+    let mut reader = crate::io::RawReader(src).take(n);
+    let mut dest = crate::io::RawWriter(dest);
+    let mut buf = [0u8; 1024 * 32]; //rustix::io::read can avoid 0-fill, but it generates wrong error...
+    while sent < n {
         match reader.read(&mut buf)? {
-            0 => return Ok(bytes_written),
-            n => {
-                target.write_all(&buf[..n])?;
-                bytes_written += n as u64;
+            0 => return Ok(sent),
+            written => {
+                dest.write_all(&buf[..written])?;
+                sent += written as u64;
             }
         }
     }
+    Ok(sent)
 }
 
 /// Return verified /dev/null