Skip to main content

harn_vm/
op_interrupt.rs

1//! Cooperative interrupt observation for blocking sync builtins.
2//!
3//! Sync builtins (including every subprocess-spawning path: the hostlib
4//! `run_command` tool family and the VM-side `process.exec`/`exec_opts`
5//! builtins) execute inline on the VM's async task. While one of them
6//! blocks — typically waiting on a child process — the interpreter's
7//! `tokio::select!` cancel/deadline race in
8//! `vm/execution.rs::execute_op_with_scope_interrupts` cannot run: the op
9//! future never yields, so scope cancellation, `deadline` expiry, and host
10//! aborts used to wait for the child to exit on its own (orphaning it on
11//! task abort / VM drop).
12//!
13//! This module closes that gap cooperatively. Before invoking a sync
14//! builtin, the VM installs the *currently armed* interrupt sources — its
15//! host cancel token (`Arc<AtomicBool>`) and the innermost deadline — into
16//! a thread-local via [`install`]. Blocking wait loops poll [`requested`]
17//! (they already poll `try_wait` every ~20ms) and, when it fires,
18//! gracefully terminate their child process tree/group (SIGTERM, then SIGKILL
19//! after [`SUBPROCESS_TERM_GRACE`]) and return. The VM then surfaces the
20//! ordinary cancellation / deadline error at the next op boundary.
21//!
22//! Trigger coverage:
23//! - **Scope / `parallel` cancellation and VM drop**: spawned-task child
24//!   VMs share the `Arc<AtomicBool>` stored in their `VmTaskHandle`;
25//!   `Vm::cancel_spawned_tasks` (also called from `Drop for Vm`) sets it,
26//!   which the blocked wait loop observes.
27//! - **Host abort**: hosts cancel a VM by setting its cancel token — same
28//!   observation path.
29//! - **`deadline` expiry**: the deadline `Instant` is captured when the
30//!   builtin starts; the wait loop compares against `Instant::now()`.
31
32use std::cell::RefCell;
33use std::collections::BTreeMap;
34use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
35use std::sync::{Arc, LazyLock, Mutex};
36use std::time::{Duration, Instant};
37
38/// Private environment marker inherited by subprocess descendants. Cleanup uses
39/// this token to rediscover escaped descendants that have reparented or moved to
40/// a different process group before the parent-edge scan runs.
41pub const PROCESS_CLEANUP_TOKEN_ENV: &str = "HARN_PROCESS_CLEANUP_TOKEN";
42
43/// How long a subprocess gets to exit after SIGTERM before the whole
44/// process group is SIGKILLed. Deliberately longer than the interpreter's
45/// 250ms async-op cancel grace (`CANCEL_GRACE_ASYNC_OP`): child processes
46/// often need to flush buffers / remove lock files on SIGTERM.
47pub const SUBPROCESS_TERM_GRACE: Duration = Duration::from_secs(2);
48#[cfg(unix)]
49const SUBPROCESS_KILL_SETTLE: Duration = Duration::from_millis(250);
50
51pub fn new_process_cleanup_token() -> String {
52    format!("harn-cleanup-{}", uuid::Uuid::now_v7().simple())
53}
54
55/// Structural evidence collected when Harn kills a child process tree.
56#[derive(Clone, Debug, Default, PartialEq, Eq)]
57pub struct ProcessCleanupReport {
58    pub root_pid: Option<u32>,
59    pub attempted_signals: Vec<i32>,
60    pub children: Vec<ProcessCleanupChild>,
61}
62
63impl ProcessCleanupReport {
64    pub fn for_signal(root_pid: Option<u32>, signal: i32) -> Self {
65        Self {
66            root_pid,
67            attempted_signals: vec![signal],
68            children: Vec::new(),
69        }
70    }
71
72    pub fn merge(&mut self, other: Self) {
73        if self.root_pid.is_none() {
74            self.root_pid = other.root_pid;
75        }
76        for signal in other.attempted_signals {
77            push_unique(&mut self.attempted_signals, signal);
78        }
79        for child in other.children {
80            self.merge_child(child);
81        }
82    }
83
84    pub fn refresh_survivor_status(&mut self) {
85        #[cfg(unix)]
86        {
87            for child in &mut self.children {
88                child.alive_after_cleanup = Some(process_exists(child.pid));
89            }
90        }
91    }
92
93    fn merge_child(&mut self, child: ProcessCleanupChild) {
94        if let Some(existing) = self
95            .children
96            .iter_mut()
97            .find(|entry| entry.pid == child.pid)
98        {
99            for signal in child.signals {
100                push_unique(&mut existing.signals, signal);
101            }
102            if existing.command_name.is_none() {
103                existing.command_name = child.command_name;
104            }
105            if child.alive_after_cleanup.is_some() {
106                existing.alive_after_cleanup = child.alive_after_cleanup;
107            }
108            return;
109        }
110        self.children.push(child);
111        self.children
112            .sort_by(|left, right| left.depth.cmp(&right.depth).then(left.pid.cmp(&right.pid)));
113    }
114}
115
116/// A descendant process Harn targeted during cleanup.
117#[derive(Clone, Debug, PartialEq, Eq)]
118pub struct ProcessCleanupChild {
119    pub pid: u32,
120    pub parent_pid: Option<u32>,
121    pub depth: u32,
122    pub command_name: Option<String>,
123    pub signals: Vec<i32>,
124    pub alive_after_cleanup: Option<bool>,
125}
126
127impl ProcessCleanupChild {
128    pub fn new(
129        pid: u32,
130        parent_pid: Option<u32>,
131        depth: u32,
132        command_name: Option<String>,
133    ) -> Self {
134        Self {
135            pid,
136            parent_pid,
137            depth,
138            command_name,
139            signals: Vec::new(),
140            alive_after_cleanup: None,
141        }
142    }
143
144    #[cfg(unix)]
145    fn with_signal(mut self, signal: i32) -> Self {
146        push_unique(&mut self.signals, signal);
147        self
148    }
149}
150
151fn push_unique<T: Copy + Eq>(values: &mut Vec<T>, value: T) {
152    if !values.contains(&value) {
153        values.push(value);
154    }
155}
156
157#[derive(Clone, Default)]
158struct OpInterrupt {
159    cancel: Option<Arc<AtomicBool>>,
160    deadline: Option<Instant>,
161}
162
163thread_local! {
164    static CURRENT: RefCell<Option<OpInterrupt>> = const { RefCell::new(None) };
165}
166
167#[derive(Clone, Debug)]
168struct ActiveProcessCleanup {
169    pid: Option<u32>,
170    cleanup_token: String,
171    owner_cancel_token: Option<Arc<AtomicBool>>,
172}
173
174static ACTIVE_PROCESS_CLEANUP_ID: AtomicU64 = AtomicU64::new(1);
175static ACTIVE_PROCESS_CLEANUPS: LazyLock<Mutex<BTreeMap<u64, ActiveProcessCleanup>>> =
176    LazyLock::new(|| Mutex::new(BTreeMap::new()));
177
178/// Registration guard for an asynchronously waited child process. The VM's
179/// sync process paths poll [`requested`] directly, but Tokio wait paths can be
180/// parked inside `wait_with_output()`/`child.wait()` and need an out-of-band
181/// cleanup hook when `harn run` is interrupted or reaches its run deadline.
182pub struct ActiveProcessCleanupGuard {
183    id: u64,
184}
185
186impl Drop for ActiveProcessCleanupGuard {
187    fn drop(&mut self) {
188        unregister_active_process_cleanup(self.id);
189    }
190}
191
192pub fn register_active_process_cleanup(
193    pid: Option<u32>,
194    cleanup_token: &str,
195    owner_cancel_token: Option<Arc<AtomicBool>>,
196) -> ActiveProcessCleanupGuard {
197    let id = ACTIVE_PROCESS_CLEANUP_ID.fetch_add(1, Ordering::SeqCst);
198    ACTIVE_PROCESS_CLEANUPS
199        .lock()
200        .expect("active process cleanup registry poisoned")
201        .insert(
202            id,
203            ActiveProcessCleanup {
204                pid,
205                cleanup_token: cleanup_token.to_string(),
206                owner_cancel_token,
207            },
208        );
209    ActiveProcessCleanupGuard { id }
210}
211
212fn unregister_active_process_cleanup(id: u64) {
213    ACTIVE_PROCESS_CLEANUPS
214        .lock()
215        .expect("active process cleanup registry poisoned")
216        .remove(&id);
217}
218
219/// Signal every actively registered async child process tree. Prefer
220/// [`signal_active_process_cleanups_for_cancel_token`] or
221/// [`signal_ownerless_active_process_cleanups`] when the caller can avoid a
222/// process-global sweep.
223pub fn signal_active_process_cleanups(signal: i32) -> ProcessCleanupReport {
224    signal_active_process_cleanups_matching(signal, |_| true)
225}
226
227pub fn signal_ownerless_active_process_cleanups(signal: i32) -> ProcessCleanupReport {
228    signal_active_process_cleanups_matching(signal, |entry| entry.owner_cancel_token.is_none())
229}
230
231pub fn signal_active_process_cleanups_for_cancel_token(
232    signal: i32,
233    cancel_token: &Arc<AtomicBool>,
234) -> ProcessCleanupReport {
235    signal_active_process_cleanups_matching(signal, |entry| {
236        entry
237            .owner_cancel_token
238            .as_ref()
239            .is_some_and(|owner| Arc::ptr_eq(owner, cancel_token))
240    })
241}
242
243#[cfg(test)]
244fn active_cleanup_tokens_for_cancel_token_for_test(cancel_token: &Arc<AtomicBool>) -> Vec<String> {
245    ACTIVE_PROCESS_CLEANUPS
246        .lock()
247        .expect("active process cleanup registry poisoned")
248        .values()
249        .filter(|entry| {
250            entry
251                .owner_cancel_token
252                .as_ref()
253                .is_some_and(|owner| Arc::ptr_eq(owner, cancel_token))
254        })
255        .map(|entry| entry.cleanup_token.clone())
256        .collect()
257}
258
259#[cfg(test)]
260fn ownerless_active_cleanup_tokens_for_test() -> Vec<String> {
261    ACTIVE_PROCESS_CLEANUPS
262        .lock()
263        .expect("active process cleanup registry poisoned")
264        .values()
265        .filter(|entry| entry.owner_cancel_token.is_none())
266        .map(|entry| entry.cleanup_token.clone())
267        .collect()
268}
269
270fn signal_active_process_cleanups_matching(
271    signal: i32,
272    matches_entry: impl Fn(&ActiveProcessCleanup) -> bool,
273) -> ProcessCleanupReport {
274    let entries = ACTIVE_PROCESS_CLEANUPS
275        .lock()
276        .expect("active process cleanup registry poisoned")
277        .values()
278        .filter(|entry| matches_entry(entry))
279        .cloned()
280        .collect::<Vec<_>>();
281    let mut report = ProcessCleanupReport::default();
282    for entry in entries {
283        if let Some(pid) = entry.pid {
284            report.merge(signal_pid_tree_group_and_token_with_report(
285                pid,
286                Some(&entry.cleanup_token),
287                signal,
288            ));
289        }
290    }
291    report
292}
293
294/// Guard returned by [`install`]. Restores the previously installed
295/// interrupt context on drop so nested builtin dispatch (child VMs running
296/// on the same thread) composes correctly.
297pub struct OpInterruptGuard {
298    // Outer Option = "guard owes a restore"; inner Option is the previous
299    // thread-local slot value (which can itself be None).
300    #[allow(clippy::option_option)]
301    prev: Option<Option<OpInterrupt>>,
302}
303
304impl Drop for OpInterruptGuard {
305    fn drop(&mut self) {
306        if let Some(prev) = self.prev.take() {
307            CURRENT.with(|slot| *slot.borrow_mut() = prev);
308        }
309    }
310}
311
312/// Install the interrupt sources a blocking builtin on this thread should
313/// observe: an optional cooperative cancel token and an optional deadline.
314/// The VM calls this around sync builtin dispatch; tests use it to simulate
315/// scope cancellation without booting a full interpreter.
316pub fn install(cancel: Option<Arc<AtomicBool>>, deadline: Option<Instant>) -> OpInterruptGuard {
317    let prev = CURRENT.with(|slot| slot.borrow_mut().replace(OpInterrupt { cancel, deadline }));
318    OpInterruptGuard { prev: Some(prev) }
319}
320
321/// Returns `true` when an interrupt context is installed on this thread.
322///
323/// This is separate from [`requested`] so blocking operations can decide
324/// whether to use a short heartbeat poll or a true indefinite wait.
325pub fn installed() -> bool {
326    CURRENT.with(|slot| slot.borrow().is_some())
327}
328
329/// Returns `true` when the interrupt context installed on this thread has
330/// fired: the cancel token is set, or the deadline has passed. Cheap enough
331/// to call from a ~20ms poll loop. Returns `false` when nothing is armed.
332pub fn requested() -> bool {
333    CURRENT.with(|slot| {
334        let ctx = slot.borrow();
335        let Some(ctx) = ctx.as_ref() else {
336            return false;
337        };
338        if ctx
339            .cancel
340            .as_ref()
341            .is_some_and(|token| token.load(Ordering::SeqCst))
342        {
343            return true;
344        }
345        ctx.deadline
346            .is_some_and(|deadline| Instant::now() >= deadline)
347    })
348}
349
350/// Put the child in its own process group (`setpgid(0, 0)`) so a later
351/// group signal reaps ordinary grandchildren too. No-op on non-Unix targets — group
352/// semantics are Unix-first; Windows callers fall back to killing the
353/// direct child handle (`TerminateProcess` via `Child::kill`).
354pub fn configure_kill_group(command: &mut std::process::Command) {
355    #[cfg(unix)]
356    {
357        use std::os::unix::process::CommandExt;
358        command.process_group(0);
359    }
360    #[cfg(not(unix))]
361    {
362        let _ = command;
363    }
364}
365
366/// Tokio-process variant of [`configure_kill_group`]. Tokio's command wrapper
367/// does not flow through `std::process::Command`, so async spawn paths must opt
368/// in separately before they rely on group/tree cleanup.
369pub fn configure_tokio_kill_group(command: &mut tokio::process::Command) {
370    #[cfg(unix)]
371    {
372        command.process_group(0);
373    }
374    #[cfg(not(unix))]
375    {
376        let _ = command;
377    }
378}
379
380/// Signal a pid and its process group. No-op on non-Unix targets.
381pub fn signal_pid_and_group(pid: u32, signal: i32) {
382    #[cfg(unix)]
383    {
384        // SAFETY: kill(2) takes a pid_t (i32 on all Unix targets) and a
385        // signal number; calling it with any valid signal is well-defined.
386        extern "C" {
387            fn kill(pid: i32, sig: i32) -> i32;
388        }
389        unsafe {
390            kill(-(pid as i32), signal);
391            kill(pid as i32, signal);
392        }
393    }
394    #[cfg(not(unix))]
395    {
396        let _ = (pid, signal);
397    }
398}
399
400/// Signal a pid, its process group, and every descendant process visible in
401/// the system process table. Descendants are signalled deepest-first so a
402/// child that escaped into its own process group (for example via `setsid`)
403/// cannot survive a timeout merely because it left the original group.
404pub fn signal_pid_tree_and_group(pid: u32, signal: i32) {
405    let _ = signal_pid_tree_and_group_with_report(pid, signal);
406}
407
408/// Signal a pid, its process group, and visible descendants, returning the
409/// structural targets observed before signaling.
410pub fn signal_pid_tree_and_group_with_report(pid: u32, signal: i32) -> ProcessCleanupReport {
411    signal_pid_tree_group_and_token_with_report(pid, None, signal)
412}
413
414/// Signal a pid, its process group, visible descendants, and any same-token
415/// process that inherited Harn's cleanup marker. The token path closes the
416/// reparented-descendant hole in pure parent-edge scanning: a child can `setsid`
417/// and outlive its direct parent, but it keeps the inherited environment unless
418/// it deliberately scrubs it.
419pub fn signal_pid_tree_group_and_token_with_report(
420    pid: u32,
421    cleanup_token: Option<&str>,
422    signal: i32,
423) -> ProcessCleanupReport {
424    #[cfg(unix)]
425    {
426        let mut report = ProcessCleanupReport::for_signal(Some(pid), signal);
427        for child in descendant_processes(pid) {
428            signal_pid_and_group(child.pid, signal);
429            report.merge_child(child.with_signal(signal));
430        }
431        if let Some(cleanup_token) = cleanup_token.filter(|token| !token.is_empty()) {
432            for child in cleanup_token_processes(cleanup_token) {
433                if child.pid == pid {
434                    continue;
435                }
436                signal_pid_and_group(child.pid, signal);
437                report.merge_child(child.with_signal(signal));
438            }
439        }
440        signal_pid_and_group(pid, signal);
441        if signal == 9 {
442            wait_for_report_children_to_exit(&report, SUBPROCESS_KILL_SETTLE);
443        }
444        report.refresh_survivor_status();
445        report
446    }
447    #[cfg(not(unix))]
448    {
449        let _ = cleanup_token;
450        ProcessCleanupReport::for_signal(Some(pid), signal)
451    }
452}
453
454#[cfg(unix)]
455fn descendant_processes(root: u32) -> Vec<ProcessCleanupChild> {
456    use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System};
457
458    let mut sys = System::new();
459    sys.refresh_processes_specifics(
460        ProcessesToUpdate::All,
461        false,
462        ProcessRefreshKind::everything(),
463    );
464    let rows = sys
465        .processes()
466        .iter()
467        .filter_map(|(pid, process)| {
468            Some((
469                pid.as_u32(),
470                process.parent()?.as_u32(),
471                command_name(process.cmd()),
472            ))
473        })
474        .collect::<Vec<_>>();
475    descendant_processes_from_parent_edges(root, &rows)
476}
477
478#[cfg(unix)]
479fn cleanup_token_processes(token: &str) -> Vec<ProcessCleanupChild> {
480    use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind};
481
482    let mut sys = System::new();
483    sys.refresh_processes_specifics(
484        ProcessesToUpdate::All,
485        false,
486        ProcessRefreshKind::nothing()
487            .with_environ(UpdateKind::Always)
488            .with_cmd(UpdateKind::Always),
489    );
490    let mut children = sys
491        .processes()
492        .iter()
493        .filter(|(_, process)| process_has_cleanup_token(process.environ(), token))
494        .map(|(pid, process)| {
495            ProcessCleanupChild::new(
496                pid.as_u32(),
497                process.parent().map(|parent| parent.as_u32()),
498                1,
499                command_name(process.cmd()),
500            )
501        })
502        .collect::<Vec<_>>();
503    children.sort_by_key(|child| child.pid);
504    children
505}
506
507#[cfg(unix)]
508fn process_has_cleanup_token(environ: &[std::ffi::OsString], token: &str) -> bool {
509    let expected = format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}");
510    environ
511        .iter()
512        .any(|entry| entry.to_string_lossy() == expected)
513}
514
515#[cfg(all(unix, test))]
516fn descendant_pids_from_parent_edges(root: u32, edges: &[(u32, u32)]) -> Vec<u32> {
517    let rows = edges
518        .iter()
519        .map(|(pid, parent)| (*pid, *parent, None))
520        .collect::<Vec<_>>();
521    descendant_processes_from_parent_edges(root, &rows)
522        .into_iter()
523        .map(|child| child.pid)
524        .collect()
525}
526
527#[cfg(unix)]
528fn descendant_processes_from_parent_edges(
529    root: u32,
530    rows: &[(u32, u32, Option<String>)],
531) -> Vec<ProcessCleanupChild> {
532    use std::collections::{HashMap, HashSet};
533
534    let mut children: HashMap<u32, Vec<u32>> = HashMap::new();
535    let mut metadata: HashMap<u32, (u32, Option<String>)> = HashMap::new();
536    for (pid, parent, command) in rows {
537        metadata.insert(*pid, (*parent, command.clone()));
538        children.entry(*parent).or_default().push(*pid);
539    }
540
541    let mut seen = HashSet::new();
542    let mut stack = vec![(root, 0usize)];
543    let mut descendants = Vec::new();
544    while let Some((pid, depth)) = stack.pop() {
545        if !seen.insert(pid) {
546            continue;
547        }
548        if pid != root {
549            descendants.push((pid, depth));
550        }
551        if let Some(kids) = children.get(&pid) {
552            for &child in kids {
553                stack.push((child, depth + 1));
554            }
555        }
556    }
557
558    descendants.sort_by(|(left_pid, left_depth), (right_pid, right_depth)| {
559        right_depth
560            .cmp(left_depth)
561            .then_with(|| left_pid.cmp(right_pid))
562    });
563    descendants
564        .into_iter()
565        .map(|(pid, depth)| {
566            let (parent_pid, command) = metadata.get(&pid).cloned().unwrap_or((root, None));
567            ProcessCleanupChild::new(pid, Some(parent_pid), depth as u32, command)
568        })
569        .collect()
570}
571
572#[cfg(unix)]
573fn command_name(command: &[std::ffi::OsString]) -> Option<String> {
574    if command.is_empty() {
575        return None;
576    }
577    std::path::Path::new(&command[0])
578        .file_name()
579        .map(|name| name.to_string_lossy().into_owned())
580        .filter(|name| !name.is_empty())
581}
582
583#[cfg(unix)]
584fn process_exists(pid: u32) -> bool {
585    extern "C" {
586        fn kill(pid: i32, sig: i32) -> i32;
587    }
588    unsafe { kill(pid as i32, 0) == 0 }
589}
590
591#[cfg(unix)]
592fn wait_for_report_children_to_exit(report: &ProcessCleanupReport, timeout: Duration) {
593    let deadline = Instant::now() + timeout;
594    while Instant::now() < deadline {
595        if report
596            .children
597            .iter()
598            .all(|child| !process_exists(child.pid))
599        {
600            return;
601        }
602        std::thread::sleep(Duration::from_millis(10));
603    }
604}
605
606/// How an interruptible child wait ended.
607pub enum ChildWait {
608    /// The child exited on its own.
609    Exited(std::process::ExitStatus),
610    /// The caller-supplied timeout elapsed; the child tree/group was killed.
611    TimedOut(ProcessCleanupReport),
612    /// [`requested`] fired; the child tree/group was SIGTERMed and, after
613    /// [`SUBPROCESS_TERM_GRACE`], SIGKILLed. Carries the reaped status when
614    /// the OS reported one.
615    Interrupted(Option<std::process::ExitStatus>, ProcessCleanupReport),
616}
617
618/// Wait for `child` while polling [`requested`] and the optional timeout.
619///
620/// Used by the VM-side `process.*` builtins (`exec`, `shell`, `exec_opts`,
621/// `spawn_captured`). The hostlib `run_command` family implements the same
622/// protocol inside its `ProcessSpawner` abstraction. Callers should have
623/// spawned the child with [`configure_kill_group`] so group signals reach
624/// ordinary grandchildren; escaped descendants are reaped by process-tree
625/// scanning on Unix.
626pub fn wait_child_interruptible(
627    child: &mut std::process::Child,
628    timeout: Option<Duration>,
629) -> std::io::Result<ChildWait> {
630    wait_child_interruptible_with_cleanup_token(child, timeout, None)
631}
632
633pub fn wait_child_interruptible_with_cleanup_token(
634    child: &mut std::process::Child,
635    timeout: Option<Duration>,
636    cleanup_token: Option<&str>,
637) -> std::io::Result<ChildWait> {
638    let deadline = timeout.map(|limit| Instant::now() + limit);
639    loop {
640        if let Some(status) = child.try_wait()? {
641            return Ok(ChildWait::Exited(status));
642        }
643        if requested() {
644            let (status, report) =
645                terminate_child_group_with_cleanup_token_report(child, cleanup_token);
646            return Ok(ChildWait::Interrupted(status, report));
647        }
648        if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
649            // Timeout keeps its historical semantics: immediate SIGKILL.
650            let mut report = child_pid(child)
651                .map(|pid| signal_pid_tree_group_and_token_with_report(pid, cleanup_token, 9))
652                .unwrap_or_default();
653            let _ = child.kill();
654            let _ = child.wait();
655            report.refresh_survivor_status();
656            return Ok(ChildWait::TimedOut(report));
657        }
658        std::thread::sleep(Duration::from_millis(20));
659    }
660}
661
662/// Gracefully terminate `child` and its process tree/group: SIGTERM, wait up to
663/// [`SUBPROCESS_TERM_GRACE`], then SIGKILL. Reaps the child and returns its
664/// exit status when available. On non-Unix targets this is a best-effort
665/// direct `Child::kill` (`TerminateProcess`), which does not reach
666/// grandchildren.
667pub fn terminate_child_group(child: &mut std::process::Child) -> Option<std::process::ExitStatus> {
668    terminate_child_group_with_report(child).0
669}
670
671/// Like [`terminate_child_group`], but also returns a structural cleanup
672/// report describing descendants observed and signalled.
673pub fn terminate_child_group_with_report(
674    child: &mut std::process::Child,
675) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
676    terminate_child_group_with_cleanup_token_report(child, None)
677}
678
679pub fn terminate_child_group_with_cleanup_token_report(
680    child: &mut std::process::Child,
681    cleanup_token: Option<&str>,
682) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
683    let mut report = child_pid(child)
684        .map(|pid| ProcessCleanupReport::for_signal(Some(pid), 15))
685        .unwrap_or_default();
686    #[cfg(not(unix))]
687    let _ = cleanup_token;
688    #[cfg(unix)]
689    {
690        if let Some(pid) = child_pid(child) {
691            const SIGTERM: i32 = 15;
692            report = signal_pid_tree_group_and_token_with_report(pid, cleanup_token, SIGTERM);
693            let grace_deadline = Instant::now() + SUBPROCESS_TERM_GRACE;
694            loop {
695                match child.try_wait() {
696                    Ok(Some(status)) => {
697                        // The direct child is gone, but SIGTERM-immune
698                        // descendants may linger — sweep the group.
699                        report.merge(signal_pid_tree_group_and_token_with_report(
700                            pid,
701                            cleanup_token,
702                            9,
703                        ));
704                        report.refresh_survivor_status();
705                        return (Some(status), report);
706                    }
707                    Ok(None) => {
708                        if Instant::now() >= grace_deadline {
709                            break;
710                        }
711                        std::thread::sleep(Duration::from_millis(20));
712                    }
713                    Err(_) => break,
714                }
715            }
716            report.merge(signal_pid_tree_group_and_token_with_report(
717                pid,
718                cleanup_token,
719                9,
720            ));
721        }
722    }
723    let _ = child.kill();
724    let status = child.wait().ok();
725    report.refresh_survivor_status();
726    (status, report)
727}
728
729fn child_pid(child: &std::process::Child) -> Option<u32> {
730    let pid = child.id();
731    (pid > 0).then_some(pid)
732}
733
734/// Collect one captured pipe from a drain thread that sends the full buffer
735/// on EOF.
736///
737/// `killed == true` (the child group was already signalled) keeps a 100ms
738/// best-effort window for partial output. Otherwise wait for EOF like
739/// `Command::output` would — but keep observing [`requested`], because a
740/// lingering grandchild that inherited the pipe can hold it open long after
741/// the direct child exited; on interrupt the group gets the same SIGTERM →
742/// grace → SIGKILL treatment.
743pub(crate) fn drain_captured_pipe(
744    rx: &std::sync::mpsc::Receiver<Vec<u8>>,
745    killed: bool,
746    child_pid: u32,
747) -> Vec<u8> {
748    use std::sync::mpsc::RecvTimeoutError;
749    if killed {
750        return rx
751            .recv_timeout(Duration::from_millis(100))
752            .unwrap_or_default();
753    }
754    loop {
755        match rx.recv_timeout(Duration::from_millis(20)) {
756            Ok(buf) => return buf,
757            Err(RecvTimeoutError::Disconnected) => return Vec::new(),
758            Err(RecvTimeoutError::Timeout) => {
759                if requested() {
760                    const SIGTERM: i32 = 15;
761                    signal_pid_tree_and_group(child_pid, SIGTERM);
762                    if let Ok(buf) = rx.recv_timeout(SUBPROCESS_TERM_GRACE) {
763                        signal_pid_tree_and_group(child_pid, 9);
764                        return buf;
765                    }
766                    signal_pid_tree_and_group(child_pid, 9);
767                    return rx
768                        .recv_timeout(Duration::from_millis(100))
769                        .unwrap_or_default();
770                }
771            }
772        }
773    }
774}
775
776/// Spawn a drain thread that reads `reader` to EOF and sends the buffer.
777pub(crate) fn spawn_pipe_drain<R: std::io::Read + Send + 'static>(
778    mut reader: R,
779) -> std::sync::mpsc::Receiver<Vec<u8>> {
780    let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
781    std::thread::spawn(move || {
782        let mut buf = Vec::new();
783        let _ = reader.read_to_end(&mut buf);
784        let _ = tx.send(buf);
785    });
786    rx
787}
788
789/// Interrupt-aware replacement for `Command::output()`: the child runs in
790/// its own kill group, stdout/stderr are captured in full, stdin is closed,
791/// and the wait polls [`requested`]. When an interrupt fires the whole
792/// group is gracefully terminated and the (signal-terminated) status is
793/// returned — the interpreter surfaces the pending cancellation / deadline
794/// error at the next op boundary.
795pub fn capture_output_interruptible(
796    command: &mut std::process::Command,
797) -> std::io::Result<std::process::Output> {
798    use std::process::Stdio;
799    command
800        .stdout(Stdio::piped())
801        .stderr(Stdio::piped())
802        .stdin(Stdio::null());
803    configure_kill_group(command);
804    let cleanup_token = new_process_cleanup_token();
805    command.env(PROCESS_CLEANUP_TOKEN_ENV, &cleanup_token);
806    let mut child = command.spawn()?;
807    let pid = child.id();
808    let rx_out = child.stdout.take().map(spawn_pipe_drain);
809    let rx_err = child.stderr.take().map(spawn_pipe_drain);
810
811    let (status, killed) = match wait_child_interruptible_with_cleanup_token(
812        &mut child,
813        None,
814        Some(&cleanup_token),
815    )? {
816        ChildWait::Exited(status) => (status, false),
817        // No timeout is armed here, but keep the arm total.
818        ChildWait::TimedOut(_) => (std::process::ExitStatus::default(), true),
819        ChildWait::Interrupted(status, _) => (status.unwrap_or_default(), true),
820    };
821    let stdout = rx_out
822        .map(|rx| drain_captured_pipe(&rx, killed, pid))
823        .unwrap_or_default();
824    let stderr = rx_err
825        .map(|rx| drain_captured_pipe(&rx, killed, pid))
826        .unwrap_or_default();
827    Ok(std::process::Output {
828        status,
829        stdout,
830        stderr,
831    })
832}
833
834#[cfg(test)]
835mod tests {
836    use super::*;
837
838    #[test]
839    fn requested_is_false_without_context() {
840        assert!(!requested());
841    }
842
843    #[test]
844    fn installed_tracks_guard_lifetime() {
845        assert!(!installed());
846        let guard = install(None, None);
847        assert!(installed());
848        drop(guard);
849        assert!(!installed());
850    }
851
852    #[test]
853    fn cancel_token_trips_requested_and_guard_restores() {
854        let token = Arc::new(AtomicBool::new(false));
855        let guard = install(Some(token.clone()), None);
856        assert!(!requested());
857        token.store(true, Ordering::SeqCst);
858        assert!(requested());
859        drop(guard);
860        assert!(!requested());
861    }
862
863    #[test]
864    fn deadline_trips_requested() {
865        let expired = Instant::now()
866            .checked_sub(Duration::from_millis(1))
867            .expect("monotonic clock supports a 1ms test lookback");
868        let _guard = install(None, Some(expired));
869        assert!(requested());
870    }
871
872    #[test]
873    fn nested_installs_restore_in_order() {
874        let outer_token = Arc::new(AtomicBool::new(true));
875        let _outer = install(Some(outer_token), None);
876        assert!(requested());
877        {
878            let _inner = install(None, None);
879            assert!(!requested());
880        }
881        assert!(requested());
882    }
883
884    #[test]
885    fn active_cleanup_owner_scopes_are_disjoint() {
886        let owner = Arc::new(AtomicBool::new(false));
887        let _owned =
888            register_active_process_cleanup(None, "owned-scope-test", Some(Arc::clone(&owner)));
889        let _ownerless = register_active_process_cleanup(None, "ownerless-scope-test", None);
890
891        assert_eq!(
892            active_cleanup_tokens_for_cancel_token_for_test(&owner),
893            vec!["owned-scope-test".to_string()]
894        );
895        assert!(
896            ownerless_active_cleanup_tokens_for_test()
897                .iter()
898                .any(|token| token == "ownerless-scope-test"),
899            "explicit ownerless fallback should remain separately discoverable"
900        );
901    }
902
903    #[test]
904    fn active_cleanup_guard_unregisters_on_drop() {
905        let owner = Arc::new(AtomicBool::new(false));
906        let token = "guard-lifetime-test";
907        let guard = register_active_process_cleanup(None, token, Some(Arc::clone(&owner)));
908
909        assert!(
910            active_cleanup_tokens_for_cancel_token_for_test(&owner)
911                .iter()
912                .any(|entry| entry == token),
913            "active cleanup must remain registered while its guard is alive"
914        );
915
916        drop(guard);
917
918        assert!(
919            !active_cleanup_tokens_for_cancel_token_for_test(&owner)
920                .iter()
921                .any(|entry| entry == token),
922            "dropping the guard must unregister the cleanup token"
923        );
924    }
925
926    #[cfg(unix)]
927    #[test]
928    fn descendant_pids_from_parent_edges_returns_deepest_first_tree_only() {
929        let edges = [
930            (20, 10),
931            (30, 20),
932            (40, 20),
933            (50, 30),
934            (60, 99),
935            (70, 60),
936            // A malformed process table cycle should not hang traversal.
937            (80, 90),
938            (90, 80),
939        ];
940
941        assert_eq!(
942            descendant_pids_from_parent_edges(10, &edges),
943            vec![50, 30, 40, 20]
944        );
945        assert_eq!(descendant_pids_from_parent_edges(99, &edges), vec![70, 60]);
946        assert_eq!(
947            descendant_pids_from_parent_edges(123, &edges),
948            Vec::<u32>::new()
949        );
950    }
951
952    #[cfg(unix)]
953    #[test]
954    fn descendant_processes_preserve_metadata_and_depth_order() {
955        let rows = [
956            (20, 10, Some("worker".to_string())),
957            (30, 20, Some("grandchild".to_string())),
958            (40, 20, None),
959            (50, 30, Some("leaf".to_string())),
960        ];
961
962        let descendants = descendant_processes_from_parent_edges(10, &rows);
963        let pids = descendants
964            .iter()
965            .map(|child| {
966                (
967                    child.pid,
968                    child.parent_pid,
969                    child.depth,
970                    child.command_name.as_deref(),
971                )
972            })
973            .collect::<Vec<_>>();
974        assert_eq!(
975            pids,
976            vec![
977                (50, Some(30), 3, Some("leaf")),
978                (30, Some(20), 2, Some("grandchild")),
979                (40, Some(20), 2, None),
980                (20, Some(10), 1, Some("worker")),
981            ]
982        );
983    }
984
985    #[cfg(unix)]
986    #[test]
987    fn command_name_keeps_only_argv0_basename() {
988        let command = vec![
989            std::ffi::OsString::from("/usr/local/bin/tool"),
990            std::ffi::OsString::from("--api-key"),
991            std::ffi::OsString::from("secret-value"),
992            std::ffi::OsString::from("plain"),
993        ];
994
995        assert_eq!(command_name(&command).as_deref(), Some("tool"));
996        assert_eq!(command_name(&[]).as_deref(), None);
997    }
998
999    #[cfg(unix)]
1000    #[test]
1001    fn process_has_cleanup_token_requires_exact_marker_entry() {
1002        let token = "tok-123";
1003        let env = vec![
1004            std::ffi::OsString::from("PATH=/usr/bin"),
1005            std::ffi::OsString::from(format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}")),
1006        ];
1007        assert!(process_has_cleanup_token(&env, token));
1008        assert!(!process_has_cleanup_token(&env, "tok"));
1009        assert!(!process_has_cleanup_token(
1010            &[std::ffi::OsString::from("OTHER=tok-123")],
1011            token
1012        ));
1013    }
1014
1015    #[cfg(unix)]
1016    #[test]
1017    fn interrupted_wait_kills_process_group() {
1018        // Child spawns a grandchild; the whole group must die on interrupt.
1019        let mut command = std::process::Command::new("sh");
1020        command.args(["-c", "sleep 30 & wait"]);
1021        configure_kill_group(&mut command);
1022        let mut child = command.spawn().expect("spawn sh");
1023        let pgid = child.id();
1024
1025        let cancel = Arc::new(AtomicBool::new(true));
1026        let _guard = install(Some(cancel), None);
1027        let started = Instant::now();
1028        let outcome = wait_child_interruptible(&mut child, None).expect("wait");
1029        assert!(matches!(outcome, ChildWait::Interrupted(_, _)));
1030        assert!(started.elapsed() < Duration::from_secs(10));
1031
1032        // kill(-pgid, 0) fails with ESRCH once every member is gone.
1033        extern "C" {
1034            fn kill(pid: i32, sig: i32) -> i32;
1035        }
1036        let group_gone = || unsafe { kill(-(pgid as i32), 0) } != 0;
1037        let deadline = Instant::now() + Duration::from_secs(5);
1038        while !group_gone() && Instant::now() < deadline {
1039            std::thread::sleep(Duration::from_millis(50));
1040        }
1041        assert!(group_gone(), "process group {pgid} survived interrupt");
1042    }
1043}