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::PressureGate;
+use crate::ebpf::poller::{DrainControl, RingBufferPoller};
 use crate::prelude::*;
 use parking_lot::Mutex;
 use std::collections::HashSet;
@@ -14,19 +15,28 @@ 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 POLL_INTERVAL_MS: u64 = 1;
 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);
+struct DrainRequest {
+    timeout: Duration,
+    ack: mpsc::Sender<Result<()>>,
+}
+
+/// SIGCONTs `pid` on drop unless the pressure monitor currently owns the
+/// whole-tree stop. The monitor then resumes the pid in its coordinated sweep.
+struct ContGuard {
+    pid: i32,
+    pressure: PressureGate,
+}
 
 impl Drop for ContGuard {
     fn drop(&mut self) {
-        // SAFETY: kill with SIGCONT has no memory effects; errors (e.g. the
-        // process already exited) are intentionally ignored.
-        unsafe {
-            libc::kill(self.0, libc::SIGCONT);
+        self.pressure.release_attach(self.pid);
+        match self.pressure.try_resume(self.pid) {
+            Ok(false) => debug!("deferring SIGCONT for pid {}", self.pid),
+            Ok(true) => {}
+            Err(error) => debug!("failed to resume pid {}: {error:#}", self.pid),
         }
     }
 }
@@ -40,24 +50,42 @@ pub(crate) struct AttachWorker {
     fatal: Arc<Mutex<Option<String>>>,
     root_pid: Arc<AtomicI32>,
     bpf: Arc<Mutex<MemtrackBpf>>,
+    pressure_drain: DrainControl,
 }
 
 impl AttachWorker {
-    pub(crate) fn start(bpf: Arc<Mutex<MemtrackBpf>>) -> Result<Self> {
+    pub(crate) fn start(bpf: Arc<Mutex<MemtrackBpf>>, pressure: PressureGate) -> Result<Self> {
         let shutdown = Arc::new(AtomicBool::new(false));
         let fatal = Arc::new(Mutex::new(None));
         let root_pid = Arc::new(AtomicI32::new(0));
-
         let (tx, rx) = mpsc::channel();
+        let (drain_tx, drain_rx) = mpsc::channel();
         let poller = bpf.lock().poll_attach_with_channel(POLL_INTERVAL_MS, tx)?;
-
+        let drain = poller.drain_control();
+        let pressure_drain = DrainControl::new(move |timeout| {
+            let (ack_tx, ack_rx) = mpsc::channel();
+            drain_tx
+                .send(DrainRequest {
+                    timeout,
+                    ack: ack_tx,
+                })
+                .context("attach worker is gone")?;
+            match ack_rx.recv_timeout(timeout) {
+                Ok(Ok(())) => Ok(()),
+                Ok(Err(error)) => Err(error),
+                Err(error) => Err(error).context("attach worker did not reach barrier"),
+            }
+        });
         let worker = Worker {
             poller,
             rx,
+            drain_rx,
             bpf: bpf.clone(),
             shutdown: shutdown.clone(),
             fatal: fatal.clone(),
             root_pid: root_pid.clone(),
+            pressure,
+            drain,
         };
 
         let handle = std::thread::spawn(move || worker.run());
@@ -68,9 +96,14 @@ impl AttachWorker {
             fatal,
             root_pid,
             bpf,
+            pressure_drain,
         })
     }
 
+    pub(crate) fn drain_control(&self) -> DrainControl {
+        self.pressure_drain.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);
@@ -106,10 +139,6 @@ impl AttachWorker {
         Ok(())
     }
 }
-
-/// Joining the worker on drop releases its `MemtrackBpf` clone so the probe
-/// links detach promptly; otherwise the thread keeps spinning on its recv loop
-/// and the links fall back to a slow serial close at process exit.
 impl Drop for AttachWorker {
     fn drop(&mut self) {
         self.shutdown.store(true, Ordering::SeqCst);
@@ -122,10 +151,13 @@ impl Drop for AttachWorker {
 struct Worker {
     poller: RingBufferPoller,
     rx: mpsc::Receiver<Vec<AttachRequest>>,
+    drain_rx: mpsc::Receiver<DrainRequest>,
     bpf: Arc<Mutex<MemtrackBpf>>,
     shutdown: Arc<AtomicBool>,
     fatal: Arc<Mutex<Option<String>>>,
     root_pid: Arc<AtomicI32>,
+    pressure: PressureGate,
+    drain: DrainControl,
 }
 
 impl Worker {
@@ -152,7 +184,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<AttachRequest> = first
@@ -164,7 +202,29 @@ 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() {
+            if let Err(error) = self.drain.drain(request.timeout) {
+                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),
@@ -191,8 +251,12 @@ impl Worker {
 
             for pid in new_pids {
                 stopped.insert(pid);
-                guards.push(ContGuard(pid as i32));
-                wait_all_stopped(pid, STOP_DEADLINE)?;
+                self.pressure.register_attach(pid as i32);
+                guards.push(ContGuard {
+                    pid: pid as i32,
+                    pressure: self.pressure.clone(),
+                });
+                wait_all_stopped(pid, STOP_DEADLINE, true)?;
             }
 
             // Every producer is stopped, so a synchronous drain is complete.

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.
@@ -88,35 +89,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 +139,13 @@ 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;
+    header->_pad[0] = 1;
+    long gate_result =
+        bpf_map_update_elem(&seen_stack_hashes, &header->hash, &header->_pad[0], 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 +158,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 +172,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/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;                                                       \
     }
 

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

@@ -0,0 +1,51 @@
+#ifndef __PRESSURE_BPF_H__
+#define __PRESSURE_BPF_H__
+
+#include <bpf/bpf_helpers.h>
+
+#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
+
+const volatile __u8 blocking_enabled = 0;
+
+/* 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 (!blocking_enabled || !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__ */