agentmux_srv/backend/blockcontroller/process_tree.rs
1// Copyright 2025, Command Line Inc.
2// Copyright 2025-2026, AgentMux Corp.
3// SPDX-License-Identifier: Apache-2.0
4
5//! Process tree traversal: BFS from a root PID to collect all descendant PIDs.
6//! Used by the sysinfo loop to aggregate CPU/memory across entire process trees.
7
8use std::collections::{HashMap, VecDeque};
9
10use sysinfo::{Pid, System};
11
12/// Maximum number of PIDs to track per block (safety cap against pathological trees).
13pub const MAX_PIDS_PER_BLOCK: usize = 64;
14
15/// Returns the root PID plus all descendant PIDs via BFS, capped at `max_pids`.
16///
17/// Requires `sys` to have been refreshed with at least a minimal `ProcessRefreshKind`
18/// (so that `Process::parent()` is populated for all processes).
19pub fn collect_descendants(sys: &System, root: Pid, max_pids: usize) -> Vec<Pid> {
20 // Build parent → children adjacency map in O(N) over all known processes.
21 let mut children: HashMap<Pid, Vec<Pid>> = HashMap::new();
22 for (pid, proc) in sys.processes() {
23 if let Some(ppid) = proc.parent() {
24 children.entry(ppid).or_default().push(*pid);
25 }
26 }
27
28 // BFS from root, stopping when we hit max_pids.
29 let mut result = Vec::with_capacity(max_pids.min(8));
30 result.push(root);
31 let mut queue = VecDeque::from([root]);
32
33 while let Some(pid) = queue.pop_front() {
34 if result.len() >= max_pids {
35 break;
36 }
37 if let Some(kids) = children.get(&pid) {
38 for &child in kids {
39 if result.len() >= max_pids {
40 break;
41 }
42 result.push(child);
43 queue.push_back(child);
44 }
45 }
46 }
47
48 result
49}