use crate::traits::ProcessStore;
use crate::tree::{ExitedProcess, ProcEvent, ProcessLink};
use crate::types::ProcessInfo;
pub fn snapshot(store: &impl ProcessStore) {
let dir = match std::fs::read_dir("/proc") {
Ok(d) => d,
Err(e) => {
eprintln!("[WARNING] proc-tree: cannot read /proc: {e}");
return;
}
};
for entry in dir.flatten() {
let name = entry.file_name();
let name_str = name.to_string_lossy();
let pid: u32 = match name_str.parse() {
Ok(p) => p,
Err(_) => continue,
};
if let Some(info) = crate::proc::parse_proc_entry(pid) {
store.insert_process(pid, info);
}
}
}
pub fn resolve(store: &impl ProcessStore, pid: u32) -> Option<ProcessInfo> {
let info = resolve_process_info(store, pid);
if let Some(ref info) = info
&& store.get_process(pid).is_none()
{
store.insert_process(pid, info.clone());
}
info
}
fn resolve_process_info(store: &impl ProcessStore, pid: u32) -> Option<ProcessInfo> {
store
.get_process(pid)
.or_else(|| crate::proc::parse_proc_entry(pid))
}
#[must_use = "call .remove(&store) on each ExitedProcess after processing related events"]
pub fn handle_events(store: &impl ProcessStore, events: &[ProcEvent]) -> Vec<ExitedProcess> {
let mut exited = Vec::new();
for event in events {
if let Some(ep) = handle_event(store, event) {
exited.push(ep);
}
}
exited
}
#[must_use = "call .remove(&store) after processing related events"]
pub fn handle_event(store: &impl ProcessStore, event: &ProcEvent) -> Option<ExitedProcess> {
match event {
ProcEvent::Fork {
child_pid,
parent_pid,
..
} => {
store.insert_process(
*child_pid,
ProcessInfo {
ppid: *parent_pid,
cmd: String::new(),
user: String::new(),
tgid: 0,
start_time_ns: 0,
},
);
}
ProcEvent::Exec { pid, .. } => {
let info = crate::proc::parse_proc_entry(*pid).unwrap_or_else(|| {
let cmd = "unknown".to_string();
ProcessInfo {
ppid: 0,
cmd,
user: "unknown".to_string(),
tgid: 0,
start_time_ns: 0,
}
});
store.insert_process(*pid, info);
}
ProcEvent::Exit { pid } => {
let children = store.children_of(*pid);
for child_pid in children {
if let Some(mut info) = store.get_process(child_pid) {
info.ppid = 1;
store.insert_process(child_pid, info);
}
}
return Some(ExitedProcess { pid: *pid });
}
}
None
}
pub fn is_descendant(store: &impl ProcessStore, pid: u32, target_cmd: &str) -> bool {
walk_ancestors(store, pid, |info| info.cmd == target_cmd)
}
fn walk_ancestors<P>(store: &impl ProcessStore, pid: u32, predicate: P) -> bool
where
P: Fn(&ProcessInfo) -> bool,
{
let mut current = pid;
let mut visited = std::collections::HashSet::new();
while let Some(info) = store.get_process(current) {
if !visited.insert(current) {
break;
}
if predicate(&info) {
return true;
}
if info.ppid == 0 || current == info.ppid {
break;
}
current = info.ppid;
}
false
}
pub fn build_chain_links(store: &impl ProcessStore, pid: u32) -> Vec<ProcessLink> {
let mut parts = Vec::new();
let mut current = pid;
let mut visited = std::collections::HashSet::new();
loop {
if !visited.insert(current) {
break;
}
match resolve_process_info(store, current) {
Some(info) => {
parts.push(ProcessLink {
pid: current,
cmd: info.cmd,
user: info.user,
});
if info.ppid == 0 || current == info.ppid {
break;
}
current = info.ppid;
}
None => {
parts.push(ProcessLink {
pid: current,
cmd: "unknown".to_string(),
user: "unknown".to_string(),
});
break;
}
}
}
parts
}
pub fn build_chain_string(store: &impl ProcessStore, pid: u32) -> String {
build_chain_links(store, pid)
.iter()
.map(|l| l.to_string())
.collect::<Vec<_>>()
.join(";")
}
pub fn children(store: &impl ProcessStore, pid: u32) -> Vec<u32> {
store.children_of(pid)
}
pub fn descendants(store: &impl ProcessStore, pid: u32) -> Vec<u32> {
let mut result = Vec::new();
let mut queue = std::collections::VecDeque::new();
queue.push_back(pid);
while let Some(current) = queue.pop_front() {
let kids = children(store, current);
for kid in kids {
result.push(kid);
queue.push_back(kid);
}
}
result
}
pub fn siblings(store: &impl ProcessStore, pid: u32) -> Vec<u32> {
let ppid = match store.get_process(pid) {
Some(info) => info.ppid,
None => return Vec::new(),
};
children(store, ppid)
.into_iter()
.filter(|&c| c != pid)
.collect()
}
pub fn find_by_cmd(store: &impl ProcessStore, target_cmd: &str) -> Vec<u32> {
find_by(
store,
|pid| {
store
.get_process(pid)
.map(|info| info.cmd)
.filter(|c| !c.is_empty())
.or_else(|| crate::proc::read_proc_comm(pid))
},
target_cmd,
)
}
pub fn find_by_user(store: &impl ProcessStore, target_user: &str) -> Vec<u32> {
find_by(
store,
|pid| {
store
.get_process(pid)
.map(|info| info.user)
.or_else(|| crate::proc::parse_proc_entry(pid).map(|info| info.user))
},
target_user,
)
}
fn find_by<F>(store: &impl ProcessStore, extract: F, target: &str) -> Vec<u32>
where
F: Fn(u32) -> Option<String>,
{
store
.all_pids()
.into_iter()
.filter(|&pid| extract(pid).as_deref() == Some(target))
.collect()
}
pub fn display(store: &impl ProcessStore, root_pid: u32) -> String {
render_tree(store, root_pid, true)
}
fn render_tree(store: &impl ProcessStore, pid: u32, is_root: bool) -> String {
let cmd = get_cmd(store, pid);
let kids = children(store, pid);
if kids.is_empty() {
return cmd;
}
let mut output = cmd;
for (i, &kid) in kids.iter().enumerate() {
let is_last = i == kids.len() - 1;
let prefix = if is_last { "└─" } else { "├─" };
let continuation = if is_last { " " } else { "│ " };
let sub = render_tree(store, kid, false);
let lines: Vec<&str> = sub.lines().collect();
if is_root && i == 0 {
output.push_str(&format!("─{}", lines[0]));
} else {
output.push('\n');
output.push_str(prefix);
output.push_str(lines[0]);
}
for line in &lines[1..] {
output.push('\n');
output.push_str(continuation);
output.push_str(line);
}
}
output
}
fn get_cmd(store: &impl ProcessStore, pid: u32) -> String {
resolve_process_info(store, pid)
.map(|info| info.cmd)
.filter(|c| !c.is_empty())
.unwrap_or_else(|| "unknown".to_string())
}
pub fn tree_len(store: &impl ProcessStore) -> usize {
store.all_pids().len()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::default_store::DefaultStore;
#[test]
fn display_single_node() {
let store = DefaultStore::new(0);
store.insert_process(
1,
ProcessInfo {
ppid: 0,
cmd: "init".into(),
user: "root".into(),
tgid: 1,
start_time_ns: 0,
},
);
assert_eq!(display(&store, 1), "init");
}
#[test]
fn display_root_with_children() {
let store = DefaultStore::new(0);
store.insert_process(
1,
ProcessInfo {
ppid: 0,
cmd: "init".into(),
user: "root".into(),
tgid: 1,
start_time_ns: 0,
},
);
store.insert_process(
100,
ProcessInfo {
ppid: 1,
cmd: "a".into(),
user: "root".into(),
tgid: 100,
start_time_ns: 0,
},
);
store.insert_process(
200,
ProcessInfo {
ppid: 1,
cmd: "b".into(),
user: "root".into(),
tgid: 200,
start_time_ns: 0,
},
);
let d = display(&store, 1);
assert!(d.starts_with("init"));
assert!(d.contains("a"));
assert!(d.contains("b"));
}
}