feat: add memory executor

Author: not-matthias

Cargo.lock

@@ -57,6 +57,8 @@ memmap2 = "0.9.5"
 nix = { version = "0.29.0", features = ["fs", "time", "user"] }
 futures = "0.3.31"
 runner-shared = { path = "crates/runner-shared" }
+memtrack = { path = "crates/memtrack", default-features = false }
+ipc-channel = "0.18"
 shellexpand = { version = "3.1.1", features = ["tilde"] }
 addr2line = "0.25"
 gimli = "0.32"

Cargo.toml

@@ -148,6 +148,7 @@ pub enum RunnerMode {
     Instrumentation,
     Simulation,
     Walltime,
+    Memory,
 }
 
 #[derive(ValueEnum, Clone, Debug, PartialEq)]

src/run/mod.rs

@@ -15,6 +15,7 @@ pub fn get_base_injected_env(
         #[allow(deprecated)]
         RunnerMode::Instrumentation | RunnerMode::Simulation => "instrumentation",
         RunnerMode::Walltime => "walltime",
+        RunnerMode::Memory => "memory",
     };
     HashMap::from([
         ("PYTHONHASHSEED", "0".into()),

src/run/runner/helpers/env.rs

@@ -10,6 +10,7 @@ pub struct RunData {
 pub enum ExecutorName {
     Valgrind,
     WallTime,
+    Memory,
 }
 
 #[allow(clippy::to_string_trait_impl)]
@@ -18,6 +19,7 @@ impl ToString for ExecutorName {
         match self {
             ExecutorName::Valgrind => "valgrind".to_string(),
             ExecutorName::WallTime => "walltime".to_string(),
+            ExecutorName::Memory => "memory".to_string(),
         }
     }
 }

src/run/runner/interfaces.rs

@@ -0,0 +1,177 @@
+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::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_sudo::wrap_with_sudo;
+use crate::run::runner::shared::fifo::RunnerFifo;
+use crate::run::runner::{ExecutorName, RunData};
+use crate::run::{check_system::SystemInfo, config::Config};
+use async_trait::async_trait;
+use ipc_channel::ipc;
+use memtrack::MemtrackIpcClient;
+use memtrack::MemtrackIpcServer;
+use runner_shared::benchmark_results::{ArtifactExt, ExecutionTimestamps};
+use runner_shared::fifo::Command as FifoCommand;
+use runner_shared::fifo::IntegrationMode;
+use std::path::Path;
+use std::process::Command;
+use std::rc::Rc;
+
+pub struct MemoryExecutor;
+
+impl MemoryExecutor {
+    fn build_memtrack_command(
+        config: &Config,
+        run_data: &RunData,
+    ) -> Result<(MemtrackIpcServer, CommandBuilder)> {
+        // 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());
+
+        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);
+
+        Ok((ipc_server, cmd_builder))
+    }
+}
+
+#[async_trait(?Send)]
+impl Executor for MemoryExecutor {
+    fn name(&self) -> ExecutorName {
+        ExecutorName::Memory
+    }
+
+    async fn setup(
+        &self,
+        _system_info: &SystemInfo,
+        _setup_cache_dir: Option<&Path>,
+    ) -> Result<()> {
+        // Validate that the codspeed-memtrack command is available
+        let memtrack_path = std::env::var("CODSPEED_MEMTRACK_BINARY")
+            .unwrap_or_else(|_| "codspeed-memtrack".to_string());
+
+        info!("Validating memtrack binary at path: {}", memtrack_path);
+        let output = Command::new(&memtrack_path).arg("--version").output()?;
+        if !output.status.success() {
+            bail!("codspeed-memtrack command is not available or failed to execute");
+        }
+
+        Ok(())
+    }
+
+    async fn run(
+        &self,
+        config: &Config,
+        _system_info: &SystemInfo,
+        run_data: &RunData,
+        _mongo_tracer: &Option<MongoTracer>,
+    ) -> Result<()> {
+        // 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 cmd = wrap_with_sudo(cmd_builder)?.build();
+        debug!("cmd: {cmd:?}");
+
+        let runner_fifo = RunnerFifo::new()?;
+        let on_process_started = async |pid| -> anyhow::Result<()> {
+            let marker_result = Self::handle_fifo(runner_fifo, pid, ipc).await?;
+
+            // Directly write to the profile folder, to avoid having to define another field
+            marker_result
+                .save_to(run_data.profile_folder.join("results"))
+                .unwrap();
+
+            Ok(())
+        };
+
+        let status = run_command_with_log_pipe_and_callback(cmd, on_process_started).await?;
+        debug!("cmd exit status: {:?}", status);
+
+        if !status.success() {
+            bail!("failed to execute memory tracker process: {status}");
+        }
+
+        Ok(())
+    }
+
+    async fn teardown(
+        &self,
+        _config: &Config,
+        _system_info: &SystemInfo,
+        _run_data: &RunData,
+    ) -> Result<()> {
+        Ok(())
+    }
+}
+
+impl MemoryExecutor {
+    async fn handle_fifo(
+        mut runner_fifo: RunnerFifo,
+        pid: u32,
+        ipc: MemtrackIpcServer,
+    ) -> anyhow::Result<ExecutionTimestamps> {
+        debug!("handle_fifo called with PID {pid}");
+
+        // Accept the IPC connection from memtrack and get the sender it sends us
+        let (_, memtrack_sender) = ipc.accept()?;
+        let ipc_client = Rc::new(MemtrackIpcClient::from_accepted(memtrack_sender));
+
+        let ipc_client_health = Rc::clone(&ipc_client);
+        let health_check = async move || {
+            // Ping memtrack via IPC to check if it's still responding
+            match ipc_client_health.ping() {
+                Ok(()) => Ok(true),
+                Err(_) => Ok(false),
+            }
+        };
+
+        let on_cmd = async move |cmd: &FifoCommand| {
+            match cmd {
+                FifoCommand::StartBenchmark => {
+                    debug!("Enabling memtrack via IPC");
+                    if let Err(e) = ipc_client.enable() {
+                        error!("Failed to enable memtrack: {e}");
+                        return Ok(FifoCommand::Err);
+                    }
+                }
+                FifoCommand::StopBenchmark => {
+                    debug!("Disabling memtrack via IPC");
+                    if let Err(e) = ipc_client.disable() {
+                        // There's a chance that memtrack has already exited here, so just log as debug
+                        debug!("Failed to disable memtrack: {e}");
+                        return Ok(FifoCommand::Err);
+                    }
+                }
+                FifoCommand::GetIntegrationMode => {
+                    return Ok(FifoCommand::IntegrationModeResponse(
+                        IntegrationMode::Analysis,
+                    ));
+                }
+                _ => {
+                    warn!("Unhandled FIFO command: {cmd:?}");
+                    return Ok(FifoCommand::Err);
+                }
+            }
+
+            Ok(FifoCommand::Ack)
+        };
+
+        let (marker_result, _) = runner_fifo
+            .handle_fifo_messages(health_check, on_cmd)
+            .await?;
+        Ok(marker_result)
+    }
+}

src/run/runner/memory/executor.rs

@@ -0,0 +1 @@
+pub mod executor;

src/run/runner/memory/mod.rs

@@ -7,6 +7,7 @@ use super::{RunnerMode, config::Config};
 mod executor;
 mod helpers;
 mod interfaces;
+mod memory;
 mod shared;
 #[cfg(test)]
 mod tests;
@@ -16,6 +17,7 @@ mod wall_time;
 use executor::Executor;
 use helpers::profile_folder::create_profile_folder;
 pub use interfaces::{ExecutorName, RunData};
+use memory::executor::MemoryExecutor;
 use valgrind::executor::ValgrindExecutor;
 use wall_time::executor::WallTimeExecutor;
 
@@ -26,6 +28,7 @@ impl Display for RunnerMode {
             RunnerMode::Instrumentation => write!(f, "instrumentation"),
             RunnerMode::Simulation => write!(f, "simulation"),
             RunnerMode::Walltime => write!(f, "walltime"),
+            RunnerMode::Memory => write!(f, "memory"),
         }
     }
 }
@@ -37,13 +40,15 @@ pub fn get_executor_from_mode(mode: &RunnerMode) -> Box<dyn Executor> {
         #[allow(deprecated)]
         RunnerMode::Instrumentation | RunnerMode::Simulation => Box::new(ValgrindExecutor),
         RunnerMode::Walltime => Box::new(WallTimeExecutor::new()),
+        RunnerMode::Memory => Box::new(MemoryExecutor),
     }
 }
 
 pub fn get_all_executors() -> Vec<Box<dyn Executor>> {
     vec![
         Box::new(ValgrindExecutor),
         Box::new(WallTimeExecutor::new()),
+        Box::new(MemoryExecutor),
     ]
 }
 

src/run/runner/mod.rs

@@ -1,7 +1,7 @@
 use crate::prelude::*;
 use anyhow::Context;
 use nix::{sys::time::TimeValLike, time::clock_gettime};
-use runner_shared::benchmark_results::MarkerResult;
+use runner_shared::benchmark_results::ExecutionTimestamps;
 use runner_shared::fifo::{Command as FifoCommand, MarkerType};
 use runner_shared::fifo::{RUNNER_ACK_FIFO, RUNNER_CTL_FIFO};
 use std::cmp::Ordering;
@@ -127,7 +127,7 @@ impl RunnerFifo {
         &mut self,
         mut health_check: impl AsyncFnMut() -> anyhow::Result<bool>,
         mut handle_cmd: impl AsyncFnMut(&FifoCommand) -> anyhow::Result<FifoCommand>,
-    ) -> anyhow::Result<(MarkerResult, FifoBenchmarkData)> {
+    ) -> anyhow::Result<(ExecutionTimestamps, FifoBenchmarkData)> {
         let mut bench_order_by_timestamp = Vec::<(u64, String)>::new();
         let mut bench_pids = HashSet::<pid_t>::new();
         let mut markers = Vec::<MarkerType>::new();
@@ -214,7 +214,7 @@ impl RunnerFifo {
             self.send_cmd(handle_cmd(&cmd).await?).await?;
         }
 
-        let marker_result = MarkerResult::new(&bench_order_by_timestamp, &markers);
+        let marker_result = ExecutionTimestamps::new(&bench_order_by_timestamp, &markers);
         let fifo_data = FifoBenchmarkData {
             integration,
             bench_pids,

src/run/runner/shared/fifo.rs

@@ -2,6 +2,7 @@ use crate::run::check_system::SystemInfo;
 use crate::run::config::Config;
 use crate::run::runner::executor::Executor;
 use crate::run::runner::interfaces::RunData;
+use crate::run::runner::memory::executor::MemoryExecutor;
 use crate::run::runner::valgrind::executor::ValgrindExecutor;
 use crate::run::{RunnerMode, runner::wall_time::executor::WallTimeExecutor};
 use rstest_reuse::{self, *};
@@ -299,3 +300,54 @@ fi
         assert!(result.is_err(), "Command should fail");
     }
 }
+
+mod memory {
+    use super::*;
+
+    async fn get_memory_executor() -> MemoryExecutor {
+        let executor = MemoryExecutor;
+        let system_info = SystemInfo::new().unwrap();
+        executor.setup(&system_info, None).await.unwrap();
+        executor
+    }
+
+    #[cfg(test)]
+    fn memory_config(command: &str) -> Config {
+        Config {
+            mode: RunnerMode::Memory,
+            command: command.to_string(),
+            ..Config::test()
+        }
+    }
+
+    #[apply(test_cases)]
+    #[tokio::test]
+    async fn test_memory_executor(#[case] cmd: &str) {
+        let (system_info, run_data, _temp_dir) = create_test_setup().await;
+        let executor = get_memory_executor().await;
+
+        let config = memory_config(cmd);
+        executor
+            .run(&config, &system_info, &run_data, &None)
+            .await
+            .unwrap();
+    }
+
+    #[apply(env_test_cases)]
+    #[tokio::test]
+    async fn test_memory_executor_with_env(#[case] env_case: (&str, &str)) {
+        let (system_info, run_data, _temp_dir) = create_test_setup().await;
+        let executor = get_memory_executor().await;
+
+        let (env_var, env_value) = env_case;
+        temp_env::async_with_vars(&[(env_var, Some(env_value))], async {
+            let cmd = env_var_validation_script(env_var, env_value);
+            let config = memory_config(&cmd);
+            executor
+                .run(&config, &system_info, &run_data, &None)
+                .await
+                .unwrap();
+        })
+        .await;
+    }
+}