use crate::traits::ProcessStore;
use crate::tree::{ExitedProcess, ProcEvent, ProcessLink};
use crate::types::ProcessInfo;
const UNKNOWN: &str = "unknown";
pub fn snapshot(store: &impl ProcessStore) -> Result<(), std::io::Error> {
let dir = std::fs::read_dir("/proc")?;
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);
}
}
Ok(())
}
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::new(
String::new(),
String::new(),
String::new(),
*parent_pid,
0,
0,
),
);
}
ProcEvent::Exec { pid, .. } => {
let info = crate::proc::parse_proc_entry(*pid).unwrap_or_else(|| {
ProcessInfo::new(
UNKNOWN.to_string(),
UNKNOWN.to_string(),
UNKNOWN.to_string(),
0,
0,
0,
)
});
store.insert_process(*pid, info);
}
ProcEvent::Exit { pid } => {
store.for_each_child(*pid, &mut |child_pid| {
if let Some(mut info) = store.get_process(child_pid) {
info = ProcessInfo::new(
info.comm().to_string(),
info.cmd().to_string(),
info.user().to_string(),
1,
info.tgid(),
info.start_time_ns(),
);
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.comm() == 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::new(
current,
info.comm().to_string(),
info.cmd().to_string(),
info.user().to_string(),
));
if info.ppid() == 0 || current == info.ppid() {
break;
}
current = info.ppid();
}
None => {
parts.push(ProcessLink::new(
current,
UNKNOWN.to_string(),
UNKNOWN.to_string(),
UNKNOWN.to_string(),
));
break;
}
}
}
parts
}
pub fn build_chain_string(store: &impl ProcessStore, pid: u32) -> String {
let links = build_chain_links(store, pid);
if links.is_empty() {
return String::new();
}
serde_json::to_string(&links).unwrap_or_default()
}
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.comm().to_string())
.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().to_string())
.or_else(|| crate::proc::parse_proc_entry(pid).map(|info| info.user().to_string()))
},
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 mut lines = sub.lines();
if let Some(first) = lines.next() {
if is_root && i == 0 {
output.push('─');
output.push_str(first);
} else {
output.push('\n');
output.push_str(prefix);
output.push_str(first);
}
for line in lines {
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().to_string())
.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::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
);
assert_eq!(display(&store, 1), "init");
}
#[test]
fn display_root_with_children() {
let store = DefaultStore::new(0);
store.insert_process(
1,
ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
);
store.insert_process(
100,
ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0),
);
store.insert_process(
200,
ProcessInfo::new("b".into(), "b".into(), "root".into(), 1, 200, 0),
);
let d = display(&store, 1);
assert!(d.starts_with("init"));
assert!(d.contains("a"));
assert!(d.contains("b"));
}
}