use std::io;
use std::process::{Command, Output, Stdio};
use std::time::{Duration, Instant};
const MAX_PROBE_OUTPUT_BYTES: usize = 8 * 1024 * 1024;
fn drain_pipe(mut pipe: impl io::Read + Send + 'static) -> std::thread::JoinHandle<Vec<u8>> {
std::thread::spawn(move || {
let mut buf = Vec::new();
let mut chunk = [0u8; 64 * 1024];
loop {
if buf.len() >= MAX_PROBE_OUTPUT_BYTES {
break;
}
let to_read = chunk.len().min(MAX_PROBE_OUTPUT_BYTES - buf.len());
match pipe.read(&mut chunk[..to_read]) {
Ok(0) | Err(_) => break,
Ok(n) => buf.extend_from_slice(&chunk[..n]),
}
}
buf
})
}
pub fn run_with_timeout(program: &str, args: &[&str], timeout: Duration) -> io::Result<Output> {
let mut child = Command::new(program)
.args(args)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()?;
let out_reader = drain_pipe(child.stdout.take().expect("stdout was piped"));
let err_reader = drain_pipe(child.stderr.take().expect("stderr was piped"));
let deadline = Instant::now() + timeout;
let status = loop {
if let Some(status) = child.try_wait()? {
break status;
}
if Instant::now() >= deadline {
let _ = child.kill();
let _ = child.wait();
let _ = out_reader.join();
let _ = err_reader.join();
return Err(io::Error::new(
io::ErrorKind::TimedOut,
format!("`{program}` timed out after {timeout:?}"),
));
}
std::thread::sleep(Duration::from_millis(50));
};
Ok(Output {
status,
stdout: out_reader.join().unwrap_or_default(),
stderr: err_reader.join().unwrap_or_default(),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_run_with_timeout_captures_stdout() {
let out = run_with_timeout("echo", &["hello"], Duration::from_secs(5)).unwrap();
assert!(out.status.success());
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "hello");
}
#[test]
fn test_run_with_timeout_kills_hung_process() {
let start = Instant::now();
let result = run_with_timeout("sleep", &["30"], Duration::from_millis(200));
assert!(result.is_err());
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::TimedOut);
assert!(
start.elapsed() < Duration::from_secs(5),
"timeout was not enforced promptly"
);
}
#[test]
fn test_run_with_timeout_nonexistent_program() {
let result = run_with_timeout("zzz_no_such_binary_xyz", &[], Duration::from_secs(5));
assert!(result.is_err());
}
#[test]
fn test_drain_pipe_caps_output_instead_of_buffering_unbounded_data() {
struct FiniteButLargeReader {
remaining: usize,
}
impl io::Read for FiniteButLargeReader {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if self.remaining == 0 {
return Ok(0);
}
let n = buf.len().min(self.remaining);
for b in &mut buf[..n] {
*b = b'x';
}
self.remaining -= n;
Ok(n)
}
}
let source = FiniteButLargeReader {
remaining: MAX_PROBE_OUTPUT_BYTES * 2,
};
let bytes = drain_pipe(source).join().unwrap();
assert!(
bytes.len() <= MAX_PROBE_OUTPUT_BYTES,
"drain_pipe buffered {} bytes, more than the {}-byte cap",
bytes.len(),
MAX_PROBE_OUTPUT_BYTES
);
}
#[test]
fn test_run_with_timeout_large_output_no_deadlock() {
let out = run_with_timeout(
"sh",
&["-c", "head -c 200000 /dev/zero"],
Duration::from_secs(10),
)
.expect("large-output command should complete, not time out");
assert!(out.status.success());
assert_eq!(out.stdout.len(), 200_000);
}
}