run.rs

use crate::{
    app::{BenchTargetFilters, PackageFilters},
    helpers::get_codspeed_target_dir,
    measurement_mode::{BuildMode, MeasurementMode},
    prelude::*,
};
use anyhow::Context;
use cargo_metadata::{Metadata, Package};
use codspeed::walltime_results::WalltimeResults;
use std::{
    io::{self, Write},
    path::{Path, PathBuf},
};

#[cfg(unix)]
use std::os::unix::process::ExitStatusExt;

struct BenchToRun {
    bench_path: PathBuf,
    bench_target_name: String,
    working_directory: PathBuf,
    package_name: String,
}

impl PackageFilters {
    /// Logic taken from [cargo::ops::Packages](https://docs.rs/cargo/0.85.0/src/cargo/ops/cargo_compile/packages.rs.html#34-42)
    pub(crate) fn packages_from_flags<'a>(
        &self,
        metadata: &'a Metadata,
    ) -> Result<Vec<&'a Package>> {
        Ok(
            match (self.workspace, self.exclude.len(), self.package.len()) {
                (false, 0, 0) => metadata.workspace_default_packages(),
                (false, 0, _) => metadata
                    .workspace_packages()
                    .into_iter()
                    .filter(|p| self.package.contains(&p.name))
                    .collect(),
                (false, _, _) => bail!("--exclude can only be used with --workspace"),
                (true, 0, _) => metadata.workspace_packages(),
                (true, _, _) => metadata
                    .workspace_packages()
                    .into_iter()
                    .filter(|p| !self.exclude.contains(&p.name))
                    .collect(),
            },
        )
    }

    fn benches_to_run(
        &self,
        metadata: &Metadata,
        bench_target_filters: BenchTargetFilters,
        codspeed_target_dir: PathBuf,
    ) -> Result<Vec<BenchToRun>> {
        let packages = self.packages_from_flags(metadata)?;

        let mut benches = vec![];
        for package in packages {
            let package_name = &package.name;
            let package_target_dir = codspeed_target_dir.join(package_name.to_string());
            let working_directory = package.manifest_path.parent().ok_or_else(|| {
                Error::msg(format!("Failed to get root dir for package {package_name}"))
            })?;
            if let io::Result::Ok(read_dir) = std::fs::read_dir(&package_target_dir) {
                for entry in read_dir {
                    let entry = entry?;
                    let bench_path = entry.path();
                    let bench_target_name = bench_path
                        .file_name()
                        .unwrap()
                        .to_str()
                        .unwrap()
                        .to_string();
                    if let Some(bench_target_filter) = &bench_target_filters.bench {
                        if !bench_target_filter.contains(&bench_target_name) {
                            continue;
                        }
                    }
                    if !bench_path.is_dir() {
                        benches.push(BenchToRun {
                            package_name: package_name.to_string(),
                            working_directory: working_directory.into(),
                            bench_path,
                            bench_target_name,
                        });
                    }
                }
            }
        }

        Ok(benches)
    }
}

pub fn run_benches(
    metadata: &Metadata,
    bench_name_filter: Option<String>,
    package_filters: PackageFilters,
    bench_target_filters: BenchTargetFilters,
    measurement_mode: MeasurementMode,
) -> Result<()> {
    let build_mode = measurement_mode.into();
    let codspeed_target_dir = get_codspeed_target_dir(metadata, build_mode);
    let workspace_root = metadata.workspace_root.as_std_path();
    if build_mode == BuildMode::Walltime {
        WalltimeResults::clear(workspace_root)?;
    }
    let benches =
        package_filters.benches_to_run(metadata, bench_target_filters, codspeed_target_dir)?;
    if benches.is_empty() {
        bail!("No benchmarks found for the {measurement_mode} mode. Run `cargo codspeed build -m {measurement_mode}` first.");
    }

    eprintln!("Collected {} benchmark suite(s) to run", benches.len());

    for bench in benches.iter() {
        let bench_target_name = &bench.bench_target_name;
        // workspace_root is needed since file! returns the path relatively to the workspace root
        // while CARGO_MANIFEST_DIR returns the path to the sub package
        let workspace_root = metadata.workspace_root.clone();
        eprintln!("Running {} {}", &bench.package_name, bench_target_name);
        let mut command = std::process::Command::new(&bench.bench_path);
        command
            .env("CODSPEED_CARGO_WORKSPACE_ROOT", workspace_root)
            .current_dir(&bench.working_directory);

        if build_mode == BuildMode::Walltime {
            command.arg("--bench"); // Walltime targets need this additional argument (inherited from running them with `cargo bench`)
        }

        if let Some(bench_name_filter) = bench_name_filter.as_ref() {
            command.arg(bench_name_filter);
        }

        command
            .status()
            .map_err(|e| anyhow!("failed to execute the benchmark process: {e}"))
            .and_then(|status| {
                if status.success() {
                    Ok(())
                } else {
                    #[cfg(unix)]
                    {
                        let code = status
                            .code()
                            .or(status.signal().map(|s| 128 + s)) // 128+N indicates that a command was interrupted by signal N (see: https://tldp.org/LDP/abs/html/exitcodes.html)
                            .unwrap_or(1);

                        eprintln!("failed to execute the benchmark process, exit code: {code}");

                        std::process::exit(code);
                    }
                    #[cfg(not(unix))]
                    {
                        bail!("failed to execute the benchmark process: {}", status)
                    }
                }
            })?;
        eprintln!("Done running {bench_target_name}");
    }
    eprintln!("Finished running {} benchmark suite(s)", benches.len());

    if build_mode == BuildMode::Walltime {
        aggregate_raw_walltime_data(workspace_root)?;
    }

    Ok(())
}

fn aggregate_raw_walltime_data(workspace_root: &Path) -> Result<()> {
    let results = WalltimeResults::collect_walltime_results(workspace_root)
        .with_context(|| {
            format!(
                "Failed to collect walltime results. This may be due to version incompatibility. \
                Ensure that your compat layer (codspeed-criterion-compat, codspeed-bencher-compat, or codspeed-divan-compat) \
                has the same major version as cargo-codspeed (currently v{}).",
                env!("CARGO_PKG_VERSION")
            )
        })?;

    if results.benchmarks().is_empty() {
        eprintln!("No walltime benchmarks found");
        return Ok(());
    }

    for bench in results.benchmarks() {
        if bench.is_invalid() {
            eprintln!(
                "Warning: Benchmark {} was possibly optimized away",
                bench.name()
            );
        }
    }

    let results_folder = std::env::var("CODSPEED_PROFILE_FOLDER")
        .map(PathBuf::from)
        .unwrap_or_else(|_| workspace_root.join("target/codspeed/profiles"))
        .join("results");
    std::fs::create_dir_all(&results_folder).context("Failed to create results folder")?;

    let results_path = results_folder.join(format!("{}.json", std::process::id()));
    let mut results_file =
        std::fs::File::create(&results_path).context("Failed to create results file")?;
    serde_json::to_writer_pretty(&results_file, &results)?;
    results_file
        .flush()
        .context("Failed to flush results file")?;
    Ok(())
}