92 lines
2.5 KiB
Rust
92 lines
2.5 KiB
Rust
use std::fs;
|
|
use std::io;
|
|
use std::thread;
|
|
use std::time::{Duration, Instant};
|
|
|
|
use nix::sys::prctl;
|
|
use rustix::io::Errno;
|
|
use rustix::process::{self, Pid, PidfdFlags, Signal, WaitOptions};
|
|
|
|
const REAP_POLL_INTERVAL: Duration = Duration::from_millis(10);
|
|
|
|
pub fn become_subreaper() -> io::Result<()> {
|
|
Ok(prctl::set_child_subreaper(true)?)
|
|
}
|
|
|
|
pub fn terminate_and_reap_all(grace_period: Duration) -> io::Result<()> {
|
|
signal_tree(process::getpid(), &[Signal::TERM, Signal::CONT]);
|
|
|
|
if reap_until(Instant::now() + grace_period)? {
|
|
return Ok(());
|
|
}
|
|
|
|
kill_and_reap_all()
|
|
}
|
|
|
|
fn signal_tree(parent: Pid, signals: &[Signal]) {
|
|
for child in read_children(parent) {
|
|
if signal_if_child_of(parent, child, signals) {
|
|
signal_tree(child, signals);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn read_children(pid: Pid) -> Vec<Pid> {
|
|
let Ok(tasks) = fs::read_dir(format!("/proc/{}/task", pid.as_raw_pid())) else {
|
|
return Vec::new();
|
|
};
|
|
|
|
tasks
|
|
.filter_map(|task| fs::read_to_string(task.ok()?.path().join("children")).ok())
|
|
.flat_map(|children| parse_pids(&children))
|
|
.collect()
|
|
}
|
|
|
|
fn parse_pids(text: &str) -> Vec<Pid> {
|
|
text.split_whitespace()
|
|
.filter_map(|n| Pid::from_raw(n.parse().ok()?))
|
|
.collect()
|
|
}
|
|
|
|
fn signal_if_child_of(parent: Pid, pid: Pid, signals: &[Signal]) -> bool {
|
|
let Ok(pidfd) = process::pidfd_open(pid, PidfdFlags::empty()) else {
|
|
return false;
|
|
};
|
|
|
|
read_parent(pid) == Some(parent)
|
|
&& signals
|
|
.iter()
|
|
.all(|&signal| process::pidfd_send_signal(&pidfd, signal).is_ok())
|
|
}
|
|
|
|
fn read_parent(pid: Pid) -> Option<Pid> {
|
|
let stat = fs::read_to_string(format!("/proc/{}/stat", pid.as_raw_pid())).ok()?;
|
|
let fields_after_comm = &stat[stat.rfind(')')? + 1..];
|
|
let ppid = fields_after_comm.split_whitespace().nth(1)?;
|
|
Pid::from_raw(ppid.parse().ok()?)
|
|
}
|
|
|
|
fn reap_until(deadline: Instant) -> io::Result<bool> {
|
|
loop {
|
|
match process::wait(WaitOptions::NOHANG) {
|
|
Ok(Some(_)) => {}
|
|
Ok(None) if Instant::now() < deadline => thread::sleep(REAP_POLL_INTERVAL),
|
|
Ok(None) => return Ok(false),
|
|
Err(Errno::CHILD) => return Ok(true),
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn kill_and_reap_all() -> io::Result<()> {
|
|
loop {
|
|
signal_tree(process::getpid(), &[Signal::KILL]);
|
|
|
|
match process::wait(WaitOptions::empty()) {
|
|
Ok(_) => {}
|
|
Err(Errno::CHILD) => return Ok(()),
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
}
|