From 64c06da533694ba2241428b1bcf75fccb2703d18 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 7 Jul 2026 19:40:28 +0200 Subject: [PATCH 1/6] feat(memtrack): add exec-mapping watcher BPF prog and inode maps --- crates/memtrack/src/ebpf/c/attach.bpf.h | 57 +++++++++++++++++++++++ crates/memtrack/src/ebpf/c/event.h | 7 +++ crates/memtrack/src/ebpf/c/memtrack.bpf.c | 9 +++- crates/memtrack/src/ebpf/events.rs | 25 ++++++++++ 4 files changed, 96 insertions(+), 2 deletions(-) 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..edee1f7e --- /dev/null +++ b/crates/memtrack/src/ebpf/c/attach.bpf.h @@ -0,0 +1,57 @@ +/* == 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 + +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->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..ff41f0be 100644 --- a/crates/memtrack/src/ebpf/c/event.h +++ b/crates/memtrack/src/ebpf/c/event.h @@ -48,4 +48,11 @@ struct event { } data; }; +/* Request from the exec-mapping watcher to the userspace attach worker */ +struct attach_request { + uint32_t pid; + 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 ad64d73b..4c93d241 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/crates/memtrack/src/ebpf/c/memtrack.bpf.c @@ -92,6 +92,11 @@ 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 == */ /* Wake the consumer only once this much unconsumed data has accumulated. @@ -111,9 +116,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; 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 8c23107a3b64b7018489ca80e850bcde8da23bff Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 8 Jul 2026 21:42:06 +0200 Subject: [PATCH 2/6] 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 | 251 ++++++++++++++++++++++ crates/memtrack/src/ebpf/memtrack.rs | 92 +++++++- crates/memtrack/src/ebpf/mod.rs | 5 +- crates/memtrack/src/ebpf/poller.rs | 98 ++++----- 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 | 63 +++--- crates/memtrack/src/session.rs | 41 ++++ 13 files changed, 852 insertions(+), 179 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 c53e2087..86eacac9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2389,6 +2389,7 @@ dependencies = [ "libc", "log", "object", + "parking_lot", "paste", "rayon", "rstest", diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index 7b76ad1a..8a2558bf 100644 --- a/crates/memtrack/Cargo.toml +++ b/crates/memtrack/Cargo.toml @@ -36,6 +36,7 @@ libbpf-rs = { version = "0.26", features = ["vendored"], optional = true } glob = "0.3.3" object = { workspace = true } rayon = "1.12" +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..d82402f2 --- /dev/null +++ b/crates/memtrack/src/ebpf/attach_worker.rs @@ -0,0 +1,251 @@ +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 let Err(e) = self.poller.drain() { + self.record_fatal(e); + 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 649d1797..9ae34dce 100644 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ b/crates/memtrack/src/ebpf/memtrack.rs @@ -575,16 +575,90 @@ impl MemtrackBpf { Ok(()) } - /// Start polling, forwarding each event over the returned channel. - 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, - crossbeam_channel::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 cb4716be..f05153d7 100644 --- a/crates/memtrack/src/ebpf/mod.rs +++ b/crates/memtrack/src/ebpf/mod.rs @@ -1,6 +1,9 @@ +mod attach_worker; mod events; mod memtrack; -mod poller; +pub(crate) mod poller; +mod proc_fs; +mod spawn; mod tracker; pub use memtrack::{MemtrackBpf, ResolvedSymbols, resolve_symbol_offsets}; diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index 6d169afe..d80a1554 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/crates/memtrack/src/ebpf/poller.rs @@ -1,77 +1,79 @@ -use anyhow::Result; -use crossbeam_channel::{Receiver, unbounded}; +use anyhow::{Context, Result}; use libbpf_rs::{MapCore, RingBufferBuilder}; -use runner_shared::artifacts::MemtrackEvent as Event; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc::{self, RecvTimeoutError, Sender}; use std::thread::JoinHandle; use std::time::Duration; -use super::events::parse_event; - -/// Polls a BPF ring buffer on a background thread and forwards events over a channel. +/// 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 { - /// Poll `rb_map` and forward each parsed event over the returned channel. - pub fn with_channel( - rb_map: &M, - poll_timeout_ms: u64, - ) -> Result<(Self, Receiver)> { - let (tx, rx) = unbounded::(); - + 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) { - let _ = tx.send(event); + if let Some(item) = parse(data) { + let _ = tx.send(item); } 0 })?; let ringbuf = builder.build()?; - let shutdown = Arc::new(AtomicBool::new(false)); - let poll_thread = std::thread::spawn({ - let shutdown = shutdown.clone(); - let timeout = Duration::from_millis(poll_timeout_ms); - move || { - while !shutdown.load(Ordering::Relaxed) { - let _ = ringbuf.poll(timeout); - - // consume() drains the buffer to empty in a single call, so - // records produced while poll() was draining are picked up - // here without paying another epoll_wait round-trip. - let _ = ringbuf.consume(); + // 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 || { + 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; + } } - - // Events may still be sitting in the ring buffer after the last - // poll; consume once more so they reach the channel before it closes. - let _ = ringbuf.consume(); } }); - Ok(( - Self { - shutdown, - poll_thread: Some(poll_thread), - }, - rx, - )) + Ok(Self { + ctl: Some(ctl), + poll_thread: Some(poll_thread), + }) } - /// Stop the polling thread and wait for it to finish draining. - 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 eb53e6e7..11ecaf04 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -1,69 +1,102 @@ -use crate::ebpf::poller::RingBufferPoller; +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 crossbeam_channel::Receiver; -use runner_shared::artifacts::MemtrackEvent as Event; +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, - poller: Option, + bpf: Arc>, + worker: Mutex>, } impl Tracker { - /// Create a new tracker instance - /// - /// This will: - /// - Initialize the BPF subsystem - /// - Bump memlock limits - /// - Attach uprobes to all libc instances - /// - Attach tracepoints for fork tracking + /// Create a new tracker. The exec-mapping watcher discovers and attaches + /// allocator probes as the tracked process tree maps executable files. pub fn new() -> Result { - let 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()?; + + let bpf = Arc::new(Mutex::new(bpf)); + let worker = AttachWorker::start(bpf.clone())?; - Ok(Self { bpf, poller: None }) + Ok(Self { + bpf, + worker: Mutex::new(Some(worker)), + }) } - pub fn attach_allocators(&mut self, libs: &[AllocatorLib]) -> Result<()> { - self.bpf.attach_allocators(libs) + /// 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); + } + + 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 of allocation events. The poller is owned by the tracker - /// and keeps running until [`Tracker::stop_polling`] is called or the tracker - /// is dropped. - pub fn track(&mut self, pid: i32) -> Result> { - self.bpf.add_tracked_pid(pid)?; - debug!("Tracking PID {pid}"); + /// Disable event tracking in the BPF program + pub fn disable_tracking(&self) -> Result<()> { + self.bpf.lock().disable_tracking() + } - let (poller, event_rx) = self.bpf.start_polling_with_channel(10)?; - self.poller = Some(poller); + /// 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() + } - Ok(event_rx) + /// 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() } - /// Stop the poll thread, draining ring-buffer stragglers. This closes the - /// event channel returned by [`track`]. - pub fn stop_polling(&mut self) { - self.poller.take(); + /// Detach all attached probes. Called explicitly at teardown because the + /// process may exit without ever dropping the tracker (the IPC thread holds + /// an Arc clone), in which case the kernel would close each link fd serially. + pub fn detach(&self) { + self.bpf.lock().detach_probes(); } /// Bump RLIMIT_MEMLOCK for kernels older than 5.11. Newer kernels account BPF @@ -84,27 +117,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() - } - - /// Detach all attached probes. Called explicitly at teardown because the - /// process may exit without ever dropping the tracker, in which case the - /// kernel would close each link fd serially at exit. - pub fn detach(&mut self) { - self.bpf.detach_probes(); - } - - /// 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 24fc75af..8ad48791 100644 --- a/crates/memtrack/src/main.rs +++ b/crates/memtrack/src/main.rs @@ -3,10 +3,9 @@ use ipc_channel::ipc; use memtrack::prelude::*; use memtrack::{MemtrackIpcMessage, Tracker, handle_ipc_message}; use runner_shared::artifacts::{ArtifactExt, MemtrackArtifact, encode_events}; -use std::os::unix::process::CommandExt; use std::path::{Path, PathBuf}; use std::process::Command; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use std::thread; #[derive(Parser)] @@ -85,41 +84,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 @@ -136,26 +132,20 @@ fn track_command( let pipeline_thread = thread::spawn(move || encode_events(event_rx, out_file, n_workers)); // 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}"); // Stop event production before draining: the child has exited, so anything // still arriving is already in the ring buffer. - let disable_result = tracker_arc - .lock() - .map_err(|_| anyhow!("tracker mutex poisoned"))? - .disable(); - if let Err(e) = disable_result { + if let Err(e) = tracker.disable_tracking() { warn!("Failed to disable tracking: {e:#}"); } - // Stopping the poller closes the event channel; without this the pipeline - // join below would block forever. + // Dropping the session drops the event poller, which does a final drain of + // the ring buffer and then closes the event channel. Without this the + // encode pipeline join below would block forever. debug!("Stopping the ring buffer poller"); - tracker_arc - .lock() - .map_err(|_| anyhow!("tracker mutex poisoned"))? - .stop_polling(); + drop(session); debug!("Waiting for the encode pipeline to finish"); let total = pipeline_thread @@ -164,20 +154,19 @@ fn track_command( info!("Wrote {total} memtrack events to disk"); + // Stop the attach worker and surface any fatal error it recorded (missed + // exec mappings mean incomplete allocator coverage). + tracker.finish()?; + // Detach probes explicitly: the IPC thread still holds an Arc clone, so the // tracker would otherwise never be dropped before process::exit and the // kernel would close every link fd serially during exit. - tracker_arc - .lock() - .map_err(|_| anyhow!("tracker mutex poisoned"))? - .detach(); + tracker.detach(); // 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 6c7ff4789f897029a3d7be1891c889ed0df5ca6e Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 8 Jul 2026 21:42:38 +0200 Subject: [PATCH 3/6] 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 86eacac9..0dfb2a92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2380,7 +2380,6 @@ dependencies = [ "clap", "crossbeam-channel", "env_logger", - "glob", "insta", "ipc-channel", "itertools 0.14.0", diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index 8a2558bf..c4311d9e 100644 --- a/crates/memtrack/Cargo.toml +++ b/crates/memtrack/Cargo.toml @@ -33,7 +33,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 } rayon = "1.12" 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 cefe5c75..89952255 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; @@ -121,44 +121,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 d29dc04232312a1665d9a0cf9734b7b076be18e3 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 8 Jul 2026 21:54:47 +0200 Subject: [PATCH 4/6] 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 4d3fbb05ecdb322a62d2b87ae018152d0a9699c1 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Thu, 9 Jul 2026 18:48:45 +0200 Subject: [PATCH 5/6] 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 49543b75..d6935beb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -95,7 +95,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 663ead6e116046910d2f6f166db2bbb9a0eed651 Mon Sep 17 00:00:00 2001 From: Guillaume Lagrange Date: Thu, 16 Jul 2026 17:22:14 +0000 Subject: [PATCH 6/6] feat(memtrack): record per-benchmark pid in execution timestamps Memory mode tracks the whole runner process tree, so between Start/StopProfiler the harness parent and sibling worker processes emit allocation events into every measured region. The CurrentBenchmark FIFO command already carries the pid that ran each benchmark; capture it alongside the uri in ExecutionTimestamps.bench_pid_by_ts (serde-default for artifacts written before this field) so the parser can scope events to the benchmarking process. --- crates/memtrack/src/ebpf/c/event.h | 8 ++ crates/memtrack/src/ebpf/c/memtrack.bpf.c | 73 ++++++++++++------- crates/memtrack/src/ebpf/events.rs | 34 +++++++++ .../src/artifacts/execution_timestamps.rs | 10 ++- .../src/artifacts/memtrack/mod.rs | 7 ++ src/executor/shared/fifo.rs | 12 ++- 6 files changed, 115 insertions(+), 29 deletions(-) diff --git a/crates/memtrack/src/ebpf/c/event.h b/crates/memtrack/src/ebpf/c/event.h index ff41f0be..a89af76e 100644 --- a/crates/memtrack/src/ebpf/c/event.h +++ b/crates/memtrack/src/ebpf/c/event.h @@ -9,6 +9,7 @@ #define EVENT_TYPE_MMAP 6 #define EVENT_TYPE_MUNMAP 7 #define EVENT_TYPE_BRK 8 +#define EVENT_TYPE_FORK 9 /* Common header shared by all event types */ struct event_header { @@ -45,6 +46,13 @@ struct event { uint64_t addr; /* address of mapping */ uint64_t size; /* size of mapping */ } mmap; + + /* Process fork event - header.pid is the child, ppid the parent. + * Lets the parser rebuild the process tree and scope allocations to a + * benchmark process and its descendants. */ + struct { + uint32_t ppid; /* parent pid that spawned header.pid */ + } fork; } data; }; diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 4c93d241..2186b017 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/crates/memtrack/src/ebpf/c/memtrack.bpf.c @@ -46,6 +46,31 @@ BPF_ARRAY_MAP(tracking_enabled, __u8, 1); /* Counter for events that couldn't be added to the ring buffer */ BPF_ARRAY_MAP(dropped_events, __u64, 1); +/* Wake the consumer only once this much unconsumed data has accumulated. + * Per-event wakeups dominate submission cost at high event rates; batching + * them behind a data watermark amortizes the wakeup to ~1 per thousand + * events. The userspace poller's poll timeout flushes the tail that never + * reaches the watermark. */ +#define WAKEUP_DATA_SIZE (64 * 1024) + +static __always_inline long wake_flags(void) { + long avail = bpf_ringbuf_query(&events, BPF_RB_AVAIL_DATA); + return avail >= WAKEUP_DATA_SIZE ? BPF_RB_FORCE_WAKEUP : BPF_RB_NO_WAKEUP; +} + +/* Helper to check if tracking is currently enabled */ +static __always_inline int is_enabled(void) { + __u32 key = 0; + __u8* enabled = bpf_map_lookup_elem(&tracking_enabled, &key); + + /* ARRAY-map lookups can't fail for a valid index; fail closed if one ever does. */ + if (!enabled) { + return 0; + } + + return *enabled; +} + /* == Code that tracks process forks and execs == */ /* Helper to check if a PID or any of its ancestors should be tracked */ @@ -87,6 +112,29 @@ int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) { bpf_map_update_elem(&pids_ppid, &child_pid, &parent_pid, BPF_ANY); // bpf_printk("auto-tracking child process: child_pid=%u", child_pid); + + /* Emit a fork event so userspace can rebuild the process tree and scope + * allocations to a process plus its descendants. Gated on is_enabled() + * like allocation events, so only forks inside a measured region reach + * the ring buffer. header.pid carries the child: current here is the + * parent, so we cannot use the SUBMIT_EVENT macro. */ + if (is_enabled()) { + struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0); + if (!e) { + __u32 zero = 0; + __u64* drops = bpf_map_lookup_elem(&dropped_events, &zero); + if (drops) { + __sync_fetch_and_add(drops, 1); + } + return 0; + } + e->header.timestamp = bpf_ktime_get_ns(); + e->header.pid = child_pid; + e->header.tid = child_pid; + e->header.event_type = EVENT_TYPE_FORK; + e->data.fork.ppid = parent_pid; + bpf_ringbuf_submit(e, wake_flags()); + } } return 0; @@ -99,31 +147,6 @@ int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) { /* == Helper functions for the allocation tracking == */ -/* Wake the consumer only once this much unconsumed data has accumulated. - * Per-event wakeups dominate submission cost at high event rates; batching - * them behind a data watermark amortizes the wakeup to ~1 per thousand - * events. The userspace poller's poll timeout flushes the tail that never - * reaches the watermark. */ -#define WAKEUP_DATA_SIZE (64 * 1024) - -static __always_inline long wake_flags(void) { - long avail = bpf_ringbuf_query(&events, BPF_RB_AVAIL_DATA); - return avail >= WAKEUP_DATA_SIZE ? BPF_RB_FORCE_WAKEUP : BPF_RB_NO_WAKEUP; -} - -/* Helper to check if tracking is currently enabled */ -static __always_inline int is_enabled(void) { - __u32 key = 0; - __u8* enabled = bpf_map_lookup_elem(&tracking_enabled, &key); - - /* ARRAY-map lookups can't fail for a valid index; fail closed if one ever does. */ - if (!enabled) { - return 0; - } - - return *enabled; -} - /* Helper to store parameter value in map for tracking between entry and return */ static __always_inline int store_param(void* map, __u64 value) { diff --git a/crates/memtrack/src/ebpf/events.rs b/crates/memtrack/src/ebpf/events.rs index f18216c6..a751b738 100644 --- a/crates/memtrack/src/ebpf/events.rs +++ b/crates/memtrack/src/ebpf/events.rs @@ -74,6 +74,12 @@ pub fn parse_event(data: &[u8]) -> Option { size: event.data.mmap.size, }, ), + EVENT_TYPE_FORK => ( + 0, + MemtrackEventKind::Fork { + ppid: event.data.fork.ppid as i32, + }, + ), unknown => { panic!("Unknown event type: {unknown}"); } @@ -185,4 +191,32 @@ mod tests { _ => panic!("Expected Malloc event kind"), } } + + #[test] + fn test_parse_fork_event() { + let mut event: bindings::event = unsafe { std::mem::zeroed() }; + event.header.event_type = bindings::EVENT_TYPE_FORK as u8; + event.header.timestamp = 12345678; + event.header.pid = 2000; // child + event.header.tid = 2000; + event.data.fork.ppid = 1000; // parent + + let bytes = unsafe { + std::slice::from_raw_parts( + &event as *const _ as *const u8, + std::mem::size_of_val(&event), + ) + }; + + let parsed = parse_event(bytes).unwrap(); + assert_eq!(parsed.pid, 2000); + assert_eq!(parsed.addr, 0); + + match parsed.kind { + MemtrackEventKind::Fork { ppid } => { + assert_eq!(ppid, 1000); + } + _ => panic!("Expected Fork event kind"), + } + } } diff --git a/crates/runner-shared/src/artifacts/execution_timestamps.rs b/crates/runner-shared/src/artifacts/execution_timestamps.rs index 80360a01..3cc29067 100644 --- a/crates/runner-shared/src/artifacts/execution_timestamps.rs +++ b/crates/runner-shared/src/artifacts/execution_timestamps.rs @@ -1,3 +1,4 @@ +use libc::pid_t; use serde::{Deserialize, Serialize}; use crate::fifo::MarkerType; @@ -6,14 +7,21 @@ use crate::fifo::MarkerType; pub struct ExecutionTimestamps { pub uri_by_ts: Vec<(u64, String)>, pub markers: Vec, + #[serde(default)] + pub bench_pid_by_ts: Vec<(u64, pid_t)>, } impl super::ArtifactExt for ExecutionTimestamps {} impl ExecutionTimestamps { - pub fn new(uri_by_ts: &[(u64, String)], markers: &[crate::fifo::MarkerType]) -> Self { + pub fn new( + uri_by_ts: &[(u64, String)], + markers: &[crate::fifo::MarkerType], + bench_pid_by_ts: &[(u64, pid_t)], + ) -> Self { Self { uri_by_ts: uri_by_ts.to_vec(), markers: markers.to_vec(), + bench_pid_by_ts: bench_pid_by_ts.to_vec(), } } } diff --git a/crates/runner-shared/src/artifacts/memtrack/mod.rs b/crates/runner-shared/src/artifacts/memtrack/mod.rs index ad843b28..898bf17d 100644 --- a/crates/runner-shared/src/artifacts/memtrack/mod.rs +++ b/crates/runner-shared/src/artifacts/memtrack/mod.rs @@ -78,6 +78,12 @@ pub enum MemtrackEventKind { Brk { size: u64, }, + /// A tracked process spawned a child. `pid` is the child, `ppid` the parent. + /// Carries no allocation; lets the parser rebuild the process tree and scope + /// allocations to a benchmark process plus its descendants. + Fork { + ppid: pid_t, + }, } pub struct MemtrackEventStream { @@ -159,6 +165,7 @@ mod tests { MemtrackEventKind::Mmap { size: 9 }, MemtrackEventKind::Munmap { size: 9 }, MemtrackEventKind::Brk { size: 9 }, + MemtrackEventKind::Fork { ppid: 1234 }, ]; for kind in kinds { diff --git a/src/executor/shared/fifo.rs b/src/executor/shared/fifo.rs index e63e4c65..d17556d9 100644 --- a/src/executor/shared/fifo.rs +++ b/src/executor/shared/fifo.rs @@ -170,6 +170,7 @@ impl RunnerFifo { std::process::ExitStatus, )> { let mut bench_order_by_timestamp = Vec::<(u64, String)>::new(); + let mut bench_pid_by_ts = Vec::<(u64, pid_t)>::new(); let mut bench_pids = HashSet::::new(); let mut markers = Vec::::new(); @@ -206,7 +207,9 @@ impl RunnerFifo { // Fall through to shared implementation for standard commands match &cmd { FifoCommand::CurrentBenchmark { pid, uri } => { - bench_order_by_timestamp.push((get_current_time(), uri.to_string())); + let ts = get_current_time(); + bench_order_by_timestamp.push((ts, uri.to_string())); + bench_pid_by_ts.push((ts, *pid)); bench_pids.insert(*pid); self.send_cmd(FifoCommand::Ack).await?; } @@ -273,8 +276,11 @@ impl RunnerFifo { debug!( "Process terminated with status: {exit_status}, stopping the command handler" ); - let marker_result = - ExecutionTimestamps::new(&bench_order_by_timestamp, &markers); + let marker_result = ExecutionTimestamps::new( + &bench_order_by_timestamp, + &markers, + &bench_pid_by_ts, + ); let fifo_data = FifoBenchmarkData { integration, bench_pids,