use anyhow::{Context, Result};
use std::process::{Child, Command, ExitStatus, Stdio};
use std::time::Duration;
pub struct ProcessGuard {
child: Child,
killed: bool,
}
impl ProcessGuard {
pub fn spawn(cmd: &mut Command) -> Result<Self> {
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
cmd.process_group(0);
}
let child = cmd.spawn().context("failed to spawn subprocess")?;
Ok(Self {
child,
killed: false,
})
}
pub fn spawn_detached(cmd: &mut Command) -> Result<Self> {
cmd.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null());
Self::spawn(cmd)
}
pub fn id(&self) -> u32 {
self.child.id()
}
pub fn child_mut(&mut self) -> &mut Child {
&mut self.child
}
pub fn wait(&mut self) -> Result<ExitStatus> {
let status = self.child.wait().context("failed to wait for subprocess")?;
self.sigterm_group();
self.killed = true;
Ok(status)
}
pub fn wait_with_timeout(&mut self, timeout: Duration) -> Result<Option<ExitStatus>> {
let pid = self.child.id();
let fired = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let fired2 = fired.clone();
let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let done2 = done.clone();
std::thread::spawn(move || {
std::thread::sleep(timeout);
if !done2.load(std::sync::atomic::Ordering::SeqCst) {
#[cfg(unix)]
unsafe {
libc::kill(-(pid as i32), libc::SIGKILL);
}
fired2.store(true, std::sync::atomic::Ordering::SeqCst);
}
});
let status = self.child.wait().context("failed to wait for subprocess")?;
done.store(true, std::sync::atomic::Ordering::SeqCst);
self.sigterm_group(); self.killed = true;
if fired.load(std::sync::atomic::Ordering::SeqCst) {
Ok(None)
} else {
Ok(Some(status))
}
}
fn sigterm_group(&self) {
#[cfg(unix)]
{
let pid = self.child.id() as i32;
unsafe {
libc::kill(-pid, libc::SIGTERM);
}
}
}
pub fn force_kill(&mut self) {
#[cfg(unix)]
{
let pid = self.child.id() as i32;
unsafe {
libc::kill(-pid, libc::SIGKILL);
}
}
let _ = self.child.kill();
self.killed = true;
}
pub fn into_child(mut self) -> Child {
self.killed = true; std::mem::replace(
&mut self.child,
Command::new("true")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.expect("failed to spawn dummy"),
)
}
}
impl Drop for ProcessGuard {
fn drop(&mut self) {
if self.killed {
return;
}
#[cfg(unix)]
{
let pid = self.child.id() as i32;
unsafe {
libc::kill(-pid, libc::SIGTERM);
}
}
let _ = self.child.kill();
let _ = self.child.wait();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn spawn_and_wait() {
let _permit = crate::test_subprocess::acquire();
let mut cmd = Command::new("echo");
cmd.arg("hello");
let mut guard = ProcessGuard::spawn(&mut cmd).unwrap();
let status = guard.wait().unwrap();
assert!(status.success());
}
#[test]
fn spawn_detached() {
let _permit = crate::test_subprocess::acquire();
let mut cmd = Command::new("echo");
cmd.arg("hello");
let mut guard = ProcessGuard::spawn_detached(&mut cmd).unwrap();
let status = guard.wait().unwrap();
assert!(status.success());
}
#[test]
fn wait_with_timeout_completes() {
let _permit = crate::test_subprocess::acquire();
let mut cmd = Command::new("echo");
cmd.arg("hello");
let mut guard = ProcessGuard::spawn(&mut cmd).unwrap();
let status = guard
.wait_with_timeout(Duration::from_secs(5))
.unwrap();
assert!(status.is_some());
assert!(status.unwrap().success());
}
#[cfg(unix)]
#[test]
fn wait_with_timeout_kills_on_expiry() {
let _permit = crate::test_subprocess::acquire();
let mut cmd = Command::new("sleep");
cmd.arg("60");
let mut guard = ProcessGuard::spawn(&mut cmd).unwrap();
let status = guard
.wait_with_timeout(Duration::from_millis(100))
.unwrap();
assert!(status.is_none());
}
#[cfg(unix)]
#[test]
fn drop_cleans_up_process_group() {
let _permit = crate::test_subprocess::acquire();
let mut cmd = Command::new("sleep");
cmd.arg("60");
let guard = ProcessGuard::spawn(&mut cmd).unwrap();
let pid = guard.id() as i32;
drop(guard);
std::thread::sleep(Duration::from_millis(100));
let result = unsafe { libc::kill(pid, 0) };
assert_eq!(result, -1);
}
}