Improved time benchmarking command - average multiple runs

time.rs

use std::process::Command;

use libc;
use clap::{App,Arg,ArgMatches};

fn commify(mut val: i64) -> String {
    if val.abs() < 1_000     { return format!("{}", val); }
    if val.abs() < 1_000_000 { return format!("{},{:03}", val/1000, val.abs() % 1000); }
    let mut vec = Vec::new();
    while val > 1000 { vec.push(val % 1000); val /= 1000; }
    return vec.iter().fold(val.to_string(), |s, v| format!("{},{:03}", s, v.abs()));
}

fn format_kib(kib: i64) -> String {
    let mib = kib / 1024;
    let gib = mib / 1024;
    if gib > 0 { return format!("{}.{:03}GiB", gib, mib % 1024); }
    if mib > 0 { return format!("{}.{:03}MiB", mib, kib % 1024); }
    return format!("{}KiB", kib);
}

fn format_us(t: u128) -> String {
    let msec = t / 1000;
    if t < 1000   { return format!("0m      {}us", t); }
    if t < 10_000 { return format!("0m    {}.{}ms", msec, t % 1000); }
    let sec  = msec / 1000;
    let min  = sec / 60;
    return format!("{}m{}.{:03}s", min, sec % 60, msec % 1000);
}

fn subtract_timeval(a: &mut libc::timeval, b: &libc::timeval, div: u128) {
    let a128 = 1_000_000_u128 * a.tv_sec as u128 + a.tv_usec as u128;
    let b128 = 1_000_000_u128 * b.tv_sec as u128 + b.tv_usec as u128;
    let rv = (a128 - b128) / div;
    a.tv_sec  = (rv / 1_000_000) as i64;
    a.tv_usec = (rv % 1_000_000) as i64;
}

enum Timestamp {
    Instant(std::time::Instant),
    Duration(std::time::Duration),
}

// clock_t looks like it should be i64 everywhere, I'm not handling i32 overflows
struct CmdResult {
    time:  Option<Timestamp>,
    usage: libc::rusage,
    rv:    Option<Result<std::process::ExitStatus, std::io::Error>>,
}
impl CmdResult {
    pub fn new() -> CmdResult {
        // WARNING: rusage is cumulative of all waited on children
        CmdResult {
            rv:    None,
            time:  None,
            usage: libc::rusage {
                ru_utime: libc::timeval { tv_sec: 0, tv_usec: 0 }, // time executing in user mode
                ru_stime: libc::timeval { tv_sec: 0, tv_usec: 0 }, // time executing in kernel mode
                ru_maxrss: 0,  // max resident set size used (in kilobytes) of the largest child, not the whole tree

                ru_majflt: 0,  // Major Faults: page faults serviced that required I/O activity
                ru_minflt: 0,  // Minor Faults: page faults serviced without any I/O activity
                               // I.e., I/O avoided because the cached page was awaiting reclamation
                               // I.e., events that would have been page faults if we had less RAM.

                ru_inblock: 0, // number of times the filesystem had to perform input
                ru_oublock: 0, // number of times the filesystem had to perform output

                ru_nvcsw: 0,   // number of voluntary context switches (e.g., waiting on I/O)
                ru_nivcsw: 0,  // number of involuntary context switches (e.g., displaced by another task)

                // Unused by current Linux (Debian 10)
                ru_ixrss: 0, ru_idrss: 0, ru_isrss: 0,
                ru_msgsnd: 0, ru_msgrcv: 0, ru_nsignals: 0,
                ru_nswap: 0,
            }
        }
    }

    pub fn user_us(&self) -> u128 { (self.usage.ru_utime.tv_usec + 1_000_000 * self.usage.ru_utime.tv_sec) as u128 }
    pub fn sys_us(&self)  -> u128 { (self.usage.ru_stime.tv_usec + 1_000_000 * self.usage.ru_stime.tv_sec) as u128 }

    pub fn update(&mut self) -> bool {
        self.time = Some(Timestamp::Instant(std::time::Instant::now()));
        0 == unsafe { libc::getrusage(libc::RUSAGE_CHILDREN, &mut self.usage) }
    }

    pub fn report(&self) {
        if let Some(Timestamp::Duration(t)) = self.time {
            println!("real    {}", format_us(t.as_micros()));
        }
        println!("user    {}", format_us(self.user_us()));
        println!("sys     {}\n", format_us(self.sys_us()));

        if self.usage.ru_maxrss  > 0 { println!("resident size       {}", format_kib(self.usage.ru_maxrss)); }
        if self.usage.ru_majflt  > 0 { println!("disk page faults    {}", commify(self.usage.ru_majflt)); }
        if self.usage.ru_minflt  > 0 { println!("cached page faults  {}", commify(self.usage.ru_minflt)); }
        if self.usage.ru_inblock > 0 { println!("blocks read         {}", commify(self.usage.ru_inblock)); }
        if self.usage.ru_oublock > 0 { println!("blocks written      {}", commify(self.usage.ru_oublock)); }
        if self.usage.ru_nvcsw   > 0 { println!("yield for I/O       {}", commify(self.usage.ru_nvcsw)); }
        if self.usage.ru_nivcsw  > 0 { println!("yield to task       {}", commify(self.usage.ru_nivcsw)); }
    }

    pub fn subtract(&mut self, other: &CmdResult, div: u32) {
        if let Some(Timestamp::Instant(b)) = self.time {
            if let Some(Timestamp::Instant(a)) = other.time {
                self.time = Some(Timestamp::Duration((b - a).checked_div(div).unwrap_or(std::time::Duration::from_secs(0))));
            } else { panic!("Can't subtract Duration from Instant"); }
        } else { panic!("Can't subtract from a Duration"); }
        subtract_timeval(&mut self.usage.ru_utime, &other.usage.ru_utime, div as u128);
        subtract_timeval(&mut self.usage.ru_stime, &other.usage.ru_stime, div as u128);
        self.usage.ru_maxrss -= other.usage.ru_maxrss; self.usage.ru_maxrss /= div as i64;
        self.usage.ru_majflt -= other.usage.ru_majflt; self.usage.ru_majflt /= div as i64;
        self.usage.ru_minflt -= other.usage.ru_minflt; self.usage.ru_minflt /= div as i64;
        self.usage.ru_inblock -= other.usage.ru_inblock; self.usage.ru_inblock /= div as i64;
        self.usage.ru_oublock -= other.usage.ru_oublock; self.usage.ru_oublock /= div as i64;
        self.usage.ru_nvcsw -= other.usage.ru_nvcsw; self.usage.ru_nvcsw /= div as i64;
        self.usage.ru_nivcsw -= other.usage.ru_nivcsw; self.usage.ru_nivcsw /= div as i64;
    }
}

fn time_it(command: &mut Command) -> CmdResult {
    let mut a = CmdResult::new();
    let mut b = CmdResult::new();
    a.update();
    let rv = command.status();
    b.update();
    b.subtract(&a, 1);
    b.rv = Some(rv);
    return b;
}

fn build_command(argv: &ArgMatches) -> Command {
    if let Some(script) = argv.value_of("SCRIPT") {
        let mut command = Command::new(argv.value_of("SHELL").unwrap_or("/usr/bin/bash"));
        command.arg("-c").arg(script);
        return command;
    }
    else if let Some(mut cmd) = argv.values_of("COMMAND") {
        let name = cmd.next();
        let args = cmd.collect::<Vec<&str>>();
        if let Some(name) = name {
            let mut command = Command::new(name);
            command.args(&args);
            return command;
        }
    }
    println!("{}", argv.usage());
    std::process::exit(1);
}
fn user_command(argv: &ArgMatches) -> String {
    if let Some(script) = argv.value_of("SCRIPT") {
        return script.to_string();
    }
    else if let Some(cmd) = argv.values_of("COMMAND") {
        return cmd.collect::<Vec<&str>>().join(" ").to_string();
    }
    println!("{}", argv.usage());
    std::process::exit(1);
}


fn main() {
    let argv = App::new("TIME")
        .version("1.0")
        .about("Benchmark a command, possibly running it multiple times")
        .arg(Arg::with_name("SECONDS").short("-s").long("--seconds").allow_hyphen_values(true).takes_value(true).help("Minimum test duration, default 5s"))
        .arg(Arg::with_name("SHELL").long("--shell").takes_value(true).help("Shell to use for -c commands"))
        .arg(Arg::with_name("quiet").short("-q").long("--quiet").help("do not show command output, even for first run"))
        .arg(Arg::with_name("verbose").short("-v").long("--verbose").help("show all command output, every run"))
        .arg(Arg::with_name("SCRIPT").short("-c").long("--command").takes_value(true).help("Pass script to a shell using `SHELL -c 'SCRIPT'`"))
        .arg(Arg::with_name("COMMAND").last(true).multiple(true).help("command and args to run"))
        .arg(Arg::with_name("ONCE").short("1").help("Run command only once"))
        .group(clap::ArgGroup::with_name("verbosity").args(&["quiet", "verbose"]))
        .get_matches();

    let mut command = build_command(&argv);
    if argv.is_present("quiet") {
        command.stdout(std::process::Stdio::null()).stderr(std::process::Stdio::null());
    }
    let mut trial = time_it(&mut command);
    match trial.rv {
        Some(Ok(rc))   => {
            if !rc.success() {
                match rc.code() {
                    Some(code) => eprintln!("{} exited with status code: {}", user_command(&argv), code),
                    None       => eprintln!("{} terminated by signal", user_command(&argv))
                }
                std::process::exit(1);
            }
        },
        Some(Err(err)) => {
            eprintln!("Fatal error running '{}': {}", user_command(&argv), err);
            std::process::exit(1);
        }
        _ => { panic!("Bad things happened running {}", user_command(&argv)); }
    }

    if !argv.is_present("ONCE") {
        if !argv.is_present("verbose") {
            command.stdout(std::process::Stdio::null()).stderr(std::process::Stdio::null());
        }

        if let Some(Timestamp::Duration(t)) = trial.time {
            let duration = argv.value_of("SECONDS").unwrap_or(&"5").parse::<f64>().unwrap_or(5.0);
            let count = if duration > 0.0 { ((1_000_000.0 * duration).round() as u128 / t.as_micros()).min(1_000_000) } else { -duration.round() as u128 };
            if count > 1 {
                println!("\nAverage of {} iterations", count);
                let mut errors = 0;
                let mut a = CmdResult::new();
                a.update();
                for _ in 1..count {
                    match command.status() {
                        Ok(rc)   => { if !rc.success() { errors += 1; } },
                        Err(err) => { eprintln!("Fatal error: {}", err); break; },
                    }
                }
                trial.update();
                trial.subtract(&a, (count-1) as u32);
                if errors > 0 { println!("Command returned error {} times out of {}", errors, count); }
            }
        }
    }

    if !argv.is_present("quiet") { println!(""); }
    trial.report();
}