fixup! refactor(walltime): port PerfRunner to Profiler trait

Author: not-matthias

src/executor/wall_time/executor.rs

@@ -13,15 +13,18 @@ use crate::executor::helpers::run_command_with_log_pipe::run_command_with_log_pi
 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::FifoBenchmarkData;
 use crate::executor::shared::fifo::RunnerFifo;
 use crate::executor::{ExecutionContext, ExecutorName, ExecutorSupport};
 use crate::instruments::mongo_tracer::MongoTracer;
 use crate::prelude::*;
 use crate::runner_mode::RunnerMode;
 use crate::system::{SupportedOs, SystemInfo};
 use async_trait::async_trait;
+use runner_shared::artifacts::ExecutionTimestamps;
 use runner_shared::fifo::Command as FifoCommand;
 use runner_shared::fifo::IntegrationMode;
+use std::cell::OnceCell;
 use std::fs::canonicalize;
 use std::io::Write;
 use std::path::Path;
@@ -79,6 +82,10 @@ impl Drop for HookScriptsGuard {
 
 pub struct WallTimeExecutor {
     profiler: Option<Box<dyn Profiler>>,
+
+    /// Stashed by [`Executor::run`] and consumed by [`Executor::teardown`] to
+    /// hand the run's outputs to [`Profiler::finalize`].
+    benchmark_state: OnceCell<(FifoBenchmarkData, ExecutionTimestamps)>,
 }
 
 impl WallTimeExecutor {
@@ -90,7 +97,10 @@ impl WallTimeExecutor {
         } else {
             None
         };
-        Self { profiler }
+        Self {
+            profiler,
+            benchmark_state: OnceCell::new(),
+        }
     }
 
     fn walltime_bench_cmd(
@@ -161,13 +171,21 @@ impl Executor for WallTimeExecutor {
         let (_env_file, _script_file, cmd_builder) =
             WallTimeExecutor::walltime_bench_cmd(&execution_context.config, execution_context)?;
 
-        let status = match self.profiler.as_mut() {
+        // Split-borrow `self` so the closure inside `run_with_profiler` can
+        // capture `benchmark_state` while we hold `&mut profiler`.
+        let Self {
+            profiler,
+            benchmark_state,
+        } = self;
+
+        let status = match profiler.as_mut() {
             Some(profiler) if execution_context.config.enable_profiler => {
                 run_with_profiler(
                     profiler.as_mut(),
                     cmd_builder,
                     &execution_context.config,
                     &execution_context.profile_folder,
+                    benchmark_state,
                 )
                 .await
             }
@@ -196,6 +214,14 @@ impl Executor for WallTimeExecutor {
     async fn teardown(&self, execution_context: &ExecutionContext) -> Result<()> {
         debug!("Copying files to the profile folder");
 
+        if let (Some(profiler), Some((fifo_data, timestamps))) =
+            (&self.profiler, self.benchmark_state.get())
+        {
+            profiler
+                .finalize(fifo_data, timestamps, &execution_context.profile_folder)
+                .await?;
+        }
+
         validate_walltime_results(
             &execution_context.profile_folder,
             execution_context.config.allow_empty,
@@ -207,12 +233,13 @@ impl Executor for WallTimeExecutor {
 
 /// Drive a single benchmark run through a [`Profiler`]: wrap the command,
 /// spawn it, dispatch FIFO commands from the integration into the profiler's
-/// hooks, and finalize once the child exits.
+/// hooks, and stash the run's outputs for [`Profiler::finalize`] in teardown.
 async fn run_with_profiler(
     profiler: &mut dyn Profiler,
     cmd_builder: CommandBuilder,
     config: &ExecutorConfig,
     profile_folder: &Path,
+    benchmark_state: &OnceCell<(FifoBenchmarkData, ExecutionTimestamps)>,
 ) -> Result<std::process::ExitStatus> {
     let wrapped = profiler
         .wrap_command(cmd_builder, config, profile_folder)
@@ -247,9 +274,7 @@ async fn run_with_profiler(
         let (timestamps, fifo_data, exit_status) =
             runner_fifo.handle_fifo_messages(&mut child, on_cmd).await?;
 
-        profiler
-            .finalize(fifo_data, timestamps, profile_folder)
-            .await?;
+        let _ = benchmark_state.set((fifo_data, timestamps));
 
         Ok(exit_status)
     })

src/executor/wall_time/profiler/mod.rs

@@ -74,8 +74,8 @@ pub trait Profiler {
     /// info) and write the unified profile metadata into `profile_folder`.
     async fn finalize(
         &self,
-        fifo_data: FifoBenchmarkData,
-        timestamps: ExecutionTimestamps,
+        fifo_data: &FifoBenchmarkData,
+        timestamps: &ExecutionTimestamps,
         profile_folder: &Path,
     ) -> anyhow::Result<()>;
 }

src/executor/wall_time/profiler/perf/mod.rs

@@ -201,8 +201,8 @@ impl Profiler for PerfProfiler {
 
     async fn finalize(
         &self,
-        fifo_data: FifoBenchmarkData,
-        timestamps: ExecutionTimestamps,
+        fifo_data: &FifoBenchmarkData,
+        timestamps: &ExecutionTimestamps,
         profile_folder: &Path,
     ) -> anyhow::Result<()> {
         let start = std::time::Instant::now();
@@ -230,9 +230,9 @@ impl Profiler for PerfProfiler {
     }
 }
 
-struct BenchmarkData {
-    fifo_data: FifoBenchmarkData,
-    marker_result: ExecutionTimestamps,
+struct BenchmarkData<'a> {
+    fifo_data: &'a FifoBenchmarkData,
+    marker_result: &'a ExecutionTimestamps,
 }
 
 #[derive(Debug)]
@@ -243,7 +243,7 @@ enum BenchmarkDataSaveError {
     FailedToHarvestJitDumps,
 }
 
-impl BenchmarkData {
+impl BenchmarkData<'_> {
     async fn save_to(
         &self,
         path: &Path,

src/executor/wall_time/profiler/samply/mod.rs

@@ -76,8 +76,8 @@ impl Profiler for SamplyProfiler {
 
     async fn finalize(
         &self,
-        fifo_data: FifoBenchmarkData,
-        timestamps: ExecutionTimestamps,
+        fifo_data: &FifoBenchmarkData,
+        timestamps: &ExecutionTimestamps,
         profile_folder: &Path,
     ) -> anyhow::Result<()> {
         let Some(integration) = fifo_data.integration.clone() else {