use std::io::{Read, Write};
use std::process::{Child, Command, ExitStatus, Output, Stdio};
use std::sync::mpsc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Grace {
Immediate,
Graceful,
}
#[cfg(not(target_os = "windows"))]
const GRACE_PERIOD: Duration = Duration::from_millis(400);
#[cfg(target_os = "windows")]
pub const DETACHED_PROCESS: u32 = 0x0000_0008;
#[cfg(target_os = "windows")]
pub const CREATE_NEW_PROCESS_GROUP: u32 = 0x0000_0200;
#[cfg(target_os = "windows")]
pub const CREATE_NO_WINDOW: u32 = 0x0800_0000;
fn is_signalable(pid: u32) -> bool {
pid > 1
}
pub async fn terminate_tree(pid: u32, grace: Grace) {
if !is_signalable(pid) {
return;
}
#[cfg(not(target_os = "windows"))]
{
if grace == Grace::Graceful {
unix_kill("-TERM", pid).await;
tokio::time::sleep(GRACE_PERIOD).await;
}
unix_kill("-KILL", pid).await;
}
#[cfg(target_os = "windows")]
{
let _ = grace; taskkill_tree(pid).await;
}
}
pub fn terminate_tree_blocking(pid: u32, grace: Grace) {
if !is_signalable(pid) {
return;
}
#[cfg(not(target_os = "windows"))]
{
if grace == Grace::Graceful {
unix_kill_blocking("-TERM", pid);
std::thread::sleep(GRACE_PERIOD);
}
unix_kill_blocking("-KILL", pid);
}
#[cfg(target_os = "windows")]
{
let _ = grace;
taskkill_tree_blocking(pid);
}
}
#[cfg(not(target_os = "windows"))]
async fn unix_kill(signal: &str, pid: u32) {
let _ = tokio::process::Command::new("kill")
.args([signal, "--", &format!("-{pid}"), &pid.to_string()])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.await;
}
#[cfg(not(target_os = "windows"))]
fn unix_kill_blocking(signal: &str, pid: u32) {
let _ = std::process::Command::new("kill")
.args([signal, "--", &format!("-{pid}"), &pid.to_string()])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status();
}
#[cfg(target_os = "windows")]
async fn taskkill_tree(pid: u32) {
let _ = tokio::process::Command::new("taskkill")
.args(["/PID", &pid.to_string(), "/T", "/F"])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.await;
}
#[cfg(target_os = "windows")]
fn taskkill_tree_blocking(pid: u32) {
let _ = std::process::Command::new("taskkill")
.args(["/PID", &pid.to_string(), "/T", "/F"])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status();
}
const POLL_INTERVAL: Duration = Duration::from_millis(10);
const READER_GRACE: Duration = Duration::from_millis(250);
const REAP_GRACE: Duration = Duration::from_millis(500);
pub fn output_with_timeout(cmd: &mut Command, timeout: Duration) -> std::io::Result<Output> {
cmd.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = cmd.spawn()?;
let stdout_rx = drain_pipe(child.stdout.take());
let stderr_rx = drain_pipe(child.stderr.take());
match wait_deadline(&mut child, timeout)? {
Some(status) => Ok(Output {
status,
stdout: collect_drained(stdout_rx),
stderr: collect_drained(stderr_rx),
}),
None => Err(timed_out(cmd, timeout)),
}
}
pub fn write_stdin_with_timeout(
cmd: &mut Command,
input: Vec<u8>,
timeout: Duration,
) -> std::io::Result<ExitStatus> {
cmd.stdin(Stdio::piped())
.stdout(Stdio::null())
.stderr(Stdio::null());
let mut child = cmd.spawn()?;
if let Some(mut stdin) = child.stdin.take() {
std::thread::spawn(move || {
let _ = stdin.write_all(&input);
});
}
match wait_deadline(&mut child, timeout)? {
Some(status) => Ok(status),
None => Err(timed_out(cmd, timeout)),
}
}
fn wait_deadline(child: &mut Child, timeout: Duration) -> std::io::Result<Option<ExitStatus>> {
let deadline = Instant::now() + timeout;
loop {
if let Some(status) = child.try_wait()? {
return Ok(Some(status));
}
if Instant::now() >= deadline {
let _ = child.kill();
let reap_deadline = Instant::now() + REAP_GRACE;
while Instant::now() < reap_deadline {
if matches!(child.try_wait(), Ok(Some(_)) | Err(_)) {
break;
}
std::thread::sleep(POLL_INTERVAL);
}
return Ok(None);
}
std::thread::sleep(POLL_INTERVAL);
}
}
fn drain_pipe<R: Read + Send + 'static>(pipe: Option<R>) -> mpsc::Receiver<Vec<u8>> {
let (tx, rx) = mpsc::channel();
if let Some(mut pipe) = pipe {
std::thread::spawn(move || {
let mut buf = [0u8; 8192];
loop {
match pipe.read(&mut buf) {
Ok(0) | Err(_) => break,
Ok(n) => {
if tx.send(buf[..n].to_vec()).is_err() {
break;
}
},
}
}
});
}
rx
}
fn collect_drained(rx: mpsc::Receiver<Vec<u8>>) -> Vec<u8> {
let mut out = Vec::new();
let deadline = Instant::now() + READER_GRACE;
while let Some(remaining) = deadline.checked_duration_since(Instant::now()) {
match rx.recv_timeout(remaining) {
Ok(chunk) => out.extend_from_slice(&chunk),
Err(_) => break,
}
}
out
}
fn timed_out(cmd: &Command, timeout: Duration) -> std::io::Error {
std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!(
"{} did not exit within {timeout:?} and was killed",
cmd.get_program().to_string_lossy()
),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pids_0_and_1_are_never_signalable() {
assert!(!is_signalable(0));
assert!(!is_signalable(1));
}
#[test]
fn real_pids_are_signalable() {
assert!(is_signalable(2));
assert!(is_signalable(1234));
assert!(is_signalable(u32::MAX));
}
#[tokio::test]
async fn terminate_tree_is_a_noop_for_unsafe_pids() {
terminate_tree(0, Grace::Immediate).await;
terminate_tree(1, Grace::Graceful).await;
terminate_tree_blocking(0, Grace::Immediate);
terminate_tree_blocking(1, Grace::Graceful);
}
#[cfg(unix)]
fn sh(script: &str) -> Command {
let mut c = Command::new("sh");
c.args(["-c", script]);
c
}
#[cfg(windows)]
fn cmd_c(script: &str) -> Command {
let mut c = Command::new("cmd");
c.args(["/C", script]);
c
}
#[cfg(unix)]
#[test]
fn output_with_timeout_captures_output() {
let out = output_with_timeout(&mut sh("echo hello"), Duration::from_secs(10)).unwrap();
assert!(out.status.success());
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "hello");
}
#[cfg(unix)]
#[test]
fn output_with_timeout_kills_a_hung_child() {
let start = Instant::now();
let err = output_with_timeout(&mut sh("sleep 30"), Duration::from_millis(200))
.expect_err("a hung child must surface as an error");
assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
assert!(
start.elapsed() < Duration::from_secs(10),
"deadline kill must not wait for the child's natural exit"
);
}
#[cfg(unix)]
#[test]
fn output_with_timeout_drains_more_than_a_pipe_buffer() {
let out = output_with_timeout(
&mut sh("head -c 1048576 /dev/zero"),
Duration::from_secs(10),
)
.unwrap();
assert!(out.status.success());
assert_eq!(out.stdout.len(), 1_048_576);
}
#[cfg(unix)]
#[test]
fn output_with_timeout_tolerates_a_grandchild_holding_the_pipe() {
let start = Instant::now();
let out = output_with_timeout(
&mut sh("echo early; sleep 5 & exit 0"),
Duration::from_secs(10),
)
.unwrap();
assert!(out.status.success());
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "early");
assert!(
start.elapsed() < Duration::from_secs(4),
"an fd-holding grandchild must not stall collection"
);
}
#[cfg(unix)]
#[test]
fn write_stdin_with_timeout_feeds_the_child() {
let status = write_stdin_with_timeout(
&mut sh(r#"input=$(cat); [ "$input" = "hello" ]"#),
b"hello".to_vec(),
Duration::from_secs(10),
)
.unwrap();
assert!(status.success());
}
#[cfg(unix)]
#[test]
fn write_stdin_with_timeout_kills_a_child_that_never_reads() {
let start = Instant::now();
let err = write_stdin_with_timeout(
&mut sh("sleep 30"),
vec![b'x'; 1_048_576],
Duration::from_millis(200),
)
.expect_err("a child that never reads must time out");
assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
assert!(start.elapsed() < Duration::from_secs(10));
}
#[cfg(windows)]
#[test]
fn output_with_timeout_captures_output() {
let out = output_with_timeout(&mut cmd_c("echo hello"), Duration::from_secs(20)).unwrap();
assert!(out.status.success());
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "hello");
}
#[cfg(windows)]
#[test]
fn output_with_timeout_kills_a_hung_child() {
let start = Instant::now();
let err = output_with_timeout(
&mut cmd_c("ping -n 30 127.0.0.1"),
Duration::from_millis(500),
)
.expect_err("a hung child must surface as an error");
assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
assert!(
start.elapsed() < Duration::from_secs(20),
"deadline kill must not wait for the child's natural exit"
);
}
#[cfg(windows)]
#[test]
fn write_stdin_with_timeout_feeds_the_child() {
let status = write_stdin_with_timeout(
&mut cmd_c("findstr hello"),
b"hello\r\n".to_vec(),
Duration::from_secs(20),
)
.unwrap();
assert!(status.success());
}
}