From f15720b04622eec63ce31f39faddd65af0825169 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 7 Jul 2026 19:40:28 +0200 Subject: [PATCH 1/8] feat(memtrack): add exec-mapping watcher BPF prog and inode maps --- crates/memtrack/src/ebpf/c/attach.bpf.h | 54 +++++++++++++++++++++++ crates/memtrack/src/ebpf/c/event.h | 8 ++++ crates/memtrack/src/ebpf/c/memtrack.bpf.c | 2 + crates/memtrack/src/ebpf/events.rs | 25 +++++++++++ 4 files changed, 89 insertions(+) create mode 100644 crates/memtrack/src/ebpf/c/attach.bpf.h diff --git a/crates/memtrack/src/ebpf/c/attach.bpf.h b/crates/memtrack/src/ebpf/c/attach.bpf.h new file mode 100644 index 00000000..b3814b8d --- /dev/null +++ b/crates/memtrack/src/ebpf/c/attach.bpf.h @@ -0,0 +1,54 @@ +/* == Exec-mapping watcher: on-demand allocator attach == */ + +#define MEMTRACK_PROT_EXEC 0x4 +#define MEMTRACK_SIGSTOP 19 + +struct inode_key { + __u64 dev; + __u64 ino; +}; + +/* (dev, ino) -> 1; populated by userspace after classify/attach */ +BPF_HASH_MAP(known_inodes, struct inode_key, __u8, 8192); +/* Requests are 24 B and rare; overflow aborts the run via the counter below */ +BPF_RINGBUF(attach_requests, 128 * 1024); +BPF_ARRAY_MAP(attach_request_dropped, __u64, 1); + +SEC("fentry/security_mmap_file") +int BPF_PROG(watch_exec_mmap, struct file* file, unsigned long prot, unsigned long flags) { + if (!file || !(prot & MEMTRACK_PROT_EXEC)) { + return 0; + } + + __u64 pid_tgid = bpf_get_current_pid_tgid(); + __u32 tgid = pid_tgid >> 32; + if (!is_tracked(tgid)) { + return 0; + } + + struct inode_key key = { + .dev = BPF_CORE_READ(file, f_inode, i_sb, s_dev), + .ino = BPF_CORE_READ(file, f_inode, i_ino), + }; + if (bpf_map_lookup_elem(&known_inodes, &key)) { + return 0; + } + + struct attach_request* req = bpf_ringbuf_reserve(&attach_requests, sizeof(*req), 0); + if (!req) { + __u32 zero = 0; + __u64* drops = bpf_map_lookup_elem(&attach_request_dropped, &zero); + if (drops) { + __sync_fetch_and_add(drops, 1); + } + return 0; + } + req->pid = tgid; + req->tid = (__u32)pid_tgid; + req->dev = key.dev; + req->ino = key.ino; + bpf_ringbuf_submit(req, 0); + + bpf_send_signal(MEMTRACK_SIGSTOP); + return 0; +} diff --git a/crates/memtrack/src/ebpf/c/event.h b/crates/memtrack/src/ebpf/c/event.h index f914f634..26635cf3 100644 --- a/crates/memtrack/src/ebpf/c/event.h +++ b/crates/memtrack/src/ebpf/c/event.h @@ -48,4 +48,12 @@ struct event { } data; }; +/* Request from the exec-mapping watcher to the userspace attach worker */ +struct attach_request { + uint32_t pid; + uint32_t tid; + uint64_t dev; /* kernel s_dev encoding: (major << 20) | minor */ + uint64_t ino; +}; + #endif /* __EVENT_H__ */ diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 06601e27..2d4b48d3 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/crates/memtrack/src/ebpf/c/memtrack.bpf.c @@ -92,6 +92,8 @@ int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) { return 0; } +#include "attach.bpf.h" + /* == Helper functions for the allocation tracking == */ /* Helper to check if tracking is currently enabled */ diff --git a/crates/memtrack/src/ebpf/events.rs b/crates/memtrack/src/ebpf/events.rs index 3ecb7a67..f18216c6 100644 --- a/crates/memtrack/src/ebpf/events.rs +++ b/crates/memtrack/src/ebpf/events.rs @@ -89,6 +89,31 @@ pub fn parse_event(data: &[u8]) -> Option { }) } +/// A request from the exec-mapping watcher to attach allocator probes. +#[derive(Debug, Clone, Copy)] +pub struct AttachRequest { + pub pid: u32, + pub dev: u64, + pub ino: u64, +} +impl AttachRequest { + /// Parse an attach request from raw ring buffer bytes. + /// + /// SAFETY: The data must be a valid `bindings::attach_request` + pub fn parse(data: &[u8]) -> Option { + if data.len() < std::mem::size_of::() { + return None; + } + + let req = unsafe { &*(data.as_ptr() as *const bindings::attach_request) }; + Some(AttachRequest { + pid: req.pid, + dev: req.dev, + ino: req.ino, + }) + } +} + #[cfg(test)] mod tests { use super::*; From 22396ecb7fa1bfb456420618a534aea13c1a05d9 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 8 Jul 2026 21:42:06 +0200 Subject: [PATCH 2/8] feat(memtrack): attach allocator probes on demand via exec-mapping watcher Instead of pre-attaching every allocator found on the system, a BPF fentry on security_mmap_file signals the first mapping of each unknown executable inode. A background worker stops the mapping processes, classifies the file by its symbols, attaches probes, and resumes them. This covers dlopen'd and statically linked allocators in spawned children, and skips libraries the benchmark never loads. The public API shrinks to Tracker (owns the attach worker) and Session (owns the spawned child and its event pipeline); ring buffers are polled through a generic RingBufferPoller with caller-supplied parsing. Fixes COD-1801 --- Cargo.lock | 1 + crates/memtrack/Cargo.toml | 1 + crates/memtrack/src/ebpf/attach_worker.rs | 250 ++++++++++++++++++++++ crates/memtrack/src/ebpf/memtrack.rs | 92 +++++++- crates/memtrack/src/ebpf/mod.rs | 7 +- crates/memtrack/src/ebpf/poller.rs | 103 ++++----- crates/memtrack/src/ebpf/proc_fs.rs | 195 +++++++++++++++++ crates/memtrack/src/ebpf/spawn.rs | 108 ++++++++++ crates/memtrack/src/ebpf/tracker.rs | 144 +++++++------ crates/memtrack/src/ipc.rs | 28 +-- crates/memtrack/src/lib.rs | 4 + crates/memtrack/src/main.rs | 40 ++-- crates/memtrack/src/session.rs | 41 ++++ 13 files changed, 837 insertions(+), 177 deletions(-) create mode 100644 crates/memtrack/src/ebpf/attach_worker.rs create mode 100644 crates/memtrack/src/ebpf/proc_fs.rs create mode 100644 crates/memtrack/src/ebpf/spawn.rs create mode 100644 crates/memtrack/src/session.rs diff --git a/Cargo.lock b/Cargo.lock index b55c08a4..b2f814d1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2388,6 +2388,7 @@ dependencies = [ "libc", "log", "object", + "parking_lot", "paste", "rstest", "runner-shared", diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index ae481941..de6214b8 100644 --- a/crates/memtrack/Cargo.toml +++ b/crates/memtrack/Cargo.toml @@ -34,6 +34,7 @@ paste = "1.0.15" libbpf-rs = { version = "0.26", features = ["vendored"], optional = true } glob = "0.3.3" object = { workspace = true } +parking_lot = "0.12" [build-dependencies] libbpf-cargo = { version = "0.26", optional = true } diff --git a/crates/memtrack/src/ebpf/attach_worker.rs b/crates/memtrack/src/ebpf/attach_worker.rs new file mode 100644 index 00000000..1eeba9de --- /dev/null +++ b/crates/memtrack/src/ebpf/attach_worker.rs @@ -0,0 +1,250 @@ +use crate::AllocatorLib; +use crate::ebpf::MemtrackBpf; +use crate::ebpf::events::AttachRequest; +use crate::ebpf::poller::RingBufferPoller; +use crate::prelude::*; +use parking_lot::Mutex; +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicI32, Ordering}; +use std::sync::mpsc::{self, RecvTimeoutError}; +use std::thread::JoinHandle; +use std::time::Duration; + +use super::proc_fs::{resolve_mapping, wait_all_stopped}; + +const STOP_DEADLINE: Duration = Duration::from_secs(1); +const POLL_INTERVAL_MS: u64 = 10; +const RECV_TIMEOUT: Duration = Duration::from_millis(100); + +/// SIGCONTs `pid` on drop, ignoring errors. Guarantees a stopped process is +/// resumed on every exit path, including panics. +struct ContGuard(i32); + +impl Drop for ContGuard { + fn drop(&mut self) { + // SAFETY: kill with SIGCONT has no memory effects; errors (e.g. the + // process already exited) are intentionally ignored. + unsafe { + libc::kill(self.0, libc::SIGCONT); + } + } +} + +/// Background worker that stops tracked processes on their first mapping of an +/// unknown executable file, classifies it, attaches allocator probes, and +/// resumes them. +pub(crate) struct AttachWorker { + handle: Option>, + shutdown: Arc, + fatal: Arc>>, + root_pid: Arc, + bpf: Arc>, +} + +impl AttachWorker { + pub(crate) fn start(bpf: Arc>) -> Result { + let shutdown = Arc::new(AtomicBool::new(false)); + let fatal = Arc::new(Mutex::new(None)); + let root_pid = Arc::new(AtomicI32::new(0)); + + let (tx, rx) = mpsc::channel(); + let poller = bpf.lock().poll_attach_with_channel(POLL_INTERVAL_MS, tx)?; + + let worker = Worker { + poller, + rx, + bpf: bpf.clone(), + shutdown: shutdown.clone(), + fatal: fatal.clone(), + root_pid: root_pid.clone(), + }; + + let handle = std::thread::spawn(move || worker.run()); + + Ok(Self { + handle: Some(handle), + shutdown, + fatal, + root_pid, + bpf, + }) + } + + /// Tell the worker which pid to SIGKILL on a fatal error. + pub(crate) fn set_root_pid(&self, pid: i32) { + self.root_pid.store(pid, Ordering::SeqCst); + } + + /// Stop the worker, join it, and surface any fatal error it recorded. + /// Fails when the attach-request ring buffer overflowed: exec mappings were + /// missed, so allocator coverage would be incomplete. + pub(crate) fn finish(mut self) -> Result<()> { + self.shutdown.store(true, Ordering::SeqCst); + if let Some(handle) = self.handle.take() + && let Err(panic) = handle.join() + { + let msg = panic + .downcast_ref::<&str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::().cloned()) + .unwrap_or_else(|| "unknown panic payload".to_string()); + bail!("attach worker thread panicked: {msg}"); + } + if let Some(err) = self.fatal.lock().take() { + bail!("{err}"); + } + + let dropped = self.bpf.lock().attach_request_dropped_count()?; + if dropped > 0 { + bail!( + "Memtrack attach-request ring buffer overflowed: {dropped} exec mappings missed, \ + aborting since allocator coverage is incomplete." + ); + } + + Ok(()) + } +} + +struct Worker { + poller: RingBufferPoller, + rx: mpsc::Receiver, + bpf: Arc>, + shutdown: Arc, + fatal: Arc>>, + root_pid: Arc, +} + +impl Worker { + fn run(self) { + let mut known: HashSet<(u64, u64)> = HashSet::new(); + + loop { + if self.shutdown.load(Ordering::SeqCst) { + // Producers are gone: a synchronous drain flushes everything. + if self.poller.drain().is_err() { + return; + } + let mut batch: Vec = self.rx.try_iter().collect(); + if batch.is_empty() { + break; + } + if let Err(e) = self.process_batch(&mut batch, &mut known) { + self.record_fatal(e); + return; + } + continue; + } + + let first = match self.rx.recv_timeout(RECV_TIMEOUT) { + Ok(req) => req, + Err(RecvTimeoutError::Timeout) => continue, + Err(RecvTimeoutError::Disconnected) => return, + }; + let mut batch: Vec = + std::iter::once(first).chain(self.rx.try_iter()).collect(); + + if let Err(e) = self.process_batch(&mut batch, &mut known) { + self.record_fatal(e); + return; + } + } + } + + /// Stop every producing pid (fixpoint, draining until no new pid appears), + /// then classify + attach for each unique `(dev, ino)`. `guards` resume every + /// stopped pid exactly once when this returns, including the error path. + fn process_batch( + &self, + batch: &mut Vec, + known: &mut HashSet<(u64, u64)>, + ) -> Result<()> { + let mut guards: Vec = Vec::new(); + let mut stopped: HashSet = HashSet::new(); + + loop { + let new_pids: HashSet = batch + .iter() + .map(|r| r.pid) + .filter(|pid| !stopped.contains(pid)) + .collect(); + + if new_pids.is_empty() { + break; + } + + for pid in new_pids { + stopped.insert(pid); + guards.push(ContGuard(pid as i32)); + wait_all_stopped(pid, STOP_DEADLINE)?; + } + + // Every producer is stopped, so a synchronous drain is complete. + self.poller.drain()?; + batch.extend(self.rx.try_iter()); + } + + let mut seen: HashSet<(u64, u64)> = HashSet::new(); + for req in batch.iter() { + let key = (req.dev, req.ino); + if !seen.insert(key) || known.contains(&key) { + continue; + } + + if self.handle_request(req)? { + self.bpf.lock().insert_known_inode(req.dev, req.ino)?; + known.insert(key); + } + } + + Ok(()) + } + + /// Returns `Ok(true)` when the inode should be marked known (attached, or + /// definitively not an allocator), `Ok(false)` when the mapping vanished + /// (retry later). An attach failure is a hard error (fail-closed). + fn handle_request(&self, req: &AttachRequest) -> Result { + let Some(mapping) = resolve_mapping(req.pid, req.dev, req.ino) else { + return Ok(false); + }; + + let Ok(lib) = AllocatorLib::from_path_static(&mapping.attach_path) else { + debug!("on-demand: not an allocator: {}", mapping.display); + return Ok(true); + }; + + let kind = lib.kind; + let mut bpf = self.bpf.lock(); + let before = bpf.probe_count(); + bpf.attach_allocator_probes(lib.kind, &lib.path)?; + if bpf.probe_count() == before { + bail!( + "{} classified at {} but zero probes attached", + kind.name(), + mapping.display + ); + } + + info!( + "on-demand attached {} probes to {}", + kind.name(), + mapping.display + ); + Ok(true) + } + + fn record_fatal(&self, err: anyhow::Error) { + let msg = format!("{err:#}"); + error!("on-demand attach worker fatal: {msg}"); + *self.fatal.lock() = Some(msg); + + let root = self.root_pid.load(Ordering::SeqCst); + if root > 0 { + // SAFETY: kill with SIGKILL has no memory effects. + unsafe { + libc::kill(root, libc::SIGKILL); + } + } + } +} diff --git a/crates/memtrack/src/ebpf/memtrack.rs b/crates/memtrack/src/ebpf/memtrack.rs index 3a1bfa26..d2d704e1 100644 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ b/crates/memtrack/src/ebpf/memtrack.rs @@ -453,16 +453,90 @@ impl MemtrackBpf { Ok(()) } - /// Start polling with an mpsc channel for events - pub fn start_polling_with_channel( + /// Poll the allocation-event ring buffer into `tx`. The returned poller + /// keeps the pipeline alive; events stop flowing when it is dropped. + pub fn poll_events_with_channel( &self, - poll_timeout_ms: u64, - ) -> Result<( - RingBufferPoller, - std::sync::mpsc::Receiver, - )> { - // Use the syscalls skeleton's ring buffer (both programs share the same one) - RingBufferPoller::with_channel(&self.skel.maps.events, poll_timeout_ms) + poll_interval_ms: u64, + tx: std::sync::mpsc::Sender, + ) -> Result { + RingBufferPoller::new( + &self.skel.maps.events, + crate::ebpf::events::parse_event, + tx, + poll_interval_ms, + ) + } +} + +// ========================================================================= +// On-demand exec-mapping watcher +// ========================================================================= + +impl MemtrackBpf { + /// Attach the exec-mapping watcher (fentry/security_mmap_file). Only used in + /// on-demand mode; the program is loaded and verified in all modes. + pub fn attach_exec_watcher(&mut self) -> Result<()> { + let link = self + .skel + .progs + .watch_exec_mmap + .attach() + .context("Failed to attach exec-mapping watcher")?; + self.probes.push(link); + Ok(()) + } + + /// Mark a (dev, ino) as classified so the watcher stops re-signalling for it. + /// The 16-byte key matches `struct inode_key { __u64 dev; __u64 ino; }` (no padding). + pub fn insert_known_inode(&self, dev: u64, ino: u64) -> Result<()> { + let mut key = [0u8; 16]; + key[..8].copy_from_slice(&dev.to_le_bytes()); + key[8..].copy_from_slice(&ino.to_le_bytes()); + self.skel + .maps + .known_inodes + .update(&key, &1u8.to_le_bytes(), libbpf_rs::MapFlags::ANY) + .context("Failed to insert known inode")?; + Ok(()) + } + + /// Number of exec-mapping requests dropped because the request ring buffer was full. + pub fn attach_request_dropped_count(&self) -> Result { + let key = 0u32; + let value = self + .skel + .maps + .attach_request_dropped + .lookup(&key.to_le_bytes(), libbpf_rs::MapFlags::ANY) + .context("Failed to read attach_request_dropped counter")? + .ok_or_else(|| anyhow!("attach_request_dropped slot 0 missing"))?; + + let bytes: [u8; 8] = value + .as_slice() + .try_into() + .map_err(|_| anyhow!("attach_request_dropped value has unexpected size"))?; + Ok(u64::from_le_bytes(bytes)) + } + + /// Number of currently-attached probes/links. + pub fn probe_count(&self) -> usize { + self.probes.len() + } + + /// Poll the exec-mapping request ring buffer into `tx`. Same contract as + /// [`Self::poll_events_with_channel`]. + pub(crate) fn poll_attach_with_channel( + &self, + poll_interval_ms: u64, + tx: std::sync::mpsc::Sender, + ) -> Result { + RingBufferPoller::new( + &self.skel.maps.attach_requests, + crate::ebpf::events::AttachRequest::parse, + tx, + poll_interval_ms, + ) } } diff --git a/crates/memtrack/src/ebpf/mod.rs b/crates/memtrack/src/ebpf/mod.rs index 0badf542..478131db 100644 --- a/crates/memtrack/src/ebpf/mod.rs +++ b/crates/memtrack/src/ebpf/mod.rs @@ -1,7 +1,10 @@ +mod attach_worker; mod events; mod memtrack; -mod poller; +pub(crate) mod poller; +mod proc_fs; +mod spawn; mod tracker; -pub use memtrack::MemtrackBpf; +pub(crate) use memtrack::MemtrackBpf; pub use tracker::Tracker; diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index 1f4ef2b4..d80a1554 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/crates/memtrack/src/ebpf/poller.rs @@ -1,88 +1,79 @@ -use anyhow::Result; +use anyhow::{Context, Result}; use libbpf_rs::{MapCore, RingBufferBuilder}; -use runner_shared::artifacts::MemtrackEvent as Event; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, mpsc}; +use std::sync::mpsc::{self, RecvTimeoutError, Sender}; use std::thread::JoinHandle; use std::time::Duration; -use super::events::parse_event; - -/// A handler function for processing ring buffer events -pub type EventHandler = Box; - -/// RingBufferPoller manages polling a BPF ring buffer in a background thread -/// and sending events to handlers +/// Polls a BPF ring buffer in a background thread, parsing raw entries with a +/// user-supplied closure and forwarding them to an mpsc channel. +/// +/// The poll thread runs until the poller is dropped, doing a final full +/// `consume()` on shutdown so no buffered entries are lost. pub struct RingBufferPoller { - shutdown: Arc, + ctl: Option>>, poll_thread: Option>, } impl RingBufferPoller { - /// Create a new RingBufferPoller for the given ring buffer map - /// - /// # Arguments - /// * `rb_map` - The BPF ring buffer map to poll - /// * `handler` - Callback function to handle each event - /// * `poll_timeout_ms` - How long to wait for events in each poll iteration - pub fn new( - rb_map: &M, - handler: EventHandler, - poll_timeout_ms: u64, - ) -> Result { + pub fn new(rb_map: &M, parse: F, tx: Sender, poll_interval_ms: u64) -> Result + where + M: MapCore, + T: Send + 'static, + F: Fn(&[u8]) -> Option + Send + 'static, + { let mut builder = RingBufferBuilder::new(); builder.add(rb_map, move |data| { - if let Some(event) = parse_event(data) { - handler(event); + if let Some(item) = parse(data) { + let _ = tx.send(item); } 0 })?; - let ringbuf = builder.build()?; - let shutdown = Arc::new(AtomicBool::new(false)); - let shutdown_clone = shutdown.clone(); + // The control channel doubles as the poll pacing: a received message is + // a drain request (acked after a full consume), a timeout is a regular + // poll tick, and disconnection is the shutdown signal. + let (ctl, ctl_rx) = mpsc::channel::>(); let poll_thread = std::thread::spawn(move || { - while !shutdown_clone.load(Ordering::Relaxed) { - let _ = ringbuf.poll(Duration::from_millis(poll_timeout_ms)); + loop { + match ctl_rx.recv_timeout(Duration::from_millis(poll_interval_ms)) { + Ok(ack) => { + let _ = ringbuf.consume(); + let _ = ack.send(()); + } + Err(RecvTimeoutError::Timeout) => { + let _ = ringbuf.poll(Duration::ZERO); + } + Err(RecvTimeoutError::Disconnected) => { + let _ = ringbuf.consume(); + break; + } + } } }); Ok(Self { - shutdown, + ctl: Some(ctl), poll_thread: Some(poll_thread), }) } - /// Create a new RingBufferPoller with an mpsc channel for events - /// - /// Returns the RingBufferPoller and the receiver end of the channel - pub fn with_channel( - rb_map: &M, - poll_timeout_ms: u64, - ) -> Result<(Self, mpsc::Receiver)> { - let (tx, rx) = mpsc::channel(); - let poller = Self::new( - rb_map, - Box::new(move |event| { - let _ = tx.send(event); - }), - poll_timeout_ms, - )?; - Ok((poller, rx)) - } - - /// Stop the polling thread and wait for it to finish - pub fn shutdown(&mut self) { - self.shutdown.store(true, Ordering::Relaxed); - if let Some(thread) = self.poll_thread.take() { - let _ = thread.join(); - } + /// Block until a full `consume()` of the ring buffer completes. When every + /// producer is stopped, all pending entries are in the channel afterwards. + pub fn drain(&self) -> Result<()> { + let (ack_tx, ack_rx) = mpsc::channel(); + let ctl = self.ctl.as_ref().context("poller already shut down")?; + ctl.send(ack_tx).context("poll thread is gone")?; + ack_rx.recv().context("poll thread died during drain")?; + Ok(()) } } impl Drop for RingBufferPoller { fn drop(&mut self) { - self.shutdown(); + drop(self.ctl.take()); + if let Some(thread) = self.poll_thread.take() { + let _ = thread.join(); + } } } diff --git a/crates/memtrack/src/ebpf/proc_fs.rs b/crates/memtrack/src/ebpf/proc_fs.rs new file mode 100644 index 00000000..1619dc4c --- /dev/null +++ b/crates/memtrack/src/ebpf/proc_fs.rs @@ -0,0 +1,195 @@ +use crate::prelude::*; +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +/// A mapping resolved from `/proc//maps` back to an attachable path. +pub(super) struct ResolvedMapping { + /// `/proc//map_files/-` — resolves uniformly for deleted + /// files and requires root (which memtrack already has for uprobes). + pub attach_path: PathBuf, + /// The pathname column, for logs. + pub display: String, +} + +/// Block until every thread of `pid` is group-stopped. +/// +/// The process state is the first non-space char after the LAST `)` in +/// `/proc//task//stat`. Success means every thread is `T`/`t`. +/// +/// - A vanished process (`/proc/` gone) is success: the stop is moot. +/// - At the deadline, threads still in uninterruptible sleep (`D`) are treated +/// as stopped — they cannot execute user code and join the group stop when the +/// syscall returns. Any thread still `R`/`S` is a hard error: leaving it +/// running breaks the drain guarantee. +pub(super) fn wait_all_stopped(pid: u32, deadline: Duration) -> Result<()> { + let start = Instant::now(); + let task_dir = format!("/proc/{pid}/task"); + + loop { + let Ok(entries) = std::fs::read_dir(&task_dir) else { + return Ok(()); + }; + + let mut all_stopped = true; + let mut running_tid: Option = None; + + for entry in entries.flatten() { + let name = entry.file_name(); + let Some(tid) = name.to_str().and_then(|s| s.parse::().ok()) else { + continue; + }; + + let Ok(stat) = std::fs::read_to_string(entry.path().join("stat")) else { + continue; + }; + let Some(state) = task_state(&stat) else { + continue; + }; + + match state { + 'T' | 't' => {} + 'D' => all_stopped = false, + _ => { + all_stopped = false; + running_tid = Some(tid); + } + } + } + + if all_stopped { + return Ok(()); + } + + if start.elapsed() >= deadline { + let Some(tid) = running_tid else { + warn!( + "pid {pid}: thread(s) in uninterruptible sleep (D) at stop deadline; treating as stopped" + ); + return Ok(()); + }; + bail!("thread {tid} of {pid} did not stop within {deadline:?}"); + } + + std::thread::sleep(Duration::from_millis(1)); + } +} + +/// The process state char: first non-space after the LAST `)` (comm can contain +/// both `)` and spaces). +fn task_state(stat: &str) -> Option { + let idx = stat.rfind(')')?; + stat[idx + 1..].trim_start().chars().next() +} + +/// Resolve the `(dev, ino)` reported by the watcher to an attachable mapping. +/// +/// `dev` is the kernel `s_dev` encoding `(major << 20) | minor`; the maps dev +/// column prints `MAJOR:MINOR` in lowercase hex (`major = dev >> 20`, +/// `minor = dev & 0xFFFFF`). Returns `None` if no row matches (the mmap failed +/// after the LSM hook, or the process exited) — the caller retries on the next +/// sighting. +pub(super) fn resolve_mapping(pid: u32, dev: u64, ino: u64) -> Option { + let maps = std::fs::read_to_string(format!("/proc/{pid}/maps")).ok()?; + let want_major = (dev >> 20) as u32; + let want_minor = (dev & 0xFFFFF) as u32; + + for line in maps.lines() { + // range perms offset dev inode pathname + let mut fields = line.split_whitespace(); + let (Some(range), Some(_perms), Some(_offset), Some(dev_col), Some(inode_col)) = ( + fields.next(), + fields.next(), + fields.next(), + fields.next(), + fields.next(), + ) else { + continue; + }; + + if inode_col.parse::().ok() != Some(ino) { + continue; + } + if parse_dev(dev_col) != Some((want_major, want_minor)) { + continue; + } + + let display = fields.next().unwrap_or("").to_string(); + return Some(ResolvedMapping { + attach_path: PathBuf::from(format!("/proc/{pid}/map_files/{range}")), + display, + }); + } + + None +} + +/// Parse a `MAJOR:MINOR` lowercase-hex dev column. +fn parse_dev(col: &str) -> Option<(u32, u32)> { + let (major, minor) = col.split_once(':')?; + Some(( + u32::from_str_radix(major, 16).ok()?, + u32::from_str_radix(minor, 16).ok()?, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn task_state_reads_char_after_last_paren() { + assert_eq!(task_state("1234 (bash) S 1 0 0"), Some('S')); + assert_eq!(task_state("1234 (weird )name) T 1"), Some('T')); + assert_eq!(task_state("1 (x) R"), Some('R')); + assert_eq!(task_state("no paren here"), None); + } + + #[test] + fn parse_dev_reads_hex_major_minor() { + assert_eq!(parse_dev("08:01"), Some((8, 1))); + assert_eq!(parse_dev("fd:00"), Some((0xfd, 0))); + assert_eq!(parse_dev("00:1a"), Some((0, 0x1a))); + assert_eq!(parse_dev("nope"), None); + assert_eq!(parse_dev("zz:01"), None); + } + + /// Round-trips the kernel dev encoding `(major << 20) | minor` through + /// resolve_mapping against our own maps (no root needed): a file-backed + /// mapping's dev/ino must resolve back to that same mapping. + #[test] + fn resolve_mapping_round_trips_self_maps() { + let pid = std::process::id(); + let maps = std::fs::read_to_string(format!("/proc/{pid}/maps")).unwrap(); + + let target = maps.lines().find_map(|line| { + let mut f = line.split_whitespace(); + let (_range, _perms, _off, dev, inode) = + (f.next()?, f.next()?, f.next()?, f.next()?, f.next()?); + let path = f.next()?; + if !path.starts_with('/') { + return None; + } + let ino: u64 = inode.parse().ok()?; + if ino == 0 { + return None; + } + let (maj, min) = parse_dev(dev)?; + let s_dev = ((maj as u64) << 20) | (min as u64); + Some((s_dev, ino, path.to_string())) + }); + + let Some((s_dev, ino, path)) = target else { + return; + }; + + let resolved = resolve_mapping(pid, s_dev, ino).expect("mapping should resolve"); + assert_eq!(resolved.display, path); + assert!( + resolved + .attach_path + .starts_with(format!("/proc/{pid}/map_files/")), + "attach_path was {:?}", + resolved.attach_path + ); + } +} diff --git a/crates/memtrack/src/ebpf/spawn.rs b/crates/memtrack/src/ebpf/spawn.rs new file mode 100644 index 00000000..f4c8dadd --- /dev/null +++ b/crates/memtrack/src/ebpf/spawn.rs @@ -0,0 +1,108 @@ +use anyhow::{Result, bail}; +use std::ffi::{OsStr, OsString}; +use std::process::{Child, Command}; + +/// Build a shell stub that stops itself with `SIGSTOP` right after the shell +/// execs, then `exec`s the real target in the same pid. +/// +/// A `pre_exec` hook that raised `SIGSTOP` cannot be used: it blocks the child +/// before `execve`, and `Command::spawn` blocks reading its exec-status pipe +/// until the child execs, so the two deadlock. The stub sidesteps that — the +/// shell execs immediately (spawn returns), then stops itself. The caller arms +/// tracking while it is stopped and resumes it; the target's own `execve` then +/// happens in the already-tracked pid, so the watcher observes its mappings. +pub fn stopped_command(program: impl AsRef, args: &[OsString]) -> Command { + let mut cmd = Command::new("sh"); + cmd.arg("-c") + .arg(r#"kill -STOP "$$"; exec "$@""#) + .arg("sh") + .arg(program.as_ref()) + .args(args); + cmd +} + +/// Re-wrap an already-configured command as a self-stopping shell stub, +/// preserving its program, args, explicit env, and working directory. +/// +/// `uid`/`gid` cannot be read back from a `Command`; apply those to the returned +/// command if needed. +pub fn wrap_stopped(command: &Command) -> Command { + let args: Vec = command.get_args().map(OsStr::to_os_string).collect(); + let mut wrapped = stopped_command(command.get_program(), &args); + + for (key, val) in command.get_envs() { + match val { + Some(val) => { + wrapped.env(key, val); + } + None => { + wrapped.env_remove(key); + } + } + } + if let Some(dir) = command.get_current_dir() { + wrapped.current_dir(dir); + } + + wrapped +} + +/// Spawn a command built by [`stopped_command`] / [`wrap_stopped`] and block +/// until the shell stub group-stops itself. +pub fn spawn_stopped(cmd: &mut Command) -> Result { + let child = cmd.spawn()?; + let pid = child.id() as i32; + + let mut status: libc::c_int = 0; + let ret = unsafe { libc::waitpid(pid, &mut status, libc::WUNTRACED) }; + if ret != pid || !libc::WIFSTOPPED(status) { + bail!("child {pid} exited before tracking was armed (status {status:#x})"); + } + + Ok(child) +} + +/// Resume a process previously stopped by [`spawn_stopped`]. +/// +/// A vanished process (`ESRCH`) is treated as success: the stop is moot. +pub fn resume(pid: i32) -> Result<()> { + let ret = unsafe { libc::kill(pid, libc::SIGCONT) }; + if ret == 0 { + return Ok(()); + } + + let err = std::io::Error::last_os_error(); + if err.raw_os_error() == Some(libc::ESRCH) { + return Ok(()); + } + bail!("failed to SIGCONT {pid}: {err}"); +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Regression for the pre_exec-SIGSTOP deadlock: this must not hang, and the + /// stub must be group-stopped before we resume it to completion. + #[test] + fn spawn_stopped_stops_then_resumes() { + let mut cmd = stopped_command("true", &[]); + let mut child = spawn_stopped(&mut cmd).expect("spawn_stopped should not hang"); + let pid = child.id() as i32; + + let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).unwrap(); + let state = stat + .rsplit(')') + .next() + .and_then(|s| s.trim_start().chars().next()) + .unwrap(); + assert!( + matches!(state, 'T' | 't'), + "expected stopped, got {state:?}" + ); + + resume(pid).expect("resume"); + let status = child.wait().expect("wait"); + assert!(status.success(), "stub-exec'd target should exit 0"); + } +} diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 9f1af587..bd2f7861 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -1,79 +1,97 @@ +use crate::ebpf::MemtrackBpf; +use crate::ebpf::attach_worker::AttachWorker; +use crate::ebpf::spawn::{resume, spawn_stopped, wrap_stopped}; use crate::prelude::*; -use crate::{AllocatorLib, ebpf::MemtrackBpf}; -use runner_shared::artifacts::MemtrackEvent as Event; -use std::sync::mpsc::{self, Receiver}; +use crate::session::Session; +use parking_lot::Mutex; +use std::os::unix::process::CommandExt; +use std::process::Command; +use std::sync::Arc; +use std::sync::mpsc; pub struct Tracker { - bpf: MemtrackBpf, + bpf: Arc>, + worker: Mutex>, } impl Tracker { - /// Create a new tracker instance + /// Create a new tracker with on-demand allocator attach. /// - /// This will: - /// - Initialize the BPF subsystem - /// - Bump memlock limits - /// - Attach uprobes to all libc instances - /// - Attach tracepoints for fork tracking + /// No allocators are pre-attached: the exec-mapping watcher discovers and + /// attaches them as the tracked tree maps executable files. pub fn new() -> Result { - let mut instance = Self::new_without_allocators()?; - - let allocators = AllocatorLib::find_all()?; - debug!("Found {} allocator instance(s)", allocators.len()); - instance.attach_allocators(&allocators)?; - - Ok(instance) - } - - pub fn new_without_allocators() -> Result { - // Bump memlock limits Self::bump_memlock_rlimit()?; let mut bpf = MemtrackBpf::new()?; bpf.attach_tracepoints()?; + bpf.attach_exec_watcher()?; - Ok(Self { bpf }) + let bpf = Arc::new(Mutex::new(bpf)); + let worker = AttachWorker::start(bpf.clone())?; + + Ok(Self { + bpf, + worker: Mutex::new(Some(worker)), + }) } - pub fn attach_allocators(&mut self, libs: &[AllocatorLib]) -> Result<()> { - for allocator in libs { - self.bpf - .attach_allocator_probes(allocator.kind, &allocator.path)?; + /// Spawn `cmd` under tracking: the target is wrapped so it stops itself + /// before exec'ing, its pid is armed while stopped, then it is resumed. + /// The watcher observes the target's own `execve` mappings — no allocation + /// escapes untracked. + /// + /// `uid_gid` drops the child's privileges (a `Command`'s uid/gid cannot be + /// read back, so it cannot be preserved through the wrap). + pub fn spawn(&self, cmd: &Command, uid_gid: Option<(u32, u32)>) -> Result { + let mut wrapped = wrap_stopped(cmd); + if let Some((uid, gid)) = uid_gid { + wrapped.uid(uid).gid(gid); } - Ok(()) + let child = spawn_stopped(&mut wrapped)?; + let pid = child.id() as i32; + self.worker + .lock() + .as_ref() + .context("tracker already finished")? + .set_root_pid(pid); + + let (tx, rx) = mpsc::channel(); + let poller = { + let mut bpf = self.bpf.lock(); + bpf.add_tracked_pid(pid)?; + bpf.poll_events_with_channel(10, tx)? + }; + resume(pid)?; + + Ok(Session::new(child, rx, poller)) } - pub fn attach_allocator(&mut self, lib: &AllocatorLib) -> Result<()> { - self.bpf.attach_allocator_probes(lib.kind, &lib.path) + /// Enable event tracking in the BPF program + pub fn enable_tracking(&self) -> Result<()> { + self.bpf.lock().enable_tracking() } - /// Start tracking allocations for a specific PID - /// - /// Returns a receiver channel that will receive allocation events. - /// The receiver will continue to produce events until the tracker is dropped. - pub fn track(&mut self, pid: i32) -> Result> { - // Add the PID to track - self.bpf.add_tracked_pid(pid)?; - debug!("Tracking PID {pid}"); - - // Start polling with channel - let (_poller, event_rx) = self.bpf.start_polling_with_channel(10)?; - - // Keep the poller alive by moving it into the channel - // When the receiver is dropped, the poller will also be dropped - let (tx, rx) = mpsc::channel(); - std::thread::spawn(move || { - // Keep poller alive - let _p = _poller; - while let Ok(event) = event_rx.recv() { - if tx.send(event).is_err() { - break; - } - } - }); - - Ok(rx) + /// Disable event tracking in the BPF program + pub fn disable_tracking(&self) -> Result<()> { + self.bpf.lock().disable_tracking() + } + + /// Number of events the kernel dropped because the ring buffer was full. + /// A non-zero value means the resulting trace is incomplete. + pub fn dropped_events_count(&self) -> Result { + self.bpf.lock().dropped_events_count() + } + + /// Stop the attach worker and surface any fatal error it recorded, + /// including missed exec mappings (incomplete allocator coverage). + pub fn finish(&self) -> Result<()> { + let worker = self + .worker + .lock() + .take() + .context("tracker already finished")?; + worker.finish() } /// Bump RLIMIT_MEMLOCK for kernels older than 5.11. Newer kernels account BPF @@ -94,20 +112,4 @@ impl Tracker { Ok(()) } - - /// Enable event tracking in the BPF program - pub fn enable(&mut self) -> anyhow::Result<()> { - self.bpf.enable_tracking() - } - - /// Disable event tracking in the BPF program - pub fn disable(&mut self) -> anyhow::Result<()> { - self.bpf.disable_tracking() - } - - /// Number of events the kernel dropped because the ring buffer was full. - /// A non-zero value means the resulting trace is incomplete. - pub fn dropped_events_count(&self) -> anyhow::Result { - self.bpf.dropped_events_count() - } } diff --git a/crates/memtrack/src/ipc.rs b/crates/memtrack/src/ipc.rs index 2b0d80d5..5d4b13d3 100644 --- a/crates/memtrack/src/ipc.rs +++ b/crates/memtrack/src/ipc.rs @@ -85,26 +85,20 @@ impl MemtrackIpcClient { /// Handle incoming IPC messages in memtrack #[cfg(feature = "ebpf")] -pub fn handle_ipc_message(msg: IpcMessage, tracker: &std::sync::Arc>) { +pub fn handle_ipc_message(msg: IpcMessage, tracker: &Tracker) { let response = match msg.command { - IpcCommand::Enable => match tracker.lock() { - Ok(mut t) => match t.enable() { - Ok(_) => { - debug!("Tracking enabled"); - IpcResponse::Ack - } - Err(_) => IpcResponse::Err, - }, + IpcCommand::Enable => match tracker.enable_tracking() { + Ok(_) => { + debug!("Tracking enabled"); + IpcResponse::Ack + } Err(_) => IpcResponse::Err, }, - IpcCommand::Disable => match tracker.lock() { - Ok(mut t) => match t.disable() { - Ok(_) => { - debug!("Tracking disabled"); - IpcResponse::Ack - } - Err(_) => IpcResponse::Err, - }, + IpcCommand::Disable => match tracker.disable_tracking() { + Ok(_) => { + debug!("Tracking disabled"); + IpcResponse::Ack + } Err(_) => IpcResponse::Err, }, IpcCommand::Ping => IpcResponse::Ack, diff --git a/crates/memtrack/src/lib.rs b/crates/memtrack/src/lib.rs index dc32092a..1186ef2d 100644 --- a/crates/memtrack/src/lib.rs +++ b/crates/memtrack/src/lib.rs @@ -3,6 +3,8 @@ mod allocators; mod ebpf; mod ipc; pub mod prelude; +#[cfg(feature = "ebpf")] +mod session; pub use allocators::{AllocatorKind, AllocatorLib}; pub use ipc::{ @@ -12,6 +14,8 @@ pub use ipc::{ #[cfg(feature = "ebpf")] pub use ebpf::*; +#[cfg(feature = "ebpf")] +pub use session::Session; #[cfg(feature = "ebpf")] pub use ipc::handle_ipc_message; diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index b14487eb..ec0d0176 100644 --- a/crates/memtrack/src/main.rs +++ b/crates/memtrack/src/main.rs @@ -3,12 +3,11 @@ use ipc_channel::ipc; use memtrack::prelude::*; use memtrack::{MemtrackIpcMessage, Tracker, handle_ipc_message}; use runner_shared::artifacts::{ArtifactExt, MemtrackArtifact, MemtrackEvent, MemtrackWriter}; -use std::os::unix::process::CommandExt; use std::path::{Path, PathBuf}; use std::process::Command; +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc::channel; -use std::sync::{Arc, Mutex}; use std::thread; use std::time::Duration; @@ -88,41 +87,38 @@ fn track_command( None }; - let tracker = Tracker::new()?; - let tracker_arc = Arc::new(Mutex::new(tracker)); + let tracker = Arc::new(Tracker::new()?); // Spawn IPC handler thread with the now-available tracker let ipc_handle = if let Some(rx) = ipc_channel { - let tracker_clone = tracker_arc.clone(); + let tracker = tracker.clone(); Some(thread::spawn(move || { while let Ok(msg) = rx.recv() { - handle_ipc_message(msg, &tracker_clone); + handle_ipc_message(msg, &tracker); } })) } else { // Without IPC, nothing toggles the tracking_enabled map, so events would // be dropped by the eBPF is_enabled() check. Enable it up front. - tracker_arc.lock().unwrap().enable()?; + tracker.enable_tracking()?; None }; - // Start the target command using bash to handle shell syntax + // Run the target command through bash to handle shell syntax. Drop + // privileges if running under sudo to avoid permission issues when the + // target accesses files owned by the original user. let mut cmd = Command::new("bash"); cmd.arg("-c").arg(cmd_string); - - // Drop privileges if running under sudo. This is required to avoid permission issues - // when the target command tries to access files or directories that the current user - // does not have permission to access. - if let Some((uid, gid)) = get_user_uid_gid() { + let uid_gid = get_user_uid_gid(); + if let Some((uid, gid)) = uid_gid { debug!("Running under sudo, dropping privileges to uid={uid}, gid={gid}"); - cmd.uid(uid).gid(gid); } - let mut child = cmd - .spawn() + let mut session = tracker + .spawn(&cmd, uid_gid) .map_err(|e| anyhow!("Failed to spawn child process: {e}"))?; - let root_pid = child.id() as i32; - let event_rx = { tracker_arc.lock().unwrap().track(root_pid)? }; + let root_pid = session.pid(); + let event_rx = session.take_events()?; debug!("Spawned child with pid {root_pid}"); // Generate output file name and create file for streaming events @@ -186,7 +182,7 @@ fn track_command( }); // Wait for the command to complete - let status = child.wait().context("Failed to wait for command")?; + let status = session.wait().context("Failed to wait for command")?; debug!("Command exited with status: {status}"); // Wait for drain thread to finish @@ -203,12 +199,12 @@ fn track_command( .join() .map_err(|_| anyhow::anyhow!("Failed to join writer thread"))??; + tracker.finish()?; + // Read the eBPF dropped-event counter after the run is complete. // A non-zero value means the ring buffer overflowed and the trace is // incomplete. - let dropped_events = tracker_arc - .lock() - .map_err(|_| anyhow!("tracker mutex poisoned"))? + let dropped_events = tracker .dropped_events_count() .context("Failed to read memtrack dropped-event counter")?; if dropped_events > 0 { diff --git a/crates/memtrack/src/session.rs b/crates/memtrack/src/session.rs new file mode 100644 index 00000000..7aec33fe --- /dev/null +++ b/crates/memtrack/src/session.rs @@ -0,0 +1,41 @@ +use crate::ebpf::poller::RingBufferPoller; +use crate::prelude::*; +use runner_shared::artifacts::MemtrackEvent; +use std::process::{Child, ExitStatus}; +use std::sync::mpsc::Receiver; + +/// A spawned, tracked process together with its event pipeline. The pipeline +/// stays alive as long as the session does; dropping it stops event delivery. +pub struct Session { + child: Child, + events: Option>, + _poller: RingBufferPoller, +} + +impl Session { + pub(crate) fn new( + child: Child, + events: Receiver, + poller: RingBufferPoller, + ) -> Self { + Self { + child, + events: Some(events), + _poller: poller, + } + } + + pub fn pid(&self) -> i32 { + self.child.id() as i32 + } + + /// Take ownership of the event receiver. Can only be taken once. + pub fn take_events(&mut self) -> Result> { + self.events.take().context("events already taken") + } + + /// Wait for the tracked process to exit. + pub fn wait(&mut self) -> Result { + Ok(self.child.wait()?) + } +} From 04025ecb909d28ac802862d0e2afbcd7dc9939cc Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 8 Jul 2026 21:42:38 +0200 Subject: [PATCH 3/8] feat(memtrack)!: remove eager allocator discovery On-demand attach supersedes the startup scan: delete the system-wide library glob, the build-dir walk, and the CODSPEED_MEMTRACK_BINARIES env var (the runner no longer resolves exec target binaries for it). BREAKING CHANGE: CODSPEED_MEMTRACK_BINARIES is no longer read Refs COD-1801 --- Cargo.lock | 1 - crates/memtrack/Cargo.toml | 1 - crates/memtrack/src/allocators/dynamic.rs | 76 --------------- crates/memtrack/src/allocators/mod.rs | 61 ------------ .../memtrack/src/allocators/static_linked.rs | 97 +------------------ src/cli/exec/mod.rs | 38 +------- 6 files changed, 3 insertions(+), 271 deletions(-) delete mode 100644 crates/memtrack/src/allocators/dynamic.rs diff --git a/Cargo.lock b/Cargo.lock index b2f814d1..228d95ca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2379,7 +2379,6 @@ dependencies = [ "bindgen", "clap", "env_logger", - "glob", "insta", "ipc-channel", "itertools 0.14.0", diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index de6214b8..4ba89714 100644 --- a/crates/memtrack/Cargo.toml +++ b/crates/memtrack/Cargo.toml @@ -32,7 +32,6 @@ static_assertions = "1.1" itertools = { workspace = true } paste = "1.0.15" libbpf-rs = { version = "0.26", features = ["vendored"], optional = true } -glob = "0.3.3" object = { workspace = true } parking_lot = "0.12" diff --git a/crates/memtrack/src/allocators/dynamic.rs b/crates/memtrack/src/allocators/dynamic.rs deleted file mode 100644 index db5b4df8..00000000 --- a/crates/memtrack/src/allocators/dynamic.rs +++ /dev/null @@ -1,76 +0,0 @@ -use crate::{AllocatorKind, AllocatorLib}; -use std::path::PathBuf; - -impl AllocatorKind { - /// Build glob patterns for finding this allocator's shared libraries. - fn search_patterns(&self) -> Vec { - const LIB_DIRS: &[&str] = &[ - // Debian, Ubuntu: multiarch paths - "/lib/*-linux-gnu", - "/usr/lib/*-linux-gnu", - // RHEL, Fedora, CentOS, Arch - "/lib*", - "/usr/lib*", - // Local installs - "/usr/local/lib*", - ]; - - let (filenames, nix_hints): (&[&str], &[&str]) = match self { - AllocatorKind::Libc => (&["libc.so.6"], &["glibc"]), - AllocatorKind::LibCpp => (&["libstdc++.so*"], &["gcc"]), - AllocatorKind::Jemalloc => (&["libjemalloc.so*"], &["jemalloc"]), - AllocatorKind::Mimalloc => (&["libmimalloc.so*"], &["mimalloc"]), - AllocatorKind::Tcmalloc => (&["libtcmalloc*.so*"], &["tcmalloc", "gperftools"]), - }; - - let mut patterns = Vec::new(); - - for dir in LIB_DIRS { - for filename in filenames { - patterns.push(format!("{dir}/{filename}")); - } - } - - for hint in nix_hints { - for filename in filenames { - patterns.push(format!("/nix/store/*{hint}*/lib/{filename}")); - } - } - - patterns - } -} - -/// Find dynamically linked allocator libraries on the system. -pub fn find_all() -> anyhow::Result> { - use std::collections::HashSet; - - let mut results = Vec::new(); - let mut seen_paths: HashSet = HashSet::new(); - - for kind in AllocatorKind::all() { - for pattern in kind.search_patterns() { - let paths = glob::glob(&pattern) - .ok() - .into_iter() - .flatten() - .filter_map(|p| p.ok()) - .filter_map(|p| p.canonicalize().ok()) - .filter(|path| { - std::fs::metadata(path) - .map(|m| m.is_file()) - .unwrap_or(false) - }) - .filter(|path| super::is_elf(path)) - .collect::>(); - - for path in paths { - if seen_paths.insert(path.clone()) { - results.push(AllocatorLib { kind: *kind, path }); - } - } - } - } - - Ok(results) -} diff --git a/crates/memtrack/src/allocators/mod.rs b/crates/memtrack/src/allocators/mod.rs index a2c4a1fb..d1a429ad 100644 --- a/crates/memtrack/src/allocators/mod.rs +++ b/crates/memtrack/src/allocators/mod.rs @@ -5,7 +5,6 @@ use std::path::PathBuf; -mod dynamic; mod static_linked; /// Represents the different allocator types we support. @@ -67,63 +66,3 @@ pub struct AllocatorLib { pub kind: AllocatorKind, pub path: PathBuf, } - -impl AllocatorLib { - pub fn find_all() -> anyhow::Result> { - let mut allocators = static_linked::find_all()?; - allocators.extend(Self::find_from_env()); - allocators.extend(dynamic::find_all()?); - Ok(allocators) - } - - /// Discover allocators from binaries listed in the `CODSPEED_MEMTRACK_BINARIES` env var. - /// - /// The variable is a platform path-list. Each path is checked for - /// a statically linked allocator via [`AllocatorLib::from_path_static`]. - /// Invalid or missing paths are silently skipped (with a debug log). - fn find_from_env() -> Vec { - let Some(raw) = std::env::var_os("CODSPEED_MEMTRACK_BINARIES") else { - return vec![]; - }; - - std::env::split_paths(&raw) - .filter(|p| !p.as_os_str().is_empty()) - .filter_map(|p| { - let path = p.as_path(); - match Self::from_path_static(path) { - Ok(alloc) => { - log::debug!( - "Found {} allocator in env-specified binary: {}", - alloc.kind.name(), - path.display() - ); - Some(alloc) - } - Err(e) => { - log::debug!("Skipping env-specified binary {}: {e}", path.display()); - None - } - } - }) - .collect() - } -} - -/// Check if a file is an ELF binary by reading its magic bytes. -fn is_elf(path: &std::path::Path) -> bool { - use std::fs; - use std::io::Read; - - let mut file = match fs::File::open(path) { - Ok(f) => f, - Err(_) => return false, - }; - - let mut magic = [0u8; 4]; - if file.read_exact(&mut magic).is_err() { - return false; - } - - // ELF magic: 0x7F 'E' 'L' 'F' - magic == [0x7F, b'E', b'L', b'F'] -} diff --git a/crates/memtrack/src/allocators/static_linked.rs b/crates/memtrack/src/allocators/static_linked.rs index 9695070d..d1891303 100644 --- a/crates/memtrack/src/allocators/static_linked.rs +++ b/crates/memtrack/src/allocators/static_linked.rs @@ -1,6 +1,6 @@ use std::collections::HashSet; use std::fs; -use std::path::{Path, PathBuf}; +use std::path::Path; use crate::allocators::{AllocatorKind, AllocatorLib}; @@ -17,79 +17,6 @@ impl AllocatorKind { } } -/// Walk upward and downward from current directory to find build directories. -/// Returns all found build directories in order of preference. -fn find_build_dirs() -> Vec { - let mut dirs = Vec::new(); - let Ok(current_dir) = std::env::current_dir() else { - return dirs; - }; - - let patterns = ["target/codspeed/analysis", "bazel-bin", "build"]; - let mut check_patterns = |dir: &Path| { - for pattern in &patterns { - let path = dir.join(pattern); - if path.is_dir() { - dirs.push(path); - } - } - }; - - // Walk upward from parent directories - // Note: We skip current_dir here since the downward walk (below) already checks it - let mut current = current_dir.clone(); - while current.pop() { - check_patterns(¤t); - } - - // Walk downward from current directory - let mut stack = vec![current_dir]; - while let Some(dir) = stack.pop() { - check_patterns(&dir); - - // Read subdirectories - let Ok(entries) = fs::read_dir(&dir) else { - continue; - }; - - for entry in entries.filter_map(Result::ok) { - let path = entry.path(); - - let Some(name) = path.file_name().and_then(|n| n.to_str()) else { - continue; - }; - - // Skip hidden dirs and common excludes - if name.starts_with('.') || matches!(name, "node_modules" | "vendor" | "venv") { - continue; - } - - // Don't recursive into dirs that we want to match. - // This can happen with `target` as it contains build dirs for statically linked crates. - if matches!(name, "target" | "bazel-bin" | "build") { - continue; - } - - if path.is_file() { - continue; - } - - stack.push(path); - } - } - - dirs -} - -fn find_binaries_in_dir(dir: &Path) -> Vec { - glob::glob(&format!("{}/**/*", dir.display())) - .into_iter() - .flatten() - .filter_map(Result::ok) - .filter(|p| p.is_file() && super::is_elf(p)) - .collect::>() -} - fn find_statically_linked_allocator(path: &Path) -> Option { use object::{Object, ObjectSymbol}; @@ -111,28 +38,6 @@ fn find_statically_linked_allocator(path: &Path) -> Option { .copied() } -pub fn find_all() -> anyhow::Result> { - let build_dirs = find_build_dirs(); - if build_dirs.is_empty() { - return Ok(vec![]); - } - - let mut allocators = Vec::new(); - for build_dir in build_dirs { - let bins = find_binaries_in_dir(&build_dir); - - for bin in bins { - let Some(kind) = find_statically_linked_allocator(&bin) else { - continue; - }; - - allocators.push(AllocatorLib { kind, path: bin }); - } - } - - Ok(allocators) -} - impl AllocatorLib { pub fn from_path_static(path: &Path) -> Result> { let kind = find_statically_linked_allocator(path).ok_or("No allocator found")?; diff --git a/src/cli/exec/mod.rs b/src/cli/exec/mod.rs index f006bddf..a63e4209 100644 --- a/src/cli/exec/mod.rs +++ b/src/cli/exec/mod.rs @@ -8,7 +8,7 @@ use crate::project_config::ProjectConfig; use crate::project_config::merger::ConfigMerger; use crate::upload::poll_results::PollResultsOptions; use clap::Args; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::path::Path; use url::Url; @@ -120,44 +120,10 @@ pub async fn run( /// Sets up the orchestrator and drives execution. Exec-harness installation is handled /// by the orchestrator when exec targets are present. pub async fn execute_config( - mut config: OrchestratorConfig, + config: OrchestratorConfig, api_client: &mut CodSpeedAPIClient, setup_cache_dir: Option<&Path>, ) -> Result<()> { - // Resolve exec target binary paths so memtrack can discover statically linked - // allocators (which may not live in known build dirs). - let memtrack_binaries: HashSet<_> = config - .targets - .iter() - .filter_map(|t| match t { - executor::BenchmarkTarget::Exec { command, .. } => command.first().cloned(), - _ => None, - }) - .filter_map(|bin| { - let result = match &config.working_directory { - Some(cwd) => which::which_in(&bin, std::env::var_os("PATH"), cwd), - None => which::which(&bin), - }; - result.ok() - }) - .collect(); - - if !memtrack_binaries.is_empty() { - let mut all_paths = memtrack_binaries; - - // Merge with any user-provided value from the parent environment. - if let Some(existing) = std::env::var_os("CODSPEED_MEMTRACK_BINARIES") { - all_paths.extend(std::env::split_paths(&existing)); - } - - let joined = - std::env::join_paths(&all_paths).expect("memtrack binary paths should be joinable"); - config.extra_env.insert( - "CODSPEED_MEMTRACK_BINARIES".into(), - joined.to_string_lossy().into_owned(), - ); - } - let orchestrator = executor::Orchestrator::new(config, api_client).await?; if !orchestrator.is_local() { From 527e6789edb953f5c1cde25e756d38572be750bb Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 8 Jul 2026 21:54:47 +0200 Subject: [PATCH 4/8] fix(memtrack): classify unprefixed jemalloc builds as jemalloc Distro libjemalloc is built without the je_ symbol prefix, so symbol classification missed it and fell through to libc++ (jemalloc exports operator new), leaving plain malloc/aligned_alloc unprobed. Match on mallocx, which keeps its name in unprefixed builds and is unique to jemalloc. --- crates/memtrack/src/allocators/static_linked.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/memtrack/src/allocators/static_linked.rs b/crates/memtrack/src/allocators/static_linked.rs index d1891303..0c2c7d79 100644 --- a/crates/memtrack/src/allocators/static_linked.rs +++ b/crates/memtrack/src/allocators/static_linked.rs @@ -10,7 +10,12 @@ impl AllocatorKind { match self { AllocatorKind::Libc => &["malloc", "free"], AllocatorKind::LibCpp => &["_Znwm", "_Znam", "_ZdlPv", "_ZdaPv"], - AllocatorKind::Jemalloc => &["_rjem_malloc", "je_malloc", "je_malloc_default"], + // "mallocx" covers unprefixed jemalloc builds (e.g. distro + // libjemalloc.so): the extended API keeps its name when the je_ + // prefix is disabled, and no other allocator exports it. + AllocatorKind::Jemalloc => { + &["_rjem_malloc", "je_malloc", "je_malloc_default", "mallocx"] + } AllocatorKind::Mimalloc => &["mi_malloc_aligned", "mi_malloc", "mi_free"], AllocatorKind::Tcmalloc => &["tc_malloc", "tc_free", "tc_version"], } From 5b7678833464cce242284d7f6bb50cf0f0b42e88 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Thu, 9 Jul 2026 18:48:45 +0200 Subject: [PATCH 5/8] test(memtrack): add dlopen/on-demand allocator tests + CI matrix --- .github/workflows/ci.yml | 2 +- .../memtrack/testdata/dlopen/dlopen_alloc.c | 34 +++++ .../memtrack/testdata/dlopen/fake_jemalloc.c | 20 +++ .../memtrack/testdata/dlopen/fake_mimalloc.c | 21 +++ .../memtrack/testdata/dlopen/thread_dlopen.c | 53 +++++++ crates/memtrack/tests/cpp_tests.rs | 28 ++-- crates/memtrack/tests/dlopen_tests.rs | 131 ++++++++++++++++++ crates/memtrack/tests/rust_tests.rs | 10 +- crates/memtrack/tests/shared.rs | 79 ++++++----- ...awn_tests__spawn_on_demand_discovery.snap} | 1 - .../spawn_tests__spawn_without_discovery.snap | 6 - crates/memtrack/tests/spawn_tests.rs | 35 +---- 12 files changed, 327 insertions(+), 93 deletions(-) create mode 100644 crates/memtrack/testdata/dlopen/dlopen_alloc.c create mode 100644 crates/memtrack/testdata/dlopen/fake_jemalloc.c create mode 100644 crates/memtrack/testdata/dlopen/fake_mimalloc.c create mode 100644 crates/memtrack/testdata/dlopen/thread_dlopen.c create mode 100644 crates/memtrack/tests/dlopen_tests.rs rename crates/memtrack/tests/snapshots/{spawn_tests__spawn_with_discovery.snap => spawn_tests__spawn_on_demand_discovery.snap} (89%) delete mode 100644 crates/memtrack/tests/snapshots/spawn_tests__spawn_without_discovery.snap diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b76aa609..fb38e88c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -96,7 +96,7 @@ jobs: # Each memtrack integration test binary runs its cases serially # (eBPF tracker can't overlap with itself in one process), so we # shard at the test-binary level to parallelize across jobs. - test: [c_tests, cpp_tests, rust_tests, spawn_tests] + test: [c_tests, cpp_tests, rust_tests, spawn_tests, dlopen_tests] steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: diff --git a/crates/memtrack/testdata/dlopen/dlopen_alloc.c b/crates/memtrack/testdata/dlopen/dlopen_alloc.c new file mode 100644 index 00000000..99b91493 --- /dev/null +++ b/crates/memtrack/testdata/dlopen/dlopen_alloc.c @@ -0,0 +1,34 @@ +#include +#include +#include + +/* dlopen an allocator lib and IMMEDIATELY (no sleep) allocate through it. The + * immediacy is the race assertion: every allocation must be captured, which + * requires the on-demand watcher to have stopped-classified-attached-resumed + * before dlopen returns. */ +int main(int argc, char** argv) { + if (argc < 2) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 1; + } + + void* handle = dlopen(argv[1], RTLD_NOW); + if (!handle) { + fprintf(stderr, "dlopen failed: %s\n", dlerror()); + return 1; + } + + void* (*mi_malloc)(size_t) = (void* (*)(size_t))dlsym(handle, "mi_malloc"); + void (*mi_free)(void*) = (void (*)(void*))dlsym(handle, "mi_free"); + if (!mi_malloc || !mi_free) { + fprintf(stderr, "dlsym failed\n"); + return 1; + } + + for (int i = 0; i < 100; i++) { + void* p = mi_malloc(4242); + mi_free(p); + } + + return 0; +} diff --git a/crates/memtrack/testdata/dlopen/fake_jemalloc.c b/crates/memtrack/testdata/dlopen/fake_jemalloc.c new file mode 100644 index 00000000..6fa72fdb --- /dev/null +++ b/crates/memtrack/testdata/dlopen/fake_jemalloc.c @@ -0,0 +1,20 @@ +#include + +/* Self-contained bump allocator; classifies as Jemalloc via the je_malloc + * symbol. See fake_mimalloc.c for why it does not call libc malloc. */ +static char arena[1 << 20]; +static size_t off = 0; + +void* je_malloc(size_t size) { + size_t aligned = (size + 15) & ~((size_t)15); + if (off + aligned > sizeof(arena)) { + return 0; + } + void* p = &arena[off]; + off += aligned; + return p; +} + +void je_free(void* p) { + (void)p; +} diff --git a/crates/memtrack/testdata/dlopen/fake_mimalloc.c b/crates/memtrack/testdata/dlopen/fake_mimalloc.c new file mode 100644 index 00000000..801f89ac --- /dev/null +++ b/crates/memtrack/testdata/dlopen/fake_mimalloc.c @@ -0,0 +1,21 @@ +#include + +/* Self-contained bump allocator so mi_malloc does not delegate to libc malloc + * (which is probed separately and would double-count). Classifies as Mimalloc + * via the mi_malloc/mi_free symbols. */ +static char arena[1 << 20]; +static size_t off = 0; + +void* mi_malloc(size_t size) { + size_t aligned = (size + 15) & ~((size_t)15); + if (off + aligned > sizeof(arena)) { + return 0; + } + void* p = &arena[off]; + off += aligned; + return p; +} + +void mi_free(void* p) { + (void)p; +} diff --git a/crates/memtrack/testdata/dlopen/thread_dlopen.c b/crates/memtrack/testdata/dlopen/thread_dlopen.c new file mode 100644 index 00000000..afdb49dd --- /dev/null +++ b/crates/memtrack/testdata/dlopen/thread_dlopen.c @@ -0,0 +1,53 @@ +#include +#include +#include +#include + +/* Two threads each dlopen a distinct allocator lib concurrently and allocate + * through it. Exercises concurrent stop-the-world: both threads trigger the + * watcher, both must be stopped, both libs classified/attached, both resumed. */ +struct job { + const char* path; + const char* alloc_sym; + const char* free_sym; + size_t size; +}; + +static void* run(void* arg) { + struct job* j = (struct job*)arg; + void* handle = dlopen(j->path, RTLD_NOW); + if (!handle) { + return NULL; + } + + void* (*alloc)(size_t) = (void* (*)(size_t))dlsym(handle, j->alloc_sym); + void (*dealloc)(void*) = (void (*)(void*))dlsym(handle, j->free_sym); + if (!alloc || !dealloc) { + return NULL; + } + + for (int i = 0; i < 100; i++) { + void* p = alloc(j->size); + dealloc(p); + } + + return NULL; +} + +int main(int argc, char** argv) { + if (argc < 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 1; + } + + struct job j1 = {argv[1], "mi_malloc", "mi_free", 4242}; + struct job j2 = {argv[2], "je_malloc", "je_free", 4243}; + + pthread_t t1, t2; + pthread_create(&t1, NULL, run, &j1); + pthread_create(&t2, NULL, run, &j2); + pthread_join(t1, NULL); + pthread_join(t2, NULL); + + return 0; +} diff --git a/crates/memtrack/tests/cpp_tests.rs b/crates/memtrack/tests/cpp_tests.rs index 2a219441..b7c5b61c 100644 --- a/crates/memtrack/tests/cpp_tests.rs +++ b/crates/memtrack/tests/cpp_tests.rs @@ -1,12 +1,30 @@ #[macro_use] mod shared; -use memtrack::AllocatorLib; use rstest::rstest; use std::path::Path; use std::process::Command; +/// A cached cmake configure pins absolute library paths from a previous +/// environment (e.g. garbage-collected nix store entries). On failure, wipe +/// the build dir and retry once from a fresh configure. +/// +/// Builds are serialized: parallel test cases share the build dir, and the +/// retry path deletes it. fn compile_cpp_project(project_dir: &Path, target: &str) -> anyhow::Result { + static BUILD_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + let _guard = BUILD_LOCK.lock().unwrap(); + + match build_cpp_target(project_dir, target) { + Ok(path) => Ok(path), + Err(_) => { + std::fs::remove_dir_all(project_dir.join("build"))?; + build_cpp_target(project_dir, target) + } + } +} + +fn build_cpp_target(project_dir: &Path, target: &str) -> anyhow::Result { let build_exists = project_dir.join("build").exists(); if !build_exists { // Configure with cmake -B build @@ -57,13 +75,7 @@ fn test_cpp_alloc_tracking(#[case] target: &str) -> Result<(), Box PathBuf { + let src = dir.join(format!("{name}.c")); + std::fs::write(&src, source).expect("write source"); + let out = dir.join(format!("{name}.so")); + let ok = Command::new("gcc") + .args(["-shared", "-fPIC", "-o"]) + .arg(&out) + .arg(&src) + .status() + .expect("run gcc") + .success(); + assert!(ok, "failed to compile {name}.so"); + out +} + +fn compile_exe(source: &str, name: &str, dir: &Path, libs: &[&str]) -> PathBuf { + let src = dir.join(format!("{name}.c")); + std::fs::write(&src, source).expect("write source"); + let out = dir.join(name); + let ok = Command::new("gcc") + .arg("-o") + .arg(&out) + .arg(&src) + .args(libs) + .status() + .expect("run gcc") + .success(); + assert!(ok, "failed to compile {name}"); + out +} + +/// dlopen -> fentry -> SIGSTOP -> resolve -> classify -> attach -> SIGCONT -> +/// first allocation captured. The dlopen'd allocator's 100 allocations must all +/// be captured, proving the attach completes before the child's first alloc. +#[test_with::env(GITHUB_ACTIONS)] +#[test_log::test] +fn test_dlopen_allocator() -> Result<(), Box> { + let dir = TempDir::new()?; + let lib = compile_shared( + include_str!("../testdata/dlopen/fake_mimalloc.c"), + "libfake_mimalloc", + dir.path(), + ); + let exe = compile_exe( + include_str!("../testdata/dlopen/dlopen_alloc.c"), + "dlopen_alloc", + dir.path(), + &["-ldl"], + ); + + let mut cmd = Command::new(&exe); + cmd.arg(&lib); + let (events, thread_handle) = shared::track_command(cmd)?; + + let malloc_addrs: HashSet = events + .iter() + .filter_map(|e| match e.kind { + MemtrackEventKind::Malloc { size: 4242 } => Some(e.addr), + _ => None, + }) + .collect(); + let malloc_count = events + .iter() + .filter(|e| matches!(e.kind, MemtrackEventKind::Malloc { size: 4242 })) + .count(); + let free_count = events + .iter() + .filter(|e| matches!(e.kind, MemtrackEventKind::Free) && malloc_addrs.contains(&e.addr)) + .count(); + + assert_eq!(malloc_count, 100, "expected 100 mi_malloc(4242) events"); + assert_eq!( + free_count, 100, + "expected 100 mi_free events for tracked addrs" + ); + + thread_handle.join().unwrap(); + Ok(()) +} + +/// Two threads dlopen distinct allocator libs concurrently. Both libs' 100 +/// allocations each must be captured, exercising concurrent stop-the-world. +#[test_with::env(GITHUB_ACTIONS)] +#[test_log::test] +fn test_thread_dlopen() -> Result<(), Box> { + let dir = TempDir::new()?; + let mimalloc = compile_shared( + include_str!("../testdata/dlopen/fake_mimalloc.c"), + "libfake_mimalloc", + dir.path(), + ); + let jemalloc = compile_shared( + include_str!("../testdata/dlopen/fake_jemalloc.c"), + "libfake_jemalloc", + dir.path(), + ); + let exe = compile_exe( + include_str!("../testdata/dlopen/thread_dlopen.c"), + "thread_dlopen", + dir.path(), + &["-ldl", "-lpthread"], + ); + + let mut cmd = Command::new(&exe); + cmd.arg(&mimalloc).arg(&jemalloc); + let (events, thread_handle) = shared::track_command(cmd)?; + + let m4242 = events + .iter() + .filter(|e| matches!(e.kind, MemtrackEventKind::Malloc { size: 4242 })) + .count(); + let m4243 = events + .iter() + .filter(|e| matches!(e.kind, MemtrackEventKind::Malloc { size: 4243 })) + .count(); + + assert_eq!(m4242, 100, "expected 100 mi_malloc(4242) events"); + assert_eq!(m4243, 100, "expected 100 je_malloc(4243) events"); + + thread_handle.join().unwrap(); + Ok(()) +} diff --git a/crates/memtrack/tests/rust_tests.rs b/crates/memtrack/tests/rust_tests.rs index 995fc2d5..ce1857aa 100644 --- a/crates/memtrack/tests/rust_tests.rs +++ b/crates/memtrack/tests/rust_tests.rs @@ -1,7 +1,6 @@ #[macro_use] mod shared; -use memtrack::AllocatorLib; use rstest::rstest; use std::path::Path; @@ -18,13 +17,8 @@ fn test_rust_alloc_tracking( let crate_path = Path::new("testdata/alloc_rust"); let binary = shared::compile_rust_binary(crate_path, "alloc_rust", features)?; - // Try to find a static allocator in the binary, then attach to it as well - // This is needed because the CWD is different, which breaks the heuristics. - let allocators = AllocatorLib::from_path_static(&binary) - .map(|a| vec![a]) - .unwrap_or_default(); - - let (events, thread_handle) = shared::track_binary_with_opts(&binary, &allocators)?; + // No extra allocators: the watcher must discover the static allocator itself. + let (events, thread_handle) = shared::track_binary(&binary)?; assert_events_with_marker!(name, &events); thread_handle.join().unwrap(); diff --git a/crates/memtrack/tests/shared.rs b/crates/memtrack/tests/shared.rs index 054b3b8e..d31bce09 100644 --- a/crates/memtrack/tests/shared.rs +++ b/crates/memtrack/tests/shared.rs @@ -1,12 +1,10 @@ #![allow(dead_code, unused)] -use anyhow::Context; +use memtrack::Tracker; use memtrack::prelude::*; -use memtrack::{AllocatorLib, Tracker}; use runner_shared::artifacts::{MemtrackEvent as Event, MemtrackEventKind}; use std::path::Path; use std::process::Command; -use std::time::Duration; type TrackResult = anyhow::Result<(Vec, std::thread::JoinHandle<()>)>; @@ -101,14 +99,28 @@ macro_rules! assert_events_with_marker { } /// Compile a Rust binary from a test crate directory. +/// +/// Each feature set builds into its own target dir: parallel test cases would +/// otherwise overwrite the same output binary and race against each other. pub fn compile_rust_binary( crate_dir: &Path, name: &str, features: &[&str], ) -> anyhow::Result { + let target_dir = match features { + [] => "target/default".to_string(), + _ => format!("target/{}", features.join("-")), + }; + let mut cmd = Command::new("cargo"); - cmd.current_dir(crate_dir) - .args(["build", "--release", "--bin", name]); + cmd.current_dir(crate_dir).args([ + "build", + "--release", + "--bin", + name, + "--target-dir", + &target_dir, + ]); if !features.is_empty() { cmd.arg("--features").arg(features.join(",")); @@ -121,51 +133,38 @@ pub fn compile_rust_binary( return Err(anyhow::anyhow!("Failed to compile Rust crate")); } - Ok(crate_dir.join(format!("target/release/{name}"))) + Ok(crate_dir.join(format!("{target_dir}/release/{name}"))) } -/// Track a spawned command, collecting all memory events. +/// Track a binary, collecting all memory events. +pub fn track_binary(binary: &Path) -> TrackResult { + track_command(Command::new(binary)) +} + +/// Track a command, collecting all memory events. /// -/// When `discover_system_allocators` is true, the tracker will scan for all -/// allocators on the system (slower). When false, only `extra_allocators` are used. -pub fn track_command( - mut command: Command, - extra_allocators: &[AllocatorLib], - discover_system_allocators: bool, -) -> TrackResult { - // IMPORTANT: Always initialize the tracker BEFORE spawning the binary, as it can take some time to - // attach to all the allocator libraries (especially when using NixOS). - let mut tracker = if discover_system_allocators { - memtrack::Tracker::new()? - } else { - memtrack::Tracker::new_without_allocators()? - }; - tracker.attach_allocators(extra_allocators)?; +/// No allocators are pre-attached: the exec-mapping watcher discovers and +/// attaches them as the tracked tree maps executable files. +pub fn track_command(command: Command) -> TrackResult { + let tracker = Tracker::new()?; + tracker.enable_tracking()?; - let child = command.spawn().context("Failed to spawn command")?; - let root_pid = child.id() as i32; + let mut session = tracker.spawn(&command, None)?; + let rx = session.take_events()?; - tracker.enable()?; - let rx = tracker.track(root_pid)?; + session.wait()?; + // Dropping the session does a final ring buffer drain and closes the + // channel, so collecting terminates without a silence timeout. + drop(session); + let events: Vec = rx.iter().collect(); - let mut events = Vec::new(); - while let Ok(event) = rx.recv_timeout(Duration::from_secs(10)) { - events.push(event); - } + tracker.finish()?; - // Drop the tracker in a new thread to not block the test - let thread_handle = std::thread::spawn(move || core::mem::drop(tracker)); + // Drop the tracker in a new thread to not block the test. + let thread_handle = std::thread::spawn(move || drop(tracker)); info!("Tracked {} events", events.len()); trace!("Events: {events:#?}"); Ok((events, thread_handle)) } - -pub fn track_binary_with_opts(binary: &Path, extra_allocators: &[AllocatorLib]) -> TrackResult { - track_command(Command::new(binary), extra_allocators, true) -} - -pub fn track_binary(binary: &Path) -> TrackResult { - track_binary_with_opts(binary, &[]) -} diff --git a/crates/memtrack/tests/snapshots/spawn_tests__spawn_with_discovery.snap b/crates/memtrack/tests/snapshots/spawn_tests__spawn_on_demand_discovery.snap similarity index 89% rename from crates/memtrack/tests/snapshots/spawn_tests__spawn_with_discovery.snap rename to crates/memtrack/tests/snapshots/spawn_tests__spawn_on_demand_discovery.snap index ab014ba8..f36dc606 100644 --- a/crates/memtrack/tests/snapshots/spawn_tests__spawn_with_discovery.snap +++ b/crates/memtrack/tests/snapshots/spawn_tests__spawn_on_demand_discovery.snap @@ -1,6 +1,5 @@ --- source: crates/memtrack/tests/spawn_tests.rs -assertion_line: 90 expression: formatted_events --- [ diff --git a/crates/memtrack/tests/snapshots/spawn_tests__spawn_without_discovery.snap b/crates/memtrack/tests/snapshots/spawn_tests__spawn_without_discovery.snap deleted file mode 100644 index f6fe99c4..00000000 --- a/crates/memtrack/tests/snapshots/spawn_tests__spawn_without_discovery.snap +++ /dev/null @@ -1,6 +0,0 @@ ---- -source: crates/memtrack/tests/spawn_tests.rs -assertion_line: 71 -expression: formatted_events ---- -[] diff --git a/crates/memtrack/tests/spawn_tests.rs b/crates/memtrack/tests/spawn_tests.rs index 8fb9f463..5511a752 100644 --- a/crates/memtrack/tests/spawn_tests.rs +++ b/crates/memtrack/tests/spawn_tests.rs @@ -1,8 +1,6 @@ #[macro_use] mod shared; -use memtrack::AllocatorLib; -use memtrack::prelude::*; use std::path::Path; use std::process::Command; @@ -10,41 +8,20 @@ fn compile_spawn_binary(name: &str) -> anyhow::Result { shared::compile_rust_binary(Path::new("testdata/spawn_wrapper"), name, &[]) } -/// Without discovering the child's static allocator, the spawned child's -/// jemalloc allocations should NOT be tracked — we expect no marker-delimited events. +/// The exec-mapping watcher discovers the statically-linked jemalloc in the +/// spawned grandchild binary automatically: no pre-attached allocators. This is +/// the COD-1801 core ask, and it exercises the watcher firing on execve. #[test_with::env(GITHUB_ACTIONS)] #[test_log::test] -fn test_spawn_without_static_allocator_discovery() -> Result<(), Box> { +fn test_spawn_static_allocator_discovery() -> Result<(), Box> { let wrapper = compile_spawn_binary("wrapper")?; let child = compile_spawn_binary("alloc_child")?; - // No allocators attached — static jemalloc in alloc_child won't be found let mut cmd = Command::new(&wrapper); cmd.arg(&child); - let (events, thread_handle) = shared::track_command(cmd, &[], false)?; + let (events, thread_handle) = shared::track_command(cmd)?; - assert_events_with_marker!("spawn_without_discovery", &events); - - thread_handle.join().unwrap(); - Ok(()) -} - -/// With the child binary's static allocator discovered (simulating -/// CODSPEED_MEMTRACK_BINARIES), allocations from the spawned child should be tracked. -#[test_with::env(GITHUB_ACTIONS)] -#[test_log::test] -fn test_spawn_with_static_allocator_discovery() -> Result<(), Box> { - let wrapper = compile_spawn_binary("wrapper")?; - let child = compile_spawn_binary("alloc_child")?; - - // Simulate what CODSPEED_MEMTRACK_BINARIES does: discover the static jemalloc - let allocator = AllocatorLib::from_path_static(&child)?; - - let mut cmd = Command::new(&wrapper); - cmd.arg(&child); - let (events, thread_handle) = shared::track_command(cmd, &[allocator], false)?; - - assert_events_with_marker!("spawn_with_discovery", &events); + assert_events_with_marker!("spawn_on_demand_discovery", &events); thread_handle.join().unwrap(); Ok(()) From ccc4c4e99262019278e3f5315a77b5ece9b54f74 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Mon, 13 Jul 2026 15:07:58 +0200 Subject: [PATCH 6/8] fixup! feat(memtrack): add exec-mapping watcher BPF prog and inode maps --- crates/memtrack/src/ebpf/c/attach.bpf.h | 1 - crates/memtrack/src/ebpf/c/event.h | 1 - 2 files changed, 2 deletions(-) diff --git a/crates/memtrack/src/ebpf/c/attach.bpf.h b/crates/memtrack/src/ebpf/c/attach.bpf.h index b3814b8d..893d115e 100644 --- a/crates/memtrack/src/ebpf/c/attach.bpf.h +++ b/crates/memtrack/src/ebpf/c/attach.bpf.h @@ -44,7 +44,6 @@ int BPF_PROG(watch_exec_mmap, struct file* file, unsigned long prot, unsigned lo return 0; } req->pid = tgid; - req->tid = (__u32)pid_tgid; req->dev = key.dev; req->ino = key.ino; bpf_ringbuf_submit(req, 0); diff --git a/crates/memtrack/src/ebpf/c/event.h b/crates/memtrack/src/ebpf/c/event.h index 26635cf3..ff41f0be 100644 --- a/crates/memtrack/src/ebpf/c/event.h +++ b/crates/memtrack/src/ebpf/c/event.h @@ -51,7 +51,6 @@ struct event { /* Request from the exec-mapping watcher to the userspace attach worker */ struct attach_request { uint32_t pid; - uint32_t tid; uint64_t dev; /* kernel s_dev encoding: (major << 20) | minor */ uint64_t ino; }; From 9de9361d78c057d75f836626a228d1894a87b428 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Mon, 13 Jul 2026 15:09:25 +0200 Subject: [PATCH 7/8] fixup! feat(memtrack): attach allocator probes on demand via exec-mapping watcher --- crates/memtrack/src/ebpf/attach_worker.rs | 3 ++- crates/memtrack/src/ebpf/memtrack.rs | 7 +------ crates/memtrack/src/ebpf/tracker.rs | 8 +++----- 3 files changed, 6 insertions(+), 12 deletions(-) diff --git a/crates/memtrack/src/ebpf/attach_worker.rs b/crates/memtrack/src/ebpf/attach_worker.rs index 1eeba9de..d82402f2 100644 --- a/crates/memtrack/src/ebpf/attach_worker.rs +++ b/crates/memtrack/src/ebpf/attach_worker.rs @@ -123,7 +123,8 @@ impl Worker { loop { if self.shutdown.load(Ordering::SeqCst) { // Producers are gone: a synchronous drain flushes everything. - if self.poller.drain().is_err() { + if let Err(e) = self.poller.drain() { + self.record_fatal(e); return; } let mut batch: Vec = self.rx.try_iter().collect(); diff --git a/crates/memtrack/src/ebpf/memtrack.rs b/crates/memtrack/src/ebpf/memtrack.rs index d2d704e1..9263492d 100644 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ b/crates/memtrack/src/ebpf/memtrack.rs @@ -154,7 +154,7 @@ macro_rules! attach_uprobe { macro_rules! attach_tracepoint { ($func:ident, $prog:ident) => { - fn $func(&mut self) -> Result<()> { + pub fn $func(&mut self) -> Result<()> { let link = self .skel .progs @@ -448,11 +448,6 @@ impl MemtrackBpf { } attach_tracepoint!(sched_fork); - pub fn attach_tracepoints(&mut self) -> Result<()> { - self.attach_sched_fork()?; - Ok(()) - } - /// Poll the allocation-event ring buffer into `tx`. The returned poller /// keeps the pipeline alive; events stop flowing when it is dropped. pub fn poll_events_with_channel( diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index bd2f7861..42f56134 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -15,15 +15,13 @@ pub struct Tracker { } impl Tracker { - /// Create a new tracker with on-demand allocator attach. - /// - /// No allocators are pre-attached: the exec-mapping watcher discovers and - /// attaches them as the tracked tree maps executable files. + /// 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::bump_memlock_rlimit()?; let mut bpf = MemtrackBpf::new()?; - bpf.attach_tracepoints()?; + bpf.attach_sched_fork()?; bpf.attach_exec_watcher()?; let bpf = Arc::new(Mutex::new(bpf)); From a0a194d8ee519a917885d2cc89f082115d34bb54 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Mon, 13 Jul 2026 16:36:42 +0200 Subject: [PATCH 8/8] fixup! feat(memtrack): add exec-mapping watcher BPF prog and inode maps --- crates/memtrack/src/ebpf/c/attach.bpf.h | 6 +++++- crates/memtrack/src/ebpf/c/memtrack.bpf.c | 7 +++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/crates/memtrack/src/ebpf/c/attach.bpf.h b/crates/memtrack/src/ebpf/c/attach.bpf.h index 893d115e..edee1f7e 100644 --- a/crates/memtrack/src/ebpf/c/attach.bpf.h +++ b/crates/memtrack/src/ebpf/c/attach.bpf.h @@ -1,4 +1,8 @@ -/* == Exec-mapping watcher: on-demand allocator attach == */ +/* == Exec-mapping watcher == + * + * On a tracked process's first exec-map of an unknown inode: stop it + * (SIGSTOP) and queue an attach request on a ring buffer. Userspace + * classifies the file, attaches allocator probes, then resumes it. */ #define MEMTRACK_PROT_EXEC 0x4 #define MEMTRACK_SIGSTOP 19 diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 2d4b48d3..1950499e 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/crates/memtrack/src/ebpf/c/memtrack.bpf.c @@ -92,6 +92,9 @@ int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) { return 0; } +/* Included as source, not a header: build.rs compiles this file as one + * translation unit into a single skeleton. Must follow is_tracked() and the + * BPF_*_MAP macros it uses. */ #include "attach.bpf.h" /* == Helper functions for the allocation tracking == */ @@ -101,9 +104,9 @@ static __always_inline int is_enabled(void) { __u32 key = 0; __u8* enabled = bpf_map_lookup_elem(&tracking_enabled, &key); - /* Default to enabled if map not initialized */ + /* ARRAY-map lookups can't fail for a valid index; fail closed if one ever does. */ if (!enabled) { - return 1; + return 0; } return *enabled;