feat(memtrack): collect RSS via rss_stat and folio-rmap reconstruction

.github/workflows/ci.yml

@@ -85,22 +85,23 @@ jobs:
         run: cargo run -- exec -m walltime --skip-upload --warmup-time 0s --max-rounds 5 -- ls -la
 
   bpf-tests:
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
+        os: [ubuntu-22.04, ubuntu-latest, ubuntu-24.04-arm, ubuntu-26.04, ubuntu-26.04-arm]
         # Each memtrack integration test binary runs its cases serially
         # (eBPF tracker can't overlap with itself in one process), so we
         # shard at the test-binary level to parallelize across jobs.
-        test: [c_tests, cpp_tests, rust_tests, spawn_tests, dlopen_tests]
+        test: [c_tests, cpp_tests, rust_tests, spawn_tests, dlopen_tests, rss_tests]
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
           lfs: true
           submodules: true
       - uses: ./.github/actions/install-rust
         with:
-          cache-key: ${{ matrix.test }}
+          cache-key: ${{ matrix.os }}-${{ matrix.test }}
       - uses: ./.github/actions/install-bpf-deps
 
       - name: Install additional allocators
@@ -109,7 +110,18 @@ jobs:
       - name: Run tests
         env:
           RUST_LOG: debug
-        run: sudo -E $(which cargo) test --lib --test ${{ matrix.test }} -- --test-threads 1 --nocapture
+        # Ubuntu 26.04 ships sudo-rs, which ignores `-E`; pass the env the
+        # rustup shims and the test gate need through `env` instead.
+        run: |
+          sudo env \
+            "HOME=$HOME" \
+            "PATH=$PATH" \
+            "CARGO_HOME=${CARGO_HOME:-$HOME/.cargo}" \
+            "RUSTUP_HOME=${RUSTUP_HOME:-$HOME/.rustup}" \
+            "CARGO_INCREMENTAL=$CARGO_INCREMENTAL" \
+            "RUST_LOG=$RUST_LOG" \
+            "GITHUB_ACTIONS=$GITHUB_ACTIONS" \
+            $(which cargo) test --lib --test ${{ matrix.test }} -- --test-threads 1 --nocapture
         working-directory: crates/memtrack
 
       # Since we ran the tests with sudo, the build artifacts will have root ownership

crates/memtrack/.clang-format

@@ -4,3 +4,4 @@ BasedOnStyle:    Google
 PointerAlignment: Left
 ColumnLimit: 100
 IndentWidth: 4
+AllowShortFunctionsOnASingleLine: None

crates/memtrack/Cargo.toml

@@ -46,7 +46,7 @@ bindgen = "0.72"
 tempfile = { workspace = true }
 rstest = { workspace = true }
 test-log = { workspace = true }
-insta = { workspace = true }
+insta = { workspace = true, features = ["json", "redactions"] }
 test-with = { workspace = true }
 
 [package.metadata.dist]

crates/memtrack/src/ebpf/c/allocator.h

@@ -5,22 +5,24 @@
 #include "utils/map_helpers.h"
 #include "utils/process_tracking.h"
 
-#define UPROBE_ARG_RET(name, arg_expr, submit_block)                                      \
-    BPF_HASH_MAP(name##_arg, __u64, __u64, 10000);                                        \
-    SEC(UPROBE_SEC)                                                                       \
-    int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_arg, arg_expr); } \
-    SEC(URETPROBE_SEC)                                                                    \
-    int uretprobe_##name(struct pt_regs* ctx) {                                           \
-        __u64* arg_ptr = take_param(&name##_arg);                                         \
-        if (!arg_ptr) {                                                                   \
-            return 0;                                                                     \
-        }                                                                                 \
-        __u64 ret_val = PT_REGS_RC(ctx);                                                  \
-        if (ret_val == 0) {                                                               \
-            return 0;                                                                     \
-        }                                                                                 \
-        __u64 arg0 = *arg_ptr;                                                            \
-        submit_block;                                                                     \
+#define UPROBE_ARG_RET(name, arg_expr, submit_block) \
+    BPF_HASH_MAP(name##_arg, __u64, __u64, 10000);   \
+    SEC(UPROBE_SEC)                                  \
+    int uprobe_##name(struct pt_regs* ctx) {         \
+        return store_param(&name##_arg, arg_expr);   \
+    }                                                \
+    SEC(URETPROBE_SEC)                               \
+    int uretprobe_##name(struct pt_regs* ctx) {      \
+        __u64* arg_ptr = take_param(&name##_arg);    \
+        if (!arg_ptr) {                              \
+            return 0;                                \
+        }                                            \
+        __u64 ret_val = PT_REGS_RC(ctx);             \
+        if (ret_val == 0) {                          \
+            return 0;                                \
+        }                                            \
+        __u64 arg0 = *arg_ptr;                       \
+        submit_block;                                \
     }
 
 #define UPROBE_RET(name, arg_expr, submit_block) \

crates/memtrack/src/ebpf/c/event.h

@@ -9,6 +9,11 @@
 #define EVENT_TYPE_MMAP 6
 #define EVENT_TYPE_MUNMAP 7
 #define EVENT_TYPE_BRK 8
+#define EVENT_TYPE_FORK 9
+#define EVENT_TYPE_EXEC 10
+#define EVENT_TYPE_EXIT 11
+#define EVENT_TYPE_RSS 12
+#define EVENT_TYPE_RMAP 13
 
 /* Common header shared by all event types */
 struct event_header {
@@ -45,6 +50,22 @@ struct event {
             uint64_t addr; /* address of mapping */
             uint64_t size; /* size of mapping */
         } mmap;
+
+        /* Process lifecycle events (fork carries the parent; exec/exit have no payload) */
+        struct {
+            uint32_t parent_pid;
+        } fork;
+
+        struct {
+            int32_t member;
+            uint64_t size;
+        } rss;
+
+        struct {
+            int32_t member; /* MM_* counter index */
+            int64_t delta;
+            uint64_t addr;
+        } rmap;
     } data;
 };
 

crates/memtrack/src/ebpf/c/main.bpf.c

@@ -8,8 +8,13 @@
 #include "allocator.h"
 #include "attach.h"
 #include "event.h"
+#include "process_tracking.bpf.h"
+#include "rmap.bpf.h"
+#include "rss.bpf.h"
 #include "utils/event_helpers.h"
+#include "utils/folio.h"
 #include "utils/map_helpers.h"
+#include "utils/mm_ownership.h"
 #include "utils/process_tracking.h"
 
 char LICENSE[] SEC("license") = "GPL";

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

@@ -0,0 +1,93 @@
+#ifndef __PROCESS_TRACKING_BPF_H__
+#define __PROCESS_TRACKING_BPF_H__
+
+#include "event.h"
+#include "utils/event_helpers.h"
+#include "utils/mm_ownership.h"
+#include "utils/process_tracking.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
+ * rss_stat out of the child's context and anon COW faults are
+ * counter-neutral, so a child that only touches inherited memory never
+ * reports its RSS on its own. EXEC and EXIT mark the points where the
+ * address space is replaced or torn down, so userland resets to zero.
+ */
+
+/* tp_btf rather than a classic tracepoint: it attaches with
+ * BPF_RAW_TRACEPOINT_OPEN, which a token can delegate, and its task_struct
+ * arguments let the child's pid resolve in the tracker's namespace. */
+SEC("tp_btf/sched_process_fork")
+int BPF_PROG(tracepoint_sched_process_fork, struct task_struct* parent, struct task_struct* child) {
+    /* copy_process assigns tgid = current->tgid for CLONE_THREAD, pid otherwise.
+     * A tid registered here would never be removed: group death untracks only
+     * the tgid. */
+    if (BPF_CORE_READ(child, pid) != BPF_CORE_READ(child, tgid)) {
+        return 0;
+    }
+
+    __u32 parent_pid = current_tgid();
+    if (!is_tracked(parent_pid)) {
+        return 0;
+    }
+
+    __u32 child_pid = task_ns_tgid(child);
+    if (!child_pid) {
+        return 0;
+    }
+    track_child(child_pid, parent_pid);
+
+    SUBMIT_EVENT_AS(child_pid, EVENT_TYPE_FORK, { e->data.fork.parent_pid = parent_pid; });
+}
+
+SEC("tracepoint/sched/sched_process_exec")
+int tracepoint_sched_process_exec(void* ctx) {
+    __u32 pid = current_tgid();
+    if (!is_tracked(pid)) {
+        return 0;
+    }
+
+    /* 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) {
+        claim_mm_owner(new_mm, pid);
+    }
+    rebind_pid_mm(pid, new_mm);
+
+    SUBMIT_EVENT_AS(pid, EVENT_TYPE_EXEC, {});
+}
+
+SEC("tracepoint/sched/sched_process_exit")
+int tracepoint_sched_process_exit(void* ctx) {
+    __u32 pid = current_tgid();
+    if (!is_tracked(pid)) {
+        return 0;
+    }
+
+    /* EXIT marks the death of the whole thread group, not of one thread: the
+     * leader can pthread_exit while workers keep running, and the last thread
+     * to exit need not be the leader. do_exit decrements signal->live before
+     * this tracepoint fires, so live == 0 identifies the dying thread group's
+     * final exit — but concurrently exiting threads can BOTH read 0, so the
+     * untrack below arbitrates: only the task that wins it emits. */
+    struct task_struct* task = bpf_get_current_task_btf();
+    if (BPF_CORE_READ(task, signal, live.counter) != 0) {
+        return 0;
+    }
+
+    /* Untrack the pid before submitting: lifetime events are gated only on
+     * is_tracked, so a stale entry would keep streaming events if the kernel
+     * reuses the pid for an unrelated process. */
+    if (!untrack_pid(pid)) {
+        return 0;
+    }
+
+    /* Drop the ownership mapping so foreign actors stop attributing to a pid
+     * the kernel may reuse. */
+    rebind_pid_mm(pid, 0);
+
+    SUBMIT_EVENT_AS(pid, EVENT_TYPE_EXIT, {});
+}
+
+#endif /* __PROCESS_TRACKING_BPF_H__ */

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

@@ -0,0 +1,113 @@
+#ifndef __RMAP_BPF_H__
+#define __RMAP_BPF_H__
+
+#include "event.h"
+#include "utils/event_helpers.h"
+#include "utils/folio.h"
+#include "utils/mm_ownership.h"
+#include "utils/process_tracking.h"
+
+static __always_inline int submit_rmap(struct vm_area_struct* vma, __s32 member, __s64 delta,
+                                       __u64 addr) {
+    __u64 mm = (__u64)BPF_CORE_READ(vma, vm_mm);
+    struct task_struct* task = bpf_get_current_task_btf();
+    __u32 pid = current_tgid();
+    __u32 owner;
+
+    if ((__u64)BPF_CORE_READ(task, mm) == mm) {
+        if (!is_tracked(pid)) {
+            return 0;
+        }
+
+        claim_mm_owner(mm, pid);
+        rebind_pid_mm(pid, mm);
+        owner = pid;
+    } else {
+        /* Foreign actor (task->mm != mm, including kthreads whose task->mm is NULL):
+         * recover the owner from the in-context registration. Fail toward dropping
+         * the event on any uncertainty about ownership. */
+        __u32* found = bpf_map_lookup_elem(&owner_by_mm, &mm);
+        if (!found) {
+            return 0;
+        }
+        owner = *found;
+        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(&mm_by_pid, &owner);
+        if (!owner_mm || *owner_mm != mm) {
+            return 0;
+        }
+    }
+
+    /* header.tid is stamped from the current task; for a foreign actor it
+     * identifies the performer, not the owning pid. */
+    SUBMIT_EVENT_AS(owner, EVENT_TYPE_RMAP, {
+        e->data.rmap.member = member;
+        e->data.rmap.delta = delta;
+        e->data.rmap.addr = addr;
+    });
+}
+
+SEC("fentry/folio_add_new_anon_rmap")
+int BPF_PROG(fentry_folio_add_new_anon_rmap, struct folio* folio, struct vm_area_struct* vma,
+             unsigned long address) {
+    return submit_rmap(vma, MM_ANONPAGES, (__s64)folio_nr_pages_est(folio), address);
+}
+
+SEC("fentry/folio_add_anon_rmap_ptes")
+int BPF_PROG(fentry_folio_add_anon_rmap_ptes, struct folio* folio, struct page* page, int nr_pages,
+             struct vm_area_struct* vma, unsigned long address) {
+    return submit_rmap(vma, MM_ANONPAGES, (__s64)nr_pages, address);
+}
+
+SEC("fentry/folio_add_anon_rmap_pmd")
+int BPF_PROG(fentry_folio_add_anon_rmap_pmd, struct folio* folio, struct page* page,
+             struct vm_area_struct* vma, unsigned long address) {
+    return submit_rmap(vma, MM_ANONPAGES, (__s64)folio_nr_pages_est(folio), address);
+}
+
+static __always_inline int submit_file_rmap(struct folio* folio, struct page* page,
+                                            struct vm_area_struct* vma, __s64 delta) {
+    return submit_rmap(vma, folio_mm_counter(folio), delta, folio_page_address(folio, page, vma));
+}
+
+SEC("fentry/folio_add_file_rmap_ptes")
+int BPF_PROG(fentry_folio_add_file_rmap_ptes, struct folio* folio, struct page* page, int nr_pages,
+             struct vm_area_struct* vma) {
+    return submit_file_rmap(folio, page, vma, (__s64)nr_pages);
+}
+
+SEC("fentry/folio_add_file_rmap_pmd")
+int BPF_PROG(fentry_folio_add_file_rmap_pmd, struct folio* folio, struct page* page,
+             struct vm_area_struct* vma) {
+    return submit_file_rmap(folio, page, vma, (__s64)folio_nr_pages_est(folio));
+}
+
+SEC("fentry/folio_add_file_rmap_pud")
+int BPF_PROG(fentry_folio_add_file_rmap_pud, struct folio* folio, struct page* page,
+             struct vm_area_struct* vma) {
+    return submit_file_rmap(folio, page, vma, (__s64)folio_nr_pages_est(folio));
+}
+
+SEC("fentry/folio_remove_rmap_ptes")
+int BPF_PROG(fentry_folio_remove_rmap_ptes, struct folio* folio, struct page* page, int nr_pages,
+             struct vm_area_struct* vma) {
+    return submit_file_rmap(folio, page, vma, -(__s64)nr_pages);
+}
+
+SEC("fentry/folio_remove_rmap_pmd")
+int BPF_PROG(fentry_folio_remove_rmap_pmd, struct folio* folio, struct page* page,
+             struct vm_area_struct* vma) {
+    return submit_file_rmap(folio, page, vma, -(__s64)folio_nr_pages_est(folio));
+}
+
+SEC("fentry/folio_remove_rmap_pud")
+int BPF_PROG(fentry_folio_remove_rmap_pud, struct folio* folio, struct page* page,
+             struct vm_area_struct* vma) {
+    return submit_file_rmap(folio, page, vma, -(__s64)folio_nr_pages_est(folio));
+}
+
+#endif /* __RMAP_BPF_H__ */