diff --git a/crates/memtrack/src/ebpf/attach_worker.rs b/crates/memtrack/src/ebpf/attach_worker.rs index 892fe6211..01b0e6db2 100644 --- a/crates/memtrack/src/ebpf/attach_worker.rs +++ b/crates/memtrack/src/ebpf/attach_worker.rs @@ -1,7 +1,8 @@ use crate::AllocatorLib; use crate::ebpf::MemtrackBpf; use crate::ebpf::events::AttachRequest; -use crate::ebpf::poller::RingBufferPoller; +use crate::ebpf::pause::StoppedProcesses; +use crate::ebpf::poller::{POLL_INTERVAL_MS, RingBufferPoller}; use crate::prelude::*; use parking_lot::Mutex; use std::collections::HashSet; @@ -14,19 +15,24 @@ use std::time::Duration; use super::proc_fs::{Resolution, 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); +/// Releases the attach worker's hold on `pid` on drop, including panics. +struct AttachHold { + stops: Arc, + pid: u32, +} -impl Drop for ContGuard { +impl AttachHold { + fn new(stops: Arc, pid: u32) -> Self { + Self { stops, pid } + } +} + +impl Drop for AttachHold { 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); + if let Err(error) = self.stops.release_attach(self.pid) { + warn!("failed to release attach stop of {}: {error:#}", self.pid); } } } @@ -49,12 +55,19 @@ impl AttachWorker { 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 (poller, stops) = { + let bpf = bpf.lock(); + ( + bpf.poll_attach_with_channel(POLL_INTERVAL_MS, tx)?, + bpf.stopped_processes(), + ) + }; let worker = Worker { poller, rx, bpf: bpf.clone(), + stops, shutdown: shutdown.clone(), fatal: fatal.clone(), root_pid: root_pid.clone(), @@ -123,6 +136,7 @@ struct Worker { poller: RingBufferPoller, rx: mpsc::Receiver>, bpf: Arc>, + stops: Arc, shutdown: Arc, fatal: Arc>>, root_pid: Arc, @@ -168,14 +182,14 @@ impl Worker { } /// 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. + /// then classify + attach for each unique `(dev, ino)`. `holds` release + /// 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 holds: Vec = Vec::new(); let mut stopped: HashSet = HashSet::new(); loop { @@ -191,8 +205,9 @@ impl Worker { for pid in new_pids { stopped.insert(pid); - guards.push(ContGuard(pid as i32)); + holds.push(AttachHold::new(self.stops.clone(), pid)); wait_all_stopped(pid, STOP_DEADLINE)?; + debug!("Stopped pid {pid} for attach"); } // Every producer is stopped, so a synchronous drain is complete. diff --git a/crates/memtrack/src/ebpf/c/attach.h b/crates/memtrack/src/ebpf/c/attach.h index e188c7d5f..2cffb604f 100644 --- a/crates/memtrack/src/ebpf/c/attach.h +++ b/crates/memtrack/src/ebpf/c/attach.h @@ -4,6 +4,7 @@ #include "event.h" #include "utils/map_helpers.h" #include "utils/process_tracking.h" +#include "utils/stopped.h" /* == Exec-mapping watcher == * @@ -12,7 +13,6 @@ * classifies the file, attaches allocator probes, then resumes it. */ #define MEMTRACK_PROT_EXEC 0x4 -#define MEMTRACK_SIGSTOP 19 struct inode_key { __u64 dev; @@ -53,12 +53,13 @@ int BPF_PROG(watch_exec_mmap, struct file* file, unsigned long prot, unsigned lo } return 0; } + /* The worker may release the pid as soon as the request is visible. */ + memtrack_stop_current(&attach_stopped, tgid); + 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/process_tracking.bpf.h b/crates/memtrack/src/ebpf/c/process_tracking.bpf.h index b82b53d6a..a39befbf5 100644 --- a/crates/memtrack/src/ebpf/c/process_tracking.bpf.h +++ b/crates/memtrack/src/ebpf/c/process_tracking.bpf.h @@ -5,6 +5,7 @@ #include "utils/event_helpers.h" #include "utils/mm_ownership.h" #include "utils/process_tracking.h" +#include "utils/stopped.h" /* FORK lets userland seed a child's RSS from its parent at fork time: the * kernel copies the mm counters during dup_mmap, but those updates fire @@ -83,9 +84,11 @@ int BPF_PROG(tracepoint_sched_process_exit) { return 0; } - /* Drop the ownership mapping so foreign actors stop attributing to a pid - * the kernel may reuse. */ + /* Drop the ownership mapping and stop records so neither foreign actors + * nor a later release act on a pid the kernel may reuse. */ rebind_pid_mm(pid, 0); + bpf_map_delete_elem(&pressure_stopped, &pid); + bpf_map_delete_elem(&attach_stopped, &pid); SUBMIT_EVENT_AS(pid, EVENT_TYPE_EXIT, {}); } diff --git a/crates/memtrack/src/ebpf/c/stack_capture.bpf.h b/crates/memtrack/src/ebpf/c/stack_capture.bpf.h index 2519767b5..92097891b 100644 --- a/crates/memtrack/src/ebpf/c/stack_capture.bpf.h +++ b/crates/memtrack/src/ebpf/c/stack_capture.bpf.h @@ -3,6 +3,7 @@ #include "event.h" #include "utils/map_helpers.h" +#include "utils/pressure.bpf.h" #include "utils/process_tracking.h" /* Emit raw stack bytes and registers once per hash for offline DWARF unwinding. @@ -18,6 +19,10 @@ const volatile __u32 stack_copy_budget = 4096; #define FNV64_OFFSET 0xcbf29ce484222325ULL #define FNV64_PRIME 0x00000100000001b3ULL +/* Map helpers reject two arguments pointing into the same ring reservation, + * so the dedup value lives in .rodata while the key stays in the record. */ +static const __u8 seen_stack_marker = 1; + struct { __uint(type, BPF_MAP_TYPE_STACK_TRACE); __uint(max_entries, 16384); @@ -88,17 +93,18 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas void* slot = bpf_ringbuf_reserve(&stacks, sizeof(struct stack_header) + stack_copy_budget, 0); if (!slot) { bump_stack_counter(MEMTRACK_STACK_COUNTER_RING_FULL); + memtrack_check_ring_pressure(&stacks, ids.tgid); return 0; } - __u64 sp = PT_REGS_SP(ctx); - __u8* payload = (__u8*)slot + sizeof(struct stack_header); - __u64 lanes[4] = { - FNV64_OFFSET ^ 0, - FNV64_OFFSET ^ 1, - FNV64_OFFSET ^ 2, - FNV64_OFFSET ^ 3, - }; + /* Keep hashing scratch in the unpublished record. Large kprobe-family BPF + * stacks may use per-CPU storage, which nested uprobes can overwrite. */ + struct stack_header* header = (struct stack_header*)slot; + __u64* lanes = &header->hash; + lanes[0] = FNV64_OFFSET ^ 0; + lanes[1] = FNV64_OFFSET ^ 1; + lanes[2] = FNV64_OFFSET ^ 2; + lanes[3] = FNV64_OFFSET ^ 3; __u32 got = 0; /* Chunked reads stop at the first unreadable stack region. @@ -106,17 +112,19 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas * so every slot access is provably in range. */ #pragma clang loop unroll(disable) for (__u32 off = 0; off + STACK_COPY_CHUNK <= stack_copy_budget; off += STACK_COPY_CHUNK) { - if (bpf_probe_read_user(payload + off, STACK_COPY_CHUNK, (void*)(sp + off)) != 0) { + if (bpf_probe_read_user((__u8*)slot + sizeof(struct stack_header) + off, STACK_COPY_CHUNK, + (void*)(PT_REGS_SP(ctx) + off)) != 0) { break; } - fnv64_hash_chunk(lanes, (const __u64*)(payload + off)); + fnv64_hash_chunk(lanes, (const __u64*)((__u8*)slot + sizeof(struct stack_header) + off)); got = off + STACK_COPY_CHUNK; } if (got == 0) { bpf_ringbuf_discard(slot, 0); bump_stack_counter(MEMTRACK_STACK_COUNTER_COPY_FAILED); + memtrack_check_ring_pressure(&stacks, ids.tgid); return 0; } @@ -135,10 +143,12 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas hash = FNV64_OFFSET; } - __u8 marker = 1; - long gate_result = bpf_map_update_elem(&seen_stack_hashes, &hash, &marker, BPF_NOEXIST); + header->hash = hash; + long gate_result = + bpf_map_update_elem(&seen_stack_hashes, &header->hash, &seen_stack_marker, BPF_NOEXIST); if (gate_result == -17) { /* -EEXIST */ bpf_ringbuf_discard(slot, 0); + memtrack_check_ring_pressure(&stacks, ids.tgid); return hash; } if (gate_result != 0) { @@ -151,11 +161,10 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas bump_stack_counter(MEMTRACK_STACK_COUNTER_STACKID_FAILED); } - struct stack_header* header = (struct stack_header*)slot; header->hash = hash; header->timestamp = bpf_ktime_get_ns(); header->stackid = stackid; - header->sp = sp; + header->sp = PT_REGS_SP(ctx); header->pid = ids.tgid; header->tid = ids.tid; header->copy_len = got; @@ -166,6 +175,7 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas fill_stack_regs(&header->regs, ctx); bpf_ringbuf_submit(slot, 0); + memtrack_check_ring_pressure(&stacks, ids.tgid); return hash; } diff --git a/crates/memtrack/src/ebpf/c/utils/pressure.bpf.h b/crates/memtrack/src/ebpf/c/utils/pressure.bpf.h new file mode 100644 index 000000000..3a70a6c58 --- /dev/null +++ b/crates/memtrack/src/ebpf/c/utils/pressure.bpf.h @@ -0,0 +1,37 @@ +#ifndef __PRESSURE_BPF_H__ +#define __PRESSURE_BPF_H__ + +#include + +#include "map_helpers.h" +#include "process_tracking.h" +#include "stopped.h" + +/* Ring pressure stop. Call only after submit/discard or a failed reserve: + * stopping with a live reservation would wedge the ring. A tracked producer + * that writes while the ring is over the watermark is stopped and recorded, + * so processes that do not write keep running. Userspace resumes the + * recorded producers once it has flushed the ring. */ + +#define MEMTRACK_PRESSURE_HEADROOM_FRAC 4 /* stop at (FRAC-1)/FRAC = 75% used */ + +static __always_inline int memtrack_ring_over_watermark(void* ring) { + __u64 size = bpf_ringbuf_query(ring, BPF_RB_RING_SIZE); + __u64 avail = bpf_ringbuf_query(ring, BPF_RB_AVAIL_DATA); + return avail >= size - size / MEMTRACK_PRESSURE_HEADROOM_FRAC; +} + +static __always_inline void memtrack_check_ring_pressure(void* ring, __u32 current_tgid) { + if (!memtrack_ring_over_watermark(ring)) { + return; + } + + /* Never stop an untracked process that happens to trigger a probe. */ + if (!is_tracked(current_tgid)) { + return; + } + + memtrack_stop_current(&pressure_stopped, current_tgid); +} + +#endif /* __PRESSURE_BPF_H__ */ diff --git a/crates/memtrack/src/ebpf/c/utils/stopped.h b/crates/memtrack/src/ebpf/c/utils/stopped.h new file mode 100644 index 000000000..2b627a852 --- /dev/null +++ b/crates/memtrack/src/ebpf/c/utils/stopped.h @@ -0,0 +1,35 @@ +#ifndef __STOPPED_H__ +#define __STOPPED_H__ + +#include + +#include "map_helpers.h" + +#define MEMTRACK_SIGCONT 18 +#define MEMTRACK_SIGSTOP 19 + +/* tgid -> 1 for every process BPF stopped, one map per reason. A process is + * recorded before its stop can take effect, and userspace resumes only + * recorded processes, so a process stopped for both reasons resumes once + * neither map holds it. Sized like tracked_pids. */ +BPF_HASH_MAP(pressure_stopped, __u32, __u8, 10000); +BPF_HASH_MAP(attach_stopped, __u32, __u8, 10000); + +/* Stop the current process and record it in `map`. SIGSTOP is queued before + * the record is written: it only takes effect on return to user mode, so any + * SIGCONT sent after userspace sees the record cancels or ends the stop. + * Requires task context with IRQs enabled, where the signal is queued + * synchronously rather than via irq_work. */ +static __always_inline void memtrack_stop_current(void* map, __u32 tgid) { + if (bpf_send_signal(MEMTRACK_SIGSTOP) != 0) { + return; + } + + __u8 marker = 1; + if (bpf_map_update_elem(map, &tgid, &marker, BPF_ANY) != 0) { + /* Unrecorded, so nothing would resume it. */ + bpf_send_signal(MEMTRACK_SIGCONT); + } +} + +#endif /* __STOPPED_H__ */ diff --git a/crates/memtrack/src/ebpf/memtrack/maps.rs b/crates/memtrack/src/ebpf/memtrack/maps.rs index 7ca16d25a..6424eee25 100644 --- a/crates/memtrack/src/ebpf/memtrack/maps.rs +++ b/crates/memtrack/src/ebpf/memtrack/maps.rs @@ -1,7 +1,9 @@ -use super::MemtrackBpf; +use super::{MemtrackBpf, Skel}; +use crate::ebpf::pause::StoppedProcesses; use crate::ebpf::stacks::StackCaptureFailureStats; use crate::prelude::*; -use libbpf_rs::MapCore; +use libbpf_rs::{MapCore, MapHandle}; +use std::sync::Arc; impl MemtrackBpf { pub fn add_tracked_pid(&mut self, pid: i32) -> Result<()> { @@ -65,6 +67,39 @@ impl MemtrackBpf { ) } + /// Owned map handles, so callers can release stopped processes without + /// holding the `MemtrackBpf` lock. + pub(crate) fn stopped_processes(&self) -> Arc { + self.stopped.clone() + } + + pub(super) fn open_stopped_processes(skel: &Skel) -> Result { + let (pressure_stopped, attach_stopped) = match skel { + Skel::Token(skel) => ( + MapHandle::try_from(&skel.maps.pressure_stopped), + MapHandle::try_from(&skel.maps.attach_stopped), + ), + Skel::Legacy(skel) => ( + MapHandle::try_from(&skel.maps.pressure_stopped), + MapHandle::try_from(&skel.maps.attach_stopped), + ), + }; + Ok(StoppedProcesses::new( + pressure_stopped.context("Failed to create handle for pressure_stopped map")?, + attach_stopped.context("Failed to create handle for attach_stopped map")?, + )) + } + + /// Callback that resumes every pressure-stopped process. + pub(super) fn on_ring_drained(&self) -> Box { + let stopped = self.stopped.clone(); + Box::new(move || { + if let Err(error) = stopped.release_pressure() { + error!("failed to release pressure-stopped producers: {error:#}"); + } + }) + } + pub fn stack_capture_stats(&self) -> Result { StackCaptureFailureStats::read(with_skel!(self, skel => &skel.maps.stack_counters)) } diff --git a/crates/memtrack/src/ebpf/memtrack/mod.rs b/crates/memtrack/src/ebpf/memtrack/mod.rs index 0c9411d43..dbe4680ad 100644 --- a/crates/memtrack/src/ebpf/memtrack/mod.rs +++ b/crates/memtrack/src/ebpf/memtrack/mod.rs @@ -126,6 +126,7 @@ pub struct MemtrackBpf { /// the same file offset as their canonical function; attaching each /// alias would double-instrument the one underlying function. attached_offsets: std::collections::HashSet<(std::path::PathBuf, usize)>, + stopped: std::sync::Arc, } impl MemtrackBpf { @@ -224,8 +225,10 @@ impl MemtrackBpf { } }; + let stopped = std::sync::Arc::new(Self::open_stopped_processes(&skel)?); Ok(Self { skel, + stopped, probes: Vec::new(), rmap, physical, @@ -235,7 +238,7 @@ impl MemtrackBpf { /// 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( + pub(crate) fn poll_events_with_channel( &self, poll_interval_ms: u64, tx: std::sync::mpsc::Sender>, @@ -245,6 +248,7 @@ impl MemtrackBpf { crate::ebpf::events::parse_event, tx, poll_interval_ms, + None, )) } @@ -278,6 +282,7 @@ impl MemtrackBpf { resolve, tx, poll_interval_ms, + Some(self.on_ring_drained()), )) } @@ -293,6 +298,7 @@ impl MemtrackBpf { crate::ebpf::events::AttachRequest::parse, tx, poll_interval_ms, + None, )) } diff --git a/crates/memtrack/src/ebpf/mod.rs b/crates/memtrack/src/ebpf/mod.rs index 2954f5d38..3aba823ac 100644 --- a/crates/memtrack/src/ebpf/mod.rs +++ b/crates/memtrack/src/ebpf/mod.rs @@ -1,6 +1,7 @@ mod attach_worker; mod events; mod memtrack; +pub(crate) mod pause; pub(crate) mod poller; mod proc_fs; mod spawn; diff --git a/crates/memtrack/src/ebpf/pause.rs b/crates/memtrack/src/ebpf/pause.rs new file mode 100644 index 000000000..a3c7a17f7 --- /dev/null +++ b/crates/memtrack/src/ebpf/pause.rs @@ -0,0 +1,80 @@ +use crate::prelude::*; +use libbpf_rs::{MapCore, MapFlags, MapHandle}; +use std::sync::atomic::{AtomicU64, Ordering::Relaxed}; + +/// Processes stopped by BPF, either for ring pressure or for an attach request. +/// +/// A single SIGCONT resumes a process no matter how many times it was stopped, +/// so a process stopped for both reasons resumes only once both are released. +/// Each side deletes its own entry before checking the other's, so two +/// concurrent releases cannot both skip the resume. +pub(crate) struct StoppedProcesses { + pressure_stopped: MapHandle, + attach_stopped: MapHandle, + // Stop counts, only used for stats. + pressure_stops: AtomicU64, + attach_stops: AtomicU64, +} + +impl StoppedProcesses { + pub(crate) fn new(pressure_stopped: MapHandle, attach_stopped: MapHandle) -> Self { + Self { + pressure_stopped, + attach_stopped, + pressure_stops: AtomicU64::new(0), + attach_stops: AtomicU64::new(0), + } + } + + /// Attach worker is done with pid. + pub(crate) fn release_attach(&self, pid: u32) -> Result<()> { + debug!("Releasing attach stop of pid {pid}"); + self.attach_stops.fetch_add(1, Relaxed); + Self::release(pid, &self.attach_stopped, &self.pressure_stopped) + } + + /// Resume every pressure-stopped producer; call once a ring is flushed. + pub(crate) fn release_pressure(&self) -> Result<()> { + // Deleting while iterating restarts hash iteration, so snapshot the keys first. + let keys: Vec> = self.pressure_stopped.keys().collect(); + if keys.is_empty() { + return Ok(()); + } + self.pressure_stops.fetch_add(keys.len() as u64, Relaxed); + for key in keys { + let pid = u32::from_le_bytes( + key.as_slice() + .try_into() + .context("Invalid pressure_stopped key size")?, + ); + debug!("Releasing pressure stop of pid {pid}"); + Self::release(pid, &self.pressure_stopped, &self.attach_stopped)?; + } + Ok(()) + } + + /// Deletes `pid` from `own`, then resumes it unless `other` still holds it. + /// A missing entry means the process already exited or was released. + fn release(pid: u32, own: &MapHandle, other: &MapHandle) -> Result<()> { + let key = pid.to_le_bytes(); + match own.delete(&key) { + Ok(()) => {} + Err(err) if err.kind() == libbpf_rs::ErrorKind::NotFound => return Ok(()), + Err(err) => return Err(err).context("Failed to delete stop entry"), + } + if other.lookup(&key, MapFlags::ANY)?.is_some() { + return Ok(()); + } + crate::ebpf::spawn::resume(pid as i32) + } +} + +impl Drop for StoppedProcesses { + fn drop(&mut self) { + debug!( + "Process stops: {} pressure, {} attach", + self.pressure_stops.get_mut(), + self.attach_stops.get_mut(), + ); + } +} diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index 9c122e068..637ed27e7 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/crates/memtrack/src/ebpf/poller.rs @@ -1,11 +1,14 @@ use anyhow::{Context, Result}; -use libbpf_rs::{MapCore, RingBufferBuilder}; +use libbpf_rs::{AsRawLibbpf, MapCore, RingBuffer, RingBufferBuilder, libbpf_sys}; use parking_lot::Mutex; use std::sync::Arc; use std::sync::mpsc::{self, RecvTimeoutError, Sender}; use std::thread::JoinHandle; use std::time::Duration; +/// Ring-buffer poll interval shared by every poller. +pub(crate) const POLL_INTERVAL_MS: u64 = 1; + /// Items buffered before a channel send. `std::sync::mpsc` allocates a block /// every 31 messages, so sending one item at a time makes that allocation /// dominate the pipeline; batching amortizes it over a whole batch. @@ -23,6 +26,21 @@ fn flush_batch(batch: &Mutex>, tx: &Sender>) { let _ = tx.send(items); } +/// `consume()` also stops at a record a producer is still writing, so retry +/// until everything reserved before the call has been consumed. Bounding by a +/// producer-position snapshot, not an empty ring, keeps producers that are +/// never stopped from starving the drain. +fn consume_all(ringbuf: &RingBuffer, ring: *mut libbpf_sys::ring) { + let target = unsafe { libbpf_sys::ring__producer_pos(ring) }; + loop { + let _ = ringbuf.consume(); + if unsafe { libbpf_sys::ring__consumer_pos(ring) } >= target { + return; + } + std::thread::sleep(Duration::from_millis(1)); + } +} + fn poll_iteration( control: std::result::Result, RecvTimeoutError>, consume: impl FnOnce(), @@ -67,6 +85,7 @@ impl RingBufferPoller { parse: F, tx: Sender>, poll_interval_ms: u64, + on_drained: Option>, ) -> Result where M: MapCore, @@ -101,17 +120,27 @@ impl RingBufferPoller { // poll tick, and disconnection is the shutdown signal. let (ctl, ctl_rx) = mpsc::channel::>(); let poll_thread = std::thread::spawn(move || { + // SAFETY: the built `RingBuffer` holds exactly the one ring added above. + let ring = + unsafe { libbpf_sys::ring_buffer__ring(ringbuf.as_libbpf_object().as_ptr(), 0) }; while poll_iteration( ctl_rx.recv_timeout(Duration::from_millis(poll_interval_ms)), - || { - let _ = ringbuf.consume(); - }, + || consume_all(&ringbuf, ring), || { let _ = ringbuf.poll(Duration::ZERO); }, &batch, &tx, - ) {} + ) { + if let Some(on_drained) = &on_drained + && unsafe { libbpf_sys::ring__avail_data_size(ring) } == 0 + { + on_drained(); + } + } + if let Some(on_drained) = &on_drained { + on_drained(); + } }); Ok(Self { @@ -163,6 +192,7 @@ impl ThreadedRingBufferPoller { resolve: R, tx: Sender>, poll_interval_ms: u64, + on_drained: Option>, ) -> Result where M: MapCore, @@ -172,7 +202,7 @@ impl ThreadedRingBufferPoller { R: Fn(T) -> U + Send + 'static, { let (parsed_tx, parsed_rx) = mpsc::channel::>(); - let ring = RingBufferPoller::new(rb_map, parse, parsed_tx, poll_interval_ms)?; + let ring = RingBufferPoller::new(rb_map, parse, parsed_tx, poll_interval_ms, on_drained)?; let resolver = std::thread::spawn(move || { for batch in parsed_rx { let resolved = batch.into_iter().map(&resolve).collect(); diff --git a/crates/memtrack/src/ebpf/proc_fs.rs b/crates/memtrack/src/ebpf/proc_fs.rs index 510fb9504..2c797c72d 100644 --- a/crates/memtrack/src/ebpf/proc_fs.rs +++ b/crates/memtrack/src/ebpf/proc_fs.rs @@ -38,7 +38,7 @@ impl ResolvedMapping { /// 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`. +/// `/proc//task//stat`. Success means every thread is `T`/`t` or exited. /// /// - A vanished process (`/proc/` gone) is success: the stop is moot. /// - At the deadline, threads still in uninterruptible sleep (`D`) are treated @@ -47,31 +47,18 @@ impl ResolvedMapping { /// 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 { + let Some(states) = task_states(pid) 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; - }; - + for (tid, state) in states { match state { - 'T' | 't' => {} + 'T' | 't' | 'Z' | 'X' => {} 'D' => all_stopped = false, _ => { all_stopped = false; @@ -105,6 +92,21 @@ fn task_state(stat: &str) -> Option { stat[idx + 1..].trim_start().chars().next() } +/// `(tid, state)` of every readable thread of `pid`; `None` once +/// `/proc//task` is gone. +fn task_states(pid: u32) -> Option> { + let entries = std::fs::read_dir(format!("/proc/{pid}/task")).ok()?; + let states = entries + .flatten() + .filter_map(|entry| { + let tid = entry.file_name().to_str()?.parse::().ok()?; + let stat = std::fs::read_to_string(entry.path().join("stat")).ok()?; + Some((tid, task_state(&stat)?)) + }) + .collect(); + Some(states) +} + /// The outcome of resolving a watcher `(dev, ino)` against `/proc//maps`. #[derive(Debug)] pub(super) enum Resolution { diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 6df8a27f6..1280733a9 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -1,4 +1,5 @@ use crate::ebpf::attach_worker::AttachWorker; +use crate::ebpf::poller::POLL_INTERVAL_MS; use crate::ebpf::spawn::{resume, spawn_stopped, wrap_stopped}; use crate::ebpf::stacks::StackCaptureFailureStats; use crate::ebpf::{BpfVariant, MemtrackBpf, OwnershipMaps}; @@ -33,6 +34,10 @@ pub struct TrackerOptions { /// Maximum bytes of user stack to copy per captured call stack. #[builder(default = 8192)] pub stack_budget: u32, + /// Event and stack ring poll interval. Larger values let the rings fill, + /// which is useful for exercising ring pressure on demand. + #[builder(default = POLL_INTERVAL_MS)] + pub poll_interval_ms: u64, } impl TrackerOptions { @@ -52,6 +57,12 @@ impl TrackerOptions { .and_then(|v| v.parse().ok()) .unwrap_or(8192), ) + .poll_interval_ms( + std::env::var("CODSPEED_MEMTRACK_POLL_INTERVAL_MS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(POLL_INTERVAL_MS), + ) .build() } } @@ -121,6 +132,7 @@ impl Tracker { /// read back, so it cannot be preserved through the wrap). pub fn spawn(&self, cmd: &Command, uid_gid: Option<(u32, u32)>) -> Result { let capture_stacks = self.options.stack_capture; + let poll_interval_ms = self.options.poll_interval_ms; let mut wrapped = wrap_stopped(cmd); if let Some((uid, gid)) = uid_gid { @@ -143,9 +155,12 @@ impl Tracker { let mut bpf = self.bpf.lock(); bpf.add_tracked_pid(pid)?; let stack_poller = capture_stacks - .then(|| bpf.poll_stacks(10, tx.clone())) + .then(|| bpf.poll_stacks(poll_interval_ms, tx.clone())) .transpose()?; - (bpf.poll_events_with_channel(10, tx.clone())?, stack_poller) + ( + bpf.poll_events_with_channel(poll_interval_ms, tx.clone())?, + stack_poller, + ) }; let perf_mapping_poller = capture_stacks .then(|| PerfMappingPoller::start(pid, tx, self.mapping_lost.clone())) @@ -174,6 +189,7 @@ impl Tracker { perf_mapping_poller, )) } + /// Enable allocator-event tracking in the BPF program. Lifetime events /// (rss_stat, rmap, fork/exec/exit) are emitted for tracked pids /// regardless of this toggle. @@ -187,9 +203,14 @@ impl Tracker { } /// Number of events the kernel dropped because the ring buffer was full. - /// A non-zero value means the resulting trace is incomplete. + /// A non-zero value means the resulting trace is incomplete. Includes + /// allocation-stack ring overflow: missing stack records make the capture + /// incomplete just like ordinary event-ring or mapping loss. pub fn dropped_events_count(&self) -> Result { - Ok(self.bpf.lock().dropped_events_count()? + self.mapping_lost.load(Ordering::Relaxed)) + let bpf = self.bpf.lock(); + Ok(bpf.dropped_events_count()? + + bpf.stack_capture_stats()?.ring_full + + self.mapping_lost.load(Ordering::Relaxed)) } /// Per-cause counts of stack captures that were skipped or truncated. diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index 7971b8c3a..5ff58cb41 100644 --- a/crates/memtrack/src/main.rs +++ b/crates/memtrack/src/main.rs @@ -136,9 +136,12 @@ fn track_command( let pipeline_thread = thread::spawn(move || encode_events(event_rx.into_iter().flatten(), out_file, n_workers)); - // Wait for the command to complete - let status = session.wait().context("Failed to wait for command")?; - debug!("Command exited with status: {status}"); + // A worker failure must not skip disabling tracking, draining, joining the + // encoder, or detaching probes. Keep the wait result until teardown is done. + let status = session.wait().context("Failed to wait for command"); + if let Ok(status) = &status { + debug!("Command exited with status: {status}"); + } // Stop allocator-event production before draining: the child has exited, // so anything still arriving is already in the ring buffer. @@ -155,18 +158,26 @@ fn track_command( debug!("Waiting for the encode pipeline to finish"); let total = pipeline_thread .join() - .map_err(|_| anyhow::anyhow!("Failed to join memtrack encode pipeline"))??; - - info!("Wrote {total} memtrack events to disk"); + .map_err(|_| anyhow::anyhow!("Failed to join memtrack encode pipeline")) + .and_then(|result| result); - // Stop the attach worker and surface any fatal error it recorded (missed - // exec mappings mean incomplete allocator coverage). - tracker.finish()?; + if let Ok(total) = &total { + info!("Wrote {total} memtrack events to disk"); + } - if tracker.stack_capture_enabled() { - let stats = tracker - .stack_capture_stats() - .context("Failed to read stack capture stats")?; + // Stop background workers after the ring pipeline has drained. Fatal + // worker errors mean the capture is incomplete. + let finish = tracker.finish(); + let stack_stats = if tracker.stack_capture_enabled() { + Some( + tracker + .stack_capture_stats() + .context("Failed to read stack capture stats"), + ) + } else { + None + }; + if let Some(Ok(stats)) = &stack_stats { debug!("stack capture stats: {stats:?}"); } @@ -175,6 +186,13 @@ fn track_command( // kernel would close every link fd serially during exit. tracker.detach(); + let status = status?; + total?; + finish?; + if let Some(stats) = stack_stats { + stats?; + } + // 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. diff --git a/crates/memtrack/testdata/alloc_storm.c b/crates/memtrack/testdata/alloc_storm.c new file mode 100644 index 000000000..d862d0fcf --- /dev/null +++ b/crates/memtrack/testdata/alloc_storm.c @@ -0,0 +1,41 @@ +// Saturates the allocator uprobes from several threads at once so the event +// ring fills faster than a slow poller can drain it. +// +// usage: alloc_storm +#include +#include +#include + +static long iterations; + +static void* storm(void* arg) { + (void)arg; + for (long i = 0; i < iterations; i++) { + volatile char* p = malloc(16); + p[0] = (char)i; + free((void*)p); + } + return NULL; +} + +int main(int argc, char** argv) { + if (argc != 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + int threads = atoi(argv[1]); + iterations = atol(argv[2]); + + pthread_t* handles = calloc((size_t)threads, sizeof(*handles)); + for (int i = 0; i < threads; i++) { + if (pthread_create(&handles[i], NULL, storm, NULL) != 0) { + perror("pthread_create"); + return 1; + } + } + for (int i = 0; i < threads; i++) { + pthread_join(handles[i], NULL); + } + free(handles); + return 0; +} diff --git a/crates/memtrack/testdata/alloc_storm_procs.c b/crates/memtrack/testdata/alloc_storm_procs.c new file mode 100644 index 000000000..fb061f8d5 --- /dev/null +++ b/crates/memtrack/testdata/alloc_storm_procs.c @@ -0,0 +1,46 @@ +// Saturates the allocator uprobes from several processes at once, so the +// event ring fills while many producers write to it concurrently. +// +// usage: alloc_storm_procs +#include +#include +#include +#include + +static void storm(long iterations) { + for (long i = 0; i < iterations; i++) { + volatile char* p = malloc(16); + p[0] = (char)i; + free((void*)p); + } +} + +int main(int argc, char** argv) { + if (argc != 3) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + int processes = atoi(argv[1]); + long iterations = atol(argv[2]); + + for (int i = 0; i < processes; i++) { + pid_t pid = fork(); + if (pid < 0) { + perror("fork"); + return 1; + } + if (pid == 0) { + storm(iterations); + _exit(0); + } + } + + int failed = 0; + for (int i = 0; i < processes; i++) { + int status; + if (wait(&status) < 0 || !WIFEXITED(status) || WEXITSTATUS(status) != 0) { + failed = 1; + } + } + return failed; +} diff --git a/crates/memtrack/tests/pressure_tests.rs b/crates/memtrack/tests/pressure_tests.rs new file mode 100644 index 000000000..546c49fdf --- /dev/null +++ b/crates/memtrack/tests/pressure_tests.rs @@ -0,0 +1,114 @@ +//! Ring-pressure pause under a deliberately slow poller: with stack capture on, +//! the stack ring fills well before the next poll tick, so writing producers +//! must stay paused until the poller has flushed the ring. + +mod shared; + +use memtrack::{Tracker, TrackerOptions}; +use std::process::Command; +use std::time::{Duration, Instant}; +use tempfile::TempDir; + +const THREADS: &str = "8"; +const PROCESSES: &str = "16"; +const ITERATIONS: &str = "400000"; +/// Long enough that the ring's 75% watermark is crossed between two polls. +const SLOW_POLL_MS: u64 = 10_000; + +struct Run { + wall: Duration, + dropped: u64, +} + +fn run_storm( + binary: &std::path::Path, + args: [&str; 2], + options: TrackerOptions, +) -> anyhow::Result { + let tracker = Tracker::with_options(options)?; + tracker.enable_tracking()?; + + let mut command = Command::new(binary); + command.args(args); + + let started = Instant::now(); + let mut session = tracker.spawn(&command, None)?; + let rx = session.take_events()?; + let status = session.wait()?; + let wall = started.elapsed(); + assert!(status.success(), "fixture failed: {status}"); + + drop(session); + let events: usize = rx.into_iter().map(|batch| batch.len()).sum(); + tracker.finish()?; + let dropped = tracker.dropped_events_count()?; + drop(tracker); + + eprintln!("wall {wall:?}, events {events}, dropped {dropped}"); + Ok(Run { wall, dropped }) +} + +#[test_with::env(GITHUB_ACTIONS)] +#[test] +fn slow_poller_pause_recovers_without_loss() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + let dir = TempDir::new()?; + let binary = shared::compile_c_source( + include_str!("../testdata/alloc_storm.c"), + "alloc_storm", + dir.path(), + ) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + eprintln!("-- baseline: fast poller"); + let baseline = run_storm( + &binary, + [THREADS, ITERATIONS], + TrackerOptions::builder().build(), + )?; + + eprintln!("-- slow poller"); + let blocked = run_storm( + &binary, + [THREADS, ITERATIONS], + TrackerOptions::builder() + .stack_capture(true) + .poll_interval_ms(SLOW_POLL_MS) + .build(), + )?; + + eprintln!( + "baseline {:?} | blocked {:?} (dropped {})", + baseline.wall, blocked.wall, blocked.dropped + ); + + assert_eq!(blocked.dropped, 0, "pressure pause lost events"); + Ok(()) +} + +/// Every writing process is stopped on its own and must be resumed: a +/// producer left stopped hangs the parent's `wait`, a missed one drops events. +#[test_with::env(GITHUB_ACTIONS)] +#[test] +fn slow_poller_pause_resumes_every_writing_process() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + let dir = TempDir::new()?; + let binary = shared::compile_c_source( + include_str!("../testdata/alloc_storm_procs.c"), + "alloc_storm_procs", + dir.path(), + ) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let blocked = run_storm( + &binary, + [PROCESSES, ITERATIONS], + TrackerOptions::builder() + .stack_capture(true) + .poll_interval_ms(SLOW_POLL_MS) + .build(), + )?; + + assert_eq!(blocked.dropped, 0, "pressure pause lost events"); + Ok(()) +} diff --git a/crates/memtrack/tests/shared.rs b/crates/memtrack/tests/shared.rs index 75b4d2105..1d01c8ec7 100644 --- a/crates/memtrack/tests/shared.rs +++ b/crates/memtrack/tests/shared.rs @@ -268,20 +268,35 @@ pub fn track_command_with_rmap_checkpoint( /// none of them can be compared across variants. type EventProfile = std::collections::BTreeMap; +/// Frees count only when they release an allocation the run saw: glibc's +/// `free` can re-enter itself on a thread's first free (lazy tcache init), so +/// one call may fire the uprobe once or twice depending on scheduling. fn event_profile(events: &[Event]) -> EventProfile { + use itertools::Itertools; + let mut profile = EventProfile::new(); - for event in events { + let mut live = std::collections::HashSet::new(); + for event in events.iter().sorted_by_key(|e| e.timestamp) { // Only allocator events are comparable across variants: RSS and // lifecycle values (sizes, pids) are per-run. - if !matches!( - event.kind, + match event.kind { + MemtrackEventKind::Free { .. } => { + if !live.remove(&event.addr) { + continue; + } + } + MemtrackEventKind::Realloc { old_addr, .. } => { + if let Some(old_addr) = old_addr { + live.remove(&old_addr); + } + live.insert(event.addr); + } MemtrackEventKind::Malloc { .. } - | MemtrackEventKind::Free { .. } - | MemtrackEventKind::Calloc { .. } - | MemtrackEventKind::Realloc { .. } - | MemtrackEventKind::AlignedAlloc { .. } - ) { - continue; + | MemtrackEventKind::Calloc { .. } + | MemtrackEventKind::AlignedAlloc { .. } => { + live.insert(event.addr); + } + _ => continue, } *profile.entry(describe_kind(&event.kind)).or_default() += 1; }