hyperfoot 0.2.0

#[cfg(target_os = "linux")]
mod cgroup;
mod measure;
mod progress;
mod report;
mod sampler;
mod stats;

use anyhow::{Context, Result};
use clap::Parser;
use std::process::{Command, Stdio};
use std::time::Duration;

use progress::Progress;
use stats::BenchResult;

/// Benchmark the resource footprint of commands (time, CPU, memory, disk I/O).
#[derive(Parser, Debug)]
#[command(
    name = "hyperfoot",
    version,
    about = "Benchmark the resource footprint of commands"
)]
struct Cli {
    /// Commands to benchmark. One command reports its stats; two or more are compared.
    #[arg(required = true, num_args = 1..)]
    commands: Vec<String>,

    /// Number of timed runs per command
    #[arg(short = 'r', long, default_value_t = 10)]
    runs: u32,

    /// Number of warmup runs (not measured) per command
    #[arg(short = 'w', long, default_value_t = 0)]
    warmup: u32,

    /// Command to run once before each timed run (e.g. cleaning build artifacts)
    #[arg(long)]
    prepare: Option<String>,

    /// Shell used to invoke each command
    #[arg(long, default_value = "sh")]
    shell: String,

    /// Write results as JSON to this path
    #[arg(long)]
    export_json: Option<String>,

    /// Write results as CSV to this path
    #[arg(long)]
    export_csv: Option<String>,

    /// Write results as a Markdown table to this path
    #[arg(long)]
    export_markdown: Option<String>,
}

fn main() -> Result<()> {
    restore_default_sigpipe();
    let cli = Cli::parse();
    run(cli)
}

/// Rust masks SIGPIPE to SIG_IGN at startup, so writes to a closed pipe
/// (e.g. `hyperfoot ... | head`) surface as an `io::Error` that `println!`
/// turns into a panic instead of just ending the process the way any C
/// command-line tool does. Restore the default disposition so a broken
/// pipe kills us quietly.
fn restore_default_sigpipe() {
    unsafe {
        libc::signal(libc::SIGPIPE, libc::SIG_DFL);
    }
}

fn run(cli: Cli) -> Result<()> {
    let mut results = Vec::with_capacity(cli.commands.len());
    let progress = Progress::new(&cli.commands);
    let total_commands = cli.commands.len();

    for (command_index, command) in cli.commands.iter().enumerate() {
        let position = (command_index, total_commands);
        for w in 0..cli.warmup {
            progress.warmup(command, position, w + 1, cli.warmup);
            run_prepare(&cli.shell, cli.prepare.as_deref())?;
            measure::measure_once(&cli.shell, command)
                .with_context(|| format!("failed to run warmup for `{command}`"))?;
        }

        let mut runs = Vec::with_capacity(cli.runs as usize);
        let mut total_wall = Duration::ZERO;
        for i in 0..cli.runs {
            let mean_so_far = (i > 0).then(|| total_wall / i);
            progress.run(command, position, i + 1, cli.runs, mean_so_far);
            run_prepare(&cli.shell, cli.prepare.as_deref())?;
            let stats = measure::measure_once(&cli.shell, command)
                .with_context(|| format!("failed to run `{command}`"))?;
            total_wall += stats.wall_time;
            runs.push(stats);
        }

        results.push(BenchResult {
            command: command.clone(),
            runs,
        });
    }
    progress.finish();

    if results.len() == 1 {
        report::print_single(&results[0]);
    } else {
        report::print_comparison(&results);
    }

    if let Some(path) = &cli.export_json {
        report::export_json(&results, path).with_context(|| format!("failed to write {path}"))?;
    }
    if let Some(path) = &cli.export_csv {
        report::export_csv(&results, path).with_context(|| format!("failed to write {path}"))?;
    }
    if let Some(path) = &cli.export_markdown {
        report::export_markdown(&results, path)
            .with_context(|| format!("failed to write {path}"))?;
    }

    Ok(())
}

fn run_prepare(shell: &str, prepare: Option<&str>) -> Result<()> {
    let Some(prepare) = prepare else {
        return Ok(());
    };
    let status = Command::new(shell)
        .arg("-c")
        .arg(prepare)
        .stdin(Stdio::null())
        .status()
        .with_context(|| format!("failed to run --prepare command `{prepare}`"))?;
    if !status.success() {
        anyhow::bail!("--prepare command `{prepare}` exited with {status}");
    }
    Ok(())
}