feat(memtrack): pause producers under ring pressure

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<StoppedProcesses>,
+    pid: u32,
+}
 
-impl Drop for ContGuard {
+impl AttachHold {
+    fn new(stops: Arc<StoppedProcesses>, 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<Vec<AttachRequest>>,
     bpf: Arc<Mutex<MemtrackBpf>>,
+    stops: Arc<StoppedProcesses>,
     shutdown: Arc<AtomicBool>,
     fatal: Arc<Mutex<Option<String>>>,
     root_pid: Arc<AtomicI32>,
@@ -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<AttachRequest>,
         known: &mut HashSet<(u64, u64)>,
     ) -> Result<()> {
-        let mut guards: Vec<ContGuard> = Vec::new();
+        let mut holds: Vec<AttachHold> = Vec::new();
         let mut stopped: HashSet<u32> = 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.

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;
 }
 

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, {});
 }

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,35 +93,38 @@ 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.
      * Loop bound is checked against stack_copy_budget (a frozen rodata constant)
      * 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;
 }
 

crates/memtrack/src/ebpf/c/utils/pressure.bpf.h

@@ -0,0 +1,37 @@
+#ifndef __PRESSURE_BPF_H__
+#define __PRESSURE_BPF_H__
+
+#include <bpf/bpf_helpers.h>
+
+#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__ */

crates/memtrack/src/ebpf/c/utils/stopped.h

@@ -0,0 +1,35 @@
+#ifndef __STOPPED_H__
+#define __STOPPED_H__
+
+#include <bpf/bpf_helpers.h>
+
+#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__ */

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<StoppedProcesses> {
+        self.stopped.clone()
+    }
+
+    pub(super) fn open_stopped_processes(skel: &Skel) -> Result<StoppedProcesses> {
+        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<dyn Fn() + Send> {
+        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> {
         StackCaptureFailureStats::read(with_skel!(self, skel => &skel.maps.stack_counters))
     }