diff --git a/crates/memtrack/src/ebpf/memtrack/mod.rs b/crates/memtrack/src/ebpf/memtrack/mod.rs index 4c32a33a..2ec1f04f 100644 --- a/crates/memtrack/src/ebpf/memtrack/mod.rs +++ b/crates/memtrack/src/ebpf/memtrack/mod.rs @@ -26,6 +26,7 @@ pub use maps::OwnershipMaps; pub use rmap::RmapSupport; use crate::bpf_token::has_delegated_bpf_token; +use crate::ebpf::TrackerOptions; /// Which attach mechanism a loaded skeleton uses for its uprobes. See /// `src/ebpf/c/utils/variant.h` for why only one of them is delegatable. @@ -119,35 +120,30 @@ pub struct MemtrackBpf { pub(super) skel: Skel, pub(super) probes: Vec, rmap: RmapSupport, + physical: bool, } impl MemtrackBpf { - /// Load the skeleton, picking the variant a BPF token is available for. - pub fn new_with_rmap(track_rmap: bool) -> Result { - let variant = if has_delegated_bpf_token() { - BpfVariant::Token - } else { - BpfVariant::Legacy - }; - Self::with_variant(variant, track_rmap) - } - - /// Load a specific variant rather than the one [`Self::new_with_rmap`] - /// would detect. Either attaches given host privileges; the token only - /// matters when `bpf()` is called from an unprivileged user namespace. - pub fn with_variant(variant: BpfVariant, track_rmap: bool) -> Result { + /// Load the skeleton, defaulting to the variant a BPF token is available for. + pub fn load(options: TrackerOptions) -> Result { + let variant = options.variant.unwrap_or_else(|| { + if has_delegated_bpf_token() { + BpfVariant::Token + } else { + BpfVariant::Legacy + } + }); + let physical = options.physical; crate::kernel::KernelBtf::ensure_available()?; let page_shift = page_shift()?; - let rmap = if track_rmap { + let rmap = if physical { RmapSupport::detect() } else { RmapSupport::Unsupported }; - // Both variants expose `rodata_data` and `progs` under the same field - // names, but as distinct generated types, so this can't be a function - // over the two. + // Both variants expose the same fields as distinct types, so this can't be a function. macro_rules! open_and_load { ($builder:expr, $skel:path) => {{ let open_object = Box::leak(Box::new(MaybeUninit::uninit())); @@ -168,9 +164,7 @@ impl MemtrackBpf { } } - // Autoload is decided before load(), so fentries whose targets - // the kernel lacks have to be turned off here or the whole - // skeleton fails to load. + // Autoload is decided before load(), so missing fentry targets must be off here. macro_rules! disable_rmap_prog { ($name:ident) => { paste::paste! { @@ -190,6 +184,10 @@ impl MemtrackBpf { RmapSupport::CoreAndPud => {} } + if !physical { + open_skel.progs.tracepoint_rss_stat.set_autoload(false); + } + $skel(Box::new( open_skel .load() @@ -211,6 +209,7 @@ impl MemtrackBpf { skel, probes: Vec::new(), rmap, + physical, }) } diff --git a/crates/memtrack/src/ebpf/memtrack/tracking.rs b/crates/memtrack/src/ebpf/memtrack/tracking.rs index 3e00ceb3..99a4fc3a 100644 --- a/crates/memtrack/src/ebpf/memtrack/tracking.rs +++ b/crates/memtrack/src/ebpf/memtrack/tracking.rs @@ -22,8 +22,10 @@ impl MemtrackBpf { 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:#}"); + if self.physical { + if let Err(e) = self.attach_rss_stat() { + warn!("Failed to attach rss_stat tracepoint, RSS collection disabled: {e:#}"); + } } // Defined here rather than as a method per group because the per-program diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index dc9bc13e..4841049c 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -16,9 +16,13 @@ pub struct TrackerOptions { /// exec-mapping watcher. #[builder(default = true)] pub allocators: bool, - /// Reconstruct per-process RSS from the folio rmap fentry hooks. + /// Track physical (resident) memory: the `rss_stat` tracepoint plus the + /// folio rmap hooks, which only attach on kernels that expose them. #[builder(default = false)] - pub rmap: bool, + pub physical: bool, + /// Uprobe attach mechanism. `None` detects it from BPF token availability. + #[builder(default, setter(strip_option))] + pub variant: Option, } impl TrackerOptions { @@ -28,7 +32,7 @@ impl TrackerOptions { 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")) + .physical(std::env::var("CODSPEED_MEMTRACK_TRACK_PHYSICAL").is_ok_and(|v| v == "1")) .build() } } @@ -40,40 +44,23 @@ pub struct Tracker { } impl Tracker { - /// Create a new tracker. The exec-mapping watcher discovers and attaches - /// allocator probes as the tracked process tree maps executable files. + /// Create a tracker configured from the environment. pub fn new() -> Result { 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 = TrackerOptions::from_env().rmap; - Self::build(MemtrackBpf::with_variant(variant, track_rmap)?, true) - } - - /// 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. - fn build(mut bpf: MemtrackBpf, allocators: bool) -> Result { Self::bump_memlock_rlimit()?; + let mut bpf = MemtrackBpf::load(options)?; bpf.attach_tracepoints()?; - if allocators { + if options.allocators { bpf.attach_exec_watcher()?; } let bpf = Arc::new(Mutex::new(bpf)); - let worker = if allocators { + let worker = if options.allocators { Some(AttachWorker::start(bpf.clone())?) } else { None @@ -82,7 +69,7 @@ impl Tracker { Ok(Self { bpf, worker: Mutex::new(worker), - allocators, + allocators: options.allocators, }) } diff --git a/crates/memtrack/tests/rss_tests.rs b/crates/memtrack/tests/rss_tests.rs index 5e45c04f..9ca76196 100644 --- a/crates/memtrack/tests/rss_tests.rs +++ b/crates/memtrack/tests/rss_tests.rs @@ -425,14 +425,10 @@ enum Reclaim { #[case::rss_stat(Reclaim::RssStat)] #[case::rmap(Reclaim::Rmap)] fn test_rss_external_reclaim(#[case] mode: Reclaim) -> Result<(), Box> { - let track: fn(Command) -> shared::TrackResult = match mode { - Reclaim::RssStat => shared::track_command, - Reclaim::Rmap => shared::track_command_with_rmap, - }; let (_report, events) = track_fixture( include_str!("../testdata/rss/madvise_extern.c"), "madvise_extern", - track, + shared::track_command_with_rmap, )?; // A = owner that faulted the file region; B = external caller, single-threaded diff --git a/crates/memtrack/tests/shared.rs b/crates/memtrack/tests/shared.rs index c9e1c626..0d27b6a4 100644 --- a/crates/memtrack/tests/shared.rs +++ b/crates/memtrack/tests/shared.rs @@ -197,28 +197,28 @@ pub fn compile_c_source( Ok(binary_path) } -/// Track a command with the default probes: no rmap, and allocators discovered -/// by the exec-mapping watcher as the tracked tree maps executables. +/// Track a command with the default probes: allocators only, discovered by the +/// exec-mapping watcher as the tracked tree maps executables. pub fn track_command(command: Command) -> TrackResult { track_command_with_opts(command, TrackerOptions::builder().build()) } /// Track a command under a specific BPF variant rather than the detected one. pub fn track_command_with_variant(command: Command, variant: BpfVariant) -> TrackResult { - track_command_with_tracker(command, Tracker::with_variant(variant)?) + track_command_with_opts(command, TrackerOptions::builder().variant(variant).build()) } -/// RSS reconstruction from the folio rmap hooks, without allocator probes. -fn rmap_only_options() -> TrackerOptions { +/// Physical-memory tracking without allocator probes. +fn physical_only_options() -> TrackerOptions { TrackerOptions::builder() .allocators(false) - .rmap(true) + .physical(true) .build() } -/// Track a command with folio rmap hooks enabled, reconstructing per-process RSS. +/// Track a command with physical-memory tracking enabled. pub fn track_command_with_rmap(command: Command) -> TrackResult { - track_command_with_opts(command, rmap_only_options()) + track_command_with_opts(command, physical_only_options()) } /// Track a command with an explicit probe selection rather than the environment's. @@ -231,7 +231,7 @@ pub fn track_command_with_opts(command: Command, options: TrackerOptions) -> Tra pub fn track_command_with_rmap_maps( command: Command, ) -> anyhow::Result<(Vec, OwnershipMaps, std::thread::JoinHandle<()>)> { - let tracker = Tracker::with_options(rmap_only_options())?; + let tracker = Tracker::with_options(physical_only_options())?; let (tracker, events, ()) = run_tracked(command, tracker, |_, _| Ok(()))?; let maps = tracker.ownership_maps()?; Ok((events, maps, std::thread::spawn(move || drop(tracker)))) @@ -247,7 +247,7 @@ pub fn track_command_with_rmap_checkpoint( ready: &Path, release: &Path, ) -> anyhow::Result<(Vec, OwnershipMaps, i32, std::thread::JoinHandle<()>)> { - let tracker = Tracker::with_options(rmap_only_options())?; + let tracker = Tracker::with_options(physical_only_options())?; let (tracker, events, (maps, root_pid)) = run_tracked(command, tracker, |tracker, pid| { let deadline = std::time::Instant::now() + std::time::Duration::from_secs(30); while !ready.exists() && std::time::Instant::now() < deadline { @@ -307,13 +307,14 @@ pub fn for_each_variant( let mut profiles: Vec<(BpfVariant, EventProfile)> = Vec::new(); for variant in [BpfVariant::Legacy, BpfVariant::Token] { - let tracker = match Tracker::with_variant(variant) { - Ok(tracker) => tracker, - Err(err) => { - eprintln!("skipping {variant:?} variant, cannot attach here: {err:#}"); - continue; - } - }; + let tracker = + match Tracker::with_options(TrackerOptions::builder().variant(variant).build()) { + Ok(tracker) => tracker, + Err(err) => { + eprintln!("skipping {variant:?} variant, cannot attach here: {err:#}"); + continue; + } + }; let (events, thread_handle) = track_command_with_tracker(workload(), tracker)?; assert_events(&events);