1use 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
38pub const PROCESS_CLEANUP_TOKEN_ENV: &str = "HARN_PROCESS_CLEANUP_TOKEN";
42
43pub const PROCESS_OWNER_TOKEN_ENV: &str = "HARN_INTERNAL_PROCESS_OWNER_TOKEN";
52
53pub const SUBPROCESS_TERM_GRACE: Duration = Duration::from_secs(2);
58#[cfg(unix)]
59const SUBPROCESS_KILL_SETTLE: Duration = Duration::from_millis(250);
60
61pub fn new_process_cleanup_token() -> String {
62 format!("harn-cleanup-{}", uuid::Uuid::now_v7().simple())
63}
64
65fn owner_process_group_journal(token: &str) -> std::path::PathBuf {
66 let digest = blake3::hash(token.as_bytes()).to_hex();
67 std::env::temp_dir().join(format!("harn-process-owner-{digest}.groups"))
68}
69
70pub fn initialize_process_owner_group_journal(token: &str) -> std::io::Result<()> {
72 #[cfg(unix)]
73 {
74 use std::os::unix::fs::OpenOptionsExt;
75
76 std::fs::OpenOptions::new()
77 .create_new(true)
78 .write(true)
79 .mode(0o600)
80 .custom_flags(libc::O_NOFOLLOW)
81 .open(owner_process_group_journal(token))
82 .map(|_| ())
83 }
84 #[cfg(not(unix))]
85 {
86 let _ = token;
87 Ok(())
88 }
89}
90
91pub fn record_current_process_owner_group(pid: u32) -> std::io::Result<()> {
98 #[cfg(unix)]
99 {
100 use std::io::Write;
101 use std::os::unix::fs::OpenOptionsExt;
102
103 let Ok(token) = std::env::var(PROCESS_OWNER_TOKEN_ENV) else {
104 return Ok(());
105 };
106 let observed_pgid = unsafe { libc::getpgid(pid as i32) };
107 let pgid = u32::try_from(observed_pgid).unwrap_or(pid);
108 let mut journal = std::fs::OpenOptions::new()
109 .create(true)
110 .append(true)
111 .mode(0o600)
112 .custom_flags(libc::O_NOFOLLOW)
113 .open(owner_process_group_journal(&token))?;
114 journal.write_all(format!("{pgid}\n").as_bytes())
115 }
116 #[cfg(not(unix))]
117 {
118 let _ = pid;
119 Ok(())
120 }
121}
122
123pub async fn record_tokio_process_owner_group(
125 child: &mut tokio::process::Child,
126 cleanup_token: &str,
127) -> std::io::Result<()> {
128 let Some(pid) = child.id() else {
129 return Ok(());
130 };
131 if let Err(error) = record_current_process_owner_group(pid) {
132 let _ = signal_pid_tree_group_and_token_with_report(pid, Some(cleanup_token), 9);
133 let _ = child.start_kill();
134 let _ = child.wait().await;
135 return Err(error);
136 }
137 Ok(())
138}
139
140#[cfg(unix)]
141fn owner_process_groups(token: &str) -> Vec<u32> {
142 let Ok(contents) = std::fs::read_to_string(owner_process_group_journal(token)) else {
143 return Vec::new();
144 };
145 let mut groups = contents
146 .lines()
147 .filter_map(|line| line.parse::<u32>().ok())
148 .collect::<Vec<_>>();
149 groups.sort_unstable();
150 groups.dedup();
151 groups
152}
153
154pub fn remove_process_owner_group_journal(token: &str) {
156 let _ = std::fs::remove_file(owner_process_group_journal(token));
157}
158
159pub fn preserve_process_owner_token(command: &mut std::process::Command) {
162 if let Some(token) = std::env::var_os(PROCESS_OWNER_TOKEN_ENV).filter(|token| !token.is_empty())
163 {
164 command.env(PROCESS_OWNER_TOKEN_ENV, token);
165 }
166}
167
168pub fn process_owner_survivors(token: &str) -> Vec<ProcessCleanupChild> {
176 #[cfg(unix)]
177 {
178 let mut survivors = cleanup_token_processes(token)
179 .into_iter()
180 .filter(|child| child.pid != std::process::id())
181 .collect::<Vec<_>>();
182 survivors.sort_by_key(|child| child.pid);
183 survivors
184 }
185 #[cfg(not(unix))]
186 {
187 let _ = token;
188 Vec::new()
189 }
190}
191
192#[derive(Clone, Debug, Default, PartialEq, Eq)]
194pub struct ProcessCleanupReport {
195 pub root_pid: Option<u32>,
196 pub attempted_signals: Vec<i32>,
197 pub children: Vec<ProcessCleanupChild>,
198}
199
200impl ProcessCleanupReport {
201 pub fn for_signal(root_pid: Option<u32>, signal: i32) -> Self {
202 Self {
203 root_pid,
204 attempted_signals: vec![signal],
205 children: Vec::new(),
206 }
207 }
208
209 pub fn merge(&mut self, other: Self) {
210 if self.root_pid.is_none() {
211 self.root_pid = other.root_pid;
212 }
213 for signal in other.attempted_signals {
214 push_unique(&mut self.attempted_signals, signal);
215 }
216 for child in other.children {
217 self.merge_child(child);
218 }
219 }
220
221 pub fn refresh_survivor_status(&mut self) {
222 #[cfg(unix)]
223 {
224 for child in &mut self.children {
225 child.alive_after_cleanup = Some(process_exists(child.pid));
226 }
227 }
228 }
229
230 fn merge_child(&mut self, child: ProcessCleanupChild) {
231 if let Some(existing) = self
232 .children
233 .iter_mut()
234 .find(|entry| entry.pid == child.pid)
235 {
236 for signal in child.signals {
237 push_unique(&mut existing.signals, signal);
238 }
239 if existing.command_name.is_none() {
240 existing.command_name = child.command_name;
241 }
242 if child.alive_after_cleanup.is_some() {
243 existing.alive_after_cleanup = child.alive_after_cleanup;
244 }
245 return;
246 }
247 self.children.push(child);
248 self.children
249 .sort_by(|left, right| left.depth.cmp(&right.depth).then(left.pid.cmp(&right.pid)));
250 }
251}
252
253#[derive(Clone, Debug, PartialEq, Eq)]
255pub struct ProcessCleanupChild {
256 pub pid: u32,
257 pub parent_pid: Option<u32>,
258 pub depth: u32,
259 pub command_name: Option<String>,
260 pub signals: Vec<i32>,
261 pub alive_after_cleanup: Option<bool>,
262}
263
264impl ProcessCleanupChild {
265 pub fn new(
266 pid: u32,
267 parent_pid: Option<u32>,
268 depth: u32,
269 command_name: Option<String>,
270 ) -> Self {
271 Self {
272 pid,
273 parent_pid,
274 depth,
275 command_name,
276 signals: Vec::new(),
277 alive_after_cleanup: None,
278 }
279 }
280
281 #[cfg(unix)]
282 fn with_signal(mut self, signal: i32) -> Self {
283 push_unique(&mut self.signals, signal);
284 self
285 }
286}
287
288fn push_unique<T: Copy + Eq>(values: &mut Vec<T>, value: T) {
289 if !values.contains(&value) {
290 values.push(value);
291 }
292}
293
294#[derive(Clone, Default)]
295struct OpInterrupt {
296 cancel: Option<Arc<AtomicBool>>,
297 deadline: Option<Instant>,
298}
299
300thread_local! {
301 static CURRENT: RefCell<Option<OpInterrupt>> = const { RefCell::new(None) };
302}
303
304#[derive(Clone, Debug)]
305struct ActiveProcessCleanup {
306 pid: Option<u32>,
307 cleanup_token: String,
308 owner_cancel_token: Option<Arc<AtomicBool>>,
309}
310
311static ACTIVE_PROCESS_CLEANUP_ID: AtomicU64 = AtomicU64::new(1);
312static ACTIVE_PROCESS_CLEANUPS: LazyLock<Mutex<BTreeMap<u64, ActiveProcessCleanup>>> =
313 LazyLock::new(|| Mutex::new(BTreeMap::new()));
314
315pub struct ActiveProcessCleanupGuard {
320 id: u64,
321}
322
323impl Drop for ActiveProcessCleanupGuard {
324 fn drop(&mut self) {
325 unregister_active_process_cleanup(self.id);
326 }
327}
328
329pub fn register_active_process_cleanup(
330 pid: Option<u32>,
331 cleanup_token: &str,
332 owner_cancel_token: Option<Arc<AtomicBool>>,
333) -> ActiveProcessCleanupGuard {
334 let id = ACTIVE_PROCESS_CLEANUP_ID.fetch_add(1, Ordering::SeqCst);
335 ACTIVE_PROCESS_CLEANUPS
336 .lock()
337 .expect("active process cleanup registry poisoned")
338 .insert(
339 id,
340 ActiveProcessCleanup {
341 pid,
342 cleanup_token: cleanup_token.to_string(),
343 owner_cancel_token,
344 },
345 );
346 ActiveProcessCleanupGuard { id }
347}
348
349fn unregister_active_process_cleanup(id: u64) {
350 ACTIVE_PROCESS_CLEANUPS
351 .lock()
352 .expect("active process cleanup registry poisoned")
353 .remove(&id);
354}
355
356pub fn signal_active_process_cleanups(signal: i32) -> ProcessCleanupReport {
361 signal_active_process_cleanups_matching(signal, |_| true)
362}
363
364pub fn signal_ownerless_active_process_cleanups(signal: i32) -> ProcessCleanupReport {
365 signal_active_process_cleanups_matching(signal, |entry| entry.owner_cancel_token.is_none())
366}
367
368pub fn signal_active_process_cleanups_for_cancel_token(
369 signal: i32,
370 cancel_token: &Arc<AtomicBool>,
371) -> ProcessCleanupReport {
372 signal_active_process_cleanups_matching(signal, |entry| {
373 entry
374 .owner_cancel_token
375 .as_ref()
376 .is_some_and(|owner| Arc::ptr_eq(owner, cancel_token))
377 })
378}
379
380#[cfg(test)]
381fn active_cleanup_tokens_for_cancel_token_for_test(cancel_token: &Arc<AtomicBool>) -> Vec<String> {
382 ACTIVE_PROCESS_CLEANUPS
383 .lock()
384 .expect("active process cleanup registry poisoned")
385 .values()
386 .filter(|entry| {
387 entry
388 .owner_cancel_token
389 .as_ref()
390 .is_some_and(|owner| Arc::ptr_eq(owner, cancel_token))
391 })
392 .map(|entry| entry.cleanup_token.clone())
393 .collect()
394}
395
396#[cfg(test)]
397fn ownerless_active_cleanup_tokens_for_test() -> Vec<String> {
398 ACTIVE_PROCESS_CLEANUPS
399 .lock()
400 .expect("active process cleanup registry poisoned")
401 .values()
402 .filter(|entry| entry.owner_cancel_token.is_none())
403 .map(|entry| entry.cleanup_token.clone())
404 .collect()
405}
406
407fn signal_active_process_cleanups_matching(
408 signal: i32,
409 matches_entry: impl Fn(&ActiveProcessCleanup) -> bool,
410) -> ProcessCleanupReport {
411 let entries = ACTIVE_PROCESS_CLEANUPS
412 .lock()
413 .expect("active process cleanup registry poisoned")
414 .values()
415 .filter(|entry| matches_entry(entry))
416 .cloned()
417 .collect::<Vec<_>>();
418 let mut report = ProcessCleanupReport::default();
419 for entry in entries {
420 if let Some(pid) = entry.pid {
421 report.merge(signal_pid_tree_group_and_token_with_report(
422 pid,
423 Some(&entry.cleanup_token),
424 signal,
425 ));
426 }
427 }
428 report
429}
430
431pub struct OpInterruptGuard {
435 #[allow(clippy::option_option)]
438 prev: Option<Option<OpInterrupt>>,
439}
440
441impl Drop for OpInterruptGuard {
442 fn drop(&mut self) {
443 if let Some(prev) = self.prev.take() {
444 CURRENT.with(|slot| *slot.borrow_mut() = prev);
445 }
446 }
447}
448
449pub fn install(cancel: Option<Arc<AtomicBool>>, deadline: Option<Instant>) -> OpInterruptGuard {
454 let prev = CURRENT.with(|slot| slot.borrow_mut().replace(OpInterrupt { cancel, deadline }));
455 OpInterruptGuard { prev: Some(prev) }
456}
457
458pub fn installed() -> bool {
463 CURRENT.with(|slot| slot.borrow().is_some())
464}
465
466pub fn requested() -> bool {
470 CURRENT.with(|slot| {
471 let ctx = slot.borrow();
472 let Some(ctx) = ctx.as_ref() else {
473 return false;
474 };
475 if ctx
476 .cancel
477 .as_ref()
478 .is_some_and(|token| token.load(Ordering::SeqCst))
479 {
480 return true;
481 }
482 ctx.deadline
483 .is_some_and(|deadline| Instant::now() >= deadline)
484 })
485}
486
487pub fn configure_kill_group(command: &mut std::process::Command) {
492 #[cfg(unix)]
493 {
494 use std::os::unix::process::CommandExt;
495 command.process_group(0);
496 }
497 #[cfg(not(unix))]
498 {
499 let _ = command;
500 }
501}
502
503pub fn configure_tokio_kill_group(command: &mut tokio::process::Command) {
507 #[cfg(unix)]
508 {
509 command.process_group(0);
510 }
511 #[cfg(not(unix))]
512 {
513 let _ = command;
514 }
515}
516
517pub fn signal_pid_and_group(pid: u32, signal: i32) {
519 #[cfg(unix)]
520 {
521 extern "C" {
524 fn kill(pid: i32, sig: i32) -> i32;
525 }
526 unsafe {
527 kill(-(pid as i32), signal);
528 kill(pid as i32, signal);
529 }
530 }
531 #[cfg(not(unix))]
532 {
533 let _ = (pid, signal);
534 }
535}
536
537pub fn signal_pid_tree_and_group(pid: u32, signal: i32) {
542 let _ = signal_pid_tree_and_group_with_report(pid, signal);
543}
544
545pub fn signal_pid_tree_and_group_with_report(pid: u32, signal: i32) -> ProcessCleanupReport {
548 signal_pid_tree_group_and_token_with_report(pid, None, signal)
549}
550
551pub fn signal_pid_tree_group_and_token_with_report(
557 pid: u32,
558 cleanup_token: Option<&str>,
559 signal: i32,
560) -> ProcessCleanupReport {
561 #[cfg(unix)]
562 {
563 let mut report = ProcessCleanupReport::for_signal(Some(pid), signal);
564 for child in descendant_processes(pid) {
565 signal_pid_and_group(child.pid, signal);
566 report.merge_child(child.with_signal(signal));
567 }
568 if let Some(cleanup_token) = cleanup_token.filter(|token| !token.is_empty()) {
569 for child in cleanup_token_processes(cleanup_token) {
570 if child.pid == pid {
571 continue;
572 }
573 signal_pid_and_group(child.pid, signal);
574 report.merge_child(child.with_signal(signal));
575 }
576 }
577 signal_pid_and_group(pid, signal);
578 if signal == 9 {
579 wait_for_report_children_to_exit(&report, SUBPROCESS_KILL_SETTLE);
580 }
581 report.refresh_survivor_status();
582 report
583 }
584 #[cfg(not(unix))]
585 {
586 let _ = cleanup_token;
587 ProcessCleanupReport::for_signal(Some(pid), signal)
588 }
589}
590
591#[cfg(unix)]
599pub fn signal_pid_tree_and_token_preserving_group_with_report(
600 pid: u32,
601 cleanup_token: Option<&str>,
602 preserved_pgid: u32,
603 signal: i32,
604) -> ProcessCleanupReport {
605 let preserved_pgid = preserved_pgid as i32;
606 let mut report = ProcessCleanupReport::for_signal(Some(pid), signal);
607 for child in descendant_processes(pid) {
608 signal_pid_preserving_group(child.pid, preserved_pgid, signal);
609 report.merge_child(child.with_signal(signal));
610 }
611 if let Some(cleanup_token) = cleanup_token.filter(|token| !token.is_empty()) {
612 for child in cleanup_token_processes(cleanup_token) {
613 if child.pid == pid {
614 continue;
615 }
616 signal_pid_preserving_group(child.pid, preserved_pgid, signal);
617 report.merge_child(child.with_signal(signal));
618 }
619 for pgid in owner_process_groups(cleanup_token) {
620 if pgid != preserved_pgid as u32 {
621 unsafe {
622 libc::kill(-(pgid as i32), signal);
623 }
624 }
625 }
626 }
627 signal_pid_preserving_group(pid, preserved_pgid, signal);
628 report
629}
630
631#[cfg(unix)]
632fn signal_pid_preserving_group(pid: u32, preserved_pgid: i32, signal: i32) {
633 let pid = pid as i32;
634 let pgid = unsafe { libc::getpgid(pid) };
635 unsafe {
636 if pgid > 0 && pgid != preserved_pgid {
637 libc::kill(-pgid, signal);
638 }
639 libc::kill(pid, signal);
640 }
641}
642
643#[cfg(unix)]
644fn descendant_processes(root: u32) -> Vec<ProcessCleanupChild> {
645 use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System};
646
647 let mut sys = System::new();
648 sys.refresh_processes_specifics(
649 ProcessesToUpdate::All,
650 false,
651 ProcessRefreshKind::everything(),
652 );
653 let rows = sys
654 .processes()
655 .iter()
656 .filter_map(|(pid, process)| {
657 Some((
658 pid.as_u32(),
659 process.parent()?.as_u32(),
660 command_name(process.cmd()),
661 ))
662 })
663 .collect::<Vec<_>>();
664 descendant_processes_from_parent_edges(root, &rows)
665}
666
667#[cfg(unix)]
668fn cleanup_token_processes(token: &str) -> Vec<ProcessCleanupChild> {
669 use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind};
670
671 let mut sys = System::new();
672 sys.refresh_processes_specifics(
673 ProcessesToUpdate::All,
674 false,
675 ProcessRefreshKind::nothing()
676 .with_environ(UpdateKind::Always)
677 .with_cmd(UpdateKind::Always),
678 );
679 let mut children = sys
680 .processes()
681 .iter()
682 .filter(|(_, process)| {
683 process_status_can_execute(process.status())
684 && process_has_cleanup_token(process.environ(), token)
685 })
686 .map(|(pid, process)| {
687 ProcessCleanupChild::new(
688 pid.as_u32(),
689 process.parent().map(|parent| parent.as_u32()),
690 1,
691 command_name(process.cmd()),
692 )
693 })
694 .collect::<Vec<_>>();
695 children.sort_by_key(|child| child.pid);
696 children
697}
698
699#[cfg(unix)]
700fn process_status_can_execute(status: sysinfo::ProcessStatus) -> bool {
701 status != sysinfo::ProcessStatus::Zombie
702}
703
704#[cfg(unix)]
705fn process_has_cleanup_token(environ: &[std::ffi::OsString], token: &str) -> bool {
706 let cleanup = format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}");
707 let owner = format!("{PROCESS_OWNER_TOKEN_ENV}={token}");
708 environ
709 .iter()
710 .any(|entry| matches!(entry.to_string_lossy().as_ref(), value if value == cleanup || value == owner))
711}
712
713#[cfg(all(unix, test))]
714fn descendant_pids_from_parent_edges(root: u32, edges: &[(u32, u32)]) -> Vec<u32> {
715 let rows = edges
716 .iter()
717 .map(|(pid, parent)| (*pid, *parent, None))
718 .collect::<Vec<_>>();
719 descendant_processes_from_parent_edges(root, &rows)
720 .into_iter()
721 .map(|child| child.pid)
722 .collect()
723}
724
725#[cfg(unix)]
726fn descendant_processes_from_parent_edges(
727 root: u32,
728 rows: &[(u32, u32, Option<String>)],
729) -> Vec<ProcessCleanupChild> {
730 use std::collections::{HashMap, HashSet};
731
732 let mut children: HashMap<u32, Vec<u32>> = HashMap::new();
733 let mut metadata: HashMap<u32, (u32, Option<String>)> = HashMap::new();
734 for (pid, parent, command) in rows {
735 metadata.insert(*pid, (*parent, command.clone()));
736 children.entry(*parent).or_default().push(*pid);
737 }
738
739 let mut seen = HashSet::new();
740 let mut stack = vec![(root, 0usize)];
741 let mut descendants = Vec::new();
742 while let Some((pid, depth)) = stack.pop() {
743 if !seen.insert(pid) {
744 continue;
745 }
746 if pid != root {
747 descendants.push((pid, depth));
748 }
749 if let Some(kids) = children.get(&pid) {
750 for &child in kids {
751 stack.push((child, depth + 1));
752 }
753 }
754 }
755
756 descendants.sort_by(|(left_pid, left_depth), (right_pid, right_depth)| {
757 right_depth
758 .cmp(left_depth)
759 .then_with(|| left_pid.cmp(right_pid))
760 });
761 descendants
762 .into_iter()
763 .map(|(pid, depth)| {
764 let (parent_pid, command) = metadata.get(&pid).cloned().unwrap_or((root, None));
765 ProcessCleanupChild::new(pid, Some(parent_pid), depth as u32, command)
766 })
767 .collect()
768}
769
770#[cfg(unix)]
771fn command_name(command: &[std::ffi::OsString]) -> Option<String> {
772 if command.is_empty() {
773 return None;
774 }
775 std::path::Path::new(&command[0])
776 .file_name()
777 .map(|name| name.to_string_lossy().into_owned())
778 .filter(|name| !name.is_empty())
779}
780
781#[cfg(unix)]
782fn process_exists(pid: u32) -> bool {
783 unsafe { libc::kill(pid as i32, 0) == 0 }
784}
785
786#[cfg(unix)]
787fn wait_for_report_children_to_exit(report: &ProcessCleanupReport, timeout: Duration) {
788 let deadline = Instant::now() + timeout;
789 while Instant::now() < deadline {
790 if report
791 .children
792 .iter()
793 .all(|child| !process_exists(child.pid))
794 {
795 return;
796 }
797 std::thread::sleep(Duration::from_millis(10));
798 }
799}
800
801pub enum ChildWait {
803 Exited(std::process::ExitStatus),
805 TimedOut(ProcessCleanupReport),
807 Interrupted(Option<std::process::ExitStatus>, ProcessCleanupReport),
811}
812
813pub fn wait_child_interruptible(
822 child: &mut std::process::Child,
823 timeout: Option<Duration>,
824) -> std::io::Result<ChildWait> {
825 wait_child_interruptible_with_cleanup_token(child, timeout, None)
826}
827
828pub fn wait_child_interruptible_with_cleanup_token(
829 child: &mut std::process::Child,
830 timeout: Option<Duration>,
831 cleanup_token: Option<&str>,
832) -> std::io::Result<ChildWait> {
833 let deadline = timeout.map(|limit| Instant::now() + limit);
834 loop {
835 if let Some(status) = child.try_wait()? {
836 return Ok(ChildWait::Exited(status));
837 }
838 if requested() {
839 let (status, report) =
840 terminate_child_group_with_cleanup_token_report(child, cleanup_token);
841 return Ok(ChildWait::Interrupted(status, report));
842 }
843 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
844 let mut report = child_pid(child)
846 .map(|pid| signal_pid_tree_group_and_token_with_report(pid, cleanup_token, 9))
847 .unwrap_or_default();
848 let _ = child.kill();
849 let _ = child.wait();
850 report.refresh_survivor_status();
851 return Ok(ChildWait::TimedOut(report));
852 }
853 std::thread::sleep(Duration::from_millis(20));
854 }
855}
856
857pub fn terminate_child_group(child: &mut std::process::Child) -> Option<std::process::ExitStatus> {
863 terminate_child_group_with_report(child).0
864}
865
866pub fn terminate_child_group_with_report(
869 child: &mut std::process::Child,
870) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
871 terminate_child_group_with_cleanup_token_report(child, None)
872}
873
874pub fn terminate_child_group_with_cleanup_token_report(
875 child: &mut std::process::Child,
876 cleanup_token: Option<&str>,
877) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
878 let mut report = child_pid(child)
879 .map(|pid| ProcessCleanupReport::for_signal(Some(pid), 15))
880 .unwrap_or_default();
881 #[cfg(not(unix))]
882 let _ = cleanup_token;
883 #[cfg(unix)]
884 {
885 if let Some(pid) = child_pid(child) {
886 const SIGTERM: i32 = 15;
887 report = signal_pid_tree_group_and_token_with_report(pid, cleanup_token, SIGTERM);
888 let grace_deadline = Instant::now() + SUBPROCESS_TERM_GRACE;
889 loop {
890 match child.try_wait() {
891 Ok(Some(status)) => {
892 report.merge(signal_pid_tree_group_and_token_with_report(
895 pid,
896 cleanup_token,
897 9,
898 ));
899 report.refresh_survivor_status();
900 return (Some(status), report);
901 }
902 Ok(None) => {
903 if Instant::now() >= grace_deadline {
904 break;
905 }
906 std::thread::sleep(Duration::from_millis(20));
907 }
908 Err(_) => break,
909 }
910 }
911 report.merge(signal_pid_tree_group_and_token_with_report(
912 pid,
913 cleanup_token,
914 9,
915 ));
916 }
917 }
918 let _ = child.kill();
919 let status = child.wait().ok();
920 report.refresh_survivor_status();
921 (status, report)
922}
923
924fn child_pid(child: &std::process::Child) -> Option<u32> {
925 let pid = child.id();
926 (pid > 0).then_some(pid)
927}
928
929pub(crate) fn drain_captured_pipe(
939 rx: &std::sync::mpsc::Receiver<Vec<u8>>,
940 killed: bool,
941 child_pid: u32,
942) -> Vec<u8> {
943 use std::sync::mpsc::RecvTimeoutError;
944 if killed {
945 return rx
946 .recv_timeout(Duration::from_millis(100))
947 .unwrap_or_default();
948 }
949 loop {
950 match rx.recv_timeout(Duration::from_millis(20)) {
951 Ok(buf) => return buf,
952 Err(RecvTimeoutError::Disconnected) => return Vec::new(),
953 Err(RecvTimeoutError::Timeout) => {
954 if requested() {
955 const SIGTERM: i32 = 15;
956 signal_pid_tree_and_group(child_pid, SIGTERM);
957 if let Ok(buf) = rx.recv_timeout(SUBPROCESS_TERM_GRACE) {
958 signal_pid_tree_and_group(child_pid, 9);
959 return buf;
960 }
961 signal_pid_tree_and_group(child_pid, 9);
962 return rx
963 .recv_timeout(Duration::from_millis(100))
964 .unwrap_or_default();
965 }
966 }
967 }
968 }
969}
970
971pub(crate) fn spawn_pipe_drain<R: std::io::Read + Send + 'static>(
973 mut reader: R,
974) -> std::sync::mpsc::Receiver<Vec<u8>> {
975 let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
976 std::thread::spawn(move || {
977 let mut buf = Vec::new();
978 let _ = reader.read_to_end(&mut buf);
979 let _ = tx.send(buf);
980 });
981 rx
982}
983
984pub fn capture_output_interruptible(
991 command: &mut std::process::Command,
992) -> std::io::Result<std::process::Output> {
993 use std::process::Stdio;
994 command
995 .stdout(Stdio::piped())
996 .stderr(Stdio::piped())
997 .stdin(Stdio::null());
998 configure_kill_group(command);
999 let cleanup_token = new_process_cleanup_token();
1000 command.env(PROCESS_CLEANUP_TOKEN_ENV, &cleanup_token);
1001 let mut child = command.spawn()?;
1002 let pid = child.id();
1003 let rx_out = child.stdout.take().map(spawn_pipe_drain);
1004 let rx_err = child.stderr.take().map(spawn_pipe_drain);
1005
1006 let (status, killed) = match wait_child_interruptible_with_cleanup_token(
1007 &mut child,
1008 None,
1009 Some(&cleanup_token),
1010 )? {
1011 ChildWait::Exited(status) => (status, false),
1012 ChildWait::TimedOut(_) => (std::process::ExitStatus::default(), true),
1014 ChildWait::Interrupted(status, _) => (status.unwrap_or_default(), true),
1015 };
1016 let stdout = rx_out
1017 .map(|rx| drain_captured_pipe(&rx, killed, pid))
1018 .unwrap_or_default();
1019 let stderr = rx_err
1020 .map(|rx| drain_captured_pipe(&rx, killed, pid))
1021 .unwrap_or_default();
1022 Ok(std::process::Output {
1023 status,
1024 stdout,
1025 stderr,
1026 })
1027}
1028
1029#[cfg(test)]
1030mod tests {
1031 use super::*;
1032
1033 #[test]
1034 fn requested_is_false_without_context() {
1035 assert!(!requested());
1036 }
1037
1038 #[test]
1039 fn installed_tracks_guard_lifetime() {
1040 assert!(!installed());
1041 let guard = install(None, None);
1042 assert!(installed());
1043 drop(guard);
1044 assert!(!installed());
1045 }
1046
1047 #[test]
1048 fn cancel_token_trips_requested_and_guard_restores() {
1049 let token = Arc::new(AtomicBool::new(false));
1050 let guard = install(Some(token.clone()), None);
1051 assert!(!requested());
1052 token.store(true, Ordering::SeqCst);
1053 assert!(requested());
1054 drop(guard);
1055 assert!(!requested());
1056 }
1057
1058 #[test]
1059 fn deadline_trips_requested() {
1060 let expired = Instant::now()
1061 .checked_sub(Duration::from_millis(1))
1062 .expect("monotonic clock supports a 1ms test lookback");
1063 let _guard = install(None, Some(expired));
1064 assert!(requested());
1065 }
1066
1067 #[test]
1068 fn nested_installs_restore_in_order() {
1069 let outer_token = Arc::new(AtomicBool::new(true));
1070 let _outer = install(Some(outer_token), None);
1071 assert!(requested());
1072 {
1073 let _inner = install(None, None);
1074 assert!(!requested());
1075 }
1076 assert!(requested());
1077 }
1078
1079 #[test]
1080 fn active_cleanup_owner_scopes_are_disjoint() {
1081 let owner = Arc::new(AtomicBool::new(false));
1082 let _owned =
1083 register_active_process_cleanup(None, "owned-scope-test", Some(Arc::clone(&owner)));
1084 let _ownerless = register_active_process_cleanup(None, "ownerless-scope-test", None);
1085
1086 assert_eq!(
1087 active_cleanup_tokens_for_cancel_token_for_test(&owner),
1088 vec!["owned-scope-test".to_string()]
1089 );
1090 assert!(
1091 ownerless_active_cleanup_tokens_for_test()
1092 .iter()
1093 .any(|token| token == "ownerless-scope-test"),
1094 "explicit ownerless fallback should remain separately discoverable"
1095 );
1096 }
1097
1098 #[test]
1099 fn active_cleanup_guard_unregisters_on_drop() {
1100 let owner = Arc::new(AtomicBool::new(false));
1101 let token = "guard-lifetime-test";
1102 let guard = register_active_process_cleanup(None, token, Some(Arc::clone(&owner)));
1103
1104 assert!(
1105 active_cleanup_tokens_for_cancel_token_for_test(&owner)
1106 .iter()
1107 .any(|entry| entry == token),
1108 "active cleanup must remain registered while its guard is alive"
1109 );
1110
1111 drop(guard);
1112
1113 assert!(
1114 !active_cleanup_tokens_for_cancel_token_for_test(&owner)
1115 .iter()
1116 .any(|entry| entry == token),
1117 "dropping the guard must unregister the cleanup token"
1118 );
1119 }
1120
1121 #[cfg(unix)]
1122 #[test]
1123 fn descendant_pids_from_parent_edges_returns_deepest_first_tree_only() {
1124 let edges = [
1125 (20, 10),
1126 (30, 20),
1127 (40, 20),
1128 (50, 30),
1129 (60, 99),
1130 (70, 60),
1131 (80, 90),
1133 (90, 80),
1134 ];
1135
1136 assert_eq!(
1137 descendant_pids_from_parent_edges(10, &edges),
1138 vec![50, 30, 40, 20]
1139 );
1140 assert_eq!(descendant_pids_from_parent_edges(99, &edges), vec![70, 60]);
1141 assert_eq!(
1142 descendant_pids_from_parent_edges(123, &edges),
1143 Vec::<u32>::new()
1144 );
1145 }
1146
1147 #[cfg(unix)]
1148 #[test]
1149 fn descendant_processes_preserve_metadata_and_depth_order() {
1150 let rows = [
1151 (20, 10, Some("worker".to_string())),
1152 (30, 20, Some("grandchild".to_string())),
1153 (40, 20, None),
1154 (50, 30, Some("leaf".to_string())),
1155 ];
1156
1157 let descendants = descendant_processes_from_parent_edges(10, &rows);
1158 let pids = descendants
1159 .iter()
1160 .map(|child| {
1161 (
1162 child.pid,
1163 child.parent_pid,
1164 child.depth,
1165 child.command_name.as_deref(),
1166 )
1167 })
1168 .collect::<Vec<_>>();
1169 assert_eq!(
1170 pids,
1171 vec![
1172 (50, Some(30), 3, Some("leaf")),
1173 (30, Some(20), 2, Some("grandchild")),
1174 (40, Some(20), 2, None),
1175 (20, Some(10), 1, Some("worker")),
1176 ]
1177 );
1178 }
1179
1180 #[cfg(unix)]
1181 #[test]
1182 fn command_name_keeps_only_argv0_basename() {
1183 let command = vec![
1184 std::ffi::OsString::from("/usr/local/bin/tool"),
1185 std::ffi::OsString::from("--api-key"),
1186 std::ffi::OsString::from("secret-value"),
1187 std::ffi::OsString::from("plain"),
1188 ];
1189
1190 assert_eq!(command_name(&command).as_deref(), Some("tool"));
1191 assert_eq!(command_name(&[]).as_deref(), None);
1192 }
1193
1194 #[cfg(unix)]
1195 #[test]
1196 fn process_has_cleanup_token_requires_exact_marker_entry() {
1197 let token = "tok-123";
1198 let env = vec![
1199 std::ffi::OsString::from("PATH=/usr/bin"),
1200 std::ffi::OsString::from(format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}")),
1201 ];
1202 assert!(process_has_cleanup_token(&env, token));
1203 assert!(!process_has_cleanup_token(&env, "tok"));
1204 assert!(!process_has_cleanup_token(
1205 &[std::ffi::OsString::from("OTHER=tok-123")],
1206 token
1207 ));
1208 }
1209
1210 #[cfg(unix)]
1211 #[test]
1212 fn process_has_cleanup_token_accepts_owner_lifetime_marker() {
1213 let token = "owner-123";
1214 let env = vec![std::ffi::OsString::from(format!(
1215 "{PROCESS_OWNER_TOKEN_ENV}={token}"
1216 ))];
1217 assert!(process_has_cleanup_token(&env, token));
1218 assert!(!process_has_cleanup_token(&env, "owner"));
1219 }
1220
1221 #[cfg(unix)]
1222 #[test]
1223 fn zombie_processes_are_not_lifetime_survivors() {
1224 assert!(!process_status_can_execute(sysinfo::ProcessStatus::Zombie));
1225 assert!(process_status_can_execute(sysinfo::ProcessStatus::Sleep));
1226 assert!(process_status_can_execute(sysinfo::ProcessStatus::Dead));
1227 }
1228
1229 #[cfg(unix)]
1230 #[test]
1231 fn owner_journal_initialization_refuses_preexisting_symlink() {
1232 let token = new_process_cleanup_token();
1233 let journal = owner_process_group_journal(&token);
1234 let target = tempfile::NamedTempFile::new().expect("create journal symlink target");
1235 std::os::unix::fs::symlink(target.path(), &journal).expect("create owner journal symlink");
1236 initialize_process_owner_group_journal(&token)
1237 .expect_err("preexisting journal symlink must fail closed");
1238 std::fs::remove_file(journal).expect("remove owner journal symlink");
1239 }
1240
1241 #[cfg(unix)]
1242 #[test]
1243 fn interrupted_wait_kills_process_group() {
1244 let mut command = std::process::Command::new("sh");
1246 command.args(["-c", "sleep 30 & wait"]);
1247 configure_kill_group(&mut command);
1248 let mut child = command.spawn().expect("spawn sh");
1249 let pgid = child.id();
1250
1251 let cancel = Arc::new(AtomicBool::new(true));
1252 let _guard = install(Some(cancel), None);
1253 let started = Instant::now();
1254 let outcome = wait_child_interruptible(&mut child, None).expect("wait");
1255 assert!(matches!(outcome, ChildWait::Interrupted(_, _)));
1256 assert!(started.elapsed() < Duration::from_secs(10));
1257
1258 extern "C" {
1260 fn kill(pid: i32, sig: i32) -> i32;
1261 }
1262 let group_gone = || unsafe { kill(-(pgid as i32), 0) } != 0;
1263 let deadline = Instant::now() + Duration::from_secs(5);
1264 while !group_gone() && Instant::now() < deadline {
1265 std::thread::sleep(Duration::from_millis(50));
1266 }
1267 assert!(group_gone(), "process group {pgid} survived interrupt");
1268 }
1269}