prock/

lib.rs

//! Fast, low-overhead CPU statistics for process trees.
//!
//! This crate provides multiple approaches for measuring CPU usage:
//!
//! 1. **Snapshot-based**: Get raw CPU time at a point in time
//! 2. **Delta-based**: Track changes between two snapshots
//! 3. **Cumulative**: Track total CPU from a baseline (good for short-lived processes)
//!
//! All approaches use direct syscalls for minimal overhead (~1-5µs per process).

#![deny(unsafe_code)] // Unsafe only in platform-specific modules

mod platform;

use std::collections::HashMap;
use std::time::{Duration, Instant};

pub use platform::{
    AllStats, CpuTime, DiskIo, ProcessInfo, build_parent_map, get_all_stats, get_children,
    get_cpu_time, get_cpu_time_with_children, get_disk_io, get_memory, get_memory_virtual,
    get_ppid, get_process_comm, get_process_environ, get_process_path, get_start_time, get_tty,
    list_all_pids, scan_all_processes,
};

/// CPU usage as a percentage (0.0 - 100.0+ for multi-core).
pub type CpuPercent = f32;

/// Statistics for a single process.
#[derive(Debug, Clone, Default)]
pub struct ProcessStats {
    pub cpu_time: CpuTime,
    pub memory_bytes: u64,
}

/// Statistics for a process tree (process + all descendants).
#[derive(Debug, Clone, Default)]
pub struct TreeStats {
    /// Total CPU time across all processes in tree
    pub cpu_time: CpuTime,
    /// Total memory across all processes in tree
    pub memory_bytes: u64,
    /// Number of processes in tree
    pub process_count: u32,
    /// List of all PIDs in tree
    pub pids: Vec<i32>,
}

/// Approach 1: Snapshot-based measurement.
///
/// Get CPU time at a point in time. Caller computes deltas.
/// Lowest overhead, most flexible.
pub fn snapshot_process(pid: i32) -> Option<ProcessStats> {
    let cpu_time = get_cpu_time(pid)?;
    let memory_bytes = platform::get_memory(pid).unwrap_or(0);
    Some(ProcessStats {
        cpu_time,
        memory_bytes,
    })
}

/// Snapshot an entire process tree.
pub fn snapshot_tree(pid: i32) -> Option<TreeStats> {
    let pids = collect_tree_pids(pid);
    if pids.is_empty() {
        return None;
    }

    let mut total_cpu = CpuTime::default();
    let mut total_memory = 0u64;

    for &p in &pids {
        if let Some(cpu) = get_cpu_time(p) {
            total_cpu = total_cpu + cpu;
        }
        total_memory += platform::get_memory(p).unwrap_or(0);
    }

    Some(TreeStats {
        cpu_time: total_cpu,
        memory_bytes: total_memory,
        process_count: pids.len() as u32,
        pids,
    })
}

/// Approach 2: Delta-based tracker.
///
/// Tracks CPU deltas between refresh calls.
/// Similar to how `top` works.
#[derive(Debug)]
pub struct DeltaTracker {
    /// Previous snapshot per PID
    snapshots: HashMap<i32, (CpuTime, Instant)>,
}

impl DeltaTracker {
    pub fn new() -> Self {
        Self {
            snapshots: HashMap::new(),
        }
    }

    /// Get CPU percentage for a process since last call.
    /// First call returns None (establishing baseline).
    pub fn get_cpu_percent(&mut self, pid: i32) -> Option<CpuPercent> {
        let now = Instant::now();
        let current = get_cpu_time(pid)?;

        let result = if let Some((prev_cpu, prev_time)) = self.snapshots.get(&pid) {
            let elapsed = now.duration_since(*prev_time);
            if elapsed.as_nanos() == 0 {
                return Some(0.0);
            }
            Some(current.cpu_percent_since(prev_cpu, elapsed))
        } else {
            None // First sample, no delta yet
        };

        self.snapshots.insert(pid, (current, now));
        result
    }

    /// Get CPU percentage for a process tree.
    ///
    /// Returns partial result with 0% CPU for processes without a baseline.
    /// Returns None only if the process tree is empty (root doesn't exist).
    pub fn get_tree_cpu_percent(&mut self, root_pid: i32) -> Option<(CpuPercent, u64, u32)> {
        let pids = collect_tree_pids(root_pid);
        if pids.is_empty() {
            return None;
        }

        let mut total_percent = 0.0f32;
        let mut total_memory = 0u64;
        let mut all_have_baseline = true;

        for &pid in &pids {
            match self.get_cpu_percent(pid) {
                Some(pct) => total_percent += pct,
                None => all_have_baseline = false,
            }
            total_memory += platform::get_memory(pid).unwrap_or(0);
        }

        if all_have_baseline {
            Some((total_percent, total_memory, pids.len() as u32))
        } else {
            // Return partial result with 0% for processes without baseline
            Some((total_percent, total_memory, pids.len() as u32))
        }
    }

    /// Remove stale entries for dead processes.
    pub fn prune_dead(&mut self) {
        self.snapshots
            .retain(|&pid, _| platform::process_exists(pid));
    }
}

impl Default for DeltaTracker {
    fn default() -> Self {
        Self::new()
    }
}

/// Approach 3: Cumulative tracker.
///
/// Tracks CPU usage from a fixed baseline (t=0).
/// Better for catching short-lived processes because the measurement
/// window grows over time (50ms, 100ms, 150ms, etc.).
#[derive(Debug)]
pub struct CumulativeTracker {
    /// Baseline CPU time per PID (captured at first observation)
    baselines: HashMap<i32, (CpuTime, Instant)>,
    /// When tracking started
    start_time: Instant,
}

impl CumulativeTracker {
    pub fn new() -> Self {
        Self {
            baselines: HashMap::new(),
            start_time: Instant::now(),
        }
    }

    /// Get CPU percentage for a process since it was first observed.
    /// Unlike DeltaTracker, this computes delta from the FIRST observation,
    /// giving a longer measurement window as time progresses.
    pub fn get_cpu_percent(&mut self, pid: i32) -> Option<CpuPercent> {
        let now = Instant::now();
        let current = get_cpu_time(pid)?;

        if let Some((baseline_cpu, baseline_time)) = self.baselines.get(&pid) {
            let elapsed = now.duration_since(*baseline_time);
            if elapsed.as_nanos() == 0 {
                return Some(0.0);
            }
            Some(current.cpu_percent_since(baseline_cpu, elapsed))
        } else {
            // First observation - store baseline
            self.baselines.insert(pid, (current, now));
            Some(0.0) // No CPU yet (just started tracking)
        }
    }

    /// Get CPU percentage for a process tree since tracking started.
    pub fn get_tree_cpu_percent(&mut self, root_pid: i32) -> Option<(CpuPercent, u64, u32)> {
        let pids = collect_tree_pids(root_pid);
        if pids.is_empty() {
            return None;
        }

        let mut total_percent = 0.0f32;
        let mut total_memory = 0u64;

        for &pid in &pids {
            if let Some(pct) = self.get_cpu_percent(pid) {
                total_percent += pct;
            }
            total_memory += platform::get_memory(pid).unwrap_or(0);
        }

        Some((total_percent, total_memory, pids.len() as u32))
    }

    /// Reset baselines for all processes (start fresh measurement).
    pub fn reset(&mut self) {
        self.baselines.clear();
        self.start_time = Instant::now();
    }

    /// Time since tracking started.
    pub fn elapsed(&self) -> Duration {
        self.start_time.elapsed()
    }
}

impl Default for CumulativeTracker {
    fn default() -> Self {
        Self::new()
    }
}

/// Collect all PIDs in a process tree (BFS traversal).
pub fn collect_tree_pids(root_pid: i32) -> Vec<i32> {
    use std::collections::HashSet;

    let mut seen = HashSet::new();
    let mut all_pids = Vec::new();
    let mut queue = vec![root_pid];

    // Check if root exists
    if !platform::process_exists(root_pid) {
        return all_pids;
    }

    while let Some(pid) = queue.pop() {
        // Use HashSet for O(1) contains check instead of Vec's O(n)
        if seen.insert(pid) {
            all_pids.push(pid);
            queue.extend(get_children(pid));
        }
    }

    all_pids
}

/// Check if `child_pid` is a descendant of `parent_pid` using pre-built map.
///
/// Returns true if child_pid == parent_pid or if there's a parent chain
/// from child_pid to parent_pid. Walks at most 50 levels to avoid infinite loops.
///
/// # Arguments
/// * `child_pid` - The potential descendant process ID
/// * `parent_pid` - The potential ancestor process ID
/// * `parent_map` - Map of pid -> parent pid (from `build_parent_map()`)
///
/// # Example
/// ```ignore
/// let parent_map = prock::build_parent_map();
/// let is_child = prock::is_descendant_of(child_pid, parent_pid, &parent_map);
/// ```
pub fn is_descendant_of<S: std::hash::BuildHasher>(
    child_pid: i32,
    parent_pid: i32,
    parent_map: &HashMap<i32, i32, S>,
) -> bool {
    if child_pid == parent_pid {
        return true;
    }

    let mut current_pid = child_pid;
    for _ in 0..50 {
        if let Some(&ppid) = parent_map.get(&current_pid) {
            if ppid == parent_pid {
                return true;
            }
            if ppid <= 1 {
                return false;
            }
            current_pid = ppid;
        } else {
            return false;
        }
    }
    false
}

/// Build a HashMap of pid -> children from a full process scan.
///
/// This is used with the full scan approach: scan all processes once,
/// then use this map to quickly find descendants of any process.
pub fn build_children_map(procs: &[ProcessInfo]) -> HashMap<i32, Vec<i32>> {
    let mut map: HashMap<i32, Vec<i32>> = HashMap::new();

    for proc in procs {
        map.entry(proc.ppid).or_default().push(proc.pid);
    }

    map
}

/// Collect all descendant PIDs from a children map (built from full scan).
///
/// This is the fast path: O(n) where n is number of descendants.
pub fn collect_descendants_from_map(
    root_pid: i32,
    children_map: &HashMap<i32, Vec<i32>>,
) -> Vec<i32> {
    let mut result = vec![root_pid];
    let mut queue = vec![root_pid];

    while let Some(pid) = queue.pop() {
        if let Some(children) = children_map.get(&pid) {
            for &child in children {
                result.push(child);
                queue.push(child);
            }
        }
    }

    result
}

/// Get stats for a process tree using the full scan approach.
///
/// This scans all processes once, builds a tree, and filters to descendants.
/// Faster than targeted approach when you have many root PIDs to check.
pub fn snapshot_tree_from_scan(root_pid: i32, procs: &[ProcessInfo]) -> Option<TreeStats> {
    // Build children map
    let children_map = build_children_map(procs);

    // Get all descendant PIDs
    let pids = collect_descendants_from_map(root_pid, &children_map);

    // Build a pid -> ProcessInfo map for quick lookup
    let proc_map: HashMap<i32, &ProcessInfo> = procs.iter().map(|p| (p.pid, p)).collect();

    // Check if root exists
    if !proc_map.contains_key(&root_pid) {
        return None;
    }

    let mut total_cpu = CpuTime::default();
    let mut total_memory = 0u64;
    let mut valid_pids = Vec::new();

    for pid in pids {
        if let Some(proc) = proc_map.get(&pid) {
            total_cpu = total_cpu + proc.cpu_time;
            total_memory += proc.memory_bytes;
            valid_pids.push(pid);
        }
    }

    Some(TreeStats {
        cpu_time: total_cpu,
        memory_bytes: total_memory,
        process_count: valid_pids.len() as u32,
        pids: valid_pids,
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_snapshot_self() {
        let pid = std::process::id() as i32;
        let stats = snapshot_process(pid);
        assert!(stats.is_some());
        let stats = stats.unwrap();
        assert!(stats.memory_bytes > 0);
    }

    #[test]
    fn test_delta_tracker() {
        let pid = std::process::id() as i32;
        let mut tracker = DeltaTracker::new();

        // First call returns None (establishing baseline)
        let first = tracker.get_cpu_percent(pid);
        assert!(first.is_none());

        // Do some work
        let mut sum = 0u64;
        for i in 0..1_000_000 {
            sum = sum.wrapping_add(i);
        }
        std::hint::black_box(sum);

        // Second call returns Some
        let second = tracker.get_cpu_percent(pid);
        assert!(second.is_some());
    }

    #[test]
    fn test_cumulative_tracker() {
        let pid = std::process::id() as i32;
        let mut tracker = CumulativeTracker::new();

        // First call returns 0 (baseline just set)
        let first = tracker.get_cpu_percent(pid);
        assert_eq!(first, Some(0.0));

        // Do some work
        let mut sum = 0u64;
        for i in 0..1_000_000 {
            sum = sum.wrapping_add(i);
        }
        std::hint::black_box(sum);

        // Second call returns cumulative from baseline
        let second = tracker.get_cpu_percent(pid);
        assert!(second.is_some());
    }

    #[test]
    fn test_collect_tree_pids() {
        let pid = std::process::id() as i32;
        let pids = collect_tree_pids(pid);
        assert!(pids.contains(&pid));
    }

    #[test]
    fn test_is_descendant_of_self() {
        let map = build_parent_map();
        let pid = std::process::id() as i32;
        assert!(is_descendant_of(pid, pid, &map));
    }

    #[test]
    fn test_is_descendant_of_parent() {
        let map = build_parent_map();
        let pid = std::process::id() as i32;
        if let Some(&ppid) = map.get(&pid) {
            assert!(is_descendant_of(pid, ppid, &map));
        }
    }

    #[test]
    fn test_is_descendant_of_unrelated() {
        let map = build_parent_map();
        let our_pid = std::process::id() as i32;
        // PID 1 should NOT be a descendant of our process
        assert!(!is_descendant_of(1, our_pid, &map));
    }

    #[test]
    fn test_is_descendant_of_empty_map() {
        let empty_map: HashMap<i32, i32> = HashMap::new();
        assert!(is_descendant_of(100, 100, &empty_map));
        assert!(!is_descendant_of(100, 200, &empty_map));
    }
}