fix(memtrack): drop rss_stat samples from recycled mm_structs

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

@@ -5,39 +5,31 @@
 #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, its mm when seeded, last
+ * in-context size}. An external (curr==0) update may only lower the counter, and
+ * only while pid_mm still binds the owner to the seeded mm: mm_id is a hash, so
+ * without both guards a stale reclaim read, a hash collision, or a recycled
+ * mm_struct slab slot could invent a peak. */
 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.
+ * pointer. pid_mm is the inverse, letting exec and exit remove an entry by
+ * value; attribution requires both to agree, so a stale mm fails closed.
  *
- * 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.
- *
- * 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 leave exec
+ * and exit unable to remove the forward entry. */
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
     __uint(max_entries, 10240);
@@ -46,6 +38,45 @@ struct {
 } mm_owner SEC(".maps");
 BPF_HASH_MAP(pid_mm, __u32, __u64, 10240);
 
+/* Rebind pid's address space; mm == 0 unbinds at exit. The mm_owner entry is only
+ * dropped while it still names pid: a live CLONE_VM sibling shares the mm and must
+ * keep its registration. */
+static __always_inline void set_pid_mm(__u32 pid, __u64 mm) {
+    __u64* cur = bpf_map_lookup_elem(&pid_mm, &pid);
+    if (cur && *cur == mm) {
+        return;
+    }
+    if (cur) {
+        __u32* owner = bpf_map_lookup_elem(&mm_owner, cur);
+        if (owner && *owner == pid) {
+            bpf_map_delete_elem(&mm_owner, cur);
+        }
+    }
+    if (mm) {
+        bpf_map_update_elem(&pid_mm, &pid, &mm, BPF_ANY);
+    } else {
+        bpf_map_delete_elem(&pid_mm, &pid);
+    }
+}
+
+/* Claim mm for pid without stealing from a live owner: CLONE_VM siblings share the
+ * mm, and overwriting would let the child's exec-time cleanup delete the entry out
+ * from under the still-live parent. */
+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);
+    }
+}
+
 #define FOLIO_MAPPING_ANON 0x1UL
 
 const volatile __u32 page_shift = 12;
@@ -73,10 +104,15 @@ 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 == 1 means current->mm is the counter's mm. pid_mm is maintained
+         * here too because the rmap hooks may not be attached. */
+        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);
+        set_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 +123,10 @@ int tracepoint_rss_stat(struct trace_event_raw_rss_stat* ctx) {
         if (cur == owner) {
             return 0;
         }
+        __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 +216,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);
+        set_pid_mm(pid, mm);
         owner = pid;
     } else {
         /* Foreign actor (task->mm != mm, including kthreads whose task->mm is NULL):
@@ -202,6 +231,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,39 +342,20 @@ 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. */
-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);
-        }
-    }
-    bpf_map_delete_elem(&pid_mm, &pid);
-}
-
 SEC("tracepoint/sched/sched_process_exec")
 int tracepoint_sched_process_exec(void* ctx) {
     __u32 pid = bpf_get_current_pid_tgid() >> 32;
     if (!is_tracked(pid)) {
         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. */
-    drop_mm_ownership(pid);
-
+    /* SUBMIT_EVENT_AS returns, so the rebind must precede it. */
     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);
-        bpf_map_update_elem(&pid_mm, &pid, &new_mm, BPF_ANY);
+        mm_owner_take(new_mm, pid);
     }
+    set_pid_mm(pid, new_mm);
 
     SUBMIT_EVENT_AS(pid, EVENT_TYPE_EXEC, {});
 }
@@ -374,7 +390,7 @@ int tracepoint_sched_process_exit(void* ctx) {
 
     /* Drop the ownership mapping so foreign actors stop attributing to a pid
      * the kernel may reuse. */
-    drop_mm_ownership(pid);
+    set_pid_mm(pid, 0);
 
     SUBMIT_EVENT_AS(pid, EVENT_TYPE_EXIT, {});
 }

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

@@ -67,6 +67,50 @@ impl MemtrackBpf {
             "dropped_events",
         )
     }
+
+    pub fn ownership_maps(&self) -> Result<OwnershipMaps> {
+        let mm_owner = entries(with_skel!(self, skel => &skel.maps.mm_owner))?;
+        let pid_mm = entries(with_skel!(self, skel => &skel.maps.pid_mm))?;
+        Ok(OwnershipMaps {
+            mm_owner: mm_owner
+                .into_iter()
+                .map(|(mm, pid)| (mm, pid as u32))
+                .collect(),
+            pid_mm: pid_mm
+                .into_iter()
+                .map(|(pid, mm)| (pid as u32, mm))
+                .collect(),
+        })
+    }
+}
+
+/// Live ownership bindings: `mm_owner` (`mm_struct` pointer -> owning pid) and its
+/// inverse `pid_mm`, which foreign-actor rmap attribution validates against.
+pub struct OwnershipMaps {
+    pub mm_owner: Vec<(u64, u32)>,
+    pub pid_mm: Vec<(u32, u64)>,
+}
+
+/// 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 entries(map: &impl MapCore) -> Result<Vec<(u64, u64)>> {
+    let mut entries = Vec::new();
+    for key in map.keys() {
+        let value = map
+            .lookup(&key, libbpf_rs::MapFlags::ANY)
+            .context("Failed to read map entry")?;
+        if let Some(value) = value {
+            entries.push((le(&key), le(&value)));
+        }
+    }
+    Ok(entries)
+}
+
+fn le(bytes: &[u8]) -> u64 {
+    bytes
+        .iter()
+        .rev()
+        .fold(0, |acc, &b| acc << 8 | u64::from(b))
 }
 
 /// Read slot 0 of a single-entry `__u64` array map.

crates/memtrack/src/ebpf/memtrack/mod.rs

@@ -22,6 +22,8 @@ mod allocator;
 mod maps;
 mod tracking;
 
+pub use maps::OwnershipMaps;
+
 use crate::bpf_token::has_delegated_bpf_token;
 
 /// Which attach mechanism a loaded skeleton uses for its uprobes. See

crates/memtrack/src/ebpf/mod.rs

@@ -6,5 +6,7 @@ mod proc_fs;
 mod spawn;
 mod tracker;
 
-pub use memtrack::{BpfVariant, MemtrackBpf, ResolvedSymbols, resolve_symbol_offsets};
+pub use memtrack::{
+    BpfVariant, MemtrackBpf, OwnershipMaps, ResolvedSymbols, resolve_symbol_offsets,
+};
 pub use tracker::{Tracker, TrackerOptions};

crates/memtrack/src/ebpf/tracker.rs

@@ -1,6 +1,6 @@
 use crate::ebpf::attach_worker::AttachWorker;
 use crate::ebpf::spawn::{resume, spawn_stopped, wrap_stopped};
-use crate::ebpf::{BpfVariant, MemtrackBpf};
+use crate::ebpf::{BpfVariant, MemtrackBpf, OwnershipMaps};
 use crate::prelude::*;
 use crate::session::Session;
 use parking_lot::Mutex;
@@ -135,6 +135,11 @@ impl Tracker {
         self.bpf.lock().dropped_events_count()
     }
 
+    /// Only meaningful while the BPF object is alive; teardown frees the maps.
+    pub fn ownership_maps(&self) -> Result<OwnershipMaps> {
+        self.bpf.lock().ownership_maps()
+    }
+
     /// 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/file_region.h

@@ -0,0 +1,55 @@
+#ifndef FILE_REGION_H
+#define FILE_REGION_H
+
+#include <fcntl.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/mman.h>
+#include <sys/syscall.h>
+#include <sys/uio.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+/* Maps len bytes of a private file mapping and faults every page in, so the
+ * pages land in the caller's RSS (MM_FILEPAGES) in-context and seed ownership.
+ * Returns the mapping, or NULL.
+ *
+ * The data file is derived from base rather than placed in /tmp, which is tmpfs
+ * on Ubuntu >= 25.04 and would account mapped file pages as shmem, not file. */
+static void* map_and_fault_file(const char* base, size_t len) {
+    char path[4096];
+    snprintf(path, sizeof(path), "%s.data-XXXXXX", base);
+    int fd = mkstemp(path);
+    if (fd < 0) return NULL;
+    unlink(path);
+    if (ftruncate(fd, len) != 0) return NULL;
+
+    void* mem = mmap(NULL, len, PROT_READ, MAP_PRIVATE, fd, 0);
+    if (mem == MAP_FAILED) return NULL;
+
+    /* volatile so the reads are not elided. 4096 is the fault stride: a smaller
+     * one would fault the same page twice, a larger one would skip pages. */
+    volatile char sink = 0;
+    for (size_t i = 0; i < len; i += 4096) sink ^= ((volatile char*)mem)[i];
+    (void)sink;
+    return mem;
+}
+
+/* Pages out [mem, mem+len) of the calling process from a forked child, and
+ * waits for it. The reclaim must run in another task's context to exercise
+ * foreign-actor attribution. Returns 0 on success. */
+static int reclaim_from_child(void* mem, size_t len) {
+    pid_t pid = fork();
+    if (pid < 0) return 1;
+    if (pid == 0) {
+        int pidfd = syscall(SYS_pidfd_open, getppid(), 0);
+        if (pidfd < 0) _exit(1);
+        struct iovec iov = {.iov_base = mem, .iov_len = len};
+        syscall(SYS_process_madvise, pidfd, &iov, 1UL, MADV_PAGEOUT, 0UL);
+        _exit(0);
+    }
+    int status;
+    return waitpid(pid, &status, 0) < 0;
+}
+
+#endif