proc_tree/
ops.rs

1//! Process tree operations: snapshot, resolve, queries, display.
2//!
3//! All functions are generic over [`ProcessStore`] so they work with any storage backend.
4
5use crate::traits::ProcessStore;
6use crate::tree::{ExitedProcess, ProcEvent, ProcessLink};
7use crate::types::ProcessInfo;
8
9const UNKNOWN: &str = "unknown";
10
11/// Snapshot all running processes from `/proc`.
12///
13/// Populates the store. Call once at startup before processing events.
14///
15/// # Errors
16///
17/// Returns an error if `/proc` cannot be read.
18///
19/// # Example
20///
21/// ```no_run
22/// use proc_tree::{DefaultStore, snapshot, ProcessStore};
23///
24/// let store = DefaultStore::new(600);
25/// snapshot(&store).expect("failed to read /proc");
26///
27/// // PID 1 should always exist on Linux
28/// assert!(store.get_process(1).is_some());
29/// ```
30pub fn snapshot(store: &impl ProcessStore) -> Result<(), std::io::Error> {
31    let dir = std::fs::read_dir("/proc")?;
32    for entry in dir.flatten() {
33        let name = entry.file_name();
34        let name_str = name.to_string_lossy();
35        let pid: u32 = match name_str.parse() {
36            Ok(p) => p,
37            Err(_) => continue,
38        };
39        if let Some(info) = crate::proc::parse_proc_entry(pid) {
40            store.insert_process(pid, info);
41        }
42    }
43    Ok(())
44}
45
46/// Resolve a PID to its process info.
47///
48/// Checks the store first, then falls back to reading `/proc` directly.
49///
50/// ```no_run
51/// use proc_tree::{DefaultStore, snapshot, resolve, ProcessStore};
52///
53/// let store = DefaultStore::new(600);
54/// snapshot(&store).expect("failed to read /proc");
55///
56/// let info = resolve(&store, 1).unwrap();
57/// assert!(!info.cmd().is_empty());
58/// ```
59pub fn resolve(store: &impl ProcessStore, pid: u32) -> Option<ProcessInfo> {
60    let info = resolve_process_info(store, pid);
61    // Populate store for future lookups if we got info from /proc
62    if let Some(ref info) = info
63        && store.get_process(pid).is_none()
64    {
65        store.insert_process(pid, info.clone());
66    }
67    info
68}
69
70/// Internal helper to resolve process info with fallback chain.
71///
72/// Tries store first, then falls back to reading `/proc` directly.
73/// Returns `None` if the process doesn't exist.
74fn resolve_process_info(store: &impl ProcessStore, pid: u32) -> Option<ProcessInfo> {
75    store
76        .get_process(pid)
77        .or_else(|| crate::proc::parse_proc_entry(pid))
78}
79
80/// Handle a batch of process lifecycle events.
81///
82/// Returns [`ExitedProcess`] handles for Exit events. The process info **stays in the store**
83/// — callers decide when to remove it via [`ExitedProcess::remove`].
84///
85/// This is critical for event-driven systems where file events (fanotify)
86/// may arrive after the proc connector exit event but still need to look
87/// up process info.
88///
89/// # Example
90///
91/// ```
92/// use proc_tree::{DefaultStore, handle_events, ProcEvent, ProcessStore};
93///
94/// let store = DefaultStore::new(0);
95///
96/// // Fork creates a process
97/// let exited = handle_events(&store, &[
98///     ProcEvent::Fork { child_pid: 200, parent_pid: 100, timestamp_ns: 0 },
99/// ]);
100/// assert!(exited.is_empty());
101///
102/// // Exit returns ExitedProcess handle — process info stays in store
103/// let exited = handle_events(&store, &[
104///     ProcEvent::Exit { pid: 200 },
105/// ]);
106/// assert_eq!(exited[0].pid, 200);
107/// assert!(store.get_process(200).is_some());
108///
109/// // Caller explicitly removes when done
110/// for ep in exited {
111///     ep.remove(&store);
112/// }
113/// assert!(store.get_process(200).is_none());
114/// ```
115#[must_use = "call .remove(&store) on each ExitedProcess after processing related events"]
116pub fn handle_events(store: &impl ProcessStore, events: &[ProcEvent]) -> Vec<ExitedProcess> {
117    let mut exited = Vec::new();
118    for event in events {
119        if let Some(ep) = handle_event(store, event) {
120            exited.push(ep);
121        }
122    }
123    exited
124}
125
126/// Handle a single process lifecycle event.
127///
128/// Returns [`ExitedProcess`] for Exit events, `None` for other events.
129/// The process info **stays in the store** — callers decide when to remove it
130/// via [`ExitedProcess::remove`].
131///
132/// # Example
133///
134/// ```
135/// use proc_tree::{DefaultStore, handle_event, ProcEvent, ProcessStore};
136///
137/// let store = DefaultStore::new(0);
138///
139/// // Create a process
140/// handle_event(&store, &ProcEvent::Fork {
141///     child_pid: 100,
142///     parent_pid: 1,
143///     timestamp_ns: 0,
144/// });
145///
146/// // Exit returns ExitedProcess handle — process info stays in store
147/// let exited = handle_event(&store, &ProcEvent::Exit { pid: 100 }).unwrap();
148/// assert_eq!(exited.pid, 100);
149/// assert!(store.get_process(100).is_some());
150///
151/// // Caller explicitly removes when done
152/// exited.remove(&store);
153/// assert!(store.get_process(100).is_none());
154/// ```
155#[must_use = "call .remove(&store) after processing related events"]
156pub fn handle_event(store: &impl ProcessStore, event: &ProcEvent) -> Option<ExitedProcess> {
157    match event {
158        ProcEvent::Fork {
159            child_pid,
160            parent_pid,
161            ..
162        } => {
163            store.insert_process(
164                *child_pid,
165                ProcessInfo::new(
166                    String::new(),
167                    String::new(),
168                    String::new(),
169                    *parent_pid,
170                    0,
171                    0,
172                ),
173            );
174        }
175        ProcEvent::Exec { pid, .. } => {
176            let info = crate::proc::parse_proc_entry(*pid).unwrap_or_else(|| {
177                ProcessInfo::new(
178                    UNKNOWN.to_string(),
179                    UNKNOWN.to_string(),
180                    UNKNOWN.to_string(),
181                    0,
182                    0,
183                    0,
184                )
185            });
186            // Keep start_time_ns from /proc, don't overwrite with event timestamp
187            store.insert_process(*pid, info);
188        }
189        ProcEvent::Exit { pid } => {
190            // Orphan children to init (PID 1)
191            store.for_each_child(*pid, &mut |child_pid| {
192                if let Some(mut info) = store.get_process(child_pid) {
193                    info = ProcessInfo::new(
194                        info.comm().to_string(),
195                        info.cmd().to_string(),
196                        info.user().to_string(),
197                        1,
198                        info.tgid(),
199                        info.start_time_ns(),
200                    );
201                    store.insert_process(child_pid, info);
202                }
203            });
204            return Some(ExitedProcess { pid: *pid });
205        }
206    }
207    None
208}
209
210/// Check if `pid` is a descendant of any process whose cmd == `target_cmd`.
211///
212/// ```
213/// use proc_tree::{DefaultStore, is_descendant, ProcessStore, ProcessInfo};
214///
215/// let store = DefaultStore::new(0);
216/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
217/// store.insert_process(100, ProcessInfo::new("sshd".into(), "sshd".into(), "root".into(), 1, 100, 0));
218/// store.insert_process(200, ProcessInfo::new("bash".into(), "bash".into(), "root".into(), 100, 200, 0));
219///
220/// assert!(is_descendant(&store, 200, "sshd"));
221/// assert!(is_descendant(&store, 200, "init"));
222/// assert!(!is_descendant(&store, 200, "nginx"));
223/// assert!(!is_descendant(&store, 1, "sshd")); // init is not a descendant of sshd
224/// ```
225pub fn is_descendant(store: &impl ProcessStore, pid: u32, target_cmd: &str) -> bool {
226    walk_ancestors(store, pid, |info| info.comm() == target_cmd)
227}
228
229/// Walk up the process tree from `pid`, calling `predicate` on each ancestor.
230///
231/// Returns `true` if the predicate matches any ancestor, `false` otherwise.
232/// Handles cycles by tracking visited PIDs.
233fn walk_ancestors<P>(store: &impl ProcessStore, pid: u32, predicate: P) -> bool
234where
235    P: Fn(&ProcessInfo) -> bool,
236{
237    let mut current = pid;
238    let mut visited = std::collections::HashSet::new();
239    while let Some(info) = store.get_process(current) {
240        if !visited.insert(current) {
241            break;
242        }
243        if predicate(&info) {
244            return true;
245        }
246        if info.ppid() == 0 || current == info.ppid() {
247            break;
248        }
249        current = info.ppid();
250    }
251    false
252}
253
254/// Build a chain of ProcessLink from the process tree.
255pub fn build_chain_links(store: &impl ProcessStore, pid: u32) -> Vec<ProcessLink> {
256    let mut parts = Vec::new();
257    let mut current = pid;
258    let mut visited = std::collections::HashSet::new();
259    loop {
260        if !visited.insert(current) {
261            break;
262        }
263        match resolve_process_info(store, current) {
264            Some(info) => {
265                parts.push(ProcessLink::new(
266                    current,
267                    info.comm().to_string(),
268                    info.cmd().to_string(),
269                    info.user().to_string(),
270                ));
271                if info.ppid() == 0 || current == info.ppid() {
272                    break;
273                }
274                current = info.ppid();
275            }
276            None => {
277                parts.push(ProcessLink::new(
278                    current,
279                    UNKNOWN.to_string(),
280                    UNKNOWN.to_string(),
281                    UNKNOWN.to_string(),
282                ));
283                break;
284            }
285        }
286    }
287    parts
288}
289
290/// Build a chain string from the process tree.
291///
292/// Format: `"102|touch|root;101|sh|root;100|openclaw|root;1|systemd|root"`
293///
294/// ```
295/// use proc_tree::{DefaultStore, build_chain_string, ProcessStore, ProcessInfo};
296///
297/// let store = DefaultStore::new(0);
298///
299/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
300/// store.insert_process(100, ProcessInfo::new("sshd".into(), "sshd".into(), "root".into(), 1, 100, 0));
301/// store.insert_process(200, ProcessInfo::new("bash".into(), "bash".into(), "root".into(), 100, 200, 0));
302///
303/// let chain = build_chain_string(&store, 200);
304/// // Returns JSON array: [{"pid":200,"comm":"bash","cmd":"bash","user":"root"}, ...]
305/// assert!(chain.contains("\"comm\":\"bash\""));
306/// ```
307pub fn build_chain_string(store: &impl ProcessStore, pid: u32) -> String {
308    let links = build_chain_links(store, pid);
309    if links.is_empty() {
310        return String::new();
311    }
312    serde_json::to_string(&links).unwrap_or_default()
313}
314
315/// Get direct children of a PID.
316///
317/// ```
318/// use proc_tree::{DefaultStore, children, ProcessStore, ProcessInfo};
319///
320/// let store = DefaultStore::new(0);
321/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
322/// store.insert_process(100, ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0));
323/// store.insert_process(200, ProcessInfo::new("b".into(), "b".into(), "root".into(), 1, 200, 0));
324/// store.insert_process(300, ProcessInfo::new("c".into(), "c".into(), "root".into(), 100, 300, 0));
325///
326/// let mut kids = children(&store, 1);
327/// kids.sort();
328/// assert_eq!(kids, vec![100, 200]);
329/// assert_eq!(children(&store, 100), vec![300]);
330/// assert!(children(&store, 999).is_empty());
331/// ```
332pub fn children(store: &impl ProcessStore, pid: u32) -> Vec<u32> {
333    store.children_of(pid)
334}
335
336/// Get all descendants of a PID (BFS traversal).
337///
338/// ```
339/// use proc_tree::{DefaultStore, descendants, ProcessStore, ProcessInfo};
340///
341/// let store = DefaultStore::new(0);
342/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
343/// store.insert_process(100, ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0));
344/// store.insert_process(200, ProcessInfo::new("b".into(), "b".into(), "root".into(), 100, 200, 0));
345/// store.insert_process(300, ProcessInfo::new("c".into(), "c".into(), "root".into(), 200, 300, 0));
346///
347/// let mut desc = descendants(&store, 1);
348/// desc.sort();
349/// assert_eq!(desc, vec![100, 200, 300]);
350/// assert_eq!(descendants(&store, 300), Vec::<u32>::new());
351/// ```
352pub fn descendants(store: &impl ProcessStore, pid: u32) -> Vec<u32> {
353    let mut result = Vec::new();
354    let mut queue = std::collections::VecDeque::new();
355    queue.push_back(pid);
356    while let Some(current) = queue.pop_front() {
357        let kids = children(store, current);
358        for kid in kids {
359            result.push(kid);
360            queue.push_back(kid);
361        }
362    }
363    result
364}
365
366/// Get siblings of a PID (processes with the same parent).
367///
368/// Excludes the given pid itself.
369///
370/// ```
371/// use proc_tree::{DefaultStore, siblings, ProcessStore, ProcessInfo};
372///
373/// let store = DefaultStore::new(0);
374/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
375/// store.insert_process(100, ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0));
376/// store.insert_process(200, ProcessInfo::new("b".into(), "b".into(), "root".into(), 1, 200, 0));
377/// store.insert_process(300, ProcessInfo::new("c".into(), "c".into(), "root".into(), 1, 300, 0));
378///
379/// let mut sibs = siblings(&store, 100);
380/// sibs.sort();
381/// assert_eq!(sibs, vec![200, 300]);
382/// assert!(siblings(&store, 1).is_empty()); // init has no siblings
383/// ```
384pub fn siblings(store: &impl ProcessStore, pid: u32) -> Vec<u32> {
385    let ppid = match store.get_process(pid) {
386        Some(info) => info.ppid(),
387        None => return Vec::new(),
388    };
389    children(store, ppid)
390        .into_iter()
391        .filter(|&c| c != pid)
392        .collect()
393}
394
395/// Find all PIDs whose cmd matches the given string.
396///
397/// ```
398/// use proc_tree::{DefaultStore, find_by_cmd, ProcessStore, ProcessInfo};
399///
400/// let store = DefaultStore::new(0);
401/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
402/// store.insert_process(100, ProcessInfo::new("sshd".into(), "sshd".into(), "root".into(), 1, 100, 0));
403/// store.insert_process(200, ProcessInfo::new("sshd".into(), "sshd".into(), "root".into(), 1, 200, 0));
404/// store.insert_process(300, ProcessInfo::new("bash".into(), "bash".into(), "root".into(), 1, 300, 0));
405///
406/// let mut sshds = find_by_cmd(&store, "sshd");
407/// sshds.sort();
408/// assert_eq!(sshds, vec![100, 200]);
409/// assert_eq!(find_by_cmd(&store, "nginx"), Vec::<u32>::new());
410/// ```
411pub fn find_by_cmd(store: &impl ProcessStore, target_cmd: &str) -> Vec<u32> {
412    find_by(
413        store,
414        |pid| {
415            store
416                .get_process(pid)
417                .map(|info| info.comm().to_string())
418                .filter(|c| !c.is_empty())
419                .or_else(|| crate::proc::read_proc_comm(pid))
420        },
421        target_cmd,
422    )
423}
424
425/// Find all PIDs whose user matches the given string.
426///
427/// ```
428/// use proc_tree::{DefaultStore, find_by_user, ProcessStore, ProcessInfo};
429///
430/// let store = DefaultStore::new(0);
431///
432/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
433/// store.insert_process(100, ProcessInfo::new("bash".into(), "bash".into(), "alice".into(), 1, 100, 0));
434///
435/// assert_eq!(find_by_user(&store, "root"), vec![1]);
436/// assert_eq!(find_by_user(&store, "alice"), vec![100]);
437/// assert_eq!(find_by_user(&store, "nobody"), Vec::<u32>::new());
438/// ```
439pub fn find_by_user(store: &impl ProcessStore, target_user: &str) -> Vec<u32> {
440    find_by(
441        store,
442        |pid| {
443            store
444                .get_process(pid)
445                .map(|info| info.user().to_string())
446                .or_else(|| crate::proc::parse_proc_entry(pid).map(|info| info.user().to_string()))
447        },
448        target_user,
449    )
450}
451
452/// Generic find function that filters PIDs by a value extractor.
453///
454/// This is a helper to reduce code duplication between `find_by_cmd` and `find_by_user`.
455fn find_by<F>(store: &impl ProcessStore, extract: F, target: &str) -> Vec<u32>
456where
457    F: Fn(u32) -> Option<String>,
458{
459    store
460        .all_pids()
461        .into_iter()
462        .filter(|&pid| extract(pid).as_deref() == Some(target))
463        .collect()
464}
465
466/// Render a pstree-style display starting from the given root PID.
467///
468/// ```
469/// use proc_tree::{DefaultStore, display, ProcessStore, ProcessInfo};
470///
471/// let store = DefaultStore::new(0);
472/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
473/// store.insert_process(100, ProcessInfo::new("sshd".into(), "sshd".into(), "root".into(), 1, 100, 0));
474/// store.insert_process(200, ProcessInfo::new("cron".into(), "cron".into(), "root".into(), 1, 200, 0));
475///
476/// let output = display(&store, 1);
477/// assert!(output.starts_with("init"));
478/// assert!(output.contains("sshd"));
479/// assert!(output.contains("cron"));
480/// ```
481pub fn display(store: &impl ProcessStore, root_pid: u32) -> String {
482    render_tree(store, root_pid, true)
483}
484
485/// Recursive tree renderer.
486///
487/// `is_root` controls the rendering style:
488/// - Root node: first child attaches with "─"
489/// - Non-root nodes: all children use tree prefixes
490fn render_tree(store: &impl ProcessStore, pid: u32, is_root: bool) -> String {
491    let cmd = get_cmd(store, pid);
492    let kids = children(store, pid);
493    if kids.is_empty() {
494        return cmd;
495    }
496    let mut output = cmd;
497    for (i, &kid) in kids.iter().enumerate() {
498        let is_last = i == kids.len() - 1;
499        let prefix = if is_last { "└─" } else { "├─" };
500        let continuation = if is_last { "  " } else { "│ " };
501        let sub = render_tree(store, kid, false);
502        let mut lines = sub.lines();
503        if let Some(first) = lines.next() {
504            if is_root && i == 0 {
505                output.push('─');
506                output.push_str(first);
507            } else {
508                output.push('\n');
509                output.push_str(prefix);
510                output.push_str(first);
511            }
512            for line in lines {
513                output.push('\n');
514                output.push_str(continuation);
515                output.push_str(line);
516            }
517        }
518    }
519    output
520}
521
522/// Get command name for a PID, with fallback chain: store -> /proc -> "unknown"
523fn get_cmd(store: &impl ProcessStore, pid: u32) -> String {
524    resolve_process_info(store, pid)
525        .map(|info| info.cmd().to_string())
526        .filter(|c| !c.is_empty())
527        .unwrap_or_else(|| UNKNOWN.to_string())
528}
529
530/// Get the number of entries in the store.
531///
532/// ```
533/// use proc_tree::{DefaultStore, tree_len, ProcessStore, ProcessInfo};
534///
535/// let store = DefaultStore::new(0);
536/// assert_eq!(tree_len(&store), 0);
537///
538/// store.insert_process(1, ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0));
539/// store.insert_process(2, ProcessInfo::new("bash".into(), "bash".into(), "root".into(), 1, 2, 0));
540/// assert_eq!(tree_len(&store), 2);
541/// ```
542pub fn tree_len(store: &impl ProcessStore) -> usize {
543    store.all_pids().len()
544}
545
546#[cfg(test)]
547mod tests {
548    use super::*;
549    use crate::default_store::DefaultStore;
550
551    #[test]
552    fn display_single_node() {
553        let store = DefaultStore::new(0);
554        store.insert_process(
555            1,
556            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
557        );
558        assert_eq!(display(&store, 1), "init");
559    }
560
561    #[test]
562    fn display_root_with_children() {
563        let store = DefaultStore::new(0);
564        store.insert_process(
565            1,
566            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
567        );
568        store.insert_process(
569            100,
570            ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0),
571        );
572        store.insert_process(
573            200,
574            ProcessInfo::new("b".into(), "b".into(), "root".into(), 1, 200, 0),
575        );
576        let d = display(&store, 1);
577        assert!(d.starts_with("init"));
578        assert!(d.contains("a"));
579        assert!(d.contains("b"));
580    }
581}
proc_tree/ops.rs

proc_tree/
ops.rs
