use anyhow::{Context, Result, anyhow};
use std::io::Read;
use std::process::{Child, Command, ExitStatus, Output, Stdio};
use std::thread::{self, JoinHandle};
use std::time::Duration;
use wait_timeout::ChildExt;
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(600);
const DEFAULT_OUTPUT_BYTES: usize = 1024 * 1024;
#[derive(Clone, Copy)]
pub(super) struct Operation<'a>(pub(super) &'a str);
#[derive(Clone, Copy)]
enum OutputStream {
Stdout,
Stderr,
}
impl OutputStream {
const fn name(self) -> &'static str {
match self {
Self::Stdout => "stdout",
Self::Stderr => "stderr",
}
}
}
#[derive(Clone, Copy)]
struct ProcessLimits {
timeout: Duration,
output_bytes: usize,
}
impl Default for ProcessLimits {
fn default() -> Self {
Self {
timeout: DEFAULT_TIMEOUT,
output_bytes: DEFAULT_OUTPUT_BYTES,
}
}
}
pub(super) fn run_command(command: &mut Command, operation: Operation<'_>) -> Result<Output> {
run_command_with_limits(command, operation, ProcessLimits::default())
}
fn run_command_with_limits(
command: &mut Command,
operation: Operation<'_>,
limits: ProcessLimits,
) -> Result<Output> {
let Operation(operation_name) = operation;
let mut child = command
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.with_context(|| format!("{operation_name}: start subprocess"))?;
let stdout_pipe = take_stdout(&mut child, operation)?;
let stderr_pipe = take_stderr(&mut child, operation)?;
let stdout_reader = spawn_bounded_reader(stdout_pipe, limits.output_bytes);
let stderr_reader = spawn_bounded_reader(stderr_pipe, limits.output_bytes);
let status = wait_for_exit(&mut child, operation, limits.timeout);
if status.is_err() {
terminate_child(&mut child, operation)?;
}
let captured_stdout = join_reader(stdout_reader, operation, OutputStream::Stdout)?;
let captured_stderr = join_reader(stderr_reader, operation, OutputStream::Stderr)?;
Ok(Output {
status: status?,
stdout: captured_stdout,
stderr: captured_stderr,
})
}
fn take_stdout(
child: &mut Child,
Operation(operation): Operation<'_>,
) -> Result<std::process::ChildStdout> {
child
.stdout
.take()
.ok_or_else(|| anyhow!("{operation}: capture subprocess stdout"))
}
fn take_stderr(
child: &mut Child,
Operation(operation): Operation<'_>,
) -> Result<std::process::ChildStderr> {
child
.stderr
.take()
.ok_or_else(|| anyhow!("{operation}: capture subprocess stderr"))
}
fn spawn_bounded_reader(
mut pipe: impl Read + Send + 'static,
output_limit: usize,
) -> JoinHandle<Result<Vec<u8>>> {
thread::spawn(move || {
let mut captured = Vec::with_capacity(output_limit.min(8192));
let mut buffer = [0_u8; 8192];
loop {
let bytes_read = pipe.read(&mut buffer).context("read subprocess output")?;
if bytes_read == 0 {
break;
}
let bytes_to_capture = output_limit.saturating_sub(captured.len()).min(bytes_read);
captured.extend(buffer.iter().take(bytes_to_capture).copied());
}
Ok(captured)
})
}
fn wait_for_exit(
child: &mut Child,
Operation(operation): Operation<'_>,
timeout: Duration,
) -> Result<ExitStatus> {
child
.wait_timeout(timeout)
.with_context(|| format!("{operation}: wait for subprocess"))?
.ok_or_else(|| {
anyhow!(
"{operation}: subprocess timed out after {}s",
timeout.as_secs()
)
})
}
fn terminate_child(child: &mut Child, Operation(operation): Operation<'_>) -> Result<()> {
if child
.try_wait()
.with_context(|| format!("{operation}: poll timed-out subprocess"))?
.is_some()
{
return Ok(());
}
child
.kill()
.with_context(|| format!("{operation}: kill timed-out subprocess"))?;
child
.wait()
.with_context(|| format!("{operation}: reap timed-out subprocess"))?;
Ok(())
}
fn join_reader(
reader: JoinHandle<Result<Vec<u8>>>,
Operation(operation): Operation<'_>,
stream: OutputStream,
) -> Result<Vec<u8>> {
let stream_name = stream.name();
reader
.join()
.map_err(|_| anyhow!("{operation}: {stream_name} reader thread panicked"))?
}
#[cfg(test)]
mod tests {
use super::{Operation, OutputStream, ProcessLimits, join_reader, run_command_with_limits};
use rstest::rstest;
use std::io::Write;
use std::process::Command;
use std::time::Duration;
#[test]
fn output_capture_is_limited_per_stream() {
let mut command = probe_command("bounded_output_probe");
let output = run_command_with_limits(
&mut command,
Operation("capture bounded output"),
ProcessLimits {
timeout: Duration::from_secs(5),
output_bytes: 256,
},
)
.expect("bounded-output probe should run");
assert!(output.status.success(), "probe failed: {output:?}");
assert_eq!(output.stdout.len(), 256);
assert_eq!(output.stderr.len(), 256);
}
#[test]
fn stalled_process_is_terminated_at_the_deadline() {
let mut command = probe_command("bounded_timeout_probe");
let error = run_command_with_limits(
&mut command,
Operation("run timeout probe"),
ProcessLimits {
timeout: Duration::from_millis(100),
output_bytes: 256,
},
)
.expect_err("stalled probe should time out");
assert!(format!("{error:#}").contains("run timeout probe: subprocess timed out"));
}
#[rstest]
#[case::stdout(OutputStream::Stdout, "read probe: stdout reader thread panicked")]
#[case::stderr(OutputStream::Stderr, "read probe: stderr reader thread panicked")]
fn reader_thread_diagnostic_names_the_stream(
#[case] stream: OutputStream,
#[case] expected: &str,
) {
let reader = std::thread::spawn(|| -> anyhow::Result<Vec<u8>> {
panic!("deliberate reader-thread failure")
});
let error = join_reader(reader, Operation("read probe"), stream)
.expect_err("panicked reader thread should return an error");
assert_eq!(format!("{error:#}"), expected);
}
fn probe_command(test_name: &str) -> Command {
let mut command = Command::new(
std::env::current_exe().expect("the integration-test executable should have a path"),
);
command.args(["--ignored", "--nocapture", test_name]);
command
}
#[test]
#[ignore = "executed as a high-output subprocess"]
fn bounded_output_probe() {
let output = [b'x'; 4096];
std::io::stdout()
.lock()
.write_all(&output)
.expect("write stdout probe data");
std::io::stderr()
.lock()
.write_all(&output)
.expect("write stderr probe data");
}
#[test]
#[ignore = "executed as a stalled subprocess"]
fn bounded_timeout_probe() {
std::thread::sleep(Duration::from_secs(30));
}
}