Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/executor/helpers/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
65 changes: 65 additions & 0 deletions src/executor/helpers/run_with_env.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
//! Forwards the current environment to a command when run with sudo.

use crate::executor::helpers::command::CommandBuilder;
use crate::prelude::*;
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))
}

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)
}
25 changes: 17 additions & 8 deletions src/executor/memory/executor.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,16 @@
use crate::executor::ExecutorName;
use crate::executor::helpers::command::CommandBuilder;
use crate::executor::helpers::env::get_base_injected_env;
use crate::executor::helpers::get_bench_command::get_bench_command;
use crate::executor::helpers::run_command_with_log_pipe::run_command_with_log_pipe_and_callback;
use crate::executor::helpers::run_with_env::wrap_with_env;
use crate::executor::helpers::run_with_sudo::wrap_with_sudo;
use crate::executor::shared::fifo::RunnerFifo;
use crate::executor::{ExecutionContext, Executor};
use crate::instruments::mongo_tracer::MongoTracer;
use crate::prelude::*;
use crate::run::check_system::SystemInfo;
use crate::runner_mode::RunnerMode;
use async_trait::async_trait;
use ipc_channel::ipc;
use memtrack::MemtrackIpcClient;
Expand All @@ -18,31 +21,37 @@ 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(
execution_context: &ExecutionContext,
) -> 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(&execution_context.config)?);
cmd_builder.arg("--output");
cmd_builder.arg(execution_context.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(&execution_context.config)?);

// Wrap command with environment forwarding
let extra_env =
get_base_injected_env(RunnerMode::Memory, &execution_context.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))
}
}

Expand Down Expand Up @@ -82,7 +91,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(execution_context.profile_folder.join("results"))?;

let (ipc, cmd_builder) = Self::build_memtrack_command(execution_context)?;
let (ipc, cmd_builder, _env_file) = Self::build_memtrack_command(execution_context)?;
let cmd = wrap_with_sudo(cmd_builder)?.build();
debug!("cmd: {cmd:?}");

Expand Down
25 changes: 25 additions & 0 deletions src/executor/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -417,4 +417,29 @@ mod memory {
)
.await;
}

#[test_log::test(tokio::test)]
async fn test_memory_executor_forwards_path() {
let custom_path = "/custom/test/path";
let current_path = std::env::var("PATH").unwrap();
let modified_path = format!("{custom_path}:{current_path}");

let cmd = format!(
r#"
if ! echo "$PATH" | grep -q "{custom_path}"; then
echo "FAIL: PATH does not contain custom path {custom_path}"
echo "Got PATH: $PATH"
exit 1
fi
"#
);
let config = memory_config(&cmd);
let (execution_context, _temp_dir) = create_test_setup(config).await;
let (_permit, _lock, executor) = get_memory_executor().await;

temp_env::async_with_vars(&[("PATH", Some(&modified_path))], async {
executor.run(&execution_context, &None).await.unwrap();
})
.await;
}
}
63 changes: 19 additions & 44 deletions src/executor/wall_time/executor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use crate::executor::helpers::get_bench_command::get_bench_command;
use crate::executor::helpers::introspected_golang;
use crate::executor::helpers::introspected_nodejs;
use crate::executor::helpers::run_command_with_log_pipe::run_command_with_log_pipe;
use crate::executor::helpers::run_with_env::wrap_with_env;
use crate::executor::helpers::run_with_sudo::wrap_with_sudo;
use crate::executor::{ExecutionContext, ExecutorName};
use crate::instruments::mongo_tracer::MongoTracer;
Expand Down Expand Up @@ -70,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>,
}
Expand All @@ -110,19 +92,10 @@ impl WallTimeExecutor {
config: &Config,
execution_context: &ExecutionContext,
) -> 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, &execution_context.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(),
Expand All @@ -132,18 +105,18 @@ 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, &execution_context.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 {
Expand All @@ -164,9 +137,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))
}
}

Expand Down
Loading