Optimize read2 streaming

Before, the whole stdout/stderr buffer would be copied unnecessarily after each data chunk was received.

Author: Kobzol

collector/src/lib.rs

@@ -177,27 +177,28 @@ fn run_command_with_output(cmd: &mut Command) -> anyhow::Result<process::Output>
         .spawn()
         .with_context(|| format!("failed to spawn process for cmd: {:?}", cmd))?;
 
-    let mut stdout = Vec::new();
-    let mut stderr = Vec::new();
-    let mut stdout_writer = std::io::LineWriter::new(std::io::stdout());
-    let mut stderr_writer = std::io::LineWriter::new(std::io::stderr());
-    read2::read2(
+    let mut stdout_writer = std::io::LineWriter::new(std::io::stdout().lock());
+    let mut stderr_writer = std::io::LineWriter::new(std::io::stderr().lock());
+
+    let mut stdout_written = 0;
+    let mut stderr_written = 0;
+    let (stdout, stderr) = read2::read2(
         child.stdout.take().unwrap(),
         child.stderr.take().unwrap(),
         &mut |is_stdout, buffer, _is_done| {
             // Send output if trace logging is enabled
             if log::log_enabled!(target: "raw_cargo_messages", log::Level::Trace) {
                 use std::io::Write;
                 if is_stdout {
-                    stdout_writer.write_all(&buffer[stdout.len()..]).unwrap();
+                    stdout_writer.write_all(&buffer[stdout_written..]).unwrap();
                 } else {
-                    stderr_writer.write_all(&buffer[stderr.len()..]).unwrap();
+                    stderr_writer.write_all(&buffer[stderr_written..]).unwrap();
                 }
             }
             if is_stdout {
-                stdout = buffer.clone();
+                stdout_written = buffer.len();
             } else {
-                stderr = buffer.clone();
+                stderr_written = buffer.len();
             }
         },
     )?;

collector/src/utils/read2.rs

@@ -31,11 +31,12 @@ mod imp {
     use std::os::unix::prelude::*;
     use std::process::{ChildStderr, ChildStdout};
 
+    /// Returns (stdout, stderr).
     pub fn read2(
         mut out_pipe: ChildStdout,
         mut err_pipe: ChildStderr,
         data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
-    ) -> io::Result<()> {
+    ) -> io::Result<(Vec<u8>, Vec<u8>)> {
         unsafe {
             libc::fcntl(out_pipe.as_raw_fd(), libc::F_SETFL, libc::O_NONBLOCK);
             libc::fcntl(err_pipe.as_raw_fd(), libc::F_SETFL, libc::O_NONBLOCK);
@@ -93,7 +94,7 @@ mod imp {
             }
             data(true, &mut out, out_done);
         }
-        Ok(())
+        Ok((out, err))
     }
 }
 
@@ -120,7 +121,7 @@ mod imp {
         out_pipe: ChildStdout,
         err_pipe: ChildStderr,
         data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
-    ) -> io::Result<()> {
+    ) -> io::Result<(Vec<u8>, Vec<u8>)> {
         let mut out = Vec::new();
         let mut err = Vec::new();
 
@@ -151,7 +152,7 @@ mod imp {
                 }
             }
 
-            Ok(())
+            Ok((out, err))
         }
     }