fix(memtrack): drop rss_stat samples from recycled mm_structs

crates/memtrack/src/ebpf/c/rss.bpf.h

@@ -5,39 +5,40 @@
 #include "utils/event_helpers.h"
 #include "utils/process_tracking.h"
 
-/* (rss_stat mm_id << 32 | member) -> {owning tgid, last in-context size}. Keyed per
- * counter so an external (curr==0) update is attributed only once that mm/member was
- * established in-context. An external event may only lower a counter: any size above
- * the last in-context value is dropped, so neither a stale/racing reclaim read nor an
- * mm_id hash collision with another task can invent a peak. LRU eviction + re-seeding
- * on the owner's next in-context event covers hash reuse, so no teardown hook needed. */
+/* (rss_stat mm_id << 32 | member) -> {owning tgid, that tgid's mm at seeding time,
+ * last in-context size}. Keyed per counter so an external (curr==0) update is
+ * attributed only once that mm/member was established in-context. An external event
+ * may only lower a counter: any size above the last in-context value is dropped, so
+ * neither a stale/racing reclaim read nor an mm_id hash collision with another task
+ * can invent a peak. mm_id is only a hash, so the pointer is stored alongside and
+ * revalidated against pid_mm on the external path: an entry whose mm the owner no
+ * longer holds describes a freed mm_struct whose slab slot (and therefore hash) has
+ * been recycled by an unrelated address space. */
 struct rss_owner {
     __u32 pid;
+    __u64 mm;
     __u64 size;
 };
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
     __uint(max_entries, 40960);
     __type(key, __u64);
     __type(value, struct rss_owner);
-} mm_to_pid SEC(".maps");
+} rss_counter_owner SEC(".maps");
 
 /* Foreign-actor rmap attribution: rmap events run by a task other than the mm's
  * owner (kswapd reclaim, another process's process_madvise, khugepaged, KSM,
  * uffd) carry no owning-pid context, so mm_owner recovers it from the mm_struct
  * pointer. pid_mm is the inverse, letting the exec and exit hooks remove an entry
  * by value.
  *
- * Lifecycle invariant: every mm_owner entry is removed when its process execs
- * (the old mm is freed mid-life) or when its thread group dies, whichever comes
- * first; LRU eviction is only a backstop. A stale entry surviving mm-pointer
- * reuse would misattribute another process's events, so ownership is only ever
- * registered from an in-context (task->mm == mm) event.
+ * Lifecycle invariant: mm_owner follows pid_mm. Registration does not steal a
+ * live CLONE_VM sibling's ownership. Moving to another mm or reaching group
+ * death removes the corresponding entry. Foreign attribution validates both
+ * maps so stale mm pointers fail closed.
  *
- * pid_mm is a plain hash on purpose: an LRU inverse could be evicted while its
- * forward twin stays lookup-hot, leaving exec/exit unable to remove the live
- * mm_owner entry. Like tracked_pids, its entries are bound to the process
- * lifecycle and removed at group death. */
+ * pid_mm must not use LRU eviction: losing the inverse binding would prevent
+ * exec and exit from removing the forward entry. */
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
     __uint(max_entries, 10240);
@@ -46,6 +47,48 @@ struct {
 } mm_owner SEC(".maps");
 BPF_HASH_MAP(pid_mm, __u32, __u64, 10240);
 
+/* Delete mm_owner[mm] only while it still names pid: the slot may have been
+ * re-registered by another process since. */
+static __always_inline void drop_mm_owner_entry(__u64 mm, __u32 pid) {
+    __u32* owner = bpf_map_lookup_elem(&mm_owner, &mm);
+    if (owner && *owner == pid) {
+        bpf_map_delete_elem(&mm_owner, &mm);
+    }
+}
+
+/* Register pid as mm's owner without stealing from a live one: CLONE_VM
+ * siblings (vfork/posix_spawn) share the mm, and overwriting would let the
+ * child's exec-time cleanup delete the entry out from under the still-live
+ * parent. Only a vacant slot or one whose owner no longer holds the mm (a
+ * recycled mm_struct) is claimed. */
+static __always_inline void mm_owner_take(__u64 mm, __u32 pid) {
+    __u32* reg = bpf_map_lookup_elem(&mm_owner, &mm);
+    if (!reg) {
+        bpf_map_update_elem(&mm_owner, &mm, &pid, BPF_ANY);
+        return;
+    }
+    if (*reg == pid) {
+        return;
+    }
+    __u64* reg_mm = bpf_map_lookup_elem(&pid_mm, reg);
+    if (!reg_mm || *reg_mm != mm) {
+        bpf_map_update_elem(&mm_owner, &mm, &pid, BPF_ANY);
+    }
+}
+
+/* Remove the previous forward binding when pid_mm changes. Exec-time faults
+ * may expose the new mm before the exec tracepoint runs. */
+static __always_inline void refresh_pid_mm(__u32 pid, __u64 mm) {
+    __u64* cur_mm = bpf_map_lookup_elem(&pid_mm, &pid);
+    if (cur_mm && *cur_mm == mm) {
+        return;
+    }
+    if (cur_mm) {
+        drop_mm_owner_entry(*cur_mm, pid);
+    }
+    bpf_map_update_elem(&pid_mm, &pid, &mm, BPF_ANY);
+}
+
 #define FOLIO_MAPPING_ANON 0x1UL
 
 const volatile __u32 page_shift = 12;
@@ -73,10 +116,17 @@ int tracepoint_rss_stat(struct trace_event_raw_rss_stat* ctx) {
             return 0;
         }
         owner = cur;
-        struct rss_owner state = {.pid = cur, .size = size};
-        bpf_map_update_elem(&mm_to_pid, &key, &state, BPF_ANY);
+        /* curr means current->mm is the mm the counter belongs to. Recording it
+         * (and keeping pid_mm current, which the rmap hooks may never do when
+         * only rss_stat is attached) is what lets the external path below tell a
+         * live mm from a recycled slab slot. */
+        struct task_struct* task = bpf_get_current_task_btf();
+        __u64 mm = (__u64)BPF_CORE_READ(task, mm);
+        struct rss_owner state = {.pid = cur, .mm = mm, .size = size};
+        bpf_map_update_elem(&rss_counter_owner, &key, &state, BPF_ANY);
+        refresh_pid_mm(cur, mm);
     } else {
-        struct rss_owner* found = bpf_map_lookup_elem(&mm_to_pid, &key);
+        struct rss_owner* found = bpf_map_lookup_elem(&rss_counter_owner, &key);
         if (!found) {
             return 0;
         }
@@ -87,6 +137,14 @@ int tracepoint_rss_stat(struct trace_event_raw_rss_stat* ctx) {
         if (cur == owner) {
             return 0;
         }
+        /* mm_id is a hash of a pointer the tracepoint never exposes, so the key
+         * survives the mm it was seeded from. Once the owner no longer holds that
+         * mm (it execed, or the entry predates a pid reuse), the hash now belongs
+         * to a recycled mm_struct and its counters describe another address space. */
+        __u64* owner_mm = bpf_map_lookup_elem(&pid_mm, &owner);
+        if (!owner_mm || *owner_mm != found->mm) {
+            return 0;
+        }
         /* An external actor may only lower a counter. A larger value is a stale
          * reclaim read or an mm_id hash collision with another task; dropping it
          * keeps the reconstructed peak identical to the in-context timeline. */
@@ -176,19 +234,8 @@ static __always_inline int submit_rmap(struct vm_area_struct* vma, __s32 member,
             return 0;
         }
 
-        /* Register ownership so foreign actors can later attribute to this pid.
-         * Both maps are validated (not just written) on every in-context event:
-         * the lookups keep the hot path cheap AND keep both entries LRU-fresh,
-         * since pid_mm is otherwise never read until exec/exit and could be
-         * evicted independently of its still-hot mm_owner twin. */
-        __u32* reg = bpf_map_lookup_elem(&mm_owner, &mm);
-        if (!reg || *reg != pid) {
-            bpf_map_update_elem(&mm_owner, &mm, &pid, BPF_ANY);
-        }
-        __u64* cur_mm = bpf_map_lookup_elem(&pid_mm, &pid);
-        if (!cur_mm || *cur_mm != mm) {
-            bpf_map_update_elem(&pid_mm, &pid, &mm, BPF_ANY);
-        }
+        mm_owner_take(mm, pid);
+        refresh_pid_mm(pid, mm);
         owner = pid;
     } else {
         /* Foreign actor (task->mm != mm, including kthreads whose task->mm is NULL):
@@ -202,6 +249,12 @@ static __always_inline int submit_rmap(struct vm_area_struct* vma, __s32 member,
         if (!is_tracked(owner)) {
             return 0;
         }
+        /* An mm_struct address may be reused while a stale owner entry remains.
+         * Accept only the current inverse binding. */
+        __u64* owner_mm = bpf_map_lookup_elem(&pid_mm, &owner);
+        if (!owner_mm || *owner_mm != mm) {
+            return 0;
+        }
     }
 
     /* header.tid is stamped from the current task; for a foreign actor it
@@ -307,17 +360,12 @@ int tracepoint_task_newtask(struct trace_event_raw_task_newtask* ctx) {
     SUBMIT_EVENT_AS(child_pid, EVENT_TYPE_FORK, { e->data.fork.parent_pid = parent_pid; });
 }
 
-/* Remove pid's ownership registration. The mm_owner value is verified against
- * pid before deleting: a stale pid_mm entry (LRU eviction skew) could otherwise
- * point at an mm since re-registered by another process, and deleting that
- * would silence a live owner's foreign attribution. */
+/* Verify the forward binding before deletion because the mm may have been
+ * re-registered by another process. */
 static __always_inline void drop_mm_ownership(__u32 pid) {
     __u64* mm = bpf_map_lookup_elem(&pid_mm, &pid);
     if (mm) {
-        __u32* owner = bpf_map_lookup_elem(&mm_owner, mm);
-        if (owner && *owner == pid) {
-            bpf_map_delete_elem(&mm_owner, mm);
-        }
+        drop_mm_owner_entry(*mm, pid);
     }
     bpf_map_delete_elem(&pid_mm, &pid);
 }
@@ -329,15 +377,13 @@ int tracepoint_sched_process_exec(void* ctx) {
         return 0;
     }
 
-    /* Maintain ownership before submitting (SUBMIT_EVENT_AS returns from the
-     * function). Exec frees the old mm long before group death, so the stale
-     * pointer must be dropped here or a reused mm_struct would be misattributed. */
+    /* SUBMIT_EVENT_AS returns, so ownership cleanup must precede it. */
     drop_mm_ownership(pid);
 
     struct task_struct* task = bpf_get_current_task_btf();
     __u64 new_mm = (__u64)BPF_CORE_READ(task, mm);
     if (new_mm) {
-        bpf_map_update_elem(&mm_owner, &new_mm, &pid, BPF_ANY);
+        mm_owner_take(new_mm, pid);
         bpf_map_update_elem(&pid_mm, &pid, &new_mm, BPF_ANY);
     }
 

crates/memtrack/src/ebpf/memtrack/maps.rs

@@ -95,4 +95,52 @@ impl MemtrackBpf {
             .map_err(|_| anyhow!("dropped_events value has unexpected size"))?;
         Ok(u64::from_le_bytes(bytes))
     }
+
+    /// Live `mm_owner` entries as `(mm_struct pointer, owning pid)`.
+    pub fn mm_owner_entries(&self) -> Result<Vec<(u64, u32)>> {
+        let entries = snapshot_entries::<8, 4>(&self.skel.maps.mm_owner, "mm_owner")?;
+        Ok(entries
+            .into_iter()
+            .map(|(mm, pid)| (u64::from_le_bytes(mm), u32::from_le_bytes(pid)))
+            .collect())
+    }
+
+    /// Live `pid_mm` entries as `(pid, mm_struct pointer)`, the inverse binding
+    /// `mm_owner` is validated against.
+    pub fn pid_mm_entries(&self) -> Result<Vec<(u32, u64)>> {
+        let entries = snapshot_entries::<4, 8>(&self.skel.maps.pid_mm, "pid_mm")?;
+        Ok(entries
+            .into_iter()
+            .map(|(pid, mm)| (u32::from_le_bytes(pid), u64::from_le_bytes(mm)))
+            .collect())
+    }
+}
+
+/// Snapshot a hash map's live entries as fixed-size key/value byte pairs.
+///
+/// Iteration is `BPF_MAP_GET_NEXT_KEY` followed by a separate lookup, so it is
+/// not atomic: a key deleted in between is skipped rather than reported.
+fn snapshot_entries<const K: usize, const V: usize>(
+    map: &impl MapCore,
+    name: &str,
+) -> Result<Vec<([u8; K], [u8; V])>> {
+    let mut entries = Vec::new();
+    for key in map.keys() {
+        let key: [u8; K] = key
+            .as_slice()
+            .try_into()
+            .map_err(|_| anyhow!("{name} key has unexpected size"))?;
+        let Some(value) = map
+            .lookup(&key, libbpf_rs::MapFlags::ANY)
+            .with_context(|| format!("Failed to read {name} entry"))?
+        else {
+            continue;
+        };
+        let value: [u8; V] = value
+            .as_slice()
+            .try_into()
+            .map_err(|_| anyhow!("{name} value has unexpected size"))?;
+        entries.push((key, value));
+    }
+    Ok(entries)
 }

crates/memtrack/src/ebpf/tracker.rs

@@ -103,6 +103,18 @@ impl Tracker {
         self.bpf.lock().dropped_events_count()
     }
 
+    /// Live `mm_owner` entries (`mm_struct` pointer -> owning pid). Only
+    /// meaningful while the BPF object is alive; teardown frees the map.
+    pub fn mm_owner_entries(&self) -> Result<Vec<(u64, u32)>> {
+        self.bpf.lock().mm_owner_entries()
+    }
+
+    /// Live `pid_mm` entries (pid -> `mm_struct` pointer). Same lifetime caveat
+    /// as [`Tracker::mm_owner_entries`].
+    pub fn pid_mm_entries(&self) -> Result<Vec<(u32, u64)>> {
+        self.bpf.lock().pid_mm_entries()
+    }
+
     /// Stop the attach worker, if any, and surface any fatal error it recorded,
     /// including missed exec mappings (incomplete allocator coverage). A tracker
     /// without an allocator watcher has no worker, so this is a no-op.

crates/memtrack/testdata/rss/exec_churn.c

@@ -0,0 +1,56 @@
+/* Repeatedly populate and replace child address spaces.
+ *
+ *   parent ── fork ──▶ child faults REGION_MIB of anonymous memory
+ *          │             │ execv(trivial image) replaces the populated mm
+ *          │             ▼
+ *          │           exits immediately
+ *          └── waitpid
+ *
+ * After the final child exits, no ownership-map entry may name these processes.
+ */
+#define _GNU_SOURCE
+#include <string.h>
+#include <sys/mman.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+#define ITERATIONS 8
+#define REGION_MIB 4
+#define EXIT_MARKER "exec-churn-exit"
+
+/* Exec into the cheapest image available. Only dropping a populated mm matters,
+ * so where the pid lands is irrelevant: distributions without /bin/true (NixOS)
+ * fall back to re-execing this fixture with a marker that returns at once. */
+static void exec_trivial(const char* self) {
+    static const char* const candidates[] = {"/bin/true", "/usr/bin/true"};
+    for (unsigned i = 0; i < sizeof(candidates) / sizeof(*candidates); i++) {
+        if (access(candidates[i], X_OK) == 0) {
+            char* args[] = {(char*)candidates[i], NULL};
+            execv(candidates[i], args);
+        }
+    }
+    char* args[] = {(char*)self, (char*)EXIT_MARKER, NULL};
+    execv(self, args);
+}
+
+int main(int argc, char** argv) {
+    if (argc > 1 && strcmp(argv[1], EXIT_MARKER) == 0) return 0;
+
+    size_t len = (size_t)REGION_MIB * 1024 * 1024;
+    for (int i = 0; i < ITERATIONS; i++) {
+        pid_t pid = fork();
+        if (pid < 0) return 1;
+        if (pid == 0) {
+            void* region =
+                mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+            if (region == MAP_FAILED) _exit(1);
+            memset(region, 0x42, len);
+            exec_trivial(argv[0]);
+            _exit(1);
+        }
+        int status;
+        if (waitpid(pid, &status, 0) < 0) return 1;
+        if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) return 1;
+    }
+    return 0;
+}