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

@@ -3,6 +3,7 @@
 
 #include "event.h"
 #include "utils/map_helpers.h"
+#include "utils/process_stop.h"
 #include "utils/process_tracking.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

@@ -4,6 +4,7 @@
 #include "event.h"
 #include "utils/event_helpers.h"
 #include "utils/mm_ownership.h"
+#include "utils/process_stop.h"
 #include "utils/process_tracking.h"
 
 /* FORK lets userland seed a child's RSS from its parent at fork time: the
@@ -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.
@@ -28,6 +29,10 @@ struct {
 /* A separate ring keeps allocation events fixed-size. */
 BPF_RINGBUF(stacks, 512 * 1024 * 1024);
 BPF_LRU_HASH_MAP(seen_stack_hashes, __u64, __u8, 262144);
+/* seen_stack_hashes is a hash set: only keys matter, every value is this
+ * marker. It lives in .rodata because map helpers reject a key and value that
+ * both point into the same ring reservation. */
+static const __u8 seen_stack_marker = 1;
 BPF_HASH_MAP(pending_stack_hash, __u64, __u64, 10000);
 BPF_ARRAY_MAP(stack_counters, __u64, MEMTRACK_STACK_COUNTER_COUNT);
 
@@ -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,38 @@
+#ifndef __PRESSURE_BPF_H__
+#define __PRESSURE_BPF_H__
+
+#include <bpf/bpf_helpers.h>
+
+#include "map_helpers.h"
+#include "process_stop.h"
+#include "process_tracking.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. */
+
+/* Stop producers once this percentage of the ring holds unconsumed data. */
+#define MEMTRACK_PRESSURE_WATERMARK_PERCENT 75
+
+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 * 100 >= size * MEMTRACK_PRESSURE_WATERMARK_PERCENT;
+}
+
+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/process_stop.h

@@ -0,0 +1,49 @@
+#ifndef __PROCESS_STOP_H__
+#define __PROCESS_STOP_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);
+/* Stops that could not be recorded because the map was full; userspace warns. */
+BPF_ARRAY_MAP(stop_record_failed, __u64, 1);
+
+/* 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) {
+        return;
+    }
+
+    __u32 zero = 0;
+    __u64* failed = bpf_map_lookup_elem(&stop_record_failed, &zero);
+    if (failed) {
+        __sync_fetch_and_add(failed, 1);
+    }
+    /* Unrecorded, so nothing would resume it for this reason. Keep it stopped
+     * if the other reason holds it: that release will resume it. */
+    if (bpf_map_lookup_elem(&pressure_stopped, &tgid) ||
+        bpf_map_lookup_elem(&attach_stopped, &tgid)) {
+        return;
+    }
+    bpf_send_signal(MEMTRACK_SIGCONT);
+}
+
+#endif /* __PROCESS_STOP_H__ */