feat: use bench_pids filters when harvesting symbols from perf file

This fixes a regression introduced in 8b37208.
We now filter pids using `bench_pids`, except for `exec-harness` integrations,
where we take all pids.

Author: GuillaumeLagrange

src/executor/shared/fifo.rs

@@ -69,6 +69,15 @@ impl GenericFifo {
 pub struct FifoBenchmarkData {
     /// Name and version of the integration
     pub integration: Option<(String, String)>,
+    pub bench_pids: HashSet<pid_t>,
+}
+
+impl FifoBenchmarkData {
+    pub fn is_exec_harness(&self) -> bool {
+        self.integration
+            .as_ref()
+            .is_some_and(|(name, _)| name == "exec-harness")
+    }
 }
 
 pub struct RunnerFifo {
@@ -254,7 +263,10 @@ impl RunnerFifo {
                     );
                     let marker_result =
                         ExecutionTimestamps::new(&bench_order_by_timestamp, &markers);
-                    let fifo_data = FifoBenchmarkData { integration };
+                    let fifo_data = FifoBenchmarkData {
+                        integration,
+                        bench_pids,
+                    };
                     return Ok((marker_result, fifo_data, exit_status));
                 }
                 Err(e) => return Err(anyhow::Error::from(e)),

src/executor/wall_time/perf/mod.rs

@@ -283,11 +283,17 @@ impl BenchmarkData {
         let path_ref = path.as_ref();
 
         debug!("Reading perf data from file for mmap extraction");
+        let pid_filter = if self.fifo_data.is_exec_harness() {
+            parse_perf_file::PidFilter::All
+        } else {
+            parse_perf_file::PidFilter::TrackedPids(self.fifo_data.bench_pids.clone())
+        };
         let MemmapRecordsOutput {
             symbols_by_pid,
             unwind_data_by_pid,
+            tracked_pids,
         } = {
-            parse_perf_file::parse_for_memmap2(perf_file_path).map_err(|e| {
+            parse_perf_file::parse_for_memmap2(perf_file_path, pid_filter).map_err(|e| {
                 error!("Failed to parse perf file: {e}");
                 BenchmarkDataSaveError::FailedToParsePerfFile
             })?
@@ -296,15 +302,14 @@ impl BenchmarkData {
         // 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 bench_pids = symbols_by_pid.keys().copied().collect();
-        debug!("Harvesting perf maps and jit dumps for pids: {bench_pids:?}");
-        harvest_perf_maps_for_pids(path_ref, &bench_pids)
+        debug!("Harvesting perf maps and jit dumps for pids: {tracked_pids:?}");
+        harvest_perf_maps_for_pids(path_ref, &tracked_pids)
             .await
             .map_err(|e| {
                 error!("Failed to harvest perf maps: {e}");
                 BenchmarkDataSaveError::FailedToHarvestPerfMaps
             })?;
-        harvest_perf_jit_for_pids(path_ref, &bench_pids)
+        harvest_perf_jit_for_pids(path_ref, &tracked_pids)
             .await
             .map_err(|e| {
                 error!("Failed to harvest jit dumps: {e}");

src/executor/wall_time/perf/parse_perf_file.rs

@@ -8,14 +8,23 @@ use linux_perf_data::linux_perf_event_reader::EventRecord;
 use linux_perf_data::linux_perf_event_reader::RecordType;
 use runner_shared::unwind_data::UnwindData;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::path::Path;
 
 pub struct MemmapRecordsOutput {
     pub symbols_by_pid: HashMap<pid_t, ProcessSymbols>,
     pub unwind_data_by_pid: HashMap<pid_t, Vec<UnwindData>>,
+    pub tracked_pids: HashSet<pid_t>,
 }
 
-pub(super) fn parse_for_memmap2<P: AsRef<Path>>(perf_file_path: P) -> Result<MemmapRecordsOutput> {
+/// Parse the perf file at `perf_file_path` and look for MMAP2 records for the given `pids`.
+/// If the pids filter is empty, all MMAP2 records will be parsed.
+///
+/// Returns process symbols and unwind data for the executable mappings found in the perf file.
+pub fn parse_for_memmap2<P: AsRef<Path>>(
+    perf_file_path: P,
+    mut pid_filter: PidFilter,
+) -> Result<MemmapRecordsOutput> {
     let mut symbols_by_pid = HashMap::<pid_t, ProcessSymbols>::new();
     let mut unwind_data_by_pid = HashMap::<pid_t, Vec<UnwindData>>::new();
 
@@ -37,85 +46,157 @@ pub(super) fn parse_for_memmap2<P: AsRef<Path>>(perf_file_path: P) -> Result<Mem
 
         // Check the type from the raw record to avoid parsing overhead since we do not care about
         // most records.
-        if record.record_type != RecordType::MMAP2 {
-            continue;
+        match record.record_type {
+            RecordType::FORK => {
+                // Process fork events to track children (and children of children) of filtered PIDs
+                let Ok(parsed_record) = record.parse() else {
+                    continue;
+                };
+
+                let EventRecord::Fork(fork_record) = parsed_record else {
+                    continue;
+                };
+
+                if pid_filter.add_child_if_parent_tracked(fork_record.ppid, fork_record.pid) {
+                    trace!(
+                        "Fork: Tracking child PID {} from parent PID {}",
+                        fork_record.pid, fork_record.ppid
+                    );
+                }
+            }
+            RecordType::MMAP2 => {
+                let Ok(parsed_record) = record.parse() else {
+                    continue;
+                };
+
+                // Should never fail since we already checked the type in the raw record
+                let EventRecord::Mmap2(mmap2_record) = parsed_record else {
+                    continue;
+                };
+
+                // Filter on pid early to avoid string allocation for unwanted records
+                if !pid_filter.should_include(mmap2_record.pid) {
+                    continue;
+                }
+
+                process_mmap2_record(mmap2_record, &mut symbols_by_pid, &mut unwind_data_by_pid);
+            }
+            _ => continue,
         }
+    }
 
-        let Ok(parsed_record) = record.parse() else {
-            continue;
-        };
-
-        // Should never fail since we already checked the type in the raw record
-        let EventRecord::Mmap2(record) = parsed_record else {
-            continue;
-        };
+    // Retrieve the set of PIDs we ended up tracking after processing all records
+    let tracked_pids: HashSet<pid_t> = match pid_filter {
+        PidFilter::All => symbols_by_pid.keys().copied().collect(),
+        PidFilter::TrackedPids(tracked) => tracked,
+    };
 
-        // Check PROT_EXEC early to avoid string allocation for non-executable mappings
-        if record.protection as i32 & libc::PROT_EXEC == 0 {
-            continue;
-        }
+    Ok(MemmapRecordsOutput {
+        symbols_by_pid,
+        unwind_data_by_pid,
+        tracked_pids,
+    })
+}
 
-        // Filter on raw bytes before allocating a String
-        let path_slice: &[u8] = &record.path.as_slice();
+/// PID filter for parsing perf records
+pub enum PidFilter {
+    /// Parse records for all PIDs
+    All,
+    /// Parse records only for specific PIDs and their children
+    TrackedPids(HashSet<pid_t>),
+}
 
-        // Skip anonymous mappings
-        if path_slice == b"//anon" {
-            continue;
+impl PidFilter {
+    /// Check if a PID should be included in parsing
+    fn should_include(&self, pid: pid_t) -> bool {
+        match self {
+            PidFilter::All => true,
+            PidFilter::TrackedPids(tracked_pids) => tracked_pids.contains(&pid),
         }
+    }
 
-        // Skip special mappings like [vdso], [heap], etc.
-        if path_slice.first() == Some(&b'[') && path_slice.last() == Some(&b']') {
-            continue;
+    /// Add a child PID to the filter if we're tracking its parent
+    /// Returns true if the child was added
+    fn add_child_if_parent_tracked(&mut self, parent_pid: pid_t, child_pid: pid_t) -> bool {
+        match self {
+            PidFilter::All => false, // Already tracking all PIDs
+            PidFilter::TrackedPids(tracked_pids) => {
+                if tracked_pids.contains(&parent_pid) {
+                    tracked_pids.insert(child_pid)
+                } else {
+                    false
+                }
+            }
         }
+    }
+}
+
+/// Process a single MMAP2 record and add it to the symbols and unwind data maps
+fn process_mmap2_record(
+    record: linux_perf_data::linux_perf_event_reader::Mmap2Record,
+    symbols_by_pid: &mut HashMap<pid_t, ProcessSymbols>,
+    unwind_data_by_pid: &mut HashMap<pid_t, Vec<UnwindData>>,
+) {
+    // Check PROT_EXEC early to avoid string allocation for non-executable mappings
+    if record.protection as i32 & libc::PROT_EXEC == 0 {
+        return;
+    }
 
-        let record_path_string = String::from_utf8_lossy(path_slice).into_owned();
-        let end_addr = record.address + record.length;
+    // Filter on raw bytes before allocating a String
+    let path_slice: &[u8] = &record.path.as_slice();
 
-        trace!(
-            "Mapping: Pid {}: {:016x}-{:016x} {:08x} {:?} (Prot {:?})",
+    // Skip anonymous mappings
+    if path_slice == b"//anon" {
+        return;
+    }
+
+    // Skip special mappings like [vdso], [heap], etc.
+    if path_slice.first() == Some(&b'[') && path_slice.last() == Some(&b']') {
+        return;
+    }
+
+    let record_path_string = String::from_utf8_lossy(path_slice).into_owned();
+    let end_addr = record.address + record.length;
+
+    trace!(
+        "Mapping: Pid {}: {:016x}-{:016x} {:08x} {:?} (Prot {:?})",
+        record.pid,
+        record.address,
+        end_addr,
+        record.page_offset,
+        record_path_string,
+        record.protection,
+    );
+    symbols_by_pid
+        .entry(record.pid)
+        .or_insert(ProcessSymbols::new(record.pid))
+        .add_mapping(
             record.pid,
+            &record_path_string,
             record.address,
             end_addr,
             record.page_offset,
-            record_path_string,
-            record.protection,
         );
-        symbols_by_pid
-            .entry(record.pid)
-            .or_insert(ProcessSymbols::new(record.pid))
-            .add_mapping(
-                record.pid,
-                &record_path_string,
-                record.address,
-                end_addr,
-                record.page_offset,
-            );
 
-        match UnwindData::new(
-            record_path_string.as_bytes(),
-            record.page_offset,
-            record.address,
-            end_addr,
-            None,
-        ) {
-            Ok(unwind_data) => {
-                unwind_data_by_pid
-                    .entry(record.pid)
-                    .or_default()
-                    .push(unwind_data);
-                trace!(
-                    "Added unwind data for {record_path_string} ({:x} - {:x})",
-                    record.address, end_addr
-                );
-            }
-            Err(error) => {
-                debug!("Failed to create unwind data for module {record_path_string}: {error}");
-            }
+    match UnwindData::new(
+        record_path_string.as_bytes(),
+        record.page_offset,
+        record.address,
+        end_addr,
+        None,
+    ) {
+        Ok(unwind_data) => {
+            unwind_data_by_pid
+                .entry(record.pid)
+                .or_default()
+                .push(unwind_data);
+            trace!(
+                "Added unwind data for {record_path_string} ({:x} - {:x})",
+                record.address, end_addr
+            );
+        }
+        Err(error) => {
+            debug!("Failed to create unwind data for module {record_path_string}: {error}");
         }
     }
-
-    Ok(MemmapRecordsOutput {
-        symbols_by_pid,
-        unwind_data_by_pid,
-    })
 }