feat: allow caching installed executor instruments on ubuntu/debian

Author: GuillaumeLagrange

src/app.rs

@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use crate::{
     api_client::CodSpeedAPIClient,
     auth,
@@ -38,6 +40,13 @@ pub struct Cli {
     #[arg(long, env = "CODSPEED_OAUTH_TOKEN", global = true, hide = true)]
     pub oauth_token: Option<String>,
 
+    /// The directory to use for caching installed tools
+    /// The runner will restore cached tools from this directory before installing them.
+    /// After successful installation, the runner will cache the installed tools to this directory.
+    /// Only supported on ubuntu and debian systems.
+    #[arg(long, env = "CODSPEED_SETUP_CACHE_DIR", global = true)]
+    pub setup_cache_dir: Option<PathBuf>,
+
     #[command(subcommand)]
     command: Commands,
 }
@@ -56,6 +65,7 @@ pub async fn run() -> Result<()> {
     let cli = Cli::parse();
     let codspeed_config = CodSpeedConfig::load_with_override(cli.oauth_token.as_deref())?;
     let api_client = CodSpeedAPIClient::try_from((&cli, &codspeed_config))?;
+    let setup_cache_dir = cli.setup_cache_dir.as_deref();
 
     match cli.command {
         Commands::Run(_) => {} // Run is responsible for its own logger initialization
@@ -65,9 +75,11 @@ pub async fn run() -> Result<()> {
     }
 
     match cli.command {
-        Commands::Run(args) => run::run(args, &api_client, &codspeed_config).await?,
+        Commands::Run(args) => {
+            run::run(args, &api_client, &codspeed_config, setup_cache_dir).await?
+        }
         Commands::Auth(args) => auth::run(args, &api_client).await?,
-        Commands::Setup => setup::setup().await?,
+        Commands::Setup => setup::setup(setup_cache_dir).await?,
     }
     Ok(())
 }

src/run/mod.rs

@@ -9,6 +9,7 @@ use instruments::mongo_tracer::{MongoTracer, install_mongodb_tracer};
 use run_environment::interfaces::{RepositoryProvider, RunEnvironment};
 use runner::get_run_data;
 use serde::Serialize;
+use std::path::Path;
 use std::path::PathBuf;
 
 pub mod check_system;
@@ -176,6 +177,7 @@ pub async fn run(
     args: RunArgs,
     api_client: &CodSpeedAPIClient,
     codspeed_config: &CodSpeedConfig,
+    setup_cache_dir: Option<&Path>,
 ) -> Result<()> {
     let output_json = args.message_format == Some(MessageFormat::Json);
     let mut config = Config::try_from(args)?;
@@ -202,7 +204,7 @@ pub async fn run(
 
     if !config.skip_setup {
         start_group!("Preparing the environment");
-        executor.setup(&system_info).await?;
+        executor.setup(&system_info, setup_cache_dir).await?;
         // TODO: refactor and move directly in the Instruments struct as a `setup` method
         if config.instruments.is_mongodb_enabled() {
             install_mongodb_tracer().await?;

src/run/runner/executor.rs

@@ -3,12 +3,13 @@ use crate::prelude::*;
 use crate::run::instruments::mongo_tracer::MongoTracer;
 use crate::run::{check_system::SystemInfo, config::Config};
 use async_trait::async_trait;
+use std::path::Path;
 
 #[async_trait(?Send)]
 pub trait Executor {
     fn name(&self) -> ExecutorName;
 
-    async fn setup(&self, _system_info: &SystemInfo) -> Result<()> {
+    async fn setup(&self, _system_info: &SystemInfo, _setup_cache_dir: Option<&Path>) -> Result<()> {
         Ok(())
     }
 

src/run/runner/helpers/apt.rs

@@ -0,0 +1,126 @@
+use super::run_with_sudo::run_with_sudo;
+use crate::prelude::*;
+use crate::run::check_system::SystemInfo;
+use std::path::Path;
+use std::process::Command;
+
+fn is_system_compatible(system_info: &SystemInfo) -> bool {
+    system_info.os == "ubuntu" || system_info.os == "debian"
+}
+
+pub fn install(system_info: &SystemInfo, packages: &[&str]) -> Result<()> {
+    if !is_system_compatible(system_info) {
+        bail!(
+            "Package installation is not supported on this system, please install necessary packages manually"
+        );
+    }
+
+    debug!("Installing packages: {packages:?}");
+
+    run_with_sudo(&["apt-get", "update"])?;
+    let mut install_cmd = vec![
+        "apt-get",
+        "install",
+        "-y",
+        "--allow-downgrades",
+        "--reinstall",
+    ];
+    install_cmd.extend_from_slice(packages);
+    run_with_sudo(&install_cmd)?;
+
+    info!("Packages installed successfully");
+    Ok(())
+}
+
+/// Restore cached tools from the cache directory to the root filesystem
+pub fn restore_from_cache(system_info: &SystemInfo, cache_dir: &Path) -> Result<()> {
+    if !is_system_compatible(system_info) {
+        warn!("Cache restore is not supported on this system, skipping");
+        return Ok(());
+    }
+
+    if !cache_dir.exists() {
+        debug!("Cache directory does not exist: {}", cache_dir.display());
+        return Ok(());
+    }
+
+    // Check if the directory has any contents
+    let has_contents = std::fs::read_dir(cache_dir)
+        .map(|mut entries| entries.next().is_some())
+        .unwrap_or(false);
+
+    if !has_contents {
+        debug!("Cache directory is empty: {}", cache_dir.display());
+        return Ok(());
+    }
+
+    info!(
+        "Restoring tools from cache directory: {}",
+        cache_dir.display()
+    );
+
+    // Use bash to properly handle glob expansion
+    let cache_dir_str = cache_dir
+        .to_str()
+        .ok_or_else(|| anyhow!("Invalid cache directory path"))?;
+
+    let copy_cmd = format!("cp -r {cache_dir_str}/* /");
+
+    let output = Command::new("sudo")
+        .args(["bash", "-c", &copy_cmd])
+        .output()
+        .context("Failed to run cache restore command")?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        error!("stderr: {stderr}");
+        bail!("Failed to restore cache");
+    }
+
+    info!("Cache restored successfully");
+    Ok(())
+}
+
+/// Save installed packages to the cache directory
+pub fn save_to_cache(system_info: &SystemInfo, cache_dir: &Path, packages: &[&str]) -> Result<()> {
+    if !is_system_compatible(system_info) {
+        warn!("Caching of installed package is not supported on this system, skipping");
+        return Ok(());
+    }
+
+    info!(
+        "Saving installed packages to cache: {}",
+        cache_dir.display()
+    );
+
+    // Create cache directory if it doesn't exist
+    std::fs::create_dir_all(cache_dir).context("Failed to create cache directory")?;
+
+    let cache_dir_str = cache_dir
+        .to_str()
+        .ok_or_else(|| anyhow!("Invalid cache directory path"))?;
+
+    // Logic taken from https://stackoverflow.com/a/59277514
+    // This shell command lists all the files outputted by the given packages and copy them to the cache directory
+    let packages_str = packages.join(" ");
+    let shell_cmd = format!(
+        "sudo dpkg -L {packages_str} | while IFS= read -r f; do if test -f \"$f\"; then echo \"$f\"; fi; done | xargs cp --parents --target-directory {cache_dir_str}",
+    );
+
+    debug!("Running cache save command: {shell_cmd}");
+
+    let output = Command::new("sh")
+        .arg("-c")
+        .arg(&shell_cmd)
+        .output()
+        .context("Failed to execute cache save command")?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        error!("stderr: {stderr}");
+        bail!("Failed to save packages to cache");
+    }
+
+    info!("Packages cached successfully");
+    Ok(())
+}

src/run/runner/helpers/mod.rs

@@ -1,7 +1,8 @@
+pub mod apt;
 pub mod env;
 pub mod get_bench_command;
 pub mod introspected_golang;
 pub mod introspected_nodejs;
 pub mod profile_folder;
 pub mod run_command_with_log_pipe;
-pub mod setup;
+pub mod run_with_sudo;

src/run/runner/helpers/setup.rs src/run/runner/helpers/run_with_sudo.rssrc/run/runner/helpers/setup.rs renamed to src/run/runner/helpers/run_with_sudo.rs

@@ -1,5 +1,4 @@
 use crate::prelude::*;
-use log::{debug, info};
 use std::process::{Command, Stdio};
 
 /// Run a command with sudo if available

src/run/runner/tests.rs

@@ -139,7 +139,7 @@ mod valgrind {
             .get_or_init(|| async {
                 let executor = ValgrindExecutor;
                 let system_info = SystemInfo::new().unwrap();
-                executor.setup(&system_info).await.unwrap();
+                executor.setup(&system_info, None).await.unwrap();
                 executor
             })
             .await
@@ -196,7 +196,7 @@ mod walltime {
             .get_or_init(|| async {
                 let executor = WallTimeExecutor::new();
                 let system_info = SystemInfo::new().unwrap();
-                executor.setup(&system_info).await.unwrap();
+                executor.setup(&system_info, None).await.unwrap();
             })
             .await;
 

src/run/runner/valgrind/executor.rs

@@ -1,4 +1,5 @@
 use async_trait::async_trait;
+use std::path::Path;
 
 use crate::prelude::*;
 use crate::run::instruments::mongo_tracer::MongoTracer;
@@ -17,8 +18,8 @@ impl Executor for ValgrindExecutor {
         ExecutorName::Valgrind
     }
 
-    async fn setup(&self, system_info: &SystemInfo) -> Result<()> {
-        install_valgrind(system_info).await?;
+    async fn setup(&self, system_info: &SystemInfo, setup_cache_dir: Option<&Path>) -> Result<()> {
+        install_valgrind(system_info, setup_cache_dir).await?;
 
         if let Err(error) = venv_compat::symlink_libpython(None) {
             warn!("Failed to symlink libpython");

src/run/runner/valgrind/setup.rs

@@ -1,7 +1,7 @@
-use crate::run::runner::helpers::setup::run_with_sudo;
+use crate::run::runner::helpers::apt;
 use crate::{VALGRIND_CODSPEED_DEB_VERSION, run::check_system::SystemInfo};
 use crate::{VALGRIND_CODSPEED_VERSION, prelude::*, run::helpers::download_file};
-use std::{env, process::Command};
+use std::{env, path::Path, process::Command};
 use url::Url;
 
 fn get_codspeed_valgrind_filename(system_info: &SystemInfo) -> Result<String> {
@@ -54,11 +54,25 @@ fn is_valgrind_installed() -> bool {
     }
 }
 
-pub async fn install_valgrind(system_info: &SystemInfo) -> Result<()> {
+pub async fn install_valgrind(
+    system_info: &SystemInfo,
+    setup_cache_dir: Option<&Path>,
+) -> Result<()> {
     if is_valgrind_installed() {
         info!("Valgrind is already installed with the correct version, skipping installation");
         return Ok(());
     }
+
+    // Try to restore from cache first
+    if let Some(cache_dir) = setup_cache_dir {
+        apt::restore_from_cache(system_info, cache_dir)?;
+
+        if is_valgrind_installed() {
+            // Valgrind has been successfully restored from cache
+            return Ok(());
+        }
+    }
+
     debug!("Installing valgrind");
     let valgrind_deb_url = format!(
         "https://github.com/CodSpeedHQ/valgrind-codspeed/releases/download/{}/{}",
@@ -67,18 +81,15 @@ pub async fn install_valgrind(system_info: &SystemInfo) -> Result<()> {
     );
     let deb_path = env::temp_dir().join("valgrind-codspeed.deb");
     download_file(&Url::parse(valgrind_deb_url.as_str()).unwrap(), &deb_path).await?;
-
-    run_with_sudo(&["apt-get", "update"])?;
-    run_with_sudo(&[
-        "apt-get",
-        "install",
-        "--allow-downgrades",
-        "-y",
-        deb_path.to_str().unwrap(),
-    ])?;
+    apt::install(system_info, &[deb_path.to_str().unwrap()])?;
 
     info!("Valgrind installation completed successfully");
 
+    // Save to cache after successful installation
+    if let Some(cache_dir) = setup_cache_dir {
+        apt::save_to_cache(system_info, cache_dir, &["valgrind"])?;
+    }
+
     Ok(())
 }
 

src/run/runner/wall_time/executor.rs

@@ -13,7 +13,8 @@ use crate::run::{check_system::SystemInfo, config::Config};
 use async_trait::async_trait;
 use std::fs::canonicalize;
 use std::io::Write;
-use std::path::{Path, PathBuf};
+use std::path::Path;
+use std::path::PathBuf;
 use std::process::Command;
 use tempfile::NamedTempFile;
 
@@ -162,9 +163,9 @@ impl Executor for WallTimeExecutor {
         ExecutorName::WallTime
     }
 
-    async fn setup(&self, _system_info: &SystemInfo) -> Result<()> {
+    async fn setup(&self, system_info: &SystemInfo, setup_cache_dir: Option<&Path>) -> Result<()> {
         if self.perf.is_some() {
-            PerfRunner::setup_environment()?;
+            PerfRunner::setup_environment(system_info, setup_cache_dir)?;
         }
 
         Ok(())