use std::process::Child;
use std::sync::{Arc, Mutex, OnceLock};
static PROCESS_REGISTRY: OnceLock<Arc<Mutex<Vec<Child>>>> = OnceLock::new();
fn get_registry() -> &'static Arc<Mutex<Vec<Child>>> {
PROCESS_REGISTRY.get_or_init(|| Arc::new(Mutex::new(Vec::new())))
}
#[must_use]
pub fn register_child(child: Child) -> usize {
get_registry().lock().map_or(0, |mut registry| {
let idx = registry.len();
registry.push(child);
idx
})
}
#[must_use]
pub fn kill_all_registered() -> usize {
get_registry().lock().map_or(0, |mut registry| {
let count = registry.len();
for child in registry.iter_mut() {
let _ = child.kill();
let _ = child.wait();
}
registry.clear();
count
})
}
pub struct ProcessGuard {
child: Option<Child>,
#[allow(dead_code)]
pid: u32,
}
impl ProcessGuard {
#[must_use]
pub fn new(child: Child) -> Self {
let pid = child.id();
Self {
child: Some(child),
pid,
}
}
pub fn wait(&mut self) -> std::io::Result<std::process::ExitStatus> {
self.child.as_mut().map_or_else(
|| Err(std::io::Error::other("Process already consumed")),
Child::wait,
)
}
pub fn wait_with_output(mut self) -> std::io::Result<std::process::Output> {
self.child.take().map_or_else(
|| Err(std::io::Error::other("Process already consumed")),
Child::wait_with_output,
)
}
#[must_use]
pub fn take(mut self) -> Option<Child> {
self.child.take()
}
#[must_use]
pub const fn pid(&self) -> u32 {
self.pid
}
}
impl Drop for ProcessGuard {
fn drop(&mut self) {
if let Some(ref mut child) = self.child {
eprintln!("[JIDOKA] Cleaning up child process {}", self.pid);
let _ = child.kill();
let _ = child.wait();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::process::Command;
#[test]
fn test_process_guard_wait() {
let child = Command::new("echo")
.arg("test")
.spawn()
.expect("Failed to spawn");
let mut guard = ProcessGuard::new(child);
let status = guard.wait().expect("Wait failed");
assert!(status.success());
}
#[test]
fn test_process_guard_wait_with_output() {
use std::process::Stdio;
let child = Command::new("echo")
.arg("hello")
.stdout(Stdio::piped())
.spawn()
.expect("Failed to spawn");
let guard = ProcessGuard::new(child);
let output = guard.wait_with_output().expect("Wait failed");
assert!(output.status.success());
assert!(String::from_utf8_lossy(&output.stdout).contains("hello"));
}
#[test]
fn test_process_guard_take() {
let child = Command::new("echo")
.arg("test")
.spawn()
.expect("Failed to spawn");
let guard = ProcessGuard::new(child);
let mut taken = guard.take().expect("Take failed");
let status = taken.wait().expect("Wait failed");
assert!(status.success());
}
#[test]
fn test_process_guard_drop_kills() {
let child = Command::new("sleep")
.arg("60")
.spawn()
.expect("Failed to spawn");
let pid = child.id();
let guard = ProcessGuard::new(child);
drop(guard);
#[cfg(unix)]
{
use std::path::Path;
std::thread::sleep(std::time::Duration::from_millis(100));
assert!(!Path::new(&format!("/proc/{pid}")).exists());
}
}
#[test]
fn test_kill_all_registered() {
let _ = kill_all_registered();
let child1 = Command::new("sleep")
.arg("60")
.spawn()
.expect("Failed to spawn");
let child2 = Command::new("sleep")
.arg("60")
.spawn()
.expect("Failed to spawn");
let _ = register_child(child1);
let _ = register_child(child2);
let count = kill_all_registered();
assert_eq!(count, 2);
let count2 = kill_all_registered();
assert_eq!(count2, 0);
}
#[test]
fn test_process_guard_pid() {
let child = Command::new("echo")
.arg("test")
.spawn()
.expect("Failed to spawn");
let expected_pid = child.id();
let guard = ProcessGuard::new(child);
assert_eq!(guard.pid(), expected_pid);
}
#[test]
fn test_process_guard_wait_already_consumed() {
let child = Command::new("echo")
.arg("test")
.spawn()
.expect("Failed to spawn");
let mut guard = ProcessGuard::new(child);
let taken = guard.child.take();
assert!(taken.is_some());
let result = guard.wait();
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("already consumed"));
}
#[test]
fn test_process_guard_wait_with_output_already_consumed() {
use std::process::Stdio;
let child = Command::new("echo")
.arg("test")
.stdout(Stdio::piped())
.spawn()
.expect("Failed to spawn");
let mut guard = ProcessGuard::new(child);
let taken = guard.child.take();
assert!(taken.is_some());
let result = guard.wait_with_output();
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("already consumed"));
}
}