1use std::cell::RefCell;
33use std::sync::atomic::{AtomicBool, Ordering};
34use std::sync::Arc;
35use std::time::{Duration, Instant};
36
37pub const PROCESS_CLEANUP_TOKEN_ENV: &str = "HARN_PROCESS_CLEANUP_TOKEN";
41
42pub const SUBPROCESS_TERM_GRACE: Duration = Duration::from_secs(2);
47#[cfg(unix)]
48const SUBPROCESS_KILL_SETTLE: Duration = Duration::from_millis(250);
49
50pub fn new_process_cleanup_token() -> String {
51 format!("harn-cleanup-{}", uuid::Uuid::now_v7().simple())
52}
53
54#[derive(Clone, Debug, Default, PartialEq, Eq)]
56pub struct ProcessCleanupReport {
57 pub root_pid: Option<u32>,
58 pub attempted_signals: Vec<i32>,
59 pub children: Vec<ProcessCleanupChild>,
60}
61
62impl ProcessCleanupReport {
63 pub fn for_signal(root_pid: Option<u32>, signal: i32) -> Self {
64 Self {
65 root_pid,
66 attempted_signals: vec![signal],
67 children: Vec::new(),
68 }
69 }
70
71 pub fn merge(&mut self, other: Self) {
72 if self.root_pid.is_none() {
73 self.root_pid = other.root_pid;
74 }
75 for signal in other.attempted_signals {
76 push_unique(&mut self.attempted_signals, signal);
77 }
78 for child in other.children {
79 self.merge_child(child);
80 }
81 }
82
83 pub fn refresh_survivor_status(&mut self) {
84 #[cfg(unix)]
85 {
86 for child in &mut self.children {
87 child.alive_after_cleanup = Some(process_exists(child.pid));
88 }
89 }
90 }
91
92 fn merge_child(&mut self, child: ProcessCleanupChild) {
93 if let Some(existing) = self
94 .children
95 .iter_mut()
96 .find(|entry| entry.pid == child.pid)
97 {
98 for signal in child.signals {
99 push_unique(&mut existing.signals, signal);
100 }
101 if existing.command_name.is_none() {
102 existing.command_name = child.command_name;
103 }
104 if child.alive_after_cleanup.is_some() {
105 existing.alive_after_cleanup = child.alive_after_cleanup;
106 }
107 return;
108 }
109 self.children.push(child);
110 self.children
111 .sort_by(|left, right| left.depth.cmp(&right.depth).then(left.pid.cmp(&right.pid)));
112 }
113}
114
115#[derive(Clone, Debug, PartialEq, Eq)]
117pub struct ProcessCleanupChild {
118 pub pid: u32,
119 pub parent_pid: Option<u32>,
120 pub depth: u32,
121 pub command_name: Option<String>,
122 pub signals: Vec<i32>,
123 pub alive_after_cleanup: Option<bool>,
124}
125
126impl ProcessCleanupChild {
127 pub fn new(
128 pid: u32,
129 parent_pid: Option<u32>,
130 depth: u32,
131 command_name: Option<String>,
132 ) -> Self {
133 Self {
134 pid,
135 parent_pid,
136 depth,
137 command_name,
138 signals: Vec::new(),
139 alive_after_cleanup: None,
140 }
141 }
142
143 #[cfg(unix)]
144 fn with_signal(mut self, signal: i32) -> Self {
145 push_unique(&mut self.signals, signal);
146 self
147 }
148}
149
150fn push_unique<T: Copy + Eq>(values: &mut Vec<T>, value: T) {
151 if !values.contains(&value) {
152 values.push(value);
153 }
154}
155
156#[derive(Clone, Default)]
157struct OpInterrupt {
158 cancel: Option<Arc<AtomicBool>>,
159 deadline: Option<Instant>,
160}
161
162thread_local! {
163 static CURRENT: RefCell<Option<OpInterrupt>> = const { RefCell::new(None) };
164}
165
166pub struct OpInterruptGuard {
170 #[allow(clippy::option_option)]
173 prev: Option<Option<OpInterrupt>>,
174}
175
176impl Drop for OpInterruptGuard {
177 fn drop(&mut self) {
178 if let Some(prev) = self.prev.take() {
179 CURRENT.with(|slot| *slot.borrow_mut() = prev);
180 }
181 }
182}
183
184pub fn install(cancel: Option<Arc<AtomicBool>>, deadline: Option<Instant>) -> OpInterruptGuard {
189 let prev = CURRENT.with(|slot| slot.borrow_mut().replace(OpInterrupt { cancel, deadline }));
190 OpInterruptGuard { prev: Some(prev) }
191}
192
193pub fn installed() -> bool {
198 CURRENT.with(|slot| slot.borrow().is_some())
199}
200
201pub fn requested() -> bool {
205 CURRENT.with(|slot| {
206 let ctx = slot.borrow();
207 let Some(ctx) = ctx.as_ref() else {
208 return false;
209 };
210 if ctx
211 .cancel
212 .as_ref()
213 .is_some_and(|token| token.load(Ordering::SeqCst))
214 {
215 return true;
216 }
217 ctx.deadline
218 .is_some_and(|deadline| Instant::now() >= deadline)
219 })
220}
221
222pub fn configure_kill_group(command: &mut std::process::Command) {
227 #[cfg(unix)]
228 {
229 use std::os::unix::process::CommandExt;
230 command.process_group(0);
231 }
232 #[cfg(not(unix))]
233 {
234 let _ = command;
235 }
236}
237
238pub fn signal_pid_and_group(pid: u32, signal: i32) {
240 #[cfg(unix)]
241 {
242 extern "C" {
245 fn kill(pid: i32, sig: i32) -> i32;
246 }
247 unsafe {
248 kill(-(pid as i32), signal);
249 kill(pid as i32, signal);
250 }
251 }
252 #[cfg(not(unix))]
253 {
254 let _ = (pid, signal);
255 }
256}
257
258pub fn signal_pid_tree_and_group(pid: u32, signal: i32) {
263 let _ = signal_pid_tree_and_group_with_report(pid, signal);
264}
265
266pub fn signal_pid_tree_and_group_with_report(pid: u32, signal: i32) -> ProcessCleanupReport {
269 signal_pid_tree_group_and_token_with_report(pid, None, signal)
270}
271
272pub fn signal_pid_tree_group_and_token_with_report(
278 pid: u32,
279 cleanup_token: Option<&str>,
280 signal: i32,
281) -> ProcessCleanupReport {
282 #[cfg(unix)]
283 {
284 let mut report = ProcessCleanupReport::for_signal(Some(pid), signal);
285 for child in descendant_processes(pid) {
286 signal_pid_and_group(child.pid, signal);
287 report.merge_child(child.with_signal(signal));
288 }
289 if let Some(cleanup_token) = cleanup_token.filter(|token| !token.is_empty()) {
290 for child in cleanup_token_processes(cleanup_token) {
291 if child.pid == pid {
292 continue;
293 }
294 signal_pid_and_group(child.pid, signal);
295 report.merge_child(child.with_signal(signal));
296 }
297 }
298 signal_pid_and_group(pid, signal);
299 if signal == 9 {
300 wait_for_report_children_to_exit(&report, SUBPROCESS_KILL_SETTLE);
301 }
302 report.refresh_survivor_status();
303 report
304 }
305 #[cfg(not(unix))]
306 {
307 let _ = cleanup_token;
308 ProcessCleanupReport::for_signal(Some(pid), signal)
309 }
310}
311
312#[cfg(unix)]
313fn descendant_processes(root: u32) -> Vec<ProcessCleanupChild> {
314 use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System};
315
316 let mut sys = System::new();
317 sys.refresh_processes_specifics(
318 ProcessesToUpdate::All,
319 false,
320 ProcessRefreshKind::everything(),
321 );
322 let rows = sys
323 .processes()
324 .iter()
325 .filter_map(|(pid, process)| {
326 Some((
327 pid.as_u32(),
328 process.parent()?.as_u32(),
329 command_name(process.cmd()),
330 ))
331 })
332 .collect::<Vec<_>>();
333 descendant_processes_from_parent_edges(root, &rows)
334}
335
336#[cfg(unix)]
337fn cleanup_token_processes(token: &str) -> Vec<ProcessCleanupChild> {
338 use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind};
339
340 let mut sys = System::new();
341 sys.refresh_processes_specifics(
342 ProcessesToUpdate::All,
343 false,
344 ProcessRefreshKind::nothing()
345 .with_environ(UpdateKind::Always)
346 .with_cmd(UpdateKind::Always),
347 );
348 let mut children = sys
349 .processes()
350 .iter()
351 .filter(|(_, process)| process_has_cleanup_token(process.environ(), token))
352 .map(|(pid, process)| {
353 ProcessCleanupChild::new(
354 pid.as_u32(),
355 process.parent().map(|parent| parent.as_u32()),
356 1,
357 command_name(process.cmd()),
358 )
359 })
360 .collect::<Vec<_>>();
361 children.sort_by_key(|child| child.pid);
362 children
363}
364
365#[cfg(unix)]
366fn process_has_cleanup_token(environ: &[std::ffi::OsString], token: &str) -> bool {
367 let expected = format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}");
368 environ
369 .iter()
370 .any(|entry| entry.to_string_lossy() == expected)
371}
372
373#[cfg(all(unix, test))]
374fn descendant_pids_from_parent_edges(root: u32, edges: &[(u32, u32)]) -> Vec<u32> {
375 let rows = edges
376 .iter()
377 .map(|(pid, parent)| (*pid, *parent, None))
378 .collect::<Vec<_>>();
379 descendant_processes_from_parent_edges(root, &rows)
380 .into_iter()
381 .map(|child| child.pid)
382 .collect()
383}
384
385#[cfg(unix)]
386fn descendant_processes_from_parent_edges(
387 root: u32,
388 rows: &[(u32, u32, Option<String>)],
389) -> Vec<ProcessCleanupChild> {
390 use std::collections::{HashMap, HashSet};
391
392 let mut children: HashMap<u32, Vec<u32>> = HashMap::new();
393 let mut metadata: HashMap<u32, (u32, Option<String>)> = HashMap::new();
394 for (pid, parent, command) in rows {
395 metadata.insert(*pid, (*parent, command.clone()));
396 children.entry(*parent).or_default().push(*pid);
397 }
398
399 let mut seen = HashSet::new();
400 let mut stack = vec![(root, 0usize)];
401 let mut descendants = Vec::new();
402 while let Some((pid, depth)) = stack.pop() {
403 if !seen.insert(pid) {
404 continue;
405 }
406 if pid != root {
407 descendants.push((pid, depth));
408 }
409 if let Some(kids) = children.get(&pid) {
410 for &child in kids {
411 stack.push((child, depth + 1));
412 }
413 }
414 }
415
416 descendants.sort_by(|(left_pid, left_depth), (right_pid, right_depth)| {
417 right_depth
418 .cmp(left_depth)
419 .then_with(|| left_pid.cmp(right_pid))
420 });
421 descendants
422 .into_iter()
423 .map(|(pid, depth)| {
424 let (parent_pid, command) = metadata.get(&pid).cloned().unwrap_or((root, None));
425 ProcessCleanupChild::new(pid, Some(parent_pid), depth as u32, command)
426 })
427 .collect()
428}
429
430#[cfg(unix)]
431fn command_name(command: &[std::ffi::OsString]) -> Option<String> {
432 if command.is_empty() {
433 return None;
434 }
435 std::path::Path::new(&command[0])
436 .file_name()
437 .map(|name| name.to_string_lossy().into_owned())
438 .filter(|name| !name.is_empty())
439}
440
441#[cfg(unix)]
442fn process_exists(pid: u32) -> bool {
443 extern "C" {
444 fn kill(pid: i32, sig: i32) -> i32;
445 }
446 unsafe { kill(pid as i32, 0) == 0 }
447}
448
449#[cfg(unix)]
450fn wait_for_report_children_to_exit(report: &ProcessCleanupReport, timeout: Duration) {
451 let deadline = Instant::now() + timeout;
452 while Instant::now() < deadline {
453 if report
454 .children
455 .iter()
456 .all(|child| !process_exists(child.pid))
457 {
458 return;
459 }
460 std::thread::sleep(Duration::from_millis(10));
461 }
462}
463
464pub enum ChildWait {
466 Exited(std::process::ExitStatus),
468 TimedOut(ProcessCleanupReport),
470 Interrupted(Option<std::process::ExitStatus>, ProcessCleanupReport),
474}
475
476pub fn wait_child_interruptible(
485 child: &mut std::process::Child,
486 timeout: Option<Duration>,
487) -> std::io::Result<ChildWait> {
488 wait_child_interruptible_with_cleanup_token(child, timeout, None)
489}
490
491pub fn wait_child_interruptible_with_cleanup_token(
492 child: &mut std::process::Child,
493 timeout: Option<Duration>,
494 cleanup_token: Option<&str>,
495) -> std::io::Result<ChildWait> {
496 let deadline = timeout.map(|limit| Instant::now() + limit);
497 loop {
498 if let Some(status) = child.try_wait()? {
499 return Ok(ChildWait::Exited(status));
500 }
501 if requested() {
502 let (status, report) =
503 terminate_child_group_with_cleanup_token_report(child, cleanup_token);
504 return Ok(ChildWait::Interrupted(status, report));
505 }
506 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
507 let mut report = child_pid(child)
509 .map(|pid| signal_pid_tree_group_and_token_with_report(pid, cleanup_token, 9))
510 .unwrap_or_default();
511 let _ = child.kill();
512 let _ = child.wait();
513 report.refresh_survivor_status();
514 return Ok(ChildWait::TimedOut(report));
515 }
516 std::thread::sleep(Duration::from_millis(20));
517 }
518}
519
520pub fn terminate_child_group(child: &mut std::process::Child) -> Option<std::process::ExitStatus> {
526 terminate_child_group_with_report(child).0
527}
528
529pub fn terminate_child_group_with_report(
532 child: &mut std::process::Child,
533) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
534 terminate_child_group_with_cleanup_token_report(child, None)
535}
536
537pub fn terminate_child_group_with_cleanup_token_report(
538 child: &mut std::process::Child,
539 cleanup_token: Option<&str>,
540) -> (Option<std::process::ExitStatus>, ProcessCleanupReport) {
541 let mut report = child_pid(child)
542 .map(|pid| ProcessCleanupReport::for_signal(Some(pid), 15))
543 .unwrap_or_default();
544 #[cfg(not(unix))]
545 let _ = cleanup_token;
546 #[cfg(unix)]
547 {
548 if let Some(pid) = child_pid(child) {
549 const SIGTERM: i32 = 15;
550 report = signal_pid_tree_group_and_token_with_report(pid, cleanup_token, SIGTERM);
551 let grace_deadline = Instant::now() + SUBPROCESS_TERM_GRACE;
552 loop {
553 match child.try_wait() {
554 Ok(Some(status)) => {
555 report.merge(signal_pid_tree_group_and_token_with_report(
558 pid,
559 cleanup_token,
560 9,
561 ));
562 report.refresh_survivor_status();
563 return (Some(status), report);
564 }
565 Ok(None) => {
566 if Instant::now() >= grace_deadline {
567 break;
568 }
569 std::thread::sleep(Duration::from_millis(20));
570 }
571 Err(_) => break,
572 }
573 }
574 report.merge(signal_pid_tree_group_and_token_with_report(
575 pid,
576 cleanup_token,
577 9,
578 ));
579 }
580 }
581 let _ = child.kill();
582 let status = child.wait().ok();
583 report.refresh_survivor_status();
584 (status, report)
585}
586
587fn child_pid(child: &std::process::Child) -> Option<u32> {
588 let pid = child.id();
589 (pid > 0).then_some(pid)
590}
591
592pub(crate) fn drain_captured_pipe(
602 rx: &std::sync::mpsc::Receiver<Vec<u8>>,
603 killed: bool,
604 child_pid: u32,
605) -> Vec<u8> {
606 use std::sync::mpsc::RecvTimeoutError;
607 if killed {
608 return rx
609 .recv_timeout(Duration::from_millis(100))
610 .unwrap_or_default();
611 }
612 loop {
613 match rx.recv_timeout(Duration::from_millis(20)) {
614 Ok(buf) => return buf,
615 Err(RecvTimeoutError::Disconnected) => return Vec::new(),
616 Err(RecvTimeoutError::Timeout) => {
617 if requested() {
618 const SIGTERM: i32 = 15;
619 signal_pid_tree_and_group(child_pid, SIGTERM);
620 if let Ok(buf) = rx.recv_timeout(SUBPROCESS_TERM_GRACE) {
621 signal_pid_tree_and_group(child_pid, 9);
622 return buf;
623 }
624 signal_pid_tree_and_group(child_pid, 9);
625 return rx
626 .recv_timeout(Duration::from_millis(100))
627 .unwrap_or_default();
628 }
629 }
630 }
631 }
632}
633
634pub(crate) fn spawn_pipe_drain<R: std::io::Read + Send + 'static>(
636 mut reader: R,
637) -> std::sync::mpsc::Receiver<Vec<u8>> {
638 let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
639 std::thread::spawn(move || {
640 let mut buf = Vec::new();
641 let _ = reader.read_to_end(&mut buf);
642 let _ = tx.send(buf);
643 });
644 rx
645}
646
647pub fn capture_output_interruptible(
654 command: &mut std::process::Command,
655) -> std::io::Result<std::process::Output> {
656 use std::process::Stdio;
657 command
658 .stdout(Stdio::piped())
659 .stderr(Stdio::piped())
660 .stdin(Stdio::null());
661 configure_kill_group(command);
662 let cleanup_token = new_process_cleanup_token();
663 command.env(PROCESS_CLEANUP_TOKEN_ENV, &cleanup_token);
664 let mut child = command.spawn()?;
665 let pid = child.id();
666 let rx_out = child.stdout.take().map(spawn_pipe_drain);
667 let rx_err = child.stderr.take().map(spawn_pipe_drain);
668
669 let (status, killed) = match wait_child_interruptible_with_cleanup_token(
670 &mut child,
671 None,
672 Some(&cleanup_token),
673 )? {
674 ChildWait::Exited(status) => (status, false),
675 ChildWait::TimedOut(_) => (std::process::ExitStatus::default(), true),
677 ChildWait::Interrupted(status, _) => (status.unwrap_or_default(), true),
678 };
679 let stdout = rx_out
680 .map(|rx| drain_captured_pipe(&rx, killed, pid))
681 .unwrap_or_default();
682 let stderr = rx_err
683 .map(|rx| drain_captured_pipe(&rx, killed, pid))
684 .unwrap_or_default();
685 Ok(std::process::Output {
686 status,
687 stdout,
688 stderr,
689 })
690}
691
692#[cfg(test)]
693mod tests {
694 use super::*;
695
696 #[test]
697 fn requested_is_false_without_context() {
698 assert!(!requested());
699 }
700
701 #[test]
702 fn installed_tracks_guard_lifetime() {
703 assert!(!installed());
704 let guard = install(None, None);
705 assert!(installed());
706 drop(guard);
707 assert!(!installed());
708 }
709
710 #[test]
711 fn cancel_token_trips_requested_and_guard_restores() {
712 let token = Arc::new(AtomicBool::new(false));
713 let guard = install(Some(token.clone()), None);
714 assert!(!requested());
715 token.store(true, Ordering::SeqCst);
716 assert!(requested());
717 drop(guard);
718 assert!(!requested());
719 }
720
721 #[test]
722 fn deadline_trips_requested() {
723 let expired = Instant::now()
724 .checked_sub(Duration::from_millis(1))
725 .expect("monotonic clock supports a 1ms test lookback");
726 let _guard = install(None, Some(expired));
727 assert!(requested());
728 }
729
730 #[test]
731 fn nested_installs_restore_in_order() {
732 let outer_token = Arc::new(AtomicBool::new(true));
733 let _outer = install(Some(outer_token), None);
734 assert!(requested());
735 {
736 let _inner = install(None, None);
737 assert!(!requested());
738 }
739 assert!(requested());
740 }
741
742 #[cfg(unix)]
743 #[test]
744 fn descendant_pids_from_parent_edges_returns_deepest_first_tree_only() {
745 let edges = [
746 (20, 10),
747 (30, 20),
748 (40, 20),
749 (50, 30),
750 (60, 99),
751 (70, 60),
752 (80, 90),
754 (90, 80),
755 ];
756
757 assert_eq!(
758 descendant_pids_from_parent_edges(10, &edges),
759 vec![50, 30, 40, 20]
760 );
761 assert_eq!(descendant_pids_from_parent_edges(99, &edges), vec![70, 60]);
762 assert_eq!(
763 descendant_pids_from_parent_edges(123, &edges),
764 Vec::<u32>::new()
765 );
766 }
767
768 #[cfg(unix)]
769 #[test]
770 fn descendant_processes_preserve_metadata_and_depth_order() {
771 let rows = [
772 (20, 10, Some("worker".to_string())),
773 (30, 20, Some("grandchild".to_string())),
774 (40, 20, None),
775 (50, 30, Some("leaf".to_string())),
776 ];
777
778 let descendants = descendant_processes_from_parent_edges(10, &rows);
779 let pids = descendants
780 .iter()
781 .map(|child| {
782 (
783 child.pid,
784 child.parent_pid,
785 child.depth,
786 child.command_name.as_deref(),
787 )
788 })
789 .collect::<Vec<_>>();
790 assert_eq!(
791 pids,
792 vec![
793 (50, Some(30), 3, Some("leaf")),
794 (30, Some(20), 2, Some("grandchild")),
795 (40, Some(20), 2, None),
796 (20, Some(10), 1, Some("worker")),
797 ]
798 );
799 }
800
801 #[cfg(unix)]
802 #[test]
803 fn command_name_keeps_only_argv0_basename() {
804 let command = vec![
805 std::ffi::OsString::from("/usr/local/bin/tool"),
806 std::ffi::OsString::from("--api-key"),
807 std::ffi::OsString::from("secret-value"),
808 std::ffi::OsString::from("plain"),
809 ];
810
811 assert_eq!(command_name(&command).as_deref(), Some("tool"));
812 assert_eq!(command_name(&[]).as_deref(), None);
813 }
814
815 #[cfg(unix)]
816 #[test]
817 fn process_has_cleanup_token_requires_exact_marker_entry() {
818 let token = "tok-123";
819 let env = vec![
820 std::ffi::OsString::from("PATH=/usr/bin"),
821 std::ffi::OsString::from(format!("{PROCESS_CLEANUP_TOKEN_ENV}={token}")),
822 ];
823 assert!(process_has_cleanup_token(&env, token));
824 assert!(!process_has_cleanup_token(&env, "tok"));
825 assert!(!process_has_cleanup_token(
826 &[std::ffi::OsString::from("OTHER=tok-123")],
827 token
828 ));
829 }
830
831 #[cfg(unix)]
832 #[test]
833 fn interrupted_wait_kills_process_group() {
834 let mut command = std::process::Command::new("sh");
836 command.args(["-c", "sleep 30 & wait"]);
837 configure_kill_group(&mut command);
838 let mut child = command.spawn().expect("spawn sh");
839 let pgid = child.id();
840
841 let cancel = Arc::new(AtomicBool::new(true));
842 let _guard = install(Some(cancel), None);
843 let started = Instant::now();
844 let outcome = wait_child_interruptible(&mut child, None).expect("wait");
845 assert!(matches!(outcome, ChildWait::Interrupted(_, _)));
846 assert!(started.elapsed() < Duration::from_secs(10));
847
848 extern "C" {
850 fn kill(pid: i32, sig: i32) -> i32;
851 }
852 let group_gone = || unsafe { kill(-(pgid as i32), 0) } != 0;
853 let deadline = Instant::now() + Duration::from_secs(5);
854 while !group_gone() && Instant::now() < deadline {
855 std::thread::sleep(Duration::from_millis(50));
856 }
857 assert!(group_gone(), "process group {pgid} survived interrupt");
858 }
859}