hyperfoot 0.2.2

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;
use std::time::Duration;

use sysinfo::{Pid, ProcessesToUpdate, System};

/// Resource peaks/totals collected by polling a process tree until it exits.
///
/// Memory and process/thread counts are peaks-of-sums observed at each poll
/// tick, so they can miss spikes shorter than `interval`. Disk I/O is
/// cumulative per process (as reported by the OS) and is captured for each
/// pid the moment it disappears from the tree, so short-lived children are
/// still counted as long as at least one poll caught them alive.
#[derive(Debug, Default, Clone, Copy)]
pub struct Sample {
    pub peak_memory_bytes: u64,
    pub max_processes: u32,
    pub max_threads: u32,
    pub disk_read_bytes: u64,
    pub disk_write_bytes: u64,
}

/// Walks `sys`'s process table to find every live descendant of `root`,
/// including `root` itself if still present.
fn collect_tree(sys: &System, root: Pid) -> Vec<Pid> {
    let mut children_of: HashMap<Pid, Vec<Pid>> = HashMap::new();
    for (pid, process) in sys.processes() {
        if let Some(parent) = process.parent() {
            children_of.entry(parent).or_default().push(*pid);
        }
    }

    let mut tree = Vec::new();
    let mut queue = vec![root];
    while let Some(pid) = queue.pop() {
        if !sys.processes().contains_key(&pid) {
            continue;
        }
        tree.push(pid);
        if let Some(kids) = children_of.get(&pid) {
            queue.extend(kids.iter().copied());
        }
    }
    tree
}

/// Polls the process tree rooted at `root` every `interval` until `stop` is
/// set, then returns the accumulated sample. Meant to run on a background
/// thread started right after the root process is spawned and joined right
/// after it's been reaped, so the last poll observes the tree just before it
/// fully drains.
pub fn watch_tree(root_pid: u32, stop: Arc<AtomicBool>, interval: Duration) -> Sample {
    let root = Pid::from(root_pid as usize);
    let mut sys = System::new();
    let mut peak_memory_bytes = 0u64;
    let mut max_processes = 0u32;
    let mut max_threads = 0u32;
    let mut retired_read = 0u64;
    let mut retired_write = 0u64;
    // Last known cumulative disk counters per pid still in the tree.
    let mut last_disk: HashMap<Pid, (u64, u64)> = HashMap::new();

    loop {
        sys.refresh_processes(ProcessesToUpdate::All, true);
        let tree = collect_tree(&sys, root);
        let tree_set: HashSet<Pid> = tree.iter().copied().collect();

        let mut memory_sum = 0u64;
        let mut thread_sum = 0u32;
        for pid in &tree {
            let Some(process) = sys.process(*pid) else {
                continue;
            };
            memory_sum += process.memory();
            // A live process always has at least one thread; an empty task
            // set here means we raced its exit while reading /proc/pid/task,
            // not that it truly had zero threads.
            thread_sum += process
                .tasks()
                .map(|t| t.len() as u32)
                .filter(|&n| n > 0)
                .unwrap_or(1);
            let disk = process.disk_usage();
            last_disk.insert(*pid, (disk.total_read_bytes, disk.total_written_bytes));
        }

        // Retire pids that dropped out of the tree since the last poll,
        // folding their last-known cumulative disk counters into the total.
        let gone: Vec<Pid> = last_disk
            .keys()
            .filter(|pid| !tree_set.contains(pid))
            .copied()
            .collect();
        for pid in gone {
            if let Some((r, w)) = last_disk.remove(&pid) {
                retired_read += r;
                retired_write += w;
            }
        }

        peak_memory_bytes = peak_memory_bytes.max(memory_sum);
        max_processes = max_processes.max(tree.len() as u32);
        max_threads = max_threads.max(thread_sum);

        if stop.load(Ordering::Relaxed) {
            break;
        }
        thread::sleep(interval);
    }

    let (still_alive_read, still_alive_write) = last_disk
        .values()
        .fold((0u64, 0u64), |acc, (r, w)| (acc.0 + r, acc.1 + w));

    Sample {
        peak_memory_bytes,
        max_processes,
        max_threads,
        disk_read_bytes: retired_read + still_alive_read,
        disk_write_bytes: retired_write + still_alive_write,
    }
}