fix(runner): forward environment in memory executor

Author: not-matthias

src/run/runner/helpers/mod.rs

@@ -6,4 +6,5 @@ pub mod introspected_golang;
 pub mod introspected_nodejs;
 pub mod profile_folder;
 pub mod run_command_with_log_pipe;
+pub mod run_with_env;
 pub mod run_with_sudo;

src/run/runner/helpers/run_with_env.rs

@@ -0,0 +1,65 @@
+//! Forwards the current environment to a command when run with sudo.
+
+use crate::prelude::*;
+use crate::run::runner::helpers::command::CommandBuilder;
+use std::collections::HashMap;
+use std::io::Write;
+use std::process::Command;
+use tempfile::NamedTempFile;
+
+/// Returns a list of exported environment variables which can be loaded with `source` in a shell.
+///
+/// Example: `declare -x outputs="out"`
+fn get_exported_system_env() -> Result<String> {
+    let output = Command::new("bash")
+        .arg("-c")
+        .arg("export")
+        .output()
+        .context("Failed to run `export`")?;
+    if !output.status.success() {
+        bail!(
+            "Failed to get system environment variables: {}",
+            String::from_utf8_lossy(&output.stderr)
+        );
+    }
+
+    String::from_utf8(output.stdout).context("Failed to parse export output as UTF-8")
+}
+
+/// Wraps a command to run with environment variables forwarded.
+///
+/// # Returns
+/// Returns a tuple of (CommandBuilder, NamedTempFile) where:
+/// - CommandBuilder is wrapped with bash to source the environment and run the original command
+/// - NamedTempFile is the environment file that must be kept alive until command execution
+pub fn wrap_with_env(
+    mut cmd_builder: CommandBuilder,
+    extra_env: &HashMap<&'static str, String>,
+) -> Result<(CommandBuilder, NamedTempFile)> {
+    let env_file = create_env_file(extra_env)?;
+
+    // Create bash command that sources the env file and runs the original command
+    let original_command = cmd_builder.as_command_line();
+    let bash_command = format!(
+        "source {} && {}",
+        env_file.path().display(),
+        original_command
+    );
+    cmd_builder.wrap("bash", ["-c", &bash_command]);
+
+    Ok((cmd_builder, env_file))
+}
+
+pub fn create_env_file(extra_env: &HashMap<&'static str, String>) -> Result<NamedTempFile> {
+    let system_env = get_exported_system_env()?;
+    let base_injected_env = extra_env
+        .iter()
+        .map(|(k, v)| format!("export {k}={v}"))
+        .collect::<Vec<_>>()
+        .join("\n");
+
+    // Create and return the environment file
+    let mut env_file = NamedTempFile::new()?;
+    env_file.write_all(format!("{system_env}\n{base_injected_env}").as_bytes())?;
+    Ok(env_file)
+}

src/run/runner/memory/executor.rs

@@ -2,11 +2,13 @@ use crate::prelude::*;
 use crate::run::instruments::mongo_tracer::MongoTracer;
 use crate::run::runner::executor::Executor;
 use crate::run::runner::helpers::command::CommandBuilder;
+use crate::run::runner::helpers::env::get_base_injected_env;
 use crate::run::runner::helpers::get_bench_command::get_bench_command;
 use crate::run::runner::helpers::run_command_with_log_pipe::run_command_with_log_pipe_and_callback;
+use crate::run::runner::helpers::run_with_env::wrap_with_env;
 use crate::run::runner::helpers::run_with_sudo::wrap_with_sudo;
 use crate::run::runner::shared::fifo::RunnerFifo;
-use crate::run::runner::{ExecutorName, RunData};
+use crate::run::runner::{ExecutorName, RunData, RunnerMode};
 use crate::run::{check_system::SystemInfo, config::Config};
 use async_trait::async_trait;
 use ipc_channel::ipc;
@@ -18,32 +20,38 @@ use runner_shared::fifo::IntegrationMode;
 use std::path::Path;
 use std::process::Command;
 use std::rc::Rc;
+use tempfile::NamedTempFile;
 
 pub struct MemoryExecutor;
 
 impl MemoryExecutor {
     fn build_memtrack_command(
         config: &Config,
         run_data: &RunData,
-    ) -> Result<(MemtrackIpcServer, CommandBuilder)> {
+    ) -> Result<(MemtrackIpcServer, CommandBuilder, NamedTempFile)> {
         // FIXME: We only support native languages for now
 
         // Find memtrack binary - check env variable or use default command name
         let memtrack_path = std::env::var("CODSPEED_MEMTRACK_BINARY")
             .unwrap_or_else(|_| "codspeed-memtrack".to_string());
 
+        // Setup memtrack IPC server
+        let (ipc_server, server_name) = ipc::IpcOneShotServer::new()?;
+
+        // Build the memtrack command
         let mut cmd_builder = CommandBuilder::new(memtrack_path);
         cmd_builder.arg("track");
-        cmd_builder.arg(get_bench_command(config)?);
         cmd_builder.arg("--output");
         cmd_builder.arg(run_data.profile_folder.join("results"));
-
-        // Setup memtrack IPC server
-        let (ipc_server, server_name) = ipc::IpcOneShotServer::new()?;
         cmd_builder.arg("--ipc-server");
         cmd_builder.arg(server_name);
+        cmd_builder.arg(get_bench_command(config)?);
+
+        // Wrap command with environment forwarding
+        let extra_env = get_base_injected_env(RunnerMode::Memory, &run_data.profile_folder);
+        let (cmd_builder, env_file) = wrap_with_env(cmd_builder, &extra_env)?;
 
-        Ok((ipc_server, cmd_builder))
+        Ok((ipc_server, cmd_builder, env_file))
     }
 }
 
@@ -85,7 +93,7 @@ impl Executor for MemoryExecutor {
         // Create the results/ directory inside the profile folder to avoid having memtrack create it with wrong permissions
         std::fs::create_dir_all(run_data.profile_folder.join("results"))?;
 
-        let (ipc, cmd_builder) = Self::build_memtrack_command(config, run_data)?;
+        let (ipc, cmd_builder, _env_file) = Self::build_memtrack_command(config, run_data)?;
         let cmd = wrap_with_sudo(cmd_builder)?.build();
         debug!("cmd: {cmd:?}");
 

src/run/runner/wall_time/executor.rs

@@ -5,11 +5,13 @@ use crate::run::RunnerMode;
 use crate::run::instruments::mongo_tracer::MongoTracer;
 use crate::run::runner::executor::Executor;
 use crate::run::runner::helpers::command::CommandBuilder;
-use crate::run::runner::helpers::env::{get_base_injected_env, is_codspeed_debug_enabled};
+use crate::run::runner::helpers::env::get_base_injected_env;
+use crate::run::runner::helpers::env::is_codspeed_debug_enabled;
 use crate::run::runner::helpers::get_bench_command::get_bench_command;
 use crate::run::runner::helpers::introspected_golang;
 use crate::run::runner::helpers::introspected_nodejs;
 use crate::run::runner::helpers::run_command_with_log_pipe::run_command_with_log_pipe;
+use crate::run::runner::helpers::run_with_env::wrap_with_env;
 use crate::run::runner::helpers::run_with_sudo::wrap_with_sudo;
 use crate::run::runner::{ExecutorName, RunData};
 use crate::run::{check_system::SystemInfo, config::Config};
@@ -69,25 +71,6 @@ impl Drop for HookScriptsGuard {
     }
 }
 
-/// Returns a list of exported environment variables which can be loaded with `source` in a shell.
-///
-/// Example: `declare -x outputs="out"`
-fn get_exported_system_env() -> Result<String> {
-    let output = Command::new("bash")
-        .arg("-c")
-        .arg("export")
-        .output()
-        .context("Failed to run `export`")?;
-    if !output.status.success() {
-        bail!(
-            "Failed to get system environment variables: {}",
-            String::from_utf8_lossy(&output.stderr)
-        );
-    }
-
-    String::from_utf8(output.stdout).context("Failed to parse export output as UTF-8")
-}
-
 pub struct WallTimeExecutor {
     perf: Option<PerfRunner>,
 }
@@ -103,19 +86,10 @@ impl WallTimeExecutor {
         config: &Config,
         run_data: &RunData,
     ) -> Result<(NamedTempFile, NamedTempFile, CommandBuilder)> {
-        let bench_cmd = get_bench_command(config)?;
-
-        let system_env = get_exported_system_env()?;
-        let base_injected_env =
-            get_base_injected_env(RunnerMode::Walltime, &run_data.profile_folder)
-                .into_iter()
-                .map(|(k, v)| format!("export {k}={v}",))
-                .collect::<Vec<_>>()
-                .join("\n");
-
+        // Build additional PATH environment variables
         let path_env = std::env::var("PATH").unwrap_or_default();
-        let path_env = format!(
-            "export PATH={}:{}:{}",
+        let path_value = format!(
+            "{}:{}:{}",
             introspected_nodejs::setup()
                 .map_err(|e| anyhow!("failed to setup NodeJS introspection. {e}"))?
                 .to_string_lossy(),
@@ -125,18 +99,17 @@ impl WallTimeExecutor {
             path_env
         );
 
-        let combined_env = format!("{system_env}\n{base_injected_env}\n{path_env}");
+        let mut extra_env = get_base_injected_env(RunnerMode::Walltime, &run_data.profile_folder);
+        extra_env.insert("PATH", path_value);
 
-        let mut env_file = NamedTempFile::new()?;
-        env_file.write_all(combined_env.as_bytes())?;
+        // We have to write the benchmark command to a script, to ensure proper formatting
+        // and to not have to manually escape everything.
+        let mut script_file = NamedTempFile::new()?;
+        script_file.write_all(get_bench_command(config)?.as_bytes())?;
 
-        let script_file = {
-            let mut file = NamedTempFile::new()?;
-            let bash_command = format!("source {} && {}", env_file.path().display(), bench_cmd);
-            debug!("Bash command: {bash_command}");
-            file.write_all(bash_command.as_bytes())?;
-            file
-        };
+        let mut bench_cmd = CommandBuilder::new("bash");
+        bench_cmd.arg(script_file.path());
+        let (mut bench_cmd, env_file) = wrap_with_env(bench_cmd, &extra_env)?;
 
         let mut cmd_builder = CommandBuilder::new("systemd-run");
         if let Some(cwd) = &config.working_directory {
@@ -157,9 +130,11 @@ impl WallTimeExecutor {
         cmd_builder.arg("--same-dir");
         cmd_builder.arg(format!("--uid={}", nix::unistd::Uid::current().as_raw()));
         cmd_builder.arg(format!("--gid={}", nix::unistd::Gid::current().as_raw()));
-        cmd_builder.args(["--", "bash"]);
-        cmd_builder.arg(script_file.path());
-        Ok((env_file, script_file, cmd_builder))
+        cmd_builder.args(["--"]);
+
+        bench_cmd.wrap_with(cmd_builder);
+
+        Ok((env_file, script_file, bench_cmd))
     }
 }