feat(runner): add experimental flag to capture allocation stacks

Author: not-matthias

src/cli/exec/mod.rs

@@ -100,6 +100,7 @@ fn build_orchestrator_config(
             .shared
             .experimental
             .experimental_disable_pythonmalloc_override,
+        memory_capture_stack: args.shared.experimental.experimental_memory_capture_stack,
     })
 }
 

src/cli/experimental.rs

@@ -26,6 +26,14 @@ pub struct ExperimentalArgs {
         value_parser = clap::builder::FalseyValueParser::new()
     )]
     pub experimental_memory_track_physical: bool,
+    /// Capture allocation call stacks in memory mode.
+    #[arg(
+        long,
+        default_value_t = false,
+        help_heading = "Experimental",
+        env = "CODSPEED_EXPERIMENTAL_MEMORY_CAPTURE_STACK"
+    )]
+    pub experimental_memory_capture_stack: bool,
 
     /// Do not set PYTHONMALLOC for simulation runs.
     #[arg(
@@ -58,6 +66,9 @@ impl ExperimentalArgs {
         if self.experimental_disable_pythonmalloc_override {
             flags.push("--experimental-disable-pythonmalloc-override");
         }
+        if self.experimental_memory_capture_stack {
+            flags.push("--experimental-memory-capture-stack");
+        }
         flags
     }
 

src/cli/run/mod.rs

@@ -81,6 +81,7 @@ impl RunArgs {
                 },
                 experimental: ExperimentalArgs {
                     experimental_fair_sched: false,
+                    experimental_memory_capture_stack: false,
                     experimental_cycle_estimation: false,
                     experimental_exclude_allocations: false,
                     experimental_memory_track_physical: false,
@@ -144,6 +145,7 @@ fn build_orchestrator_config(
             .shared
             .experimental
             .experimental_disable_pythonmalloc_override,
+        memory_capture_stack: args.shared.experimental.experimental_memory_capture_stack,
     })
 }
 

src/executor/config.rs

@@ -102,6 +102,9 @@ pub struct OrchestratorConfig {
     pub memory_track_physical: bool,
     /// Do not set PYTHONMALLOC for simulation runs.
     pub disable_pythonmalloc_override: bool,
+    /// Capture allocation call stacks in memory mode, so allocations can be
+    /// attributed to the code that made them.
+    pub memory_capture_stack: bool,
 }
 
 /// Per-execution configuration passed to executors.
@@ -149,6 +152,12 @@ pub struct ExecutorConfig {
     pub memory_track_physical: bool,
     /// Do not set PYTHONMALLOC for simulation runs.
     pub disable_pythonmalloc_override: bool,
+    /// Capture allocation call stacks in memory mode, so allocations can be
+    /// attributed to the code that made them.
+    ///
+    /// Only read by the memory executor, which is Linux-only.
+    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
+    pub memory_capture_stack: bool,
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -223,6 +232,7 @@ impl OrchestratorConfig {
             simulation_track_subprocess: self.simulation_track_subprocess,
             memory_track_physical: self.memory_track_physical,
             disable_pythonmalloc_override: self.disable_pythonmalloc_override,
+            memory_capture_stack: self.memory_capture_stack,
         }
     }
 }
@@ -260,6 +270,7 @@ impl OrchestratorConfig {
             simulation_track_subprocess: false,
             memory_track_physical: false,
             disable_pythonmalloc_override: false,
+            memory_capture_stack: false,
         }
     }
 }

src/executor/memory/executor.rs

@@ -74,6 +74,9 @@ impl MemoryExecutor {
         cmd_builder.arg("--ipc-server");
         cmd_builder.arg(server_name);
         cmd_builder.arg(bench_command);
+        if execution_context.config.memory_capture_stack {
+            cmd_builder.env("CODSPEED_MEMTRACK_CAPTURE_STACKS", "1");
+        }
 
         // Set working directory if specified
         if let Some(cwd) = &execution_context.config.working_directory {