pfx_run/
ending.rs

1pub fn parent(stat: &str) -> Option<u32> {
2    stat.rsplit_once(')')?
3        .1
4        .split_whitespace()
5        .nth(1)?
6        .parse()
7        .ok()
8}
9
10#[cfg(target_os = "linux")]
11pub use linux::{Process, below, terminate};
12
13#[cfg(target_os = "linux")]
14mod linux {
15    use super::parent;
16    use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
17
18    pub struct Process(OwnedFd);
19
20    impl Process {
21        pub fn end(&self) {
22            unsafe {
23                libc::syscall(
24                    libc::SYS_pidfd_send_signal,
25                    self.0.as_raw_fd(),
26                    libc::SIGTERM,
27                    std::ptr::null::<libc::siginfo_t>(),
28                    0u32,
29                );
30            }
31        }
32    }
33
34    pub fn terminate(child: &mut std::process::Child) {
35        if let Ok(pid) = libc::pid_t::try_from(child.id()) {
36            unsafe {
37                libc::kill(pid, libc::SIGTERM);
38            }
39        }
40    }
41
42    fn open(pid: u32) -> Option<Process> {
43        let fd = unsafe { libc::syscall(libc::SYS_pidfd_open, pid as libc::pid_t, 0u32) };
44        (fd >= 0).then(|| Process(unsafe { OwnedFd::from_raw_fd(fd as i32) }))
45    }
46
47    pub fn below(root: u32) -> Vec<Process> {
48        let parents: Vec<(u32, u32)> = std::fs::read_dir("/proc")
49            .into_iter()
50            .flatten()
51            .flatten()
52            .filter_map(|entry| {
53                let pid = entry.file_name().to_str()?.parse().ok()?;
54                let stat = std::fs::read_to_string(entry.path().join("stat")).ok()?;
55                Some((pid, parent(&stat)?))
56            })
57            .collect();
58        let mut tree = vec![root];
59        let mut at = 0;
60        while at < tree.len() {
61            let from = tree[at];
62            tree.extend(parents.iter().filter(|(_, p)| *p == from).map(|(c, _)| *c));
63            at += 1;
64        }
65        tree.into_iter().skip(1).filter_map(open).collect()
66    }
67}
68
69#[cfg(not(target_os = "linux"))]
70pub struct Process;
71
72#[cfg(not(target_os = "linux"))]
73impl Process {
74    pub fn end(&self) {}
75}
76
77#[cfg(not(target_os = "linux"))]
78pub fn below(_root: u32) -> Vec<Process> {
79    Vec::new()
80}
81
82#[cfg(not(target_os = "linux"))]
83pub fn terminate(child: &mut std::process::Child) {
84    let _ = child.kill();
85}
86
87use std::sync::atomic::{AtomicI32, Ordering};
88
89static STOP: AtomicI32 = AtomicI32::new(0);
90#[cfg(target_os = "linux")]
91static QUIT: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
92
93#[cfg(target_os = "linux")]
94static CHILDREN: std::sync::Mutex<Vec<i32>> = std::sync::Mutex::new(Vec::new());
95
96pub struct Tracked {
97    #[cfg(target_os = "linux")]
98    pid: i32,
99}
100
101impl Drop for Tracked {
102    fn drop(&mut self) {
103        #[cfg(target_os = "linux")]
104        {
105            let mut kids = CHILDREN.lock().unwrap_or_else(|poison| poison.into_inner());
106            if let Some(at) = kids.iter().position(|pid| *pid == self.pid) {
107                kids.remove(at);
108            }
109        }
110    }
111}
112
113pub fn track(child: &std::process::Child) -> Tracked {
114    #[cfg(target_os = "linux")]
115    {
116        let pid = rustix::process::Pid::from_child(child).as_raw_pid();
117        CHILDREN
118            .lock()
119            .unwrap_or_else(|poison| poison.into_inner())
120            .push(pid);
121        Tracked { pid }
122    }
123    #[cfg(not(target_os = "linux"))]
124    {
125        let _ = child;
126        Tracked {}
127    }
128}
129
130pub fn stopped() -> Option<i32> {
131    let signal = STOP.load(Ordering::SeqCst);
132    (signal != 0).then_some(signal)
133}
134
135pub fn stop_like(signal: i32) -> ! {
136    #[cfg(target_os = "linux")]
137    {
138        if QUIT.swap(true, Ordering::SeqCst) {
139            loop {
140                std::thread::sleep(std::time::Duration::from_secs(60));
141            }
142        }
143        let _ = signal_hook::low_level::emulate_default_handler(signal);
144        signal_hook::low_level::exit(128 + signal);
145    }
146    #[cfg(not(target_os = "linux"))]
147    std::process::exit(128 + signal);
148}
149
150pub fn install() -> Result<(), String> {
151    #[cfg(not(target_os = "linux"))]
152    {
153        return Ok(());
154    }
155    #[cfg(target_os = "linux")]
156    {
157        static ONCE: std::sync::Once = std::sync::Once::new();
158        let mut result: Result<(), String> = Ok(());
159        ONCE.call_once(|| {
160            result = arm();
161        });
162        result
163    }
164}
165
166#[cfg(target_os = "linux")]
167fn arm() -> Result<(), String> {
168    let mut signals = signal_hook::iterator::Signals::new([
169        signal_hook::consts::SIGTERM,
170        signal_hook::consts::SIGINT,
171        signal_hook::consts::SIGHUP,
172    ])
173    .map_err(|error| format!("signals: {error}"))?;
174    std::thread::Builder::new()
175        .name("pfx-signals".into())
176        .spawn(move || {
177            if let Some(signal) = signals.forever().next() {
178                STOP.store(signal, Ordering::SeqCst);
179                pass_term();
180                stop_like(signal);
181            }
182        })
183        .map_err(|error| format!("signals: {error}"))?;
184    Ok(())
185}
186
187#[cfg(target_os = "linux")]
188fn pass_term() {
189    let pids = CHILDREN
190        .lock()
191        .map(|kids| kids.clone())
192        .unwrap_or_else(|poison| poison.into_inner().clone());
193    for pid in &pids {
194        if let Some(pid) = rustix::process::Pid::from_raw(*pid) {
195            let _ = rustix::process::kill_process(pid, rustix::process::Signal::TERM);
196        }
197    }
198    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
199    while std::time::Instant::now() < deadline {
200        let mut alive = false;
201        for pid in &pids {
202            let Some(pid) = rustix::process::Pid::from_raw(*pid) else {
203                continue;
204            };
205            if matches!(
206                rustix::process::waitpid(Some(pid), rustix::process::WaitOptions::NOHANG),
207                Ok(None)
208            ) {
209                alive = true;
210            }
211        }
212        if !alive {
213            break;
214        }
215        std::thread::sleep(std::time::Duration::from_millis(20));
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn a_parent_is_read_past_a_name_with_spaces_and_brackets() {
225        assert_eq!(parent("4242 (papp (x) y) S 17 4242 17 0 -1"), Some(17));
226        assert_eq!(parent("4242 (render (x) y) S 17 4242 17 0 -1"), Some(17));
227        assert_eq!(parent("1 (systemd) S 0 1 1 0 -1"), Some(0));
228        assert_eq!(parent("garbage"), None);
229    }
230
231    #[cfg(target_os = "linux")]
232    #[test]
233    fn the_processes_below_a_wrapper_end_and_the_wrapper_lives_to_report_it() {
234        use std::time::{Duration, Instant};
235        let mut wrapper = std::process::Command::new("sh")
236            .args(["-c", "sleep 60 & wait $!"])
237            .spawn()
238            .unwrap();
239        let asked = Instant::now();
240        let mut found = below(wrapper.id());
241        while found.is_empty() && asked.elapsed() < Duration::from_secs(5) {
242            std::thread::sleep(Duration::from_millis(20));
243            found = below(wrapper.id());
244        }
245        assert_eq!(found.len(), 1);
246        found.iter().for_each(Process::end);
247        let status = wrapper.wait().unwrap();
248        assert!(asked.elapsed() < Duration::from_secs(30));
249        assert_eq!(status.code(), Some(128 + 15));
250    }
251
252    #[cfg(target_os = "linux")]
253    #[test]
254    fn a_child_is_asked_to_end_with_term() {
255        use std::os::unix::process::ExitStatusExt;
256        let mut child = std::process::Command::new("sleep")
257            .arg("60")
258            .spawn()
259            .unwrap();
260        terminate(&mut child);
261        assert_eq!(child.wait().unwrap().signal(), Some(libc::SIGTERM));
262    }
263
264    #[cfg(target_os = "linux")]
265    #[test]
266    fn a_tracked_child_is_asked_to_end() {
267        let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
268            .join("../../tmp")
269            .join(format!("tracked-{}", std::process::id()));
270        std::fs::create_dir_all(&dir).unwrap();
271        let got = dir.join("got");
272        let ready = dir.join("ready");
273        let mut child = std::process::Command::new("/usr/bin/python3")
274            .arg("-c")
275            .arg(concat!(
276                "import signal,sys,time\n",
277                "got,ready=sys.argv[1],sys.argv[2]\n",
278                "def stop(signum,frame):\n",
279                "    open(got,'w').write('term')\n",
280                "    sys.exit(0)\n",
281                "signal.signal(signal.SIGTERM, stop)\n",
282                "open(ready,'w').write('ready')\n",
283                "while True:\n",
284                "    time.sleep(30)\n",
285            ))
286            .arg(&got)
287            .arg(&ready)
288            .spawn()
289            .unwrap();
290        let tracked = track(&child);
291        let asked = std::time::Instant::now();
292        while !ready.exists() && asked.elapsed() < std::time::Duration::from_secs(5) {
293            std::thread::sleep(std::time::Duration::from_millis(20));
294        }
295        assert!(ready.exists());
296        pass_term();
297        drop(tracked);
298        let _ = child.try_wait();
299        assert_eq!(std::fs::read_to_string(&got).unwrap(), "term");
300        std::fs::remove_dir_all(dir).unwrap();
301    }
302}
pfx_run/ending.rs

pfx_run/
ending.rs
