git.delta.rocks / jrsonnet / refs/commits / f78972ccc432

1use std::{2	ffi::OsString,3	mem,4	process::{Command, Stdio},5	time::Instant,6};78use anyhow::{Result, bail};9use nix::{libc, sys::wait::WaitStatus, unistd::Pid};1011#[derive(Debug, Clone)]12pub struct Stats {13	pub min: f64,14	pub max: f64,15	pub mean: f64,16	pub stddev: f64,17}1819impl Stats {20	fn from_samples(samples: &[f64]) -> Self {21		let n = samples.len() as f64;22		let mean = samples.iter().sum::<f64>() / n;23		let var = if samples.len() > 1 {24			samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (n - 1.0)25		} else {26			0.027		};28		Self {29			min: samples.iter().copied().fold(f64::INFINITY, f64::min),30			max: samples.iter().copied().fold(f64::NEG_INFINITY, f64::max),31			mean,32			stddev: var.sqrt(),33		}34	}35}3637#[derive(Debug, Clone)]38pub struct BenchResult {39	pub runs: u32,40	/// Wall-clock time per run, seconds.41	pub wall_secs: Stats,42	/// Peak resident set per run, KiB (Linux `ru_maxrss`).43	pub max_rss_kib: Stats,44}4546pub struct BenchOpts<'a> {47	pub program: &'a OsString,48	pub args: &'a [OsString],49	pub runs: u32,50	pub warmup: u32,51	pub output: bool,52}5354pub fn bench(opts: BenchOpts<'_>) -> Result<BenchResult> {55	if opts.runs == 0 {56		bail!("runs must be >= 1");57	}5859	for _ in 0..opts.warmup {60		run_once(opts.program, opts.args, opts.output)?;61	}6263	let mut wall = Vec::with_capacity(opts.runs as usize);64	let mut rss = Vec::with_capacity(opts.runs as usize);65	for _ in 0..opts.runs {66		let s = run_once(opts.program, opts.args, opts.output)?;67		wall.push(s.wall_secs);68		rss.push(s.max_rss_kib as f64);69	}7071	Ok(BenchResult {72		runs: opts.runs,73		wall_secs: Stats::from_samples(&wall),74		max_rss_kib: Stats::from_samples(&rss),75	})76}7778struct Sample {79	wall_secs: f64,80	max_rss_kib: i64,81}8283fn run_once(program: &OsString, args: &[OsString], silent: bool) -> Result<Sample> {84	let mut cmd = Command::new(program);85	cmd.args(args);86	if silent {87		cmd.stdout(Stdio::null()).stderr(Stdio::null());88	}8990	let start = Instant::now();91	let child = cmd.spawn()?;92	let pid = child.id() as libc::pid_t;93	// We'll reap via wait4 ourselves; don't let std touch this handle again.94	mem::forget(child);9596	let mut status: libc::c_int = 0;97	let mut ru: libc::rusage = unsafe { mem::zeroed() };98	let waited = unsafe { libc::wait4(pid, &raw mut status, 0, &raw mut ru) };99	let elapsed = start.elapsed();100	if waited < 0 {101		return Err(std::io::Error::last_os_error().into());102	}103104	match WaitStatus::from_raw(Pid::from_raw(pid), status)? {105		WaitStatus::Exited(_, 0) => {}106		WaitStatus::Exited(_, code) => bail!("child exited with code {code}"),107		WaitStatus::Signaled(_, sig, _) => bail!("child killed by signal {sig:?}"),108		other => bail!("unexpected wait status: {other:?}"),109	}110111	Ok(Sample {112		wall_secs: elapsed.as_secs_f64(),113		max_rss_kib: ru.ru_maxrss,114	})115}116117#[cfg(target_os = "linux")]118pub fn bench_cmd(args: &[String], runs: u32, warmup: u32, output: bool) -> Result<()> {119	let program = std::ffi::OsString::from(&args[0]);120	let rest: Vec<std::ffi::OsString> = args[1..].iter().map(Into::into).collect();121	let r = bench(BenchOpts {122		program: &program,123		args: &rest,124		runs,125		warmup,126		output,127	})?;128	eprintln!(129		"runs: {}  wall: {:.3}s ± {:.3}s  [{:.3}..{:.3}]",130		r.runs, r.wall_secs.mean, r.wall_secs.stddev, r.wall_secs.min, r.wall_secs.max,131	);132	eprintln!(133		"           max_rss: {} ± {} KiB  [{}..{}]",134		r.max_rss_kib.mean as i64,135		r.max_rss_kib.stddev as i64,136		r.max_rss_kib.min as i64,137		r.max_rss_kib.max as i64,138	);139	Ok(())140}