126 lines
3.5 KiB
Rust
126 lines
3.5 KiB
Rust
use std::collections::HashSet;
|
|
use std::fs;
|
|
use std::io;
|
|
use std::os::fd::OwnedFd;
|
|
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);
|
|
const MAX_WALK_PASSES: usize = 16;
|
|
|
|
pub fn become_subreaper() -> io::Result<()> {
|
|
Ok(prctl::set_child_subreaper(true)?)
|
|
}
|
|
|
|
pub fn terminate_and_reap_all(grace_period: Duration) -> io::Result<()> {
|
|
signal_all_children(&[Signal::TERM, Signal::CONT]);
|
|
|
|
if reap_until(Instant::now() + grace_period)? {
|
|
return Ok(());
|
|
}
|
|
|
|
kill_and_reap_all()
|
|
}
|
|
|
|
// Not exact: a process that forks during the last pass can miss the signals.
|
|
// That is good enough, because every survivor receives SIGKILL after the grace period.
|
|
fn signal_all_children(signals: &[Signal]) {
|
|
let mut signaled = HashSet::new();
|
|
for _ in 0..MAX_WALK_PASSES {
|
|
let root = process::getpid();
|
|
if !signal_new_children(root, read_children(root), signals, &mut signaled) {
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
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_new_children(
|
|
parent: Pid,
|
|
children: Vec<Pid>,
|
|
signals: &[Signal],
|
|
signaled: &mut HashSet<Pid>,
|
|
) -> bool {
|
|
let mut has_signaled_new = false;
|
|
|
|
for child in children {
|
|
if signaled.contains(&child) {
|
|
continue;
|
|
}
|
|
let Some(pidfd) = open_if_child_of(parent, child) else {
|
|
continue;
|
|
};
|
|
|
|
let children_before_signal = read_children(child);
|
|
send_signals(&pidfd, signals);
|
|
signaled.insert(child);
|
|
has_signaled_new = true;
|
|
|
|
signal_new_children(child, children_before_signal, signals, signaled);
|
|
}
|
|
|
|
has_signaled_new
|
|
}
|
|
|
|
fn open_if_child_of(parent: Pid, pid: Pid) -> Option<OwnedFd> {
|
|
let pidfd = process::pidfd_open(pid, PidfdFlags::empty()).ok()?;
|
|
(read_parent(pid) == Some(parent)).then_some(pidfd)
|
|
}
|
|
|
|
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 send_signals(pidfd: &OwnedFd, signals: &[Signal]) {
|
|
for &signal in signals {
|
|
let _ = process::pidfd_send_signal(pidfd, signal);
|
|
}
|
|
}
|
|
|
|
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_all_children(&[Signal::KILL]);
|
|
|
|
match process::wait(WaitOptions::empty()) {
|
|
Ok(_) => {}
|
|
Err(Errno::CHILD) => return Ok(()),
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
}
|