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) {
497 #[cfg(unix)]
498 {
499 use std::os::unix::process::CommandExt;
500 unsafe {
505 command.pre_exec(start_kill_session);
506 }
507 }
508 #[cfg(not(unix))]
509 {
510 let _ = command;
511 }
512}
513
514pub fn configure_tokio_kill_group(command: &mut tokio::process::Command) {
518 #[cfg(unix)]
519 {
520 unsafe {
523 command.pre_exec(start_kill_session);
524 }
525 }
526 #[cfg(not(unix))]
527 {
528 let _ = command;
529 }
530}
531
532#[cfg(unix)]
533fn start_kill_session() -> std::io::Result<()> {
534 if unsafe { libc::setsid() } == -1 {
535 return Err(std::io::Error::last_os_error());
536 }
537 Ok(())
538}
539
540pub fn signal_pid_and_group(pid: u32, signal: i32) {
542 #[cfg(unix)]
543 {
544 extern "C" {
547 fn kill(pid: i32, sig: i32) -> i32;
548 }
549 unsafe {
550 kill(-(pid as i32), signal);
551 kill(pid as i32, signal);
552 }
553 }
554 #[cfg(not(unix))]
555 {
556 let _ = (pid, signal);
557 }
558}
559
560pub fn signal_pid_tree_and_group(pid: u32, signal: i32) {
565 let _ = signal_pid_tree_and_group_with_report(pid, signal);
566}
567
568pub fn signal_pid_tree_and_group_with_report(pid: u32, signal: i32) -> ProcessCleanupReport {
571 signal_pid_tree_group_and_token_with_report(pid, None, signal)
572}
573
574pub fn signal_pid_tree_group_and_token_with_report(
580 pid: u32,
581 cleanup_token: Option<&str>,
582 signal: i32,
583) -> ProcessCleanupReport {
584 #[cfg(unix)]
585 {
586 let mut report = ProcessCleanupReport::for_signal(Some(pid), signal);
587 for child in descendant_processes(pid) {
588 signal_pid_and_group(child.pid, signal);
589 report.merge_child(child.with_signal(signal));
590 }
591 if let Some(cleanup_token) = cleanup_token.filter(|token| !token.is_empty()) {
592 for child in cleanup_token_processes(cleanup_token) {
593 if child.pid == pid {
594 continue;
595 }
596 signal_pid_and_group(child.pid, signal);
597 report.merge_child(child.with_signal(signal));
598 }
599 }
600 signal_pid_and_group(pid, signal);
601 if signal == 9 {
602 wait_for_report_children_to_exit(&report, SUBPROCESS_KILL_SETTLE);
603 }
604 report.refresh_survivor_status();
605 report
606 }
607 #[cfg(not(unix))]
608 {
609 let _ = cleanup_token;
610 ProcessCleanupReport::for_signal(Some(pid), signal)
611 }
612}
613
614#[cfg(unix)]
622pub fn signal_pid_tree_and_token_preserving_group_with_report(
623 pid: u32,
624 cleanup_token: Option<&str>,
625 preserved_pgid: u32,
626 signal: i32,
627) -> ProcessCleanupReport {
628 let preserved_pgid = preserved_pgid as i32;
629 let mut report = ProcessCleanupReport::for_signal(Some(pid), signal);
630 for child in descendant_processes(pid) {
631 signal_pid_preserving_group(child.pid, preserved_pgid, signal);
632 report.merge_child(child.with_signal(signal));
633 }
634 if let Some(cleanup_token) = cleanup_token.filter(|token| !token.is_empty()) {
635 for child in cleanup_token_processes(cleanup_token) {
636 if child.pid == pid {
637 continue;
638 }
639 signal_pid_preserving_group(child.pid, preserved_pgid, signal);
640 report.merge_child(child.with_signal(signal));
641 }
642 for pgid in owner_process_groups(cleanup_token) {
643 if pgid != preserved_pgid as u32 {
644 unsafe {
645 libc::kill(-(pgid as i32), signal);
646 }
647 }
648 }
649 }
650 signal_pid_preserving_group(pid, preserved_pgid, signal);
651 report
652}
653
654#[cfg(unix)]
655fn signal_pid_preserving_group(pid: u32, preserved_pgid: i32, signal: i32) {
656 let pid = pid as i32;
657 let pgid = unsafe { libc::getpgid(pid) };
658 unsafe {
659 if pgid > 0 && pgid != preserved_pgid {
660 libc::kill(-pgid, signal);
661 }
662 libc::kill(pid, signal);
663 }
664}
665
666#[cfg(unix)]
667fn descendant_processes(root: u32) -> Vec<ProcessCleanupChild> {
668 use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System};
669
670 let mut sys = System::new();
671 sys.refresh_processes_specifics(
672 ProcessesToUpdate::All,
673 false,
674 ProcessRefreshKind::everything(),
675 );
676 let rows = sys
677 .processes()
678 .iter()
679 .filter_map(|(pid, process)| {
680 Some((
681 pid.as_u32(),
682 process.parent()?.as_u32(),
683 command_name(process.cmd()),
684 ))
685 })
686 .collect::<Vec<_>>();
687 descendant_processes_from_parent_edges(root, &rows)
688}
689
690#[cfg(unix)]
691fn cleanup_token_processes(token: &str) -> Vec<ProcessCleanupChild> {
692 use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind};
693
694 let mut sys = System::new();
695 sys.refresh_processes_specifics(
696 ProcessesToUpdate::All,
697 false,
698 ProcessRefreshKind::nothing()
699 .with_environ(UpdateKind::Always)
700 .with_cmd(UpdateKind::Always),
701 );
702 let mut children = sys
703 .processes()
704 .iter()
705 .filter(|(_, process)| {
706 process_status_can_execute(process.status())
707 && process_has_cleanup_token(process.environ(), token)
708 })
709 .map(|(pid, process)| {
710 ProcessCleanupChild::new(
711 pid.as_u32(),
712 process.parent().map(|parent| parent.as_u32()),
713 1,
714 command_name(process.cmd()),
715 )
716 })
717 .collect::<Vec<_>>();
718 children.sort_by_key(|child| child.pid);
719 children
720}
721
722#[cfg(unix)]
723fn process_status_can_execute(status: sysinfo::ProcessStatus) -> bool {
724 status != sysinfo::ProcessStatus::Zombie
725}
726
727#[cfg(unix)]
728fn process_has_cleanup_token(environ: &[std::ffi::OsString], token: &str) -> bool {
729 let cleanup = format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}");
730 let owner = format!("{PROCESS_OWNER_TOKEN_ENV}={token}");
731 environ
732 .iter()
733 .any(|entry| matches!(entry.to_string_lossy().as_ref(), value if value == cleanup || value == owner))
734}
735
736#[cfg(all(unix, test))]
737fn descendant_pids_from_parent_edges(root: u32, edges: &[(u32, u32)]) -> Vec<u32> {
738 let rows = edges
739 .iter()
740 .map(|(pid, parent)| (*pid, *parent, None))
741 .collect::<Vec<_>>();
742 descendant_processes_from_parent_edges(root, &rows)
743 .into_iter()
744 .map(|child| child.pid)
745 .collect()
746}
747
748#[cfg(unix)]
749fn descendant_processes_from_parent_edges(
750 root: u32,
751 rows: &[(u32, u32, Option<String>)],
752) -> Vec<ProcessCleanupChild> {
753 use std::collections::{HashMap, HashSet};
754
755 let mut children: HashMap<u32, Vec<u32>> = HashMap::new();
756 let mut metadata: HashMap<u32, (u32, Option<String>)> = HashMap::new();
757 for (pid, parent, command) in rows {
758 metadata.insert(*pid, (*parent, command.clone()));
759 children.entry(*parent).or_default().push(*pid);
760 }
761
762 let mut seen = HashSet::new();
763 let mut stack = vec![(root, 0usize)];
764 let mut descendants = Vec::new();
765 while let Some((pid, depth)) = stack.pop() {
766 if !seen.insert(pid) {
767 continue;
768 }
769 if pid != root {
770 descendants.push((pid, depth));
771 }
772 if let Some(kids) = children.get(&pid) {
773 for &child in kids {
774 stack.push((child, depth + 1));
775 }
776 }
777 }
778
779 descendants.sort_by(|(left_pid, left_depth), (right_pid, right_depth)| {
780 right_depth
781 .cmp(left_depth)
782 .then_with(|| left_pid.cmp(right_pid))
783 });
784 descendants
785 .into_iter()
786 .map(|(pid, depth)| {
787 let (parent_pid, command) = metadata.get(&pid).cloned().unwrap_or((root, None));
788 ProcessCleanupChild::new(pid, Some(parent_pid), depth as u32, command)
789 })
790 .collect()
791}
792
793#[cfg(unix)]
794fn command_name(command: &[std::ffi::OsString]) -> Option<String> {
795 if command.is_empty() {
796 return None;
797 }
798 std::path::Path::new(&command[0])
799 .file_name()
800 .map(|name| name.to_string_lossy().into_owned())
801 .filter(|name| !name.is_empty())
802}
803
804#[cfg(unix)]
805fn process_exists(pid: u32) -> bool {
806 unsafe { libc::kill(pid as i32, 0) == 0 }
807}
808
809#[cfg(unix)]
810fn wait_for_report_children_to_exit(report: &ProcessCleanupReport, timeout: Duration) {
811 let deadline = Instant::now() + timeout;
812 while Instant::now() < deadline {
813 if report
814 .children
815 .iter()
816 .all(|child| !process_exists(child.pid))
817 {
818 return;
819 }
820 std::thread::sleep(Duration::from_millis(10));
821 }
822}
823
824pub enum ChildWait {
826 Exited(std::process::ExitStatus),
828 TimedOut(ProcessCleanupReport),
830 Interrupted(Option<std::process::ExitStatus>, ProcessCleanupReport),
834}
835
836pub fn wait_child_interruptible(
845 child: &mut std::process::Child,
846 timeout: Option<Duration>,
847) -> std::io::Result<ChildWait> {
848 wait_child_interruptible_with_cleanup_token(child, timeout, None)
849}
850
851pub fn wait_child_interruptible_with_cleanup_token(
852 child: &mut std::process::Child,
853 timeout: Option<Duration>,
854 cleanup_token: Option<&str>,
855) -> std::io::Result<ChildWait> {
856 let deadline = timeout.map(|limit| Instant::now() + limit);
857 loop {
858 if let Some(status) = child.try_wait()? {
859 return Ok(ChildWait::Exited(status));
860 }
861 if requested() {
862 let (status, report) =
863 terminate_child_group_with_cleanup_token_report(child, cleanup_token);
864 return Ok(ChildWait::Interrupted(status, report));
865 }
866 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
867 let mut report = child_pid(child)
869 .map(|pid| signal_pid_tree_group_and_token_with_report(pid, cleanup_token, 9))
870 .unwrap_or_default();
871 let _ = child.kill();
872 let _ = child.wait();
873 report.refresh_survivor_status();
874 return Ok(ChildWait::TimedOut(report));
875 }
876 std::thread::sleep(Duration::from_millis(20));
877 }
878}
879
880pub fn terminate_child_group(child: &mut std::process::Child) -> Option<std::process::ExitStatus> {
886 terminate_child_group_with_report(child).0
887}
888
889pub fn terminate_child_group_with_report(
892 child: &mut std::process::Child,
893) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
894 terminate_child_group_with_cleanup_token_report(child, None)
895}
896
897pub fn terminate_child_group_with_cleanup_token_report(
898 child: &mut std::process::Child,
899 cleanup_token: Option<&str>,
900) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
901 let mut report = child_pid(child)
902 .map(|pid| ProcessCleanupReport::for_signal(Some(pid), 15))
903 .unwrap_or_default();
904 #[cfg(not(unix))]
905 let _ = cleanup_token;
906 #[cfg(unix)]
907 {
908 if let Some(pid) = child_pid(child) {
909 const SIGTERM: i32 = 15;
910 report = signal_pid_tree_group_and_token_with_report(pid, cleanup_token, SIGTERM);
911 let grace_deadline = Instant::now() + SUBPROCESS_TERM_GRACE;
912 loop {
913 match child.try_wait() {
914 Ok(Some(status)) => {
915 report.merge(signal_pid_tree_group_and_token_with_report(
918 pid,
919 cleanup_token,
920 9,
921 ));
922 report.refresh_survivor_status();
923 return (Some(status), report);
924 }
925 Ok(None) => {
926 if Instant::now() >= grace_deadline {
927 break;
928 }
929 std::thread::sleep(Duration::from_millis(20));
930 }
931 Err(_) => break,
932 }
933 }
934 report.merge(signal_pid_tree_group_and_token_with_report(
935 pid,
936 cleanup_token,
937 9,
938 ));
939 }
940 }
941 let _ = child.kill();
942 let status = child.wait().ok();
943 report.refresh_survivor_status();
944 (status, report)
945}
946
947fn child_pid(child: &std::process::Child) -> Option<u32> {
948 let pid = child.id();
949 (pid > 0).then_some(pid)
950}
951
952pub(crate) fn drain_captured_pipe(
962 rx: &std::sync::mpsc::Receiver<Vec<u8>>,
963 killed: bool,
964 child_pid: u32,
965) -> Vec<u8> {
966 use std::sync::mpsc::RecvTimeoutError;
967 if killed {
968 return rx
969 .recv_timeout(Duration::from_millis(100))
970 .unwrap_or_default();
971 }
972 loop {
973 match rx.recv_timeout(Duration::from_millis(20)) {
974 Ok(buf) => return buf,
975 Err(RecvTimeoutError::Disconnected) => return Vec::new(),
976 Err(RecvTimeoutError::Timeout) => {
977 if requested() {
978 const SIGTERM: i32 = 15;
979 signal_pid_tree_and_group(child_pid, SIGTERM);
980 if let Ok(buf) = rx.recv_timeout(SUBPROCESS_TERM_GRACE) {
981 signal_pid_tree_and_group(child_pid, 9);
982 return buf;
983 }
984 signal_pid_tree_and_group(child_pid, 9);
985 return rx
986 .recv_timeout(Duration::from_millis(100))
987 .unwrap_or_default();
988 }
989 }
990 }
991 }
992}
993
994pub(crate) fn spawn_pipe_drain<R: std::io::Read + Send + 'static>(
996 mut reader: R,
997) -> std::sync::mpsc::Receiver<Vec<u8>> {
998 let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
999 std::thread::spawn(move || {
1000 let mut buf = Vec::new();
1001 let _ = reader.read_to_end(&mut buf);
1002 let _ = tx.send(buf);
1003 });
1004 rx
1005}
1006
1007pub fn capture_output_interruptible(
1014 command: &mut std::process::Command,
1015) -> std::io::Result<std::process::Output> {
1016 use std::process::Stdio;
1017 command
1018 .stdout(Stdio::piped())
1019 .stderr(Stdio::piped())
1020 .stdin(Stdio::null());
1021 configure_kill_group(command);
1022 let cleanup_token = new_process_cleanup_token();
1023 command.env(PROCESS_CLEANUP_TOKEN_ENV, &cleanup_token);
1024 let mut child = command.spawn()?;
1025 let pid = child.id();
1026 let rx_out = child.stdout.take().map(spawn_pipe_drain);
1027 let rx_err = child.stderr.take().map(spawn_pipe_drain);
1028
1029 let (status, killed) = match wait_child_interruptible_with_cleanup_token(
1030 &mut child,
1031 None,
1032 Some(&cleanup_token),
1033 )? {
1034 ChildWait::Exited(status) => (status, false),
1035 ChildWait::TimedOut(_) => (std::process::ExitStatus::default(), true),
1037 ChildWait::Interrupted(status, _) => (status.unwrap_or_default(), true),
1038 };
1039 let stdout = rx_out
1040 .map(|rx| drain_captured_pipe(&rx, killed, pid))
1041 .unwrap_or_default();
1042 let stderr = rx_err
1043 .map(|rx| drain_captured_pipe(&rx, killed, pid))
1044 .unwrap_or_default();
1045 Ok(std::process::Output {
1046 status,
1047 stdout,
1048 stderr,
1049 })
1050}
1051
1052#[cfg(test)]
1053mod tests {
1054 use super::*;
1055
1056 #[test]
1057 fn requested_is_false_without_context() {
1058 assert!(!requested());
1059 }
1060
1061 #[test]
1062 fn installed_tracks_guard_lifetime() {
1063 assert!(!installed());
1064 let guard = install(None, None);
1065 assert!(installed());
1066 drop(guard);
1067 assert!(!installed());
1068 }
1069
1070 #[test]
1071 fn cancel_token_trips_requested_and_guard_restores() {
1072 let token = Arc::new(AtomicBool::new(false));
1073 let guard = install(Some(token.clone()), None);
1074 assert!(!requested());
1075 token.store(true, Ordering::SeqCst);
1076 assert!(requested());
1077 drop(guard);
1078 assert!(!requested());
1079 }
1080
1081 #[test]
1082 fn deadline_trips_requested() {
1083 let expired = Instant::now()
1084 .checked_sub(Duration::from_millis(1))
1085 .expect("monotonic clock supports a 1ms test lookback");
1086 let _guard = install(None, Some(expired));
1087 assert!(requested());
1088 }
1089
1090 #[test]
1091 fn nested_installs_restore_in_order() {
1092 let outer_token = Arc::new(AtomicBool::new(true));
1093 let _outer = install(Some(outer_token), None);
1094 assert!(requested());
1095 {
1096 let _inner = install(None, None);
1097 assert!(!requested());
1098 }
1099 assert!(requested());
1100 }
1101
1102 #[test]
1103 fn active_cleanup_owner_scopes_are_disjoint() {
1104 let owner = Arc::new(AtomicBool::new(false));
1105 let _owned =
1106 register_active_process_cleanup(None, "owned-scope-test", Some(Arc::clone(&owner)));
1107 let _ownerless = register_active_process_cleanup(None, "ownerless-scope-test", None);
1108
1109 assert_eq!(
1110 active_cleanup_tokens_for_cancel_token_for_test(&owner),
1111 vec!["owned-scope-test".to_string()]
1112 );
1113 assert!(
1114 ownerless_active_cleanup_tokens_for_test()
1115 .iter()
1116 .any(|token| token == "ownerless-scope-test"),
1117 "explicit ownerless fallback should remain separately discoverable"
1118 );
1119 }
1120
1121 #[test]
1122 fn active_cleanup_guard_unregisters_on_drop() {
1123 let owner = Arc::new(AtomicBool::new(false));
1124 let token = "guard-lifetime-test";
1125 let guard = register_active_process_cleanup(None, token, Some(Arc::clone(&owner)));
1126
1127 assert!(
1128 active_cleanup_tokens_for_cancel_token_for_test(&owner)
1129 .iter()
1130 .any(|entry| entry == token),
1131 "active cleanup must remain registered while its guard is alive"
1132 );
1133
1134 drop(guard);
1135
1136 assert!(
1137 !active_cleanup_tokens_for_cancel_token_for_test(&owner)
1138 .iter()
1139 .any(|entry| entry == token),
1140 "dropping the guard must unregister the cleanup token"
1141 );
1142 }
1143
1144 #[cfg(unix)]
1145 #[test]
1146 fn descendant_pids_from_parent_edges_returns_deepest_first_tree_only() {
1147 let edges = [
1148 (20, 10),
1149 (30, 20),
1150 (40, 20),
1151 (50, 30),
1152 (60, 99),
1153 (70, 60),
1154 (80, 90),
1156 (90, 80),
1157 ];
1158
1159 assert_eq!(
1160 descendant_pids_from_parent_edges(10, &edges),
1161 vec![50, 30, 40, 20]
1162 );
1163 assert_eq!(descendant_pids_from_parent_edges(99, &edges), vec![70, 60]);
1164 assert_eq!(
1165 descendant_pids_from_parent_edges(123, &edges),
1166 Vec::<u32>::new()
1167 );
1168 }
1169
1170 #[cfg(unix)]
1171 #[test]
1172 fn descendant_processes_preserve_metadata_and_depth_order() {
1173 let rows = [
1174 (20, 10, Some("worker".to_string())),
1175 (30, 20, Some("grandchild".to_string())),
1176 (40, 20, None),
1177 (50, 30, Some("leaf".to_string())),
1178 ];
1179
1180 let descendants = descendant_processes_from_parent_edges(10, &rows);
1181 let pids = descendants
1182 .iter()
1183 .map(|child| {
1184 (
1185 child.pid,
1186 child.parent_pid,
1187 child.depth,
1188 child.command_name.as_deref(),
1189 )
1190 })
1191 .collect::<Vec<_>>();
1192 assert_eq!(
1193 pids,
1194 vec![
1195 (50, Some(30), 3, Some("leaf")),
1196 (30, Some(20), 2, Some("grandchild")),
1197 (40, Some(20), 2, None),
1198 (20, Some(10), 1, Some("worker")),
1199 ]
1200 );
1201 }
1202
1203 #[cfg(unix)]
1204 #[test]
1205 fn command_name_keeps_only_argv0_basename() {
1206 let command = vec![
1207 std::ffi::OsString::from("/usr/local/bin/tool"),
1208 std::ffi::OsString::from("--api-key"),
1209 std::ffi::OsString::from("secret-value"),
1210 std::ffi::OsString::from("plain"),
1211 ];
1212
1213 assert_eq!(command_name(&command).as_deref(), Some("tool"));
1214 assert_eq!(command_name(&[]).as_deref(), None);
1215 }
1216
1217 #[cfg(unix)]
1218 #[test]
1219 fn process_has_cleanup_token_requires_exact_marker_entry() {
1220 let token = "tok-123";
1221 let env = vec![
1222 std::ffi::OsString::from("PATH=/usr/bin"),
1223 std::ffi::OsString::from(format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}")),
1224 ];
1225 assert!(process_has_cleanup_token(&env, token));
1226 assert!(!process_has_cleanup_token(&env, "tok"));
1227 assert!(!process_has_cleanup_token(
1228 &[std::ffi::OsString::from("OTHER=tok-123")],
1229 token
1230 ));
1231 }
1232
1233 #[cfg(unix)]
1234 #[test]
1235 fn process_has_cleanup_token_accepts_owner_lifetime_marker() {
1236 let token = "owner-123";
1237 let env = vec![std::ffi::OsString::from(format!(
1238 "{PROCESS_OWNER_TOKEN_ENV}={token}"
1239 ))];
1240 assert!(process_has_cleanup_token(&env, token));
1241 assert!(!process_has_cleanup_token(&env, "owner"));
1242 }
1243
1244 #[cfg(unix)]
1245 #[test]
1246 fn zombie_processes_are_not_lifetime_survivors() {
1247 assert!(!process_status_can_execute(sysinfo::ProcessStatus::Zombie));
1248 assert!(process_status_can_execute(sysinfo::ProcessStatus::Sleep));
1249 assert!(process_status_can_execute(sysinfo::ProcessStatus::Dead));
1250 }
1251
1252 #[cfg(unix)]
1253 #[test]
1254 fn owner_journal_initialization_refuses_preexisting_symlink() {
1255 let token = new_process_cleanup_token();
1256 let journal = owner_process_group_journal(&token);
1257 let target = tempfile::NamedTempFile::new().expect("create journal symlink target");
1258 std::os::unix::fs::symlink(target.path(), &journal).expect("create owner journal symlink");
1259 initialize_process_owner_group_journal(&token)
1260 .expect_err("preexisting journal symlink must fail closed");
1261 std::fs::remove_file(journal).expect("remove owner journal symlink");
1262 }
1263
1264 #[cfg(unix)]
1265 #[test]
1266 fn interrupted_wait_kills_process_group() {
1267 let mut command = std::process::Command::new("sh");
1269 command.args(["-c", "sleep 30 & wait"]);
1270 configure_kill_group(&mut command);
1271 let mut child = command.spawn().expect("spawn sh");
1272 let pgid = child.id();
1273
1274 let cancel = Arc::new(AtomicBool::new(true));
1275 let _guard = install(Some(cancel), None);
1276 let started = Instant::now();
1277 let outcome = wait_child_interruptible(&mut child, None).expect("wait");
1278 assert!(matches!(outcome, ChildWait::Interrupted(_, _)));
1279 assert!(started.elapsed() < Duration::from_secs(10));
1280
1281 extern "C" {
1283 fn kill(pid: i32, sig: i32) -> i32;
1284 }
1285 let group_gone = || unsafe { kill(-(pgid as i32), 0) } != 0;
1286 let deadline = Instant::now() + Duration::from_secs(5);
1287 while !group_gone() && Instant::now() < deadline {
1288 std::thread::sleep(Duration::from_millis(50));
1289 }
1290 assert!(group_gone(), "process group {pgid} survived interrupt");
1291 }
1292}