Skip to main content

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}