use std::cell::RefCell;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
pub const SUBPROCESS_TERM_GRACE: Duration = Duration::from_secs(2);
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ProcessCleanupReport {
pub root_pid: Option<u32>,
pub attempted_signals: Vec<i32>,
pub children: Vec<ProcessCleanupChild>,
}
impl ProcessCleanupReport {
pub fn for_signal(root_pid: Option<u32>, signal: i32) -> Self {
Self {
root_pid,
attempted_signals: vec![signal],
children: Vec::new(),
}
}
pub fn merge(&mut self, other: Self) {
if self.root_pid.is_none() {
self.root_pid = other.root_pid;
}
for signal in other.attempted_signals {
push_unique(&mut self.attempted_signals, signal);
}
for child in other.children {
self.merge_child(child);
}
}
pub fn refresh_survivor_status(&mut self) {
#[cfg(unix)]
{
for child in &mut self.children {
child.alive_after_cleanup = Some(process_exists(child.pid));
}
}
}
fn merge_child(&mut self, child: ProcessCleanupChild) {
if let Some(existing) = self
.children
.iter_mut()
.find(|entry| entry.pid == child.pid)
{
for signal in child.signals {
push_unique(&mut existing.signals, signal);
}
if existing.command_name.is_none() {
existing.command_name = child.command_name;
}
if child.alive_after_cleanup.is_some() {
existing.alive_after_cleanup = child.alive_after_cleanup;
}
return;
}
self.children.push(child);
self.children
.sort_by(|left, right| left.depth.cmp(&right.depth).then(left.pid.cmp(&right.pid)));
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ProcessCleanupChild {
pub pid: u32,
pub parent_pid: Option<u32>,
pub depth: u32,
pub command_name: Option<String>,
pub signals: Vec<i32>,
pub alive_after_cleanup: Option<bool>,
}
impl ProcessCleanupChild {
pub fn new(
pid: u32,
parent_pid: Option<u32>,
depth: u32,
command_name: Option<String>,
) -> Self {
Self {
pid,
parent_pid,
depth,
command_name,
signals: Vec::new(),
alive_after_cleanup: None,
}
}
#[cfg(unix)]
fn with_signal(mut self, signal: i32) -> Self {
push_unique(&mut self.signals, signal);
self
}
}
fn push_unique<T: Copy + Eq>(values: &mut Vec<T>, value: T) {
if !values.contains(&value) {
values.push(value);
}
}
#[derive(Clone, Default)]
struct OpInterrupt {
cancel: Option<Arc<AtomicBool>>,
deadline: Option<Instant>,
}
thread_local! {
static CURRENT: RefCell<Option<OpInterrupt>> = const { RefCell::new(None) };
}
pub struct OpInterruptGuard {
#[allow(clippy::option_option)]
prev: Option<Option<OpInterrupt>>,
}
impl Drop for OpInterruptGuard {
fn drop(&mut self) {
if let Some(prev) = self.prev.take() {
CURRENT.with(|slot| *slot.borrow_mut() = prev);
}
}
}
pub fn install(cancel: Option<Arc<AtomicBool>>, deadline: Option<Instant>) -> OpInterruptGuard {
let prev = CURRENT.with(|slot| slot.borrow_mut().replace(OpInterrupt { cancel, deadline }));
OpInterruptGuard { prev: Some(prev) }
}
pub fn installed() -> bool {
CURRENT.with(|slot| slot.borrow().is_some())
}
pub fn requested() -> bool {
CURRENT.with(|slot| {
let ctx = slot.borrow();
let Some(ctx) = ctx.as_ref() else {
return false;
};
if ctx
.cancel
.as_ref()
.is_some_and(|token| token.load(Ordering::SeqCst))
{
return true;
}
ctx.deadline
.is_some_and(|deadline| Instant::now() >= deadline)
})
}
pub fn configure_kill_group(command: &mut std::process::Command) {
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
command.process_group(0);
}
#[cfg(not(unix))]
{
let _ = command;
}
}
pub fn signal_pid_and_group(pid: u32, signal: i32) {
#[cfg(unix)]
{
extern "C" {
fn kill(pid: i32, sig: i32) -> i32;
}
unsafe {
kill(-(pid as i32), signal);
kill(pid as i32, signal);
}
}
#[cfg(not(unix))]
{
let _ = (pid, signal);
}
}
pub fn signal_pid_tree_and_group(pid: u32, signal: i32) {
let _ = signal_pid_tree_and_group_with_report(pid, signal);
}
pub fn signal_pid_tree_and_group_with_report(pid: u32, signal: i32) -> ProcessCleanupReport {
#[cfg(unix)]
{
let mut report = ProcessCleanupReport::for_signal(Some(pid), signal);
for child in descendant_processes(pid) {
signal_pid_and_group(child.pid, signal);
report.merge_child(child.with_signal(signal));
}
signal_pid_and_group(pid, signal);
report.refresh_survivor_status();
report
}
#[cfg(not(unix))]
{
ProcessCleanupReport::for_signal(Some(pid), signal)
}
}
#[cfg(unix)]
fn descendant_processes(root: u32) -> Vec<ProcessCleanupChild> {
use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System};
let mut sys = System::new();
sys.refresh_processes_specifics(
ProcessesToUpdate::All,
false,
ProcessRefreshKind::everything(),
);
let rows = sys
.processes()
.iter()
.filter_map(|(pid, process)| {
Some((
pid.as_u32(),
process.parent()?.as_u32(),
command_name(process.cmd()),
))
})
.collect::<Vec<_>>();
descendant_processes_from_parent_edges(root, &rows)
}
#[cfg(all(unix, test))]
fn descendant_pids_from_parent_edges(root: u32, edges: &[(u32, u32)]) -> Vec<u32> {
let rows = edges
.iter()
.map(|(pid, parent)| (*pid, *parent, None))
.collect::<Vec<_>>();
descendant_processes_from_parent_edges(root, &rows)
.into_iter()
.map(|child| child.pid)
.collect()
}
#[cfg(unix)]
fn descendant_processes_from_parent_edges(
root: u32,
rows: &[(u32, u32, Option<String>)],
) -> Vec<ProcessCleanupChild> {
use std::collections::{HashMap, HashSet};
let mut children: HashMap<u32, Vec<u32>> = HashMap::new();
let mut metadata: HashMap<u32, (u32, Option<String>)> = HashMap::new();
for (pid, parent, command) in rows {
metadata.insert(*pid, (*parent, command.clone()));
children.entry(*parent).or_default().push(*pid);
}
let mut seen = HashSet::new();
let mut stack = vec![(root, 0usize)];
let mut descendants = Vec::new();
while let Some((pid, depth)) = stack.pop() {
if !seen.insert(pid) {
continue;
}
if pid != root {
descendants.push((pid, depth));
}
if let Some(kids) = children.get(&pid) {
for &child in kids {
stack.push((child, depth + 1));
}
}
}
descendants.sort_by(|(left_pid, left_depth), (right_pid, right_depth)| {
right_depth
.cmp(left_depth)
.then_with(|| left_pid.cmp(right_pid))
});
descendants
.into_iter()
.map(|(pid, depth)| {
let (parent_pid, command) = metadata.get(&pid).cloned().unwrap_or((root, None));
ProcessCleanupChild::new(pid, Some(parent_pid), depth as u32, command)
})
.collect()
}
#[cfg(unix)]
fn command_name(command: &[std::ffi::OsString]) -> Option<String> {
if command.is_empty() {
return None;
}
std::path::Path::new(&command[0])
.file_name()
.map(|name| name.to_string_lossy().into_owned())
.filter(|name| !name.is_empty())
}
#[cfg(unix)]
fn process_exists(pid: u32) -> bool {
extern "C" {
fn kill(pid: i32, sig: i32) -> i32;
}
unsafe { kill(pid as i32, 0) == 0 }
}
pub enum ChildWait {
Exited(std::process::ExitStatus),
TimedOut(ProcessCleanupReport),
Interrupted(Option<std::process::ExitStatus>, ProcessCleanupReport),
}
pub fn wait_child_interruptible(
child: &mut std::process::Child,
timeout: Option<Duration>,
) -> std::io::Result<ChildWait> {
let deadline = timeout.map(|limit| Instant::now() + limit);
loop {
if let Some(status) = child.try_wait()? {
return Ok(ChildWait::Exited(status));
}
if requested() {
let (status, report) = terminate_child_group_with_report(child);
return Ok(ChildWait::Interrupted(status, report));
}
if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
let mut report = child_pid(child)
.map(|pid| signal_pid_tree_and_group_with_report(pid, 9))
.unwrap_or_default();
let _ = child.kill();
let _ = child.wait();
report.refresh_survivor_status();
return Ok(ChildWait::TimedOut(report));
}
std::thread::sleep(Duration::from_millis(20));
}
}
pub fn terminate_child_group(child: &mut std::process::Child) -> Option<std::process::ExitStatus> {
terminate_child_group_with_report(child).0
}
pub fn terminate_child_group_with_report(
child: &mut std::process::Child,
) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
let mut report = child_pid(child)
.map(|pid| ProcessCleanupReport::for_signal(Some(pid), 15))
.unwrap_or_default();
#[cfg(unix)]
{
if let Some(pid) = child_pid(child) {
const SIGTERM: i32 = 15;
report = signal_pid_tree_and_group_with_report(pid, SIGTERM);
let grace_deadline = Instant::now() + SUBPROCESS_TERM_GRACE;
loop {
match child.try_wait() {
Ok(Some(status)) => {
report.merge(signal_pid_tree_and_group_with_report(pid, 9));
report.refresh_survivor_status();
return (Some(status), report);
}
Ok(None) => {
if Instant::now() >= grace_deadline {
break;
}
std::thread::sleep(Duration::from_millis(20));
}
Err(_) => break,
}
}
report.merge(signal_pid_tree_and_group_with_report(pid, 9));
}
}
let _ = child.kill();
let status = child.wait().ok();
report.refresh_survivor_status();
(status, report)
}
fn child_pid(child: &std::process::Child) -> Option<u32> {
let pid = child.id();
(pid > 0).then_some(pid)
}
pub(crate) fn drain_captured_pipe(
rx: &std::sync::mpsc::Receiver<Vec<u8>>,
killed: bool,
child_pid: u32,
) -> Vec<u8> {
use std::sync::mpsc::RecvTimeoutError;
if killed {
return rx
.recv_timeout(Duration::from_millis(100))
.unwrap_or_default();
}
loop {
match rx.recv_timeout(Duration::from_millis(20)) {
Ok(buf) => return buf,
Err(RecvTimeoutError::Disconnected) => return Vec::new(),
Err(RecvTimeoutError::Timeout) => {
if requested() {
const SIGTERM: i32 = 15;
signal_pid_tree_and_group(child_pid, SIGTERM);
if let Ok(buf) = rx.recv_timeout(SUBPROCESS_TERM_GRACE) {
signal_pid_tree_and_group(child_pid, 9);
return buf;
}
signal_pid_tree_and_group(child_pid, 9);
return rx
.recv_timeout(Duration::from_millis(100))
.unwrap_or_default();
}
}
}
}
}
pub(crate) fn spawn_pipe_drain<R: std::io::Read + Send + 'static>(
mut reader: R,
) -> std::sync::mpsc::Receiver<Vec<u8>> {
let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
std::thread::spawn(move || {
let mut buf = Vec::new();
let _ = reader.read_to_end(&mut buf);
let _ = tx.send(buf);
});
rx
}
pub fn capture_output_interruptible(
command: &mut std::process::Command,
) -> std::io::Result<std::process::Output> {
use std::process::Stdio;
command
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.stdin(Stdio::null());
configure_kill_group(command);
let mut child = command.spawn()?;
let pid = child.id();
let rx_out = child.stdout.take().map(spawn_pipe_drain);
let rx_err = child.stderr.take().map(spawn_pipe_drain);
let (status, killed) = match wait_child_interruptible(&mut child, None)? {
ChildWait::Exited(status) => (status, false),
ChildWait::TimedOut(_) => (std::process::ExitStatus::default(), true),
ChildWait::Interrupted(status, _) => (status.unwrap_or_default(), true),
};
let stdout = rx_out
.map(|rx| drain_captured_pipe(&rx, killed, pid))
.unwrap_or_default();
let stderr = rx_err
.map(|rx| drain_captured_pipe(&rx, killed, pid))
.unwrap_or_default();
Ok(std::process::Output {
status,
stdout,
stderr,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn requested_is_false_without_context() {
assert!(!requested());
}
#[test]
fn installed_tracks_guard_lifetime() {
assert!(!installed());
let guard = install(None, None);
assert!(installed());
drop(guard);
assert!(!installed());
}
#[test]
fn cancel_token_trips_requested_and_guard_restores() {
let token = Arc::new(AtomicBool::new(false));
let guard = install(Some(token.clone()), None);
assert!(!requested());
token.store(true, Ordering::SeqCst);
assert!(requested());
drop(guard);
assert!(!requested());
}
#[test]
fn deadline_trips_requested() {
let expired = Instant::now()
.checked_sub(Duration::from_millis(1))
.expect("monotonic clock supports a 1ms test lookback");
let _guard = install(None, Some(expired));
assert!(requested());
}
#[test]
fn nested_installs_restore_in_order() {
let outer_token = Arc::new(AtomicBool::new(true));
let _outer = install(Some(outer_token), None);
assert!(requested());
{
let _inner = install(None, None);
assert!(!requested());
}
assert!(requested());
}
#[cfg(unix)]
#[test]
fn descendant_pids_from_parent_edges_returns_deepest_first_tree_only() {
let edges = [
(20, 10),
(30, 20),
(40, 20),
(50, 30),
(60, 99),
(70, 60),
(80, 90),
(90, 80),
];
assert_eq!(
descendant_pids_from_parent_edges(10, &edges),
vec![50, 30, 40, 20]
);
assert_eq!(descendant_pids_from_parent_edges(99, &edges), vec![70, 60]);
assert_eq!(
descendant_pids_from_parent_edges(123, &edges),
Vec::<u32>::new()
);
}
#[cfg(unix)]
#[test]
fn descendant_processes_preserve_metadata_and_depth_order() {
let rows = [
(20, 10, Some("worker".to_string())),
(30, 20, Some("grandchild".to_string())),
(40, 20, None),
(50, 30, Some("leaf".to_string())),
];
let descendants = descendant_processes_from_parent_edges(10, &rows);
let pids = descendants
.iter()
.map(|child| {
(
child.pid,
child.parent_pid,
child.depth,
child.command_name.as_deref(),
)
})
.collect::<Vec<_>>();
assert_eq!(
pids,
vec![
(50, Some(30), 3, Some("leaf")),
(30, Some(20), 2, Some("grandchild")),
(40, Some(20), 2, None),
(20, Some(10), 1, Some("worker")),
]
);
}
#[cfg(unix)]
#[test]
fn command_name_keeps_only_argv0_basename() {
let command = vec![
std::ffi::OsString::from("/usr/local/bin/tool"),
std::ffi::OsString::from("--api-key"),
std::ffi::OsString::from("secret-value"),
std::ffi::OsString::from("plain"),
];
assert_eq!(command_name(&command).as_deref(), Some("tool"));
assert_eq!(command_name(&[]).as_deref(), None);
}
#[cfg(unix)]
#[test]
fn interrupted_wait_kills_process_group() {
let mut command = std::process::Command::new("sh");
command.args(["-c", "sleep 30 & wait"]);
configure_kill_group(&mut command);
let mut child = command.spawn().expect("spawn sh");
let pgid = child.id();
let cancel = Arc::new(AtomicBool::new(true));
let _guard = install(Some(cancel), None);
let started = Instant::now();
let outcome = wait_child_interruptible(&mut child, None).expect("wait");
assert!(matches!(outcome, ChildWait::Interrupted(_, _)));
assert!(started.elapsed() < Duration::from_secs(10));
extern "C" {
fn kill(pid: i32, sig: i32) -> i32;
}
let group_gone = || unsafe { kill(-(pgid as i32), 0) } != 0;
let deadline = Instant::now() + Duration::from_secs(5);
while !group_gone() && Instant::now() < deadline {
std::thread::sleep(Duration::from_millis(50));
}
assert!(group_gone(), "process group {pgid} survived interrupt");
}
}