hyperfoot 0.2.2

use std::io;
use std::process::{Command, Stdio};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;
use std::time::{Duration, Instant};

use crate::sampler::watch_tree;

#[cfg(target_os = "linux")]
use crate::cgroup::CGroup;

/// How the resource numbers in a [`RunStats`] were obtained.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Accounting {
    /// Exact, kernel-maintained totals from a Linux cgroup v2 subtree —
    /// unaffected by how short-lived individual child processes were.
    Cgroup,
    /// Peaks/totals observed by polling the process tree. Can undercount
    /// processes that start and exit between polls.
    Sampled,
}

#[derive(Debug, Clone, Copy)]
pub struct RunStats {
    pub wall_time: Duration,
    pub cpu_time: Duration,
    pub peak_memory_bytes: u64,
    pub disk_read_bytes: u64,
    pub disk_write_bytes: u64,
    pub max_processes: u32,
    pub max_threads: u32,
    pub accounting: Accounting,
}

const DEFAULT_SAMPLE_INTERVAL: Duration = Duration::from_millis(20);

/// Runs `cmd` under `shell -c` once and measures its whole process tree's
/// resource footprint.
pub fn measure_once(shell: &str, cmd: &str) -> io::Result<RunStats> {
    #[cfg(target_os = "linux")]
    {
        if let Some(cgroup) = CGroup::try_create() {
            return measure_with_cgroup(shell, cmd, cgroup);
        }
    }
    measure_sampled_only(shell, cmd)
}

fn reap_with_rusage(pid: i32) -> Duration {
    let mut status: i32 = 0;
    let mut usage: libc::rusage = unsafe { std::mem::zeroed() };
    unsafe {
        libc::wait4(pid, &mut status, 0, &mut usage);
    }
    // ru_utime/ru_stime already include every reaped descendant's usage:
    // the kernel folds a child's RUSAGE_SELF + RUSAGE_CHILDREN into its
    // parent's RUSAGE_CHILDREN whenever the parent reaps it, so this is an
    // exact total for the whole tree provided each generation waits on its
    // own children (true for shells and build tools; orphaned daemons that
    // detach from their parent are the one case this can't see).
    seconds_and_micros_to_duration(usage.ru_utime.tv_sec, usage.ru_utime.tv_usec as i64)
        + seconds_and_micros_to_duration(usage.ru_stime.tv_sec, usage.ru_stime.tv_usec as i64)
}

fn seconds_and_micros_to_duration(secs: i64, micros: i64) -> Duration {
    Duration::from_secs(secs.max(0) as u64) + Duration::from_micros(micros.max(0) as u64)
}

fn measure_sampled_only(shell: &str, cmd: &str) -> io::Result<RunStats> {
    let start = Instant::now();
    let child = Command::new(shell)
        .arg("-c")
        .arg(cmd)
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()?;
    let pid = child.id();

    let stop = Arc::new(AtomicBool::new(false));
    let stop_clone = stop.clone();
    let sampler = thread::spawn(move || watch_tree(pid, stop_clone, DEFAULT_SAMPLE_INTERVAL));

    let cpu_time = reap_with_rusage(pid as i32);
    let wall_time = start.elapsed();
    stop.store(true, Ordering::Relaxed);
    let sample = sampler.join().expect("sampler thread panicked");

    Ok(RunStats {
        wall_time,
        cpu_time,
        peak_memory_bytes: sample.peak_memory_bytes,
        disk_read_bytes: sample.disk_read_bytes,
        disk_write_bytes: sample.disk_write_bytes,
        max_processes: sample.max_processes,
        max_threads: sample.max_threads,
        accounting: Accounting::Sampled,
    })
}

#[cfg(target_os = "linux")]
fn measure_with_cgroup(shell: &str, cmd: &str, cgroup: CGroup) -> io::Result<RunStats> {
    use std::os::unix::process::CommandExt;

    let procs_path = cgroup.procs_path();
    let start = Instant::now();

    let mut command = Command::new(shell);
    command
        .arg("-c")
        .arg(cmd)
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null());
    // Safety: runs after fork, before exec, while the child is still
    // single-threaded, so allocation/locking hazards from Rust's runtime
    // don't apply here. Joining our own freshly-created (writable) cgroup
    // needs no special privileges under cgroup v2.
    unsafe {
        command.pre_exec(move || {
            std::fs::write(&procs_path, std::process::id().to_string())?;
            Ok(())
        });
    }
    let child = command.spawn()?;
    let pid = child.id();

    let stop = Arc::new(AtomicBool::new(false));
    let stop_clone = stop.clone();
    let sampler = thread::spawn(move || watch_tree(pid, stop_clone, DEFAULT_SAMPLE_INTERVAL));

    let cpu_time = reap_with_rusage(pid as i32);
    let wall_time = start.elapsed();
    stop.store(true, Ordering::Relaxed);
    let sample = sampler.join().expect("sampler thread panicked");

    cgroup.wait_drain();
    let cgroup_stats = cgroup.read_stats();
    cgroup.cleanup();

    let accounting = if cgroup_stats.peak_memory_bytes.is_some() {
        Accounting::Cgroup
    } else {
        Accounting::Sampled
    };

    Ok(RunStats {
        wall_time,
        cpu_time,
        peak_memory_bytes: cgroup_stats
            .peak_memory_bytes
            .unwrap_or(sample.peak_memory_bytes),
        disk_read_bytes: cgroup_stats
            .disk_read_bytes
            .unwrap_or(sample.disk_read_bytes),
        disk_write_bytes: cgroup_stats
            .disk_write_bytes
            .unwrap_or(sample.disk_write_bytes),
        max_processes: sample.max_processes,
        max_threads: sample.max_threads,
        accounting,
    })
}