diff --git a/examples/progress.rs b/examples/progress.rs new file mode 100644 index 0000000..d110513 --- /dev/null +++ b/examples/progress.rs @@ -0,0 +1,8 @@ +use cmd_lib::{run_cmd, CmdResult}; + +#[cmd_lib::main] +fn main() -> CmdResult { + run_cmd!(dd if=/dev/urandom of=/dev/null bs=1M status=progress)?; + + Ok(()) +} diff --git a/src/child.rs b/src/child.rs index a54a2b7..95b7d29 100644 --- a/src/child.rs +++ b/src/child.rs @@ -342,22 +342,72 @@ struct StderrThread { impl StderrThread { fn new(cmd: &str, file: &str, line: u32, stderr: Option, capture: bool) -> Self { if let Some(stderr) = stderr { + let file_ = file.to_owned(); let thread = std::thread::spawn(move || { - let mut output = String::new(); - BufReader::new(stderr) - .lines() - .map_while(Result::ok) - .for_each(|line| { - if !capture { - info!("{line}"); - } else { - if !output.is_empty() { - output.push('\n'); + if capture { + let mut output = String::new(); + BufReader::new(stderr) + .lines() + .map_while(Result::ok) + .for_each(|line| { + if !capture { + info!("{line}"); + } else { + if !output.is_empty() { + output.push('\n'); + } + output.push_str(&line); } - output.push_str(&line); + }); + return output; + } + + // Log output one line at a time, including progress output separated by CR + let mut reader = BufReader::new(stderr); + let mut buffer = vec![]; + loop { + // Unconditionally try to read more data, since the BufReader buffer is empty + let result = match reader.fill_buf() { + Ok(buffer) => buffer, + Err(error) => { + warn!("Error reading from child process: {error:?} at {file_}:{line}"); + break; + } + }; + // Add the result onto our own buffer + buffer.extend(result); + // Empty the BufReader + let read_len = result.len(); + reader.consume(read_len); + + // Log output. Take whole “lines” at every LF or CR (for progress bars etc), + // but leave any incomplete lines in our buffer so we can try to complete them. + while let Some(offset) = buffer.iter().position(|&b| b == b'\n' || b == b'\r') { + let line = &buffer[..offset]; + let line = str::from_utf8(line).map_err(|_| line); + match line { + Ok(string) => info!("{string}"), + Err(bytes) => info!("{bytes:?}"), } - }); - output + buffer = buffer.split_off(offset + 1); + } + + if read_len == 0 { + break; + } + } + + // Log any remaining incomplete line + if !buffer.is_empty() { + let line = &buffer; + let line = str::from_utf8(line).map_err(|_| line); + match line { + Ok(string) => info!("{string}"), + Err(bytes) => info!("{bytes:?}"), + } + } + + "".to_owned() }); Self { cmd: cmd.into(),