Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/memtrack/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
10 changes: 10 additions & 0 deletions crates/memtrack/src/ebpf/memtrack/tracking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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:#}");
}
Expand Down
2 changes: 1 addition & 1 deletion crates/memtrack/src/ebpf/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
52 changes: 37 additions & 15 deletions crates/memtrack/src/ebpf/tracker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Mutex<MemtrackBpf>>,
Expand All @@ -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<Self> {
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> {
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<Self> {
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> {
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.
Expand Down
15 changes: 15 additions & 0 deletions crates/memtrack/testdata/malloc_mmap.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>

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;
}
49 changes: 49 additions & 0 deletions crates/memtrack/tests/c_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn std::error::Error>> {
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(())
}
37 changes: 25 additions & 12 deletions crates/memtrack/tests/shared.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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<String> = $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)
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down
Loading