feat: run python with perf jit dump

Author: not-matthias

src/run/mod.rs

@@ -133,7 +133,7 @@ pub struct RunArgs {
     pub command: Vec<String>,
 }
 
-#[derive(ValueEnum, Clone, Debug, Serialize)]
+#[derive(ValueEnum, Clone, Debug, Serialize, PartialEq)]
 #[serde(rename_all = "lowercase")]
 pub enum RunnerMode {
     Instrumentation,

src/run/runner/helpers/env.rs

@@ -8,6 +8,14 @@ pub fn get_base_injected_env(
 ) -> HashMap<&'static str, String> {
     HashMap::from([
         ("PYTHONHASHSEED", "0".into()),
+        (
+            "PYTHON_PERF_JIT_SUPPORT",
+            if mode == RunnerMode::Walltime {
+                "1".into()
+            } else {
+                "0".into()
+            },
+        ),
         ("ARCH", ARCH.into()),
         ("CODSPEED_ENV", "runner".into()),
         ("CODSPEED_RUNNER_MODE", mode.to_string()),

src/run/runner/wall_time/perf/jit_dump.rs

@@ -0,0 +1,102 @@
+use crate::{prelude::*, run::runner::wall_time::perf::unwind_data::UnwindData};
+use linux_perf_data::jitdump::{JitDumpReader, JitDumpRecord};
+use std::{
+    collections::HashSet,
+    path::{Path, PathBuf},
+};
+
+struct JitDump {
+    path: PathBuf,
+}
+
+impl JitDump {
+    pub fn new(path: PathBuf) -> Self {
+        Self { path }
+    }
+
+    /// Parses the JIT dump file and converts it into a list of `UnwindData`.
+    ///
+    /// The JIT dump file contains synthetic `eh_frame` data for jitted functions. This can be parsed and
+    /// then converted to `UnwindData` which is used for stack unwinding.
+    ///
+    /// See: https://github.com/python/cpython/blob/main/Python/perf_jit_trampoline.c
+    pub fn into_unwind_data(self) -> Result<Vec<UnwindData>> {
+        let file = std::fs::File::open(self.path)?;
+
+        let mut jit_unwind_data = Vec::new();
+        let mut current_unwind_info: Option<(Vec<u8>, Vec<u8>)> = None;
+
+        let mut reader = JitDumpReader::new(file)?;
+        while let Some(raw_record) = reader.next_record()? {
+            // The first recording is always the unwind info, followed by the code load event
+            // (see `perf_map_jit_write_entry` in https://github.com/python/cpython/blob/9743d069bd53e9d3a8f09df899ec1c906a79da24/Python/perf_jit_trampoline.c#L1163C13-L1163C37)
+            match raw_record.parse()? {
+                JitDumpRecord::CodeLoad(record) => {
+                    let name = record.function_name.as_slice();
+                    let name = String::from_utf8_lossy(&name);
+
+                    let avma_start = record.vma;
+                    let code_size = record.code_bytes.len() as u64;
+                    let avma_end = avma_start + code_size;
+
+                    let Some((eh_frame, eh_frame_hdr)) = current_unwind_info.take() else {
+                        warn!("No unwind info available for JIT code load: {name}");
+                        continue;
+                    };
+
+                    jit_unwind_data.push(UnwindData {
+                        path: format!("jit_{name}"),
+                        avma_range: avma_start..avma_end,
+                        base_avma: avma_start,
+                        eh_frame_hdr,
+                        eh_frame_hdr_svma: 0..0, // TODO: Do we need this?
+                        eh_frame,
+                        eh_frame_svma: 0..0, // TODO: Do we need this?
+                    });
+                }
+                JitDumpRecord::CodeUnwindingInfo(record) => {
+                    // Store unwind info for the next code loads
+                    current_unwind_info = Some((
+                        record.eh_frame.as_slice().to_vec(),
+                        record.eh_frame_hdr.as_slice().to_vec(),
+                    ));
+                }
+                JitDumpRecord::CodeMove(_) => {
+                    panic!("CodeMove not implemented yet");
+                }
+                _ => {
+                    warn!("Unhandled JIT dump record: {raw_record:?}");
+                }
+            }
+        }
+
+        Ok(jit_unwind_data)
+    }
+}
+
+/// Converts all the `jit-<pid>.dump` into unwind data and copies it to the profile folder.
+pub async fn harvest_perf_jit_for_pids(profile_folder: &Path, pids: &HashSet<i32>) -> Result<()> {
+    for pid in pids {
+        let name = format!("jit-{pid}.dump");
+        let path = PathBuf::from("/tmp").join(&name);
+
+        if !path.exists() {
+            continue;
+        }
+        debug!("Found JIT dump file: {path:?}");
+
+        let unwind_data = match JitDump::new(path).into_unwind_data() {
+            Ok(unwind_data) => unwind_data,
+            Err(error) => {
+                warn!("Failed to convert jit dump into unwind data: {error:?}");
+                continue;
+            }
+        };
+
+        for module in unwind_data {
+            module.save_to(profile_folder, *pid as _)?;
+        }
+    }
+
+    Ok(())
+}

src/run/runner/wall_time/perf/mod.rs

@@ -8,6 +8,7 @@ use crate::run::runner::helpers::run_command_with_log_pipe::run_command_with_log
 use crate::run::runner::helpers::setup::run_with_sudo;
 use crate::run::runner::valgrind::helpers::ignored_objects_path::get_objects_path_to_ignore;
 use crate::run::runner::valgrind::helpers::perf_maps::harvest_perf_maps_for_pids;
+use crate::run::runner::wall_time::perf::jit_dump::harvest_perf_jit_for_pids;
 use anyhow::Context;
 use fifo::{PerfFifo, RunnerFifo};
 use futures::stream::FuturesUnordered;
@@ -22,6 +23,7 @@ use std::{cell::OnceCell, collections::HashMap, process::ExitStatus};
 use tempfile::TempDir;
 use unwind_data::UnwindData;
 
+mod jit_dump;
 mod metadata;
 mod setup;
 mod shared;
@@ -195,15 +197,16 @@ impl PerfRunner {
         // Harvest the perf maps generated by python. This will copy the perf
         // maps from /tmp to the profile folder. We have to write our own perf
         // maps to these files AFTERWARDS, otherwise it'll be overwritten!
-        let perf_map_pids = futures::future::try_join_all(copy_tasks)
+        let bench_pids = futures::future::try_join_all(copy_tasks)
             .await?
             .into_iter()
             .filter_map(|result| {
                 debug!("Copy task result: {result:?}");
                 result.ok()
             })
             .collect::<HashSet<_>>();
-        harvest_perf_maps_for_pids(profile_folder, &perf_map_pids).await?;
+        harvest_perf_maps_for_pids(profile_folder, &bench_pids).await?;
+        harvest_perf_jit_for_pids(profile_folder, &bench_pids).await?;
 
         // Append perf maps, unwind info and other metadata
         if let Err(BenchmarkDataSaveError::MissingIntegration) = bench_data.save_to(profile_folder)