Optionally log Cargo's stderr/stdout to stderr/stdout

Mostly intended for debugging.

Author: Mark-Simulacrum

collector/src/lib.rs

@@ -6,8 +6,10 @@ use std::fmt;
 use std::process::{self, Command};
 
 pub mod api;
+mod read2;
 pub mod self_profile;
 
+use process::Stdio;
 pub use self_profile::{QueryData, SelfProfile};
 
 #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Deserialize)]
@@ -135,20 +137,47 @@ pub fn run_command(cmd: &mut Command) -> anyhow::Result<()> {
 
 pub fn command_output(cmd: &mut Command) -> anyhow::Result<process::Output> {
     log::trace!("running: {:?}", cmd);
-    let output = cmd.output()?;
-    if !output.status.success() {
+    let mut child = cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).spawn()?;
+
+    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(
+        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: "collector_raw_cargo", log::level::Trace) {
+                use std::io::Write;
+                if is_stdout {
+                    stdout_writer.write_all(&buffer[stdout.len()..]).unwrap();
+                } else {
+                    stderr_writer.write_all(&buffer[stderr.len()..]).unwrap();
+                }
+            }
+            if is_stdout {
+                stdout = buffer.clone();
+            } else {
+                stderr = buffer.clone();
+            }
+        },
+    )?;
+
+    let status = child.wait()?;
+    if !status.success() {
         return Err(anyhow::anyhow!(
             "expected success, got {}\n\nstderr={}\n\n stdout={}",
-            output.status,
-            String::from_utf8_lossy(&output.stderr),
-            String::from_utf8_lossy(&output.stdout)
+            status,
+            String::from_utf8_lossy(&stderr),
+            String::from_utf8_lossy(&stdout)
         ));
-    } else {
-        //        log::trace!(
-        //            "stderr={}\n\nstdout={}",
-        //            String::from_utf8_lossy(&output.stderr),
-        //            String::from_utf8_lossy(&output.stdout),
-        //        );
     }
+
+    let output = process::Output {
+        status,
+        stdout: stdout,
+        stderr: stderr,
+    };
     Ok(output)
 }

collector/src/read2.rs

@@ -0,0 +1,74 @@
+// This is copied from
+// https://github.com/rust-lang/rust/blob/master/src/tools/compiletest/src/read2.rs
+// and falls under the MIT or Apache 2.0 licenses.
+
+use std::io;
+use std::io::prelude::*;
+use std::mem;
+use std::os::unix::prelude::*;
+use std::process::{ChildStderr, ChildStdout};
+
+pub fn read2(
+    mut out_pipe: ChildStdout,
+    mut err_pipe: ChildStderr,
+    data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
+) -> io::Result<()> {
+    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);
+    }
+
+    let mut out_done = false;
+    let mut err_done = false;
+    let mut out = Vec::new();
+    let mut err = Vec::new();
+
+    let mut fds: [libc::pollfd; 2] = unsafe { mem::zeroed() };
+    fds[0].fd = out_pipe.as_raw_fd();
+    fds[0].events = libc::POLLIN;
+    fds[1].fd = err_pipe.as_raw_fd();
+    fds[1].events = libc::POLLIN;
+    let mut nfds = 2;
+    let mut errfd = 1;
+
+    while nfds > 0 {
+        // wait for either pipe to become readable using `select`
+        let r = unsafe { libc::poll(fds.as_mut_ptr(), nfds, -1) };
+        if r == -1 {
+            let err = io::Error::last_os_error();
+            if err.kind() == io::ErrorKind::Interrupted {
+                continue;
+            }
+            return Err(err);
+        }
+
+        // Read as much as we can from each pipe, ignoring EWOULDBLOCK or
+        // EAGAIN. If we hit EOF, then this will happen because the underlying
+        // reader will return Ok(0), in which case we'll see `Ok` ourselves. In
+        // this case we flip the other fd back into blocking mode and read
+        // whatever's leftover on that file descriptor.
+        let handle = |res: io::Result<_>| match res {
+            Ok(_) => Ok(true),
+            Err(e) => {
+                if e.kind() == io::ErrorKind::WouldBlock {
+                    Ok(false)
+                } else {
+                    Err(e)
+                }
+            }
+        };
+        if !err_done && fds[errfd].revents != 0 && handle(err_pipe.read_to_end(&mut err))? {
+            err_done = true;
+            nfds -= 1;
+        }
+        data(false, &mut err, err_done);
+        if !out_done && fds[0].revents != 0 && handle(out_pipe.read_to_end(&mut out))? {
+            out_done = true;
+            fds[0].fd = err_pipe.as_raw_fd();
+            errfd = 0;
+            nfds -= 1;
+        }
+        data(true, &mut out, out_done);
+    }
+    Ok(())
+}