use std::future::Future;
use std::io;
use std::path::{Path, PathBuf};
use std::process::{ExitStatus, Stdio};
pub trait BenchRunner {
fn run_benches(
&self,
argv: &[String],
env: &[(String, String)],
) -> impl Future<Output = io::Result<EngineStatus>>;
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct EngineStatus {
pub success: bool,
pub code: Option<i32>,
}
impl EngineStatus {
fn from_exit(status: ExitStatus) -> Self {
Self {
success: status.success(),
code: status.code(),
}
}
}
#[derive(Clone, Debug)]
pub struct CommandOutput {
pub status: ExitStatus,
pub stdout: String,
}
#[derive(Clone, Debug, Default)]
pub struct TokioBenchRunner {
dir: Option<PathBuf>,
detach_from_launcher_toolchain: bool,
}
impl TokioBenchRunner {
#[must_use]
pub fn in_dir(dir: impl Into<PathBuf>) -> Self {
Self {
dir: Some(dir.into()),
detach_from_launcher_toolchain: false,
}
}
#[must_use]
pub fn in_worktree(dir: impl Into<PathBuf>) -> Self {
Self {
dir: Some(dir.into()),
detach_from_launcher_toolchain: true,
}
}
fn command(&self, argv: &[String], env: &[(String, String)]) -> tokio::process::Command {
let mut command = engine_command(argv);
if let Some(dir) = self.dir.as_deref() {
command.current_dir(dir);
}
if self.detach_from_launcher_toolchain {
detach_from_launcher_toolchain(&mut command);
}
for (name, value) in env {
command.env(name, value);
}
command
}
}
impl BenchRunner for TokioBenchRunner {
async fn run_benches(
&self,
argv: &[String],
env: &[(String, String)],
) -> io::Result<EngineStatus> {
Ok(EngineStatus::from_exit(
self.command(argv, env).status().await?,
))
}
}
const LAUNCHER_TOOLCHAIN_VARS: [&str; 2] = ["RUSTUP_TOOLCHAIN", "CARGO"];
fn detach_from_launcher_toolchain(command: &mut tokio::process::Command) {
for name in LAUNCHER_TOOLCHAIN_VARS {
command.env_remove(name);
}
}
pub async fn capture(program: &str, args: &[&str]) -> io::Result<CommandOutput> {
let mut command = tokio::process::Command::new(program);
command.args(args);
capture_output(command).await
}
pub async fn capture_in_worktree(
program: &str,
args: &[&str],
dir: &Path,
) -> io::Result<CommandOutput> {
let mut command = tokio::process::Command::new(program);
command.args(args).current_dir(dir);
detach_from_launcher_toolchain(&mut command);
capture_output(command).await
}
async fn capture_output(mut command: tokio::process::Command) -> io::Result<CommandOutput> {
let output = command.stdin(Stdio::null()).output().await?;
Ok(CommandOutput {
status: output.status,
stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
})
}
fn engine_command(argv: &[String]) -> tokio::process::Command {
let (program, args) = argv
.split_first()
.expect("engine argv is non-empty; build_bench_argv rejects empty commands");
let mut command = tokio::process::Command::new(program);
command.args(args);
command
}
#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
use std::collections::HashMap;
use std::ffi::OsString;
use std::path::PathBuf;
use super::{LAUNCHER_TOOLCHAIN_VARS, TokioBenchRunner, engine_command};
fn env_overrides(runner: &TokioBenchRunner) -> HashMap<OsString, Option<OsString>> {
let argv = vec!["cargo".to_owned(), "bench".to_owned()];
runner
.command(&argv, &[])
.as_std()
.get_envs()
.map(|(name, value)| (name.to_owned(), value.map(ToOwned::to_owned)))
.collect()
}
#[test]
fn in_dir_sets_the_working_directory() {
let runner = TokioBenchRunner::in_dir("some/worktree");
assert_eq!(runner.dir, Some(PathBuf::from("some/worktree")));
assert_eq!(TokioBenchRunner::default().dir, None);
}
#[test]
fn in_worktree_runs_in_the_directory_like_in_dir() {
let runner = TokioBenchRunner::in_worktree("some/worktree");
assert_eq!(runner.dir, Some(PathBuf::from("some/worktree")));
}
#[test]
#[cfg_attr(all(miri, windows), ignore)]
fn in_worktree_detaches_the_child_from_the_launcher_toolchain() {
let overrides = env_overrides(&TokioBenchRunner::in_worktree("some/worktree"));
for name in LAUNCHER_TOOLCHAIN_VARS {
assert_eq!(
overrides.get(&OsString::from(name)),
Some(&None),
"{name} must be removed for a historical checkout: {overrides:?}"
);
}
}
#[test]
fn in_dir_keeps_the_launcher_toolchain() {
let overrides = env_overrides(&TokioBenchRunner::in_dir("some/workspace"));
assert!(overrides.is_empty(), "{overrides:?}");
let overrides = env_overrides(&TokioBenchRunner::default());
assert!(overrides.is_empty(), "{overrides:?}");
}
#[test]
#[cfg_attr(all(miri, windows), ignore)]
fn injected_environment_wins_over_the_detachment() {
let runner = TokioBenchRunner::in_worktree("some/worktree");
let argv = vec!["cargo".to_owned(), "bench".to_owned()];
let env = vec![("CARGO".to_owned(), "/custom/cargo".to_owned())];
let command = runner.command(&argv, &env);
let overrides: HashMap<OsString, Option<OsString>> = command
.as_std()
.get_envs()
.map(|(name, value)| (name.to_owned(), value.map(ToOwned::to_owned)))
.collect();
assert_eq!(
overrides.get(&OsString::from("CARGO")),
Some(&Some(OsString::from("/custom/cargo"))),
"{overrides:?}"
);
}
#[test]
fn engine_command_runs_argv_directly() {
let argv = vec!["echo".to_owned(), "hi there".to_owned()];
let command = engine_command(&argv);
let std = command.as_std();
let program = std.get_program().to_string_lossy().into_owned();
let args: Vec<String> = std
.get_args()
.map(|arg| arg.to_string_lossy().into_owned())
.collect();
assert_eq!(program, "echo");
assert_eq!(args, vec!["hi there".to_owned()]);
}
}