From e15887036f40c9d59b51d7868a2e30013b1f19ea Mon Sep 17 00:00:00 2001 From: not-matthias Date: Thu, 23 Jul 2026 20:03:05 +0200 Subject: [PATCH] feat: add --track-allocators toggle for memory mode Add a --track-allocators flag (default on, env CODSPEED_TRACK_ALLOCATORS) to the memtrack track subcommand. When disabled, memtrack skips the allocator uprobe machinery (exec watcher + attach worker) and only emits coarse mmap/munmap/brk events, reducing overhead on allocation-heavy programs. The mmap/munmap/brk syscall tracepoints are now always attached in every memory run. The runner does not add a CLI flag for this: it relies on the CODSPEED_TRACK_ALLOCATORS environment variable being inherited by the memtrack subprocess, keeping the runner decoupled from the installed memtrack version. Standalone memtrack can still use the CLI flag. --- Cargo.lock | 21 ++++++++ crates/memtrack/Cargo.toml | 1 + crates/memtrack/src/ebpf/memtrack/tracking.rs | 10 ++++ crates/memtrack/src/ebpf/mod.rs | 2 +- crates/memtrack/src/ebpf/tracker.rs | 52 +++++++++++++------ crates/memtrack/testdata/malloc_mmap.c | 15 ++++++ crates/memtrack/tests/c_tests.rs | 49 +++++++++++++++++ crates/memtrack/tests/shared.rs | 37 ++++++++----- 8 files changed, 159 insertions(+), 28 deletions(-) create mode 100644 crates/memtrack/testdata/malloc_mmap.c diff --git a/Cargo.lock b/Cargo.lock index 350a9e81f..8d0a92783 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2411,6 +2411,7 @@ dependencies = [ "tempfile", "test-log", "test-with", + "typed-builder", "vmlinux", ] @@ -4897,6 +4898,26 @@ version = "2.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" +[[package]] +name = "typed-builder" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31aa81521b70f94402501d848ccc0ecaa8f93c8eb6999eb9747e72287757ffda" +dependencies = [ + "typed-builder-macro", +] + +[[package]] +name = "typed-builder-macro" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "076a02dc54dd46795c2e9c8282ed40bcfb1e22747e955de9389a1de28190fb26" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "typenum" version = "1.20.0" diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index d79e49eae..4912802da 100644 --- a/crates/memtrack/Cargo.toml +++ b/crates/memtrack/Cargo.toml @@ -36,6 +36,7 @@ libbpf-rs = { version = "0.26", features = ["vendored"], optional = true } object = { workspace = true } rayon = "1.12" parking_lot = "0.12" +typed-builder = "0.23.2" [build-dependencies] libbpf-cargo = { version = "0.26", optional = true } diff --git a/crates/memtrack/src/ebpf/memtrack/tracking.rs b/crates/memtrack/src/ebpf/memtrack/tracking.rs index 406c103f2..a1ce98585 100644 --- a/crates/memtrack/src/ebpf/memtrack/tracking.rs +++ b/crates/memtrack/src/ebpf/memtrack/tracking.rs @@ -7,6 +7,11 @@ impl MemtrackBpf { attach_tracepoint!(task_newtask); attach_tracepoint!(sched_process_exec); attach_tracepoint!(sched_process_exit); + attach_tracepoint!(sys_enter_mmap); + attach_tracepoint!(sys_exit_mmap); + attach_tracepoint!(sys_enter_munmap); + attach_tracepoint!(sys_enter_brk); + attach_tracepoint!(sys_exit_brk); fn rmap_target_enabled(&self, target: &str) -> bool { !self.btf_disabled_rmap_targets.contains(&target) @@ -16,6 +21,11 @@ impl MemtrackBpf { self.attach_task_newtask()?; self.attach_sched_process_exec()?; self.attach_sched_process_exit()?; + self.attach_sys_enter_mmap()?; + self.attach_sys_exit_mmap()?; + self.attach_sys_enter_munmap()?; + self.attach_sys_enter_brk()?; + self.attach_sys_exit_brk()?; if let Err(e) = self.attach_rss_stat() { warn!("Failed to attach rss_stat tracepoint, RSS collection disabled: {e:#}"); } diff --git a/crates/memtrack/src/ebpf/mod.rs b/crates/memtrack/src/ebpf/mod.rs index cb6632e10..d44ca957f 100644 --- a/crates/memtrack/src/ebpf/mod.rs +++ b/crates/memtrack/src/ebpf/mod.rs @@ -7,4 +7,4 @@ mod spawn; mod tracker; pub use memtrack::{BpfVariant, MemtrackBpf, ResolvedSymbols, resolve_symbol_offsets}; -pub use tracker::Tracker; +pub use tracker::{Tracker, TrackerOptions}; diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 5fe065e66..d40079f47 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -8,6 +8,33 @@ use std::os::unix::process::CommandExt; use std::process::Command; use std::sync::Arc; use std::sync::mpsc; +use typed_builder::TypedBuilder; + +/// Which probes a tracker attaches. [`Tracker::new`] resolves these from the +/// environment, [`Tracker::with_options`] takes them directly. +#[derive(Debug, Clone, Copy, TypedBuilder)] +pub struct TrackerOptions { + /// Attach allocator uprobes (malloc/free/calloc/...) through the + /// exec-mapping watcher. The coarse mmap/munmap/brk tracepoints fire either + /// way, so disabling this still yields memory events. + #[builder(default = true)] + pub allocators: bool, + /// Reconstruct per-process RSS from the folio rmap fentry hooks. + #[builder(default = false)] + pub rmap: bool, +} + +impl TrackerOptions { + fn from_env() -> Self { + Self::builder() + .allocators(!matches!( + std::env::var("CODSPEED_MEMTRACK_TRACK_ALLOCATORS").as_deref(), + Ok("0") | Ok("false") + )) + .rmap(std::env::var("CODSPEED_MEMTRACK_TRACK_RMAP").is_ok_and(|v| v == "1")) + .build() + } +} pub struct Tracker { bpf: Arc>, @@ -18,29 +45,24 @@ impl Tracker { /// Create a new tracker. The exec-mapping watcher discovers and attaches /// allocator probes as the tracked process tree maps executable files. pub fn new() -> Result { - let track_rmap = Self::track_rmap_from_env(); - Self::build(MemtrackBpf::new_with_rmap(track_rmap)?, true) + Self::with_options(TrackerOptions::from_env()) + } + + /// Create a tracker from an explicit probe selection rather than the environment. + pub fn with_options(options: TrackerOptions) -> Result { + Self::build( + MemtrackBpf::new_with_rmap(options.rmap)?, + options.allocators, + ) } /// Like [`Tracker::new`], but pinned to a specific BPF variant instead of /// the detected one. pub fn with_variant(variant: BpfVariant) -> Result { - let track_rmap = Self::track_rmap_from_env(); + let track_rmap = TrackerOptions::from_env().rmap; Self::build(MemtrackBpf::with_variant(variant, track_rmap)?, true) } - fn track_rmap_from_env() -> bool { - std::env::var("CODSPEED_MEMTRACK_TRACK_RMAP").is_ok_and(|v| v == "1") - } - - /// Track per-process RSS via the rss_stat tracepoint and folio rmap fentry - /// hooks without attaching any allocator probes. No exec-mapping watcher - /// runs, so `spawn` arms no on-demand allocator attachment — the tracker - /// observes only tracepoints and (when `track_rmap`) the rmap fentries. - pub fn new_without_allocators_with_rmap(track_rmap: bool) -> Result { - Self::build(MemtrackBpf::new_with_rmap(track_rmap)?, false) - } - /// Build a tracker: attach lifetime tracepoints (and rmap fentries when the /// skeleton was opened for them), plus, when `allocators` is set, the /// exec-mapping watcher and the on-demand allocator-attach worker. diff --git a/crates/memtrack/testdata/malloc_mmap.c b/crates/memtrack/testdata/malloc_mmap.c new file mode 100644 index 000000000..26fb2b5b8 --- /dev/null +++ b/crates/memtrack/testdata/malloc_mmap.c @@ -0,0 +1,15 @@ +#include +#include +#include + +int main(void) { + void *p = malloc(1 << 20); + memset(p, 1, 1 << 20); + free(p); + + void *m = mmap(NULL, 4 << 20, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + memset(m, 1, 4 << 20); + munmap(m, 4 << 20); + return 0; +} diff --git a/crates/memtrack/tests/c_tests.rs b/crates/memtrack/tests/c_tests.rs index b67eaa9d7..db5fe6439 100644 --- a/crates/memtrack/tests/c_tests.rs +++ b/crates/memtrack/tests/c_tests.rs @@ -76,3 +76,52 @@ fn test_allocation_tracking( Ok(()) } + +#[test_with::env(GITHUB_ACTIONS)] +#[test_log::test] +fn test_track_allocators_disabled_skips_allocations() -> Result<(), Box> { + use runner_shared::artifacts::MemtrackEventKind; + + let temp_dir = TempDir::new()?; + let binary = shared::compile_c_source( + include_str!("../testdata/malloc_mmap.c"), + "malloc_mmap", + temp_dir.path(), + )?; + + let (events, thread_handle) = shared::track_command_with_opts( + std::process::Command::new(&binary), + memtrack::TrackerOptions::builder() + .allocators(false) + .build(), + )?; + + let has_mmap = events + .iter() + .any(|e| matches!(e.kind, MemtrackEventKind::Mmap { .. })); + assert!( + has_mmap, + "expected at least one Mmap event with allocators disabled" + ); + + let alloc_events: Vec<_> = events + .iter() + .filter(|e| { + matches!( + e.kind, + MemtrackEventKind::Malloc { .. } + | MemtrackEventKind::Free + | MemtrackEventKind::Calloc { .. } + | MemtrackEventKind::Realloc { .. } + | MemtrackEventKind::AlignedAlloc { .. } + ) + }) + .collect(); + assert!( + alloc_events.is_empty(), + "expected no allocation events with allocators disabled, got {alloc_events:?}" + ); + + thread_handle.join().unwrap(); + Ok(()) +} diff --git a/crates/memtrack/tests/shared.rs b/crates/memtrack/tests/shared.rs index 627df1597..0dda7e851 100644 --- a/crates/memtrack/tests/shared.rs +++ b/crates/memtrack/tests/shared.rs @@ -1,7 +1,7 @@ #![allow(dead_code, unused)] use memtrack::prelude::*; -use memtrack::{BpfVariant, Tracker}; +use memtrack::{BpfVariant, Tracker, TrackerOptions}; use runner_shared::artifacts::{MemtrackEvent as Event, MemtrackEventKind}; use std::path::Path; use std::process::Command; @@ -21,17 +21,19 @@ macro_rules! assert_events_snapshot { use runner_shared::artifacts::MemtrackEventKind; use std::mem::discriminant; + // Keep only allocator events. mmap/munmap/brk sizes reflect allocator + // arena reservations that vary per run, so including them here would + // make these snapshots nondeterministic. let formatted_events: Vec = $events .iter() .filter(|e| { - // Allocation snapshots track only allocator events; RSS and - // process-lifecycle events are asserted by dedicated tests. - !matches!( + matches!( e.kind, - MemtrackEventKind::Rss { .. } - | MemtrackEventKind::Fork { .. } - | MemtrackEventKind::Exec - | MemtrackEventKind::Exit + MemtrackEventKind::Malloc { .. } + | MemtrackEventKind::Free + | MemtrackEventKind::Calloc { .. } + | MemtrackEventKind::Realloc { .. } + | MemtrackEventKind::AlignedAlloc { .. } ) }) .sorted_by_key(|e| e.timestamp) @@ -184,10 +186,10 @@ pub fn compile_c_source( Ok(binary_path) } -/// Track a command, collecting all memory events. No allocators are pre-attached: -/// the exec-mapping watcher discovers them as the tracked tree maps executables. +/// Track a command with the default probes: no rmap, and allocators discovered +/// by the exec-mapping watcher as the tracked tree maps executables. pub fn track_command(command: Command) -> TrackResult { - track_command_with_tracker(command, Tracker::new()?) + track_command_with_opts(command, TrackerOptions::builder().build()) } /// Track a command under a specific BPF variant rather than the detected one. @@ -197,7 +199,18 @@ pub fn track_command_with_variant(command: Command, variant: BpfVariant) -> Trac /// Track a command with folio rmap hooks enabled, reconstructing per-process RSS. pub fn track_command_with_rmap(command: Command) -> TrackResult { - track_command_with_tracker(command, Tracker::new_without_allocators_with_rmap(true)?) + track_command_with_opts( + command, + TrackerOptions::builder() + .allocators(false) + .rmap(true) + .build(), + ) +} + +/// Track a command with an explicit probe selection rather than the environment's. +pub fn track_command_with_opts(command: Command, options: TrackerOptions) -> TrackResult { + track_command_with_tracker(command, Tracker::with_options(options)?) } /// How many events of each kind-and-size a run saw. Addresses, timestamps, pids