feat: add pre- and post-benchmark scripts

Author: not-matthias

src/run/runner/wall_time/perf/mod.rs

@@ -31,6 +31,52 @@ pub mod unwind_data;
 
 const PERF_DATA_PREFIX: &str = "perf.data.";
 
+struct EnvGuard;
+
+impl EnvGuard {
+    fn execute_script_from_env(script_env_var: &str) -> anyhow::Result<()> {
+        let Ok(script_path) = std::env::var(script_env_var) else {
+            debug!("Couldn't find {script_env_var}, skipping script execution");
+            return Ok(());
+        };
+
+        if script_path.is_empty() {
+            return Ok(());
+        }
+
+        let path = std::path::Path::new(&script_path);
+        if !path.exists() || !path.is_file() {
+            warn!("Script not found or not a file: {}", script_path);
+            return Ok(());
+        }
+
+        let output = Command::new("bash").args([&script_path]).output()?;
+        if !output.status.success() {
+            info!("stdout: {}", String::from_utf8_lossy(&output.stdout));
+            error!("stderr: {}", String::from_utf8_lossy(&output.stderr));
+            bail!("Failed to execute script: {}", script_path);
+        }
+
+        Ok(())
+    }
+
+    pub fn setup() -> Self {
+        if let Err(e) = Self::execute_script_from_env("CODSPEED_PRE_STARTUP_SCRIPT") {
+            warn!("Failed to execute pre-startup script: {}", e);
+            println!("asdf: {e}");
+        }
+        Self
+    }
+}
+
+impl Drop for EnvGuard {
+    fn drop(&mut self) {
+        if let Err(e) = Self::execute_script_from_env("CODSPEED_POST_CLEANUP_SCRIPT") {
+            warn!("Failed to execute post-cleanup script: {}", e);
+        }
+    }
+}
+
 pub struct PerfRunner {
     perf_dir: TempDir,
     benchmark_data: OnceCell<BenchmarkData>,
@@ -126,7 +172,11 @@ impl PerfRunner {
 
             Ok(())
         };
-        run_command_with_log_pipe_and_callback(cmd, on_process_started).await
+
+        {
+            let _guard = EnvGuard::setup();
+            run_command_with_log_pipe_and_callback(cmd, on_process_started).await
+        }
     }
 
     pub async fn save_files_to(&self, profile_folder: &PathBuf) -> anyhow::Result<()> {
@@ -389,3 +439,82 @@ impl BenchmarkData {
         self.bench_order_by_pid.values().map(|v| v.len()).sum()
     }
 }
+#[cfg(test)]
+mod tests {
+    use tempfile::NamedTempFile;
+
+    use super::*;
+    use std::{
+        io::{Read, Write},
+        os::unix::fs::PermissionsExt,
+    };
+
+    fn with_env<F>(vars: &[(&str, &str)], mut f: F)
+    where
+        F: FnMut(),
+    {
+        let original_vars: Vec<(&str, std::result::Result<String, std::env::VarError>)> =
+            vars.iter().map(|(k, _)| (*k, std::env::var(*k))).collect();
+
+        for (k, v) in vars {
+            std::env::set_var(k, v);
+        }
+
+        f();
+
+        for (k, v) in original_vars {
+            if let Ok(val) = v {
+                std::env::set_var(k, val);
+            } else {
+                std::env::remove_var(k);
+            }
+        }
+    }
+
+    #[test]
+    fn test_env_guard_no_crash() {
+        fn create_run_script(content: &str) -> anyhow::Result<NamedTempFile> {
+            let rwx = std::fs::Permissions::from_mode(0o777);
+            let mut script_file = tempfile::Builder::new()
+                .suffix(".sh")
+                .permissions(rwx)
+                .keep(true)
+                .tempfile()?;
+            script_file.write_all(content.as_bytes())?;
+
+            Ok(script_file)
+        }
+
+        let mut tmp_dst = tempfile::NamedTempFile::new().unwrap();
+
+        let pre_script = create_run_script(&format!(
+            "#!/usr/bin/env bash\necho \"pre\" >> {}",
+            tmp_dst.path().display()
+        ))
+        .unwrap();
+        let post_script = create_run_script(&format!(
+            "#!/usr/bin/env bash\necho \"post\" >> {}",
+            tmp_dst.path().display()
+        ))
+        .unwrap();
+
+        let env_vars = [
+            (
+                "CODSPEED_PRE_STARTUP_SCRIPT",
+                &*pre_script.path().to_string_lossy(),
+            ),
+            (
+                "CODSPEED_POST_CLEANUP_SCRIPT",
+                &post_script.path().to_string_lossy(),
+            ),
+        ];
+
+        with_env(&env_vars, || {
+            let _guard = EnvGuard::setup();
+        });
+
+        let mut result = String::new();
+        tmp_dst.read_to_string(&mut result).unwrap();
+        assert_eq!(result, "pre\npost\n");
+    }
+}