proc_tree/ops.rs
1//! Process tree operations: snapshot, resolve, queries, display.
2//!
3//! All functions are generic over [`TreeStore`] and [`CacheStore`] so they
4//! work with any storage backend.
5
6use crate::cache::ProcInfo;
7use crate::traits::{CacheStore, TreeStore};
8use crate::tree::{ProcEvent, ProcessLink};
9use crate::types::PidNode;
10
11/// Snapshot all running processes from `/proc`.
12///
13/// Populates both the tree and cache. Call once at startup before
14/// processing events.
15///
16/// ```no_run
17/// use proc_tree::{DefaultTree, DefaultCache, snapshot, TreeStore};
18///
19/// let tree = DefaultTree::new(65536, 600);
20/// let cache = DefaultCache::new(65536, 600);
21/// snapshot(&tree, &cache);
22///
23/// // PID 1 should always exist on Linux
24/// assert!(tree.get_node(1).is_some());
25/// ```
26pub fn snapshot(tree: &impl TreeStore, cache: &impl CacheStore) {
27 let dir = match std::fs::read_dir("/proc") {
28 Ok(d) => d,
29 Err(e) => {
30 eprintln!("[WARNING] proc-tree: cannot read /proc: {e}");
31 return;
32 }
33 };
34 for entry in dir.flatten() {
35 let name = entry.file_name();
36 let name_str = name.to_string_lossy();
37 let pid: u32 = match name_str.parse() {
38 Ok(p) => p,
39 Err(_) => continue,
40 };
41 if let Some((node, info)) = crate::proc::parse_proc_entry(pid) {
42 tree.insert_node(pid, node);
43 cache.insert_info(pid, info);
44 }
45 }
46}
47
48/// Resolve a PID to its process info.
49///
50/// Checks the cache first, then falls back to reading `/proc` directly.
51///
52/// ```no_run
53/// use proc_tree::{DefaultTree, DefaultCache, snapshot, resolve, TreeStore};
54///
55/// let tree = DefaultTree::new(65536, 600);
56/// let cache = DefaultCache::new(65536, 600);
57/// snapshot(&tree, &cache);
58///
59/// let info = resolve(&cache, 1).unwrap();
60/// assert!(!info.cmd.is_empty());
61/// ```
62pub fn resolve(cache: &impl CacheStore, pid: u32) -> Option<ProcInfo> {
63 // Try cache first
64 if let Some(info) = cache.get_info(pid) {
65 return Some(info);
66 }
67 // Fallback: read /proc directly
68 let cmd = crate::proc::read_proc_comm(pid)?;
69 let (user, ppid, tgid) =
70 crate::proc::read_proc_status_fields(pid).unwrap_or_else(|| ("unknown".to_string(), 0, 0));
71 let start_time_ns = crate::proc::read_proc_start_time_ns(pid);
72 let info = ProcInfo {
73 cmd,
74 user,
75 ppid,
76 tgid,
77 start_time_ns,
78 };
79 // Populate cache for future lookups
80 cache.insert_info(pid, info.clone());
81 Some(info)
82}
83
84/// Handle a batch of process lifecycle events.
85///
86/// ```
87/// use proc_tree::{DefaultTree, DefaultCache, handle_events, ProcEvent, CacheStore, TreeStore};
88///
89/// let tree = DefaultTree::new(100, 0);
90/// let cache = DefaultCache::new(100, 0);
91///
92/// handle_events(&tree, &cache, &[
93/// ProcEvent::Fork { child_pid: 200, parent_pid: 100, timestamp_ns: 0 },
94/// ]);
95///
96/// let node = tree.get_node(200).unwrap();
97/// assert_eq!(node.ppid, 100);
98/// ```
99pub fn handle_events(tree: &impl TreeStore, cache: &impl CacheStore, events: &[ProcEvent]) {
100 for event in events {
101 handle_event(tree, cache, event);
102 }
103}
104
105/// Handle a single process lifecycle event.
106pub fn handle_event(tree: &impl TreeStore, cache: &impl CacheStore, event: &ProcEvent) {
107 match event {
108 ProcEvent::Fork {
109 child_pid,
110 parent_pid,
111 ..
112 } => {
113 tree.insert_node(
114 *child_pid,
115 PidNode {
116 ppid: *parent_pid,
117 cmd: String::new(),
118 },
119 );
120 }
121 ProcEvent::Exec { pid, timestamp_ns } => {
122 let cmd = crate::proc::read_proc_comm(*pid).unwrap_or_else(|| "unknown".to_string());
123 let (user, ppid, tgid) = crate::proc::read_proc_status_fields(*pid)
124 .unwrap_or_else(|| ("unknown".to_string(), 0, 0));
125 tree.insert_node(
126 *pid,
127 PidNode {
128 ppid,
129 cmd: cmd.clone(),
130 },
131 );
132 cache.insert_info(
133 *pid,
134 ProcInfo {
135 cmd,
136 user,
137 ppid,
138 tgid,
139 start_time_ns: *timestamp_ns,
140 },
141 );
142 }
143 ProcEvent::Exit { .. } => {
144 // Keep the node — still valid for historical chain lookups
145 }
146 }
147}
148
149/// Check if `pid` is a descendant of any process whose cmd == `target_cmd`.
150///
151/// ```
152/// use proc_tree::{DefaultTree, is_descendant, TreeStore, PidNode};
153///
154/// let tree = DefaultTree::new(100, 0);
155/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
156/// tree.insert_node(100, PidNode { ppid: 1, cmd: "sshd".into() });
157/// tree.insert_node(200, PidNode { ppid: 100, cmd: "bash".into() });
158///
159/// assert!(is_descendant(&tree, 200, "sshd"));
160/// assert!(is_descendant(&tree, 200, "init"));
161/// assert!(!is_descendant(&tree, 200, "nginx"));
162/// assert!(!is_descendant(&tree, 1, "sshd")); // init is not a descendant of sshd
163/// ```
164pub fn is_descendant(tree: &impl TreeStore, pid: u32, target_cmd: &str) -> bool {
165 let mut current = pid;
166 let mut visited = std::collections::HashSet::new();
167 while let Some(node) = tree.get_node(current) {
168 if !visited.insert(current) {
169 break;
170 }
171 if node.cmd == target_cmd {
172 return true;
173 }
174 if node.ppid == 0 || current == node.ppid {
175 break;
176 }
177 current = node.ppid;
178 }
179 false
180}
181
182/// Build a chain of ProcessLink from the process tree.
183pub fn build_chain_links(
184 tree: &impl TreeStore,
185 cache: &impl CacheStore,
186 pid: u32,
187) -> Vec<ProcessLink> {
188 let mut parts = Vec::new();
189 let mut current = pid;
190 let mut visited = std::collections::HashSet::new();
191 loop {
192 if !visited.insert(current) {
193 break;
194 }
195 let (ppid, cmd, user) = if let Some(node) = tree.get_node(current) {
196 let user = cache
197 .get_info(current)
198 .map(|info| info.user)
199 .unwrap_or_else(|| "unknown".to_string());
200 (node.ppid, node.cmd, user)
201 } else {
202 match crate::proc::read_proc_status_fields(current) {
203 Some((u, p, _)) => {
204 let c = crate::proc::read_proc_comm(current)
205 .unwrap_or_else(|| "unknown".to_string());
206 (p, c, u)
207 }
208 None => {
209 parts.push(ProcessLink {
210 pid: current,
211 cmd: "unknown".to_string(),
212 user: "unknown".to_string(),
213 });
214 break;
215 }
216 }
217 };
218 parts.push(ProcessLink {
219 pid: current,
220 cmd,
221 user,
222 });
223 if ppid == 0 || current == ppid {
224 break;
225 }
226 current = ppid;
227 }
228 parts
229}
230
231/// Build a chain string from the process tree.
232///
233/// Format: `"102|touch|root;101|sh|root;100|openclaw|root;1|systemd|root"`
234///
235/// ```
236/// use proc_tree::{DefaultTree, DefaultCache, build_chain_string, TreeStore, CacheStore, PidNode, ProcInfo};
237///
238/// let tree = DefaultTree::new(100, 0);
239/// let cache = DefaultCache::new(100, 0);
240///
241/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
242/// cache.insert_info(1, ProcInfo { cmd: "init".into(), user: "root".into(), ppid: 0, tgid: 1, start_time_ns: 0 });
243/// tree.insert_node(100, PidNode { ppid: 1, cmd: "sshd".into() });
244/// cache.insert_info(100, ProcInfo { cmd: "sshd".into(), user: "root".into(), ppid: 1, tgid: 100, start_time_ns: 0 });
245/// tree.insert_node(200, PidNode { ppid: 100, cmd: "bash".into() });
246/// cache.insert_info(200, ProcInfo { cmd: "bash".into(), user: "root".into(), ppid: 100, tgid: 200, start_time_ns: 0 });
247///
248/// let chain = build_chain_string(&tree, &cache, 200);
249/// assert_eq!(chain, "200|bash|root;100|sshd|root;1|init|root");
250/// ```
251pub fn build_chain_string(tree: &impl TreeStore, cache: &impl CacheStore, pid: u32) -> String {
252 build_chain_links(tree, cache, pid)
253 .iter()
254 .map(|l| l.to_string())
255 .collect::<Vec<_>>()
256 .join(";")
257}
258
259/// Get direct children of a PID.
260///
261/// ```
262/// use proc_tree::{DefaultTree, children, TreeStore, PidNode};
263///
264/// let tree = DefaultTree::new(100, 0);
265/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
266/// tree.insert_node(100, PidNode { ppid: 1, cmd: "a".into() });
267/// tree.insert_node(200, PidNode { ppid: 1, cmd: "b".into() });
268/// tree.insert_node(300, PidNode { ppid: 100, cmd: "c".into() });
269///
270/// let mut kids = children(&tree, 1);
271/// kids.sort();
272/// assert_eq!(kids, vec![100, 200]);
273/// assert_eq!(children(&tree, 100), vec![300]);
274/// assert!(children(&tree, 999).is_empty());
275/// ```
276pub fn children(tree: &impl TreeStore, pid: u32) -> Vec<u32> {
277 tree.all_pids()
278 .into_iter()
279 .filter(|&p| tree.get_node(p).map(|n| n.ppid == pid).unwrap_or(false))
280 .collect()
281}
282
283/// Get all descendants of a PID (BFS traversal).
284///
285/// ```
286/// use proc_tree::{DefaultTree, descendants, TreeStore, PidNode};
287///
288/// let tree = DefaultTree::new(100, 0);
289/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
290/// tree.insert_node(100, PidNode { ppid: 1, cmd: "a".into() });
291/// tree.insert_node(200, PidNode { ppid: 100, cmd: "b".into() });
292/// tree.insert_node(300, PidNode { ppid: 200, cmd: "c".into() });
293///
294/// let mut desc = descendants(&tree, 1);
295/// desc.sort();
296/// assert_eq!(desc, vec![100, 200, 300]);
297/// assert_eq!(descendants(&tree, 300), Vec::<u32>::new());
298/// ```
299pub fn descendants(tree: &impl TreeStore, pid: u32) -> Vec<u32> {
300 let mut result = Vec::new();
301 let mut queue = std::collections::VecDeque::new();
302 queue.push_back(pid);
303 while let Some(current) = queue.pop_front() {
304 let kids = children(tree, current);
305 for kid in kids {
306 result.push(kid);
307 queue.push_back(kid);
308 }
309 }
310 result
311}
312
313/// Get siblings of a PID (processes with the same parent).
314///
315/// Excludes the given pid itself.
316///
317/// ```
318/// use proc_tree::{DefaultTree, siblings, TreeStore, PidNode};
319///
320/// let tree = DefaultTree::new(100, 0);
321/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
322/// tree.insert_node(100, PidNode { ppid: 1, cmd: "a".into() });
323/// tree.insert_node(200, PidNode { ppid: 1, cmd: "b".into() });
324/// tree.insert_node(300, PidNode { ppid: 1, cmd: "c".into() });
325///
326/// let mut sibs = siblings(&tree, 100);
327/// sibs.sort();
328/// assert_eq!(sibs, vec![200, 300]);
329/// assert!(siblings(&tree, 1).is_empty()); // init has no siblings
330/// ```
331pub fn siblings(tree: &impl TreeStore, pid: u32) -> Vec<u32> {
332 let ppid = match tree.get_node(pid) {
333 Some(node) => node.ppid,
334 None => return Vec::new(),
335 };
336 children(tree, ppid)
337 .into_iter()
338 .filter(|&c| c != pid)
339 .collect()
340}
341
342/// Find all PIDs whose cmd matches the given string.
343///
344/// ```
345/// use proc_tree::{DefaultTree, find_by_cmd, TreeStore, PidNode};
346///
347/// let tree = DefaultTree::new(100, 0);
348/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
349/// tree.insert_node(100, PidNode { ppid: 1, cmd: "sshd".into() });
350/// tree.insert_node(200, PidNode { ppid: 1, cmd: "sshd".into() });
351/// tree.insert_node(300, PidNode { ppid: 1, cmd: "bash".into() });
352///
353/// let mut sshds = find_by_cmd(&tree, "sshd");
354/// sshds.sort();
355/// assert_eq!(sshds, vec![100, 200]);
356/// assert_eq!(find_by_cmd(&tree, "nginx"), Vec::<u32>::new());
357/// ```
358pub fn find_by_cmd(tree: &impl TreeStore, target_cmd: &str) -> Vec<u32> {
359 tree.all_pids()
360 .into_iter()
361 .filter(|&pid| {
362 let cmd = tree
363 .get_node(pid)
364 .map(|n| n.cmd)
365 .filter(|c| !c.is_empty())
366 .or_else(|| crate::proc::read_proc_comm(pid));
367 cmd.as_deref() == Some(target_cmd)
368 })
369 .collect()
370}
371
372/// Find all PIDs whose user matches the given string.
373///
374/// ```
375/// use proc_tree::{DefaultTree, DefaultCache, find_by_user, TreeStore, CacheStore, PidNode, ProcInfo};
376///
377/// let tree = DefaultTree::new(100, 0);
378/// let cache = DefaultCache::new(100, 0);
379///
380/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
381/// cache.insert_info(1, ProcInfo { cmd: "init".into(), user: "root".into(), ppid: 0, tgid: 1, start_time_ns: 0 });
382/// tree.insert_node(100, PidNode { ppid: 1, cmd: "bash".into() });
383/// cache.insert_info(100, ProcInfo { cmd: "bash".into(), user: "alice".into(), ppid: 1, tgid: 100, start_time_ns: 0 });
384///
385/// assert_eq!(find_by_user(&tree, &cache, "root"), vec![1]);
386/// assert_eq!(find_by_user(&tree, &cache, "alice"), vec![100]);
387/// assert_eq!(find_by_user(&tree, &cache, "nobody"), Vec::<u32>::new());
388/// ```
389pub fn find_by_user(tree: &impl TreeStore, cache: &impl CacheStore, target_user: &str) -> Vec<u32> {
390 tree.all_pids()
391 .into_iter()
392 .filter(|&pid| {
393 let user = cache
394 .get_info(pid)
395 .map(|info| info.user)
396 .or_else(|| crate::proc::read_proc_status_fields(pid).map(|(u, _, _)| u));
397 user.as_deref() == Some(target_user)
398 })
399 .collect()
400}
401
402/// Render a pstree-style display starting from the given root PID.
403///
404/// ```
405/// use proc_tree::{DefaultTree, display, TreeStore, PidNode};
406///
407/// let tree = DefaultTree::new(100, 0);
408/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
409/// tree.insert_node(100, PidNode { ppid: 1, cmd: "sshd".into() });
410/// tree.insert_node(200, PidNode { ppid: 1, cmd: "cron".into() });
411///
412/// let output = display(&tree, 1);
413/// assert!(output.starts_with("init"));
414/// assert!(output.contains("sshd"));
415/// assert!(output.contains("cron"));
416/// ```
417pub fn display(tree: &impl TreeStore, root_pid: u32) -> String {
418 display_inner(tree, root_pid, true)
419}
420
421/// Recursive display helper.
422///
423/// `is_root` controls whether the first child is attached with `"─"` (root
424/// case) or with the standard tree prefix (non-root case).
425fn display_inner(tree: &impl TreeStore, pid: u32, is_root: bool) -> String {
426 let cmd = tree
427 .get_node(pid)
428 .map(|n| n.cmd)
429 .filter(|c| !c.is_empty())
430 .or_else(|| crate::proc::read_proc_comm(pid))
431 .unwrap_or_else(|| "unknown".to_string());
432 let kids = children(tree, pid);
433 if kids.is_empty() {
434 return cmd;
435 }
436 let mut output = cmd;
437 for (i, &kid) in kids.iter().enumerate() {
438 let is_last = i == kids.len() - 1;
439 let prefix = if is_last { "└─" } else { "├─" };
440 let continuation = if is_last { " " } else { "│ " };
441 let sub = display_inner(tree, kid, false);
442 let lines: Vec<&str> = sub.lines().collect();
443 if i == 0 && is_root {
444 output.push_str(&format!("─{}", lines[0]));
445 } else {
446 if i > 0 || is_root {
447 output.push('\n');
448 }
449 output.push_str(prefix);
450 output.push_str(lines[0]);
451 }
452 for line in &lines[1..] {
453 output.push('\n');
454 output.push_str(continuation);
455 output.push_str(line);
456 }
457 }
458 output
459}
460
461/// Get the number of entries in the tree.
462///
463/// ```
464/// use proc_tree::{DefaultTree, tree_len, TreeStore, PidNode};
465///
466/// let tree = DefaultTree::new(100, 0);
467/// assert_eq!(tree_len(&tree), 0);
468///
469/// tree.insert_node(1, PidNode { ppid: 0, cmd: "init".into() });
470/// tree.insert_node(2, PidNode { ppid: 1, cmd: "bash".into() });
471/// assert_eq!(tree_len(&tree), 2);
472/// ```
473pub fn tree_len(tree: &impl TreeStore) -> u64 {
474 tree.all_pids().len() as u64
475}
476
477#[cfg(test)]
478mod tests {
479 use super::*;
480 use crate::default_store::DefaultTree;
481
482 #[test]
483 fn display_single_node() {
484 let tree = DefaultTree::new(100, 0);
485 tree.insert_node(
486 1,
487 PidNode {
488 ppid: 0,
489 cmd: "init".into(),
490 },
491 );
492 assert_eq!(display(&tree, 1), "init");
493 }
494
495 #[test]
496 fn display_root_with_children() {
497 let tree = DefaultTree::new(100, 0);
498 tree.insert_node(
499 1,
500 PidNode {
501 ppid: 0,
502 cmd: "init".into(),
503 },
504 );
505 tree.insert_node(
506 100,
507 PidNode {
508 ppid: 1,
509 cmd: "a".into(),
510 },
511 );
512 tree.insert_node(
513 200,
514 PidNode {
515 ppid: 1,
516 cmd: "b".into(),
517 },
518 );
519 let d = display(&tree, 1);
520 assert!(d.starts_with("init"));
521 assert!(d.contains("a"));
522 assert!(d.contains("b"));
523 }
524}