diff --git a/crates/memtrack/src/ebpf/attach_worker.rs b/crates/memtrack/src/ebpf/attach_worker.rs index 892fe6211..276cbd89c 100644 --- a/crates/memtrack/src/ebpf/attach_worker.rs +++ b/crates/memtrack/src/ebpf/attach_worker.rs @@ -1,7 +1,7 @@ use crate::AllocatorLib; use crate::ebpf::MemtrackBpf; use crate::ebpf::events::AttachRequest; -use crate::ebpf::poller::RingBufferPoller; +use crate::ebpf::poller::{DrainRequest, POLL_INTERVAL_MS, PollerHandle, RingBufferPoller}; use crate::prelude::*; use parking_lot::Mutex; use std::collections::HashSet; @@ -14,7 +14,6 @@ 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 @@ -40,6 +39,7 @@ pub(crate) struct AttachWorker { fatal: Arc>>, root_pid: Arc, bpf: Arc>, + poller_handle: PollerHandle, } impl AttachWorker { @@ -49,11 +49,13 @@ impl AttachWorker { let root_pid = Arc::new(AtomicI32::new(0)); let (tx, rx) = mpsc::channel(); + let (poller_handle, drain_rx) = PollerHandle::channel(); let poller = bpf.lock().poll_attach_with_channel(POLL_INTERVAL_MS, tx)?; let worker = Worker { poller, rx, + drain_rx, bpf: bpf.clone(), shutdown: shutdown.clone(), fatal: fatal.clone(), @@ -68,9 +70,14 @@ impl AttachWorker { fatal, root_pid, bpf, + poller_handle, }) } + pub(crate) fn handle(&self) -> PollerHandle { + self.poller_handle.clone() + } + /// 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); @@ -80,6 +87,7 @@ impl AttachWorker { /// 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.poller_handle.close(); self.shutdown.store(true, Ordering::SeqCst); if let Some(handle) = self.handle.take() && let Err(panic) = handle.join() @@ -112,6 +120,7 @@ impl AttachWorker { /// and the links fall back to a slow serial close at process exit. impl Drop for AttachWorker { fn drop(&mut self) { + self.poller_handle.close(); self.shutdown.store(true, Ordering::SeqCst); if let Some(handle) = self.handle.take() { let _ = handle.join(); @@ -122,6 +131,7 @@ impl Drop for AttachWorker { struct Worker { poller: RingBufferPoller, rx: mpsc::Receiver>, + drain_rx: mpsc::Receiver, bpf: Arc>, shutdown: Arc, fatal: Arc>>, @@ -152,7 +162,13 @@ impl Worker { let first = match self.rx.recv_timeout(RECV_TIMEOUT) { Ok(reqs) => reqs, - Err(RecvTimeoutError::Timeout) => continue, + Err(RecvTimeoutError::Timeout) => { + if let Err(error) = self.service_barriers(&mut known) { + self.record_fatal(error); + return; + } + continue; + } Err(RecvTimeoutError::Disconnected) => return, }; let mut batch: Vec = first @@ -164,9 +180,34 @@ impl Worker { self.record_fatal(e); return; } + if let Err(e) = self.service_barriers(&mut known) { + self.record_fatal(e); + return; + } } } + fn service_barriers(&self, known: &mut HashSet<(u64, u64)>) -> Result<()> { + for request in self.drain_rx.try_iter() { + let remaining = request + .deadline + .saturating_duration_since(std::time::Instant::now()); + if let Err(error) = self.poller.handle().drain(remaining) { + let _ = request.ack.send(Err(error)); + continue; + } + let mut batch = self.rx.try_iter().flatten().collect(); + if let Err(error) = self.process_batch(&mut batch, known) { + let _ = request + .ack + .send(Err(anyhow::anyhow!("attach processing failed: {error:#}"))); + return Err(error); + } + let _ = request.ack.send(Ok(())); + } + Ok(()) + } + /// 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. 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/event_helpers.h b/crates/memtrack/src/ebpf/c/utils/event_helpers.h index 10dab1368..5284474bb 100644 --- a/crates/memtrack/src/ebpf/c/utils/event_helpers.h +++ b/crates/memtrack/src/ebpf/c/utils/event_helpers.h @@ -4,6 +4,7 @@ #include "../event.h" #include "../stack_capture.bpf.h" #include "map_helpers.h" +#include "pressure.bpf.h" #include "process_tracking.h" BPF_RINGBUF(events, 256 * 1024 * 1024); @@ -61,6 +62,7 @@ static __always_inline __u64* take_param(void* map) { if (drops) { \ __sync_fetch_and_add(drops, 1); \ } \ + memtrack_check_ring_pressure(&events, ids.tgid); \ return 0; \ } \ \ @@ -72,6 +74,7 @@ static __always_inline __u64* take_param(void* map) { fill_data; \ \ bpf_ringbuf_submit(e, wake_flags()); \ + memtrack_check_ring_pressure(&events, ids.tgid); \ return 0; \ } 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..6aae31dd6 --- /dev/null +++ b/crates/memtrack/src/ebpf/c/utils/pressure.bpf.h @@ -0,0 +1,49 @@ +#ifndef __PRESSURE_BPF_H__ +#define __PRESSURE_BPF_H__ + +#include + +#include "map_helpers.h" +#include "process_tracking.h" + +/* Ring pressure latch. Call only after submit/discard or a failed reserve: + * stopping with a live reservation would wedge the ring. The first + * BPF_NOEXIST insert wins the episode; userspace clears it after quiescence + * and drain. Only tracked producers may receive SIGSTOP. */ + +#ifndef MEMTRACK_SIGSTOP +#define MEMTRACK_SIGSTOP 19 +#endif + +/* Key 0 -> episode timestamp; absent means idle. A hash insert is the + * portable atomic latch for the supported 5.11 kernel floor. */ +BPF_HASH_MAP(pressure_since, __u32, __u64, 1); + +#define MEMTRACK_PRESSURE_HEADROOM_FRAC 4 /* latch 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; +} + +/* Check one ring against its watermark, latch a durable episode timestamp, + * and stop the current producer if it won the episode and is tracked. */ +static __always_inline void memtrack_check_ring_pressure(void* ring, __u32 current_tgid) { + if (!memtrack_ring_over_watermark(ring)) { + return; + } + + __u32 key = 0; + __u64 now = bpf_ktime_get_ns(); + if (bpf_map_update_elem(&pressure_since, &key, &now, BPF_NOEXIST) != 0) { + return; /* The episode is already latched. */ + } + + /* Never stop a foreign RSS/rmap producer sharing the event ring. */ + if (is_tracked(current_tgid)) { + bpf_send_signal(MEMTRACK_SIGSTOP); + } +} + +#endif /* __PRESSURE_BPF_H__ */ diff --git a/crates/memtrack/src/ebpf/memtrack/maps.rs b/crates/memtrack/src/ebpf/memtrack/maps.rs index 7ca16d25a..2461de83f 100644 --- a/crates/memtrack/src/ebpf/memtrack/maps.rs +++ b/crates/memtrack/src/ebpf/memtrack/maps.rs @@ -65,6 +65,50 @@ impl MemtrackBpf { ) } + /// The first over-watermark timestamp latched by BPF, or zero while idle. + pub fn pressure_since(&self) -> Result { + let key = 0u32; + let value = with_skel!(self, skel => skel.maps.pressure_since.lookup( + &key.to_le_bytes(), + libbpf_rs::MapFlags::ANY, + )) + .context("Failed to read pressure latch")?; + let Some(value) = value else { + return Ok(0); + }; + let bytes: [u8; 8] = value + .as_slice() + .try_into() + .map_err(|_| anyhow!("pressure latch has unexpected size"))?; + Ok(u64::from_le_bytes(bytes)) + } + + /// Clear the BPF pressure latch after the tree is quiesced and rings drain. + pub fn clear_pressure(&self) -> Result<()> { + let key = 0u32.to_le_bytes(); + let present = with_skel!(self, skel => skel.maps.pressure_since.lookup( + &key, + libbpf_rs::MapFlags::ANY, + )) + .context("Failed to inspect pressure latch")?; + if present.is_some() { + with_skel!(self, skel => skel.maps.pressure_since.delete(&key)) + .context("Failed to clear pressure latch")?; + } + Ok(()) + } + + /// PIDs currently present in the BPF tracked set. + pub fn tracked_pids(&self) -> Result> { + Ok(with_skel!(self, skel => { + skel.maps + .tracked_pids + .keys() + .map(|key| le(&key) as u32) + .collect() + })) + } + pub fn stack_capture_stats(&self) -> Result { StackCaptureFailureStats::read(with_skel!(self, skel => &skel.maps.stack_counters)) } 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..1428ab6fa --- /dev/null +++ b/crates/memtrack/src/ebpf/pause.rs @@ -0,0 +1,265 @@ +use super::proc_fs::ProcessTree; +use crate::ebpf::MemtrackBpf; +use crate::ebpf::poller::PollerHandle; +use crate::prelude::*; +use parking_lot::Mutex; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc::RecvTimeoutError; +use std::thread::JoinHandle; +use std::time::Duration; + +const OBSERVE_INTERVAL: Duration = Duration::from_millis(1); +const DRAIN_SLICE: Duration = Duration::from_millis(100); +const PAUSE_DEADLINE: Duration = Duration::from_secs(30); + +pub(crate) struct PressureWorker { + handle: Option>, + shutdown: Arc, + fatal: Arc>>, +} + +impl PressureWorker { + pub(crate) fn start(bpf: Arc>, pollers: Vec) -> Self { + let shutdown = Arc::new(AtomicBool::new(false)); + let fatal = Arc::new(Mutex::new(None)); + let worker_shutdown = shutdown.clone(); + let worker_fatal = fatal.clone(); + let handle = std::thread::spawn(move || { + monitor_loop(bpf, pollers, worker_shutdown, worker_fatal); + }); + Self { + handle: Some(handle), + shutdown, + fatal, + } + } + + 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!("pressure worker thread panicked: {msg}"); + } + if let Some(err) = self.fatal.lock().take() { + bail!("{err}"); + } + Ok(()) + } +} + +impl Drop for PressureWorker { + fn drop(&mut self) { + self.shutdown.store(true, Ordering::SeqCst); + if let Some(handle) = self.handle.take() { + let _ = handle.join(); + } + } +} + +fn record_fatal(fatal: &Mutex>, error: impl std::fmt::Display) { + let mut guard = fatal.lock(); + if guard.is_none() { + *guard = Some(error.to_string()); + } +} + +fn monitor_loop( + bpf: Arc>, + pollers: Vec, + shutdown: Arc, + fatal: Arc>>, +) { + let mut episodes = 0u64; + while !shutdown.load(Ordering::Acquire) { + let pressure = bpf.lock().pressure_since(); + let pressure_since = match pressure { + Ok(value) => value, + Err(error) => { + let message = format!("failed to read memtrack pressure latch: {error:#}"); + error!("{message}; terminating tracked tree"); + if let Ok(members) = scan_tracked_tree(&bpf) { + let _ = members.signal_all(libc::SIGKILL); + } + let _ = bpf.lock().clear_pressure(); + record_fatal(&fatal, message); + shutdown.store(true, Ordering::Release); + break; + } + }; + if pressure_since == 0 { + sleep_or_shutdown(&shutdown); + continue; + } + debug!("memtrack pressure episode begin: latch={pressure_since}ns"); + episodes += 1; + match on_high_pressure(&bpf, pressure_since, &pollers, &shutdown) { + Ok(()) => {} + Err(_error) if shutdown.load(Ordering::Acquire) => break, + Err(error) => { + record_fatal(&fatal, format!("{error:#}")); + error!("memtrack pressure monitor failed: {error:#}"); + break; + } + } + } + debug!("memtrack pressure worker exiting after {episodes} episodes"); +} + +fn on_high_pressure( + bpf: &Arc>, + pressure_since: u64, + pollers: &[PollerHandle], + shutdown: &AtomicBool, +) -> Result<()> { + let mut owned = ProcessTree::default(); + let result = (|| { + stop_tree(bpf, pressure_since, shutdown, &mut owned)?; + let stopped_at = monotonic_ns(); + drain_pollers(pollers, pressure_since, shutdown)?; + let drained_at = monotonic_ns(); + if shutdown.load(Ordering::Acquire) { + bail!("pressure monitor cancelled after stopping tracked tree"); + } + // Clear only after all producers are stopped and every ring has acked a + // synchronous drain, so resumed producers start from an empty ring. + bpf.lock().clear_pressure()?; + owned.signal_all(libc::SIGCONT)?; + let resumed_at = monotonic_ns(); + debug!( + "memtrack pressure episode recovered: {} pids, stop {}us, drain {}us, resume {}us, total {}us", + owned.len(), + stopped_at.saturating_sub(pressure_since) / 1_000, + drained_at.saturating_sub(stopped_at) / 1_000, + resumed_at.saturating_sub(drained_at) / 1_000, + resumed_at.saturating_sub(pressure_since) / 1_000, + ); + Ok(()) + })(); + + if result.is_ok() { + return Ok(()); + } + let error = result.unwrap_err(); + let timed_out = elapsed_since(pressure_since) >= PAUSE_DEADLINE; + if timed_out { + error!( + "memtrack pressure pause timed out after {}ms; terminating tracked tree", + elapsed_since(pressure_since).as_millis() + ); + } + let can_resume = + shutdown.load(Ordering::Acquire) && !timed_out && bpf.lock().clear_pressure().is_ok(); + if !can_resume || owned.signal_all(libc::SIGCONT).is_err() { + if owned.is_empty() { + if let Ok(scanned) = scan_tracked_tree(bpf) { + owned = scanned; + } + } + let _ = owned.signal_all(libc::SIGKILL); + } + let _ = bpf.lock().clear_pressure(); + Err(error) +} + +fn stop_tree( + bpf: &Arc>, + pressure_since: u64, + shutdown: &AtomicBool, + owned: &mut ProcessTree, +) -> Result<()> { + let mut clean_passes = 0u8; + loop { + if shutdown.load(Ordering::Acquire) { + bail!("pressure monitor cancelled while stopping tracked tree"); + } + check_deadline(pressure_since)?; + // Pin every pid before signalling, so cleanup owns all stopped processes. + let grew = owned.absorb(scan_tracked_tree(bpf)?); + owned.signal_all(libc::SIGSTOP)?; + // A child forked between the /proc scan and the SIGSTOP is missed by + // that pass, so the tree counts as stopped only after two passes in a + // row that find no new process and every process stopped. + if owned.all_stopped() && !grew { + clean_passes = clean_passes.saturating_add(1); + if clean_passes >= 2 { + return Ok(()); + } + } else { + clean_passes = 0; + } + sleep_or_shutdown(shutdown); + } +} + +fn drain_pollers( + pollers: &[PollerHandle], + pressure_since: u64, + shutdown: &AtomicBool, +) -> Result<()> { + for poller in pollers { + loop { + if shutdown.load(Ordering::Acquire) { + bail!("pressure monitor cancelled while draining ring buffers"); + } + let remaining = PAUSE_DEADLINE.saturating_sub(elapsed_since(pressure_since)); + if remaining.is_zero() { + bail!("memtrack pressure pause timeout"); + } + match poller.drain(remaining.min(DRAIN_SLICE)) { + Ok(()) => break, + Err(error) + if error.downcast_ref::() + == Some(&RecvTimeoutError::Timeout) => + { + sleep_or_shutdown(shutdown) + } + Err(error) => return Err(error), + } + } + } + Ok(()) +} + +fn scan_tracked_tree(bpf: &Arc>) -> Result { + let roots = bpf.lock().tracked_pids()?; + // The BPF map keys are u32; convert to pid_t at the kernel boundary. + let roots: Vec = roots.iter().map(|&pid| pid as libc::pid_t).collect(); + ProcessTree::scan(&roots) +} + +fn check_deadline(pressure_since: u64) -> Result<()> { + if elapsed_since(pressure_since) >= PAUSE_DEADLINE { + bail!("memtrack pressure pause timeout"); + } + Ok(()) +} + +fn elapsed_since(pressure_since: u64) -> Duration { + Duration::from_nanos(monotonic_ns().saturating_sub(pressure_since)) +} + +fn monotonic_ns() -> u64 { + let mut timestamp = libc::timespec { + tv_sec: 0, + tv_nsec: 0, + }; + if unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC, &mut timestamp) } != 0 { + return 0; + } + (timestamp.tv_sec as u64) + .saturating_mul(1_000_000_000) + .saturating_add(timestamp.tv_nsec as u64) +} + +fn sleep_or_shutdown(shutdown: &AtomicBool) { + if !shutdown.load(Ordering::Acquire) { + std::thread::sleep(OBSERVE_INTERVAL); + } +} diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index 9c122e068..e074aae09 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/crates/memtrack/src/ebpf/poller.rs @@ -1,10 +1,82 @@ use anyhow::{Context, Result}; -use libbpf_rs::{MapCore, RingBufferBuilder}; +use libbpf_rs::{AsRawLibbpf, MapCore, RingBufferBuilder, libbpf_sys}; use parking_lot::Mutex; use std::sync::Arc; -use std::sync::mpsc::{self, RecvTimeoutError, Sender}; +use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender}; use std::thread::JoinHandle; -use std::time::Duration; +use std::time::{Duration, Instant}; + +/// Ring-buffer poll interval shared by every poller. +pub(crate) const POLL_INTERVAL_MS: u64 = 1; + +/// A drain request routed through a channel to the poller's worker. `deadline` +/// is absolute and bounds queue wait together with every drain stage. +pub(crate) struct DrainRequest { + pub(crate) deadline: Instant, + pub(crate) ack: Sender>, +} + +/// Cloneable handle for requesting a timeout-bounded pipeline drain from the +/// poller that created it via [`PollerHandle::channel`]. +#[derive(Clone)] +pub(crate) struct PollerHandle { + request: Arc>>>, +} + +impl PollerHandle { + /// Pairs a handle with the receiver that its owner's worker polls. + pub(crate) fn channel() -> (Self, Receiver) { + let (tx, rx) = mpsc::channel::(); + ( + Self { + request: Arc::new(Mutex::new(Some(tx))), + }, + rx, + ) + } + + /// Revokes the shared sender, disabling every clone of this handle. The + /// owner's `Drop` calls this so surviving clones can neither delay shutdown + /// nor keep the worker alive. + pub(crate) fn close(&self) { + self.request.lock().take(); + } + + /// Requests a drain bounded by an absolute deadline that spans queue wait + /// and the upstream drain. A ring stopped at an uncommitted reservation is + /// retried within the remaining budget. + pub(crate) fn drain(&self, timeout: Duration) -> Result<()> { + let deadline = Instant::now() + timeout; + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + let (ack_tx, ack_rx) = mpsc::channel(); + let ctl = self + .request + .lock() + .as_ref() + .cloned() + .context("poller already shut down")?; + ctl.send(DrainRequest { + deadline, + ack: ack_tx, + }) + .context("poll thread is gone")?; + match ack_rx.recv_timeout(remaining) { + Ok(Err(error)) + if error.downcast_ref::() + == Some(&RecvTimeoutError::Timeout) + && Instant::now() < deadline => + { + std::thread::sleep(Duration::from_millis(1)); + } + Ok(result) => return result, + Err(error) => { + return Err(error).context("poller did not drain before deadline"); + } + } + } + } +} /// 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 @@ -24,18 +96,17 @@ fn flush_batch(batch: &Mutex>, tx: &Sender>) { } fn poll_iteration( - control: std::result::Result, RecvTimeoutError>, - consume: impl FnOnce(), + control: std::result::Result, + consume: impl FnOnce() -> Result<()>, poll: impl FnOnce(), batch: &Mutex>, tx: &Sender>, ) -> bool { match control { - Ok(ack) => { - consume(); - // `drain` promises pending entries are in the channel before returning. + Ok(request) => { + let result = consume(); flush_batch(batch, tx); - let _ = ack.send(()); + let _ = request.ack.send(result); true } Err(RecvTimeoutError::Timeout) => { @@ -44,7 +115,7 @@ fn poll_iteration( true } Err(RecvTimeoutError::Disconnected) => { - consume(); + let _ = consume(); flush_batch(batch, tx); false } @@ -57,7 +128,7 @@ fn poll_iteration( /// 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 { - ctl: Option>>, + handle: PollerHandle, poll_thread: Option>, } @@ -96,15 +167,26 @@ impl RingBufferPoller { })?; let ringbuf = builder.build()?; - // The control channel doubles as the poll pacing: a received message is + // The request 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 (handle, request_rx) = PollerHandle::channel(); let poll_thread = std::thread::spawn(move || { while poll_iteration( - ctl_rx.recv_timeout(Duration::from_millis(poll_interval_ms)), + request_rx.recv_timeout(Duration::from_millis(poll_interval_ms)), || { - let _ = ringbuf.consume(); + ringbuf.consume()?; + // consume() stops at an uncommitted reservation, not only + // at an empty ring. Such a stop must not acknowledge drain. + let available = unsafe { + let ring = + libbpf_sys::ring_buffer__ring(ringbuf.as_libbpf_object().as_ptr(), 0); + libbpf_sys::ring__avail_data_size(ring) + }; + if available != 0 { + return Err(RecvTimeoutError::Timeout.into()); + } + Ok(()) }, || { let _ = ringbuf.poll(Duration::ZERO); @@ -115,31 +197,68 @@ impl RingBufferPoller { }); Ok(Self { - ctl: Some(ctl), + handle, poll_thread: Some(poll_thread), }) } + pub(crate) fn handle(&self) -> PollerHandle { + self.handle.clone() + } + /// 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(()) + self.handle().drain(Duration::from_secs(30)) } } impl Drop for RingBufferPoller { fn drop(&mut self) { - drop(self.ctl.take()); + self.handle.close(); if let Some(thread) = self.poll_thread.take() { let _ = thread.join(); } } } +fn resolve_batches( + parsed_rx: mpsc::Receiver>, + request_rx: mpsc::Receiver, + ring: PollerHandle, + tx: Sender>, + resolve: impl Fn(T) -> U, +) { + let resolve_batch = |batch: Vec| { + let _ = tx.send(batch.into_iter().map(&resolve).collect()); + }; + loop { + let disconnected = match parsed_rx.recv_timeout(Duration::from_millis(1)) { + Ok(batch) => { + resolve_batch(batch); + false + } + Err(RecvTimeoutError::Timeout) => false, + Err(RecvTimeoutError::Disconnected) => true, + }; + for request in request_rx.try_iter() { + let remaining = request.deadline.saturating_duration_since(Instant::now()); + if let Err(error) = ring.drain(remaining) { + let _ = request.ack.send(Err(error)); + continue; + } + // Data preceding the barrier can arrive during resolution. + for batch in parsed_rx.try_iter() { + resolve_batch(batch); + } + let _ = request.ack.send(Ok(())); + } + if disconnected { + return; + } + } +} + /// A [`RingBufferPoller`] whose parsed items need a further, potentially /// expensive step (e.g. a BPF map lookup, which is a syscall) before they are /// forwarded on `tx`. That step runs on a dedicated resolver thread instead @@ -150,6 +269,7 @@ pub struct ThreadedRingBufferPoller { // The resolver then drains parsed items and can be joined safely. ring: Option, resolver: Option>, + handle: PollerHandle, } impl ThreadedRingBufferPoller { @@ -173,23 +293,28 @@ impl ThreadedRingBufferPoller { { let (parsed_tx, parsed_rx) = mpsc::channel::>(); let ring = RingBufferPoller::new(rb_map, parse, parsed_tx, poll_interval_ms)?; + let ring_handle = ring.handle(); + let (handle, request_rx) = PollerHandle::channel(); let resolver = std::thread::spawn(move || { - for batch in parsed_rx { - let resolved = batch.into_iter().map(&resolve).collect(); - let _ = tx.send(resolved); - } + resolve_batches(parsed_rx, request_rx, ring_handle, tx, resolve); }); Ok(Self { ring: Some(ring), resolver: Some(resolver), + handle, }) } + + pub(crate) fn handle(&self) -> PollerHandle { + self.handle.clone() + } } impl Drop for ThreadedRingBufferPoller { fn drop(&mut self) { drop(self.ring.take()); + self.handle.close(); if let Some(resolver) = self.resolver.take() { let _ = resolver.join(); } @@ -233,8 +358,14 @@ mod tests { let consumed = Cell::new(false); assert!(poll_iteration( - Ok(ack_tx), - || consumed.set(true), + Ok(DrainRequest { + deadline: Instant::now(), + ack: ack_tx, + }), + || { + consumed.set(true); + Ok(()) + }, || unreachable!(), &batch, &tx, @@ -242,7 +373,7 @@ mod tests { assert!(consumed.get()); assert_eq!(rx.recv().unwrap(), expected); - ack_rx.recv().unwrap(); + ack_rx.recv().unwrap().unwrap(); } #[test] @@ -254,7 +385,10 @@ mod tests { assert!(!poll_iteration( Err(RecvTimeoutError::Disconnected), - || consumed.set(true), + || { + consumed.set(true); + Ok(()) + }, || unreachable!(), &batch, &tx, @@ -263,4 +397,52 @@ mod tests { assert!(consumed.get()); assert_eq!(rx.recv().unwrap(), expected); } + + #[test] + fn resolver_drain_waits_for_batches_queued_during_resolution() { + let timeout = Duration::from_secs(5); + let (input_tx, input_rx) = mpsc::channel(); + let (request_tx, request_rx) = mpsc::channel::(); + let (ring_handle, ring_rx) = PollerHandle::channel(); + let (output_tx, output_rx) = mpsc::channel(); + let (started_tx, started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let ring_stub = std::thread::spawn(move || { + for request in ring_rx { + let _ = request.ack.send(Ok(())); + } + }); + let worker = std::thread::spawn(move || { + resolve_batches(input_rx, request_rx, ring_handle, output_tx, |value| { + started_tx.send(value).unwrap(); + release_rx.recv_timeout(timeout).unwrap(); + value + }); + }); + + input_tx.send(vec![1]).unwrap(); + assert_eq!(started_rx.recv_timeout(timeout).unwrap(), 1); + input_tx.send(vec![2]).unwrap(); + let (ack_tx, ack_rx) = mpsc::channel(); + request_tx + .send(DrainRequest { + deadline: Instant::now() + timeout, + ack: ack_tx, + }) + .unwrap(); + release_tx.send(()).unwrap(); + assert_eq!(started_rx.recv_timeout(timeout).unwrap(), 2); + let premature_ack = ack_rx.try_recv().is_ok(); + release_tx.send(()).unwrap(); + if !premature_ack { + ack_rx.recv_timeout(timeout).unwrap().unwrap(); + } + drop(input_tx); + worker.join().unwrap(); + drop(request_tx); + ring_stub.join().unwrap(); + + assert!(!premature_ack, "drain completed before the queued batch"); + assert_eq!(output_rx.into_iter().flatten().collect::>(), [1, 2]); + } } diff --git a/crates/memtrack/src/ebpf/proc_fs.rs b/crates/memtrack/src/ebpf/proc_fs.rs index 510fb9504..8cf11f295 100644 --- a/crates/memtrack/src/ebpf/proc_fs.rs +++ b/crates/memtrack/src/ebpf/proc_fs.rs @@ -1,7 +1,12 @@ use crate::prelude::*; +use std::collections::hash_map::Entry; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; use std::path::PathBuf; use std::time::{Duration, Instant}; +type ProcessKey = (libc::pid_t, u64); + /// A mapping resolved from `/proc//maps` back to an attachable path. #[derive(Debug)] pub(super) struct ResolvedMapping { @@ -38,7 +43,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 +52,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 +97,193 @@ 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) +} + +/// A process pinned by a pidfd, so signals never reach a reused pid. +#[derive(Debug)] +pub(super) struct Process { + pid: libc::pid_t, + starttime: u64, + pidfd: OwnedFd, +} + +impl Process { + fn key(&self) -> ProcessKey { + (self.pid, self.starttime) + } + + /// ESRCH (process gone) is success. + fn signal(&self, signal: libc::c_int) -> Result<()> { + let ret = unsafe { + libc::syscall( + libc::SYS_pidfd_send_signal, + self.pidfd.as_raw_fd(), + signal, + std::ptr::null::(), + 0, + ) + }; + if ret != 0 { + let error = std::io::Error::last_os_error(); + if error.raw_os_error() == Some(libc::ESRCH) { + return Ok(()); + } + return Err(error.into()); + } + Ok(()) + } + + /// Non-blocking; a vanished process counts as stopped. + fn is_stopped(&self) -> bool { + let Some(states) = task_states(self.pid as u32) else { + return true; + }; + // `D` threads cannot run user code and join the group stop on syscall return. + states + .iter() + .all(|&(_, state)| matches!(state, 'T' | 't' | 'Z' | 'X' | 'D')) + } +} + +#[derive(Debug, Default)] +pub(super) struct ProcessTree(HashMap); + +impl ProcessTree { + /// Pin `roots` and all their descendants. + pub(super) fn scan(roots: &[libc::pid_t]) -> Result { + let entries = std::fs::read_dir("/proc").context("failed to scan /proc")?; + let mut processes = HashMap::::new(); + for entry in entries.flatten() { + let Some(pid) = entry + .file_name() + .to_str() + .and_then(|name| name.parse().ok()) + else { + continue; + }; + let Ok(stat) = std::fs::read_to_string(entry.path().join("stat")) else { + continue; + }; + let Some((ppid, starttime)) = parse_stat(&stat) else { + continue; + }; + processes.insert(pid, (ppid, starttime)); + } + + let mut children = HashMap::>::new(); + for (&pid, &(ppid, _)) in &processes { + children.entry(ppid).or_default().push(pid); + } + let mut queue: VecDeque = roots.iter().copied().collect(); + let mut seen = HashSet::new(); + let mut tree = HashMap::new(); + while let Some(pid) = queue.pop_front() { + if !seen.insert(pid) { + continue; + } + let Some(&(_, starttime)) = processes.get(&pid) else { + continue; + }; + let Some(pidfd) = pidfd_open(pid)? else { + continue; + }; + // Recheck identity after opening the pidfd; a reused numeric pid is + // never admitted to the owned tree. + let Some(stat) = std::fs::read_to_string(format!("/proc/{pid}/stat")) + .ok() + .and_then(|value| parse_stat(&value)) + else { + continue; + }; + if stat.1 != starttime { + continue; + } + let process = Process { + pid, + starttime, + pidfd, + }; + tree.insert(process.key(), process); + if let Some(descendants) = children.get(&pid) { + queue.extend(descendants.iter().copied()); + } + } + Ok(Self(tree)) + } + + /// Move processes not already present into `self`; `true` if any were added. + pub(super) fn absorb(&mut self, other: Self) -> bool { + let mut grew = false; + for (key, process) in other.0 { + if let Entry::Vacant(entry) = self.0.entry(key) { + entry.insert(process); + grew = true; + } + } + grew + } + + /// Signal every process, continuing past failures; returns the first error. + pub(super) fn signal_all(&self, signal: libc::c_int) -> Result<()> { + let mut first_error = None; + for process in self.0.values() { + if let Err(error) = process.signal(signal) + && first_error.is_none() + { + first_error = Some(error); + } + } + first_error.map_or(Ok(()), Err) + } + + pub(super) fn all_stopped(&self) -> bool { + self.0.values().all(Process::is_stopped) + } + + pub(super) fn len(&self) -> usize { + self.0.len() + } + + pub(super) fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +/// `(ppid, starttime)` from `/proc//stat`, parsed after the LAST `)`. +fn parse_stat(stat: &str) -> Option<(libc::pid_t, u64)> { + let close = stat.rfind(')')?; + let fields: Vec<&str> = stat[close + 1..].split_whitespace().collect(); + let ppid = fields.get(1)?.parse().ok()?; + let starttime = fields.get(19)?.parse().ok()?; + Some((ppid, starttime)) +} + +fn pidfd_open(pid: libc::pid_t) -> Result> { + let fd = unsafe { libc::syscall(libc::SYS_pidfd_open, pid, 0) }; + if fd < 0 { + let error = std::io::Error::last_os_error(); + if error.raw_os_error() == Some(libc::ESRCH) { + return Ok(None); + } + return Err(error.into()); + } + // SAFETY: pidfd_open returned a fresh owned descriptor. + Ok(Some(unsafe { OwnedFd::from_raw_fd(fd as i32) })) +} + /// The outcome of resolving a watcher `(dev, ino)` against `/proc//maps`. #[derive(Debug)] pub(super) enum Resolution { @@ -202,6 +381,36 @@ mod tests { assert_eq!(task_state("no paren here"), None); } + #[test] + fn parses_proc_stat_starttime_boundary() { + let mut stat = "123 (name) with ) chars) S 42".to_owned(); + for _ in 0..17 { + stat.push_str(" 1"); + } + stat.push_str(" 99"); + let (ppid, starttime) = parse_stat(&stat).unwrap(); + assert_eq!(ppid, 42); + assert_eq!(starttime, 99); + } + + #[test] + fn scan_tree_includes_descendant() { + let mut child = std::process::Command::new("sh") + .args(["-c", "sleep 10"]) + .spawn() + .unwrap(); + let tree = ProcessTree::scan(&[std::process::id() as libc::pid_t]).unwrap(); + assert!( + tree.0 + .values() + .any(|process| process.pid == child.id() as libc::pid_t) + ); + unsafe { + libc::kill(child.id() as libc::pid_t, libc::SIGKILL); + } + let _ = child.wait(); + } + #[test] fn parse_dev_reads_hex_major_minor() { assert_eq!(parse_dev("08:01"), Some((8, 1))); diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 6df8a27f6..828d3a568 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -1,4 +1,6 @@ use crate::ebpf::attach_worker::AttachWorker; +use crate::ebpf::pause::PressureWorker; +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 +35,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 +58,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() } } @@ -64,7 +76,7 @@ impl Default for TrackerOptions { pub struct Tracker { bpf: Arc>, - worker: Mutex>, + attach_worker: Mutex>, options: TrackerOptions, /// Number of native perf mapping records lost due to ring-buffer overflow. mapping_lost: Arc, @@ -97,7 +109,7 @@ impl Tracker { } let bpf = Arc::new(Mutex::new(bpf)); - let worker = if options.allocators { + let attach_worker = if options.allocators { Some(AttachWorker::start(bpf.clone())?) } else { None @@ -105,7 +117,7 @@ impl Tracker { Ok(Self { bpf, - worker: Mutex::new(worker), + attach_worker: Mutex::new(attach_worker), options, mapping_lost: Arc::new(AtomicU64::new(0)), }) @@ -121,6 +133,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 { @@ -131,7 +144,7 @@ impl Tracker { let pid = child.id() as i32; let setup = (|| -> Result<_> { - match self.worker.lock().as_ref() { + match self.attach_worker.lock().as_ref() { Some(worker) => worker.set_root_pid(pid), // No watcher to arm means exec mappings would be missed. None if self.options.allocators => bail!("tracker already finished"), @@ -143,9 +156,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())) @@ -161,6 +177,16 @@ impl Tracker { } }; + let mut pollers = vec![]; + if let Some(worker) = self.attach_worker.lock().as_ref() { + pollers.push(worker.handle()); + } + pollers.push(poller.handle()); + if let Some(stack_poller) = &stack_poller { + pollers.push(stack_poller.handle()); + } + let pressure = PressureWorker::start(self.bpf.clone(), pollers); + if let Err(error) = resume(pid) { kill_and_wait(&mut child); return Err(error); @@ -172,8 +198,10 @@ impl Tracker { poller, stack_poller, perf_mapping_poller, + pressure, )) } + /// 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 +215,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. @@ -209,7 +242,7 @@ impl Tracker { /// Stop the attach worker, if any, and surface any fatal error it recorded, /// including missed exec mappings (incomplete allocator coverage). pub fn finish(&self) -> Result<()> { - match self.worker.lock().take() { + match self.attach_worker.lock().take() { Some(worker) => worker.finish(), None => Ok(()), } diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index 7971b8c3a..3a4303c90 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. @@ -146,27 +149,35 @@ fn track_command( warn!("Failed to disable tracking: {e:#}"); } - // Dropping the session drops the event poller, which does a final drain of + // Finishing 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"); - drop(session); + let pressure = session.finish(); 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,14 @@ fn track_command( // kernel would close every link fd serially during exit. tracker.detach(); + let status = status?; + total?; + finish?; + pressure?; + 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/src/session.rs b/crates/memtrack/src/session.rs index 0fdcb8b01..7de7b42ef 100644 --- a/crates/memtrack/src/session.rs +++ b/crates/memtrack/src/session.rs @@ -1,3 +1,4 @@ +use crate::ebpf::pause::PressureWorker; use crate::ebpf::poller::{RingBufferPoller, ThreadedRingBufferPoller}; use crate::perf_mappings::PerfMappingPoller; use crate::prelude::*; @@ -8,6 +9,8 @@ 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 { + // Joined before the pollers it drains are dropped. + pressure: PressureWorker, child: Child, events: Option>>, @@ -28,8 +31,10 @@ impl Session { poller: RingBufferPoller, stack_poller: Option, perf_mapping_poller: Option, + pressure: PressureWorker, ) -> Self { Self { + pressure, child, events: Some(events), _poller: poller, @@ -51,4 +56,9 @@ impl Session { pub fn wait(&mut self) -> Result { Ok(self.child.wait()?) } + + /// Stop pressure recovery and surface its failure, then tear down the pipeline. + pub fn finish(self) -> Result<()> { + self.pressure.finish() + } } 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/tests/pressure_tests.rs b/crates/memtrack/tests/pressure_tests.rs new file mode 100644 index 000000000..fce66204b --- /dev/null +++ b/crates/memtrack/tests/pressure_tests.rs @@ -0,0 +1,75 @@ +//! Ring-pressure pause under a deliberately slow poller: the event ring fills +//! well before the next poll tick, so the tree must be paused and drained. + +mod shared; + +use memtrack::{Tracker, TrackerOptions}; +use std::process::Command; +use std::time::{Duration, Instant}; +use tempfile::TempDir; + +const THREADS: &str = "8"; +const ITERATIONS: &str = "400000"; +/// Long enough that the ring's 75% watermark is crossed between two polls. +const SLOW_POLL_MS: u64 = 1000; + +struct Run { + wall: Duration, + dropped: u64, +} + +fn run_storm(binary: &std::path::Path, options: TrackerOptions) -> anyhow::Result { + let tracker = Tracker::with_options(options)?; + tracker.enable_tracking()?; + + let mut command = Command::new(binary); + command.args([THREADS, ITERATIONS]); + + 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}"); + + session.finish()?; + 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, TrackerOptions::builder().build())?; + + eprintln!("-- slow poller"); + let blocked = run_storm( + &binary, + TrackerOptions::builder() + .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(()) +}