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 "rmap_legacy.bpf.h"
+#include "rss.bpf.h"
 #include "utils/event_helpers.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,97 @@
+#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.
+ */
+
+#define CLONE_THREAD 0x00010000
+
+SEC("tp_btf/task_newtask")
+int BPF_PROG(tracepoint_task_newtask, struct task_struct* child, __u64 clone_flags) {
+    if (clone_flags & CLONE_THREAD) {
+        return 0;
+    }
+
+    __u32 parent_pid = current_tgid();
+    if (!is_tracked(parent_pid)) {
+        return 0;
+    }
+
+    /* Register the child here rather than on sched_process_fork: that
+     * tracepoint fires for CLONE_THREAD too and carries only raw task pids,
+     * which would fill the tracking maps with thread tids that no exit path
+     * removes (group death deletes only the tgid). task_newtask fires before
+     * wake_up_new_task, so registration precedes any event from the child.
+     * The BTF-typed variant is used because the child's pid must be read in
+     * the tracker's namespace, which the tracepoint's raw pid field cannot
+     * give. */
+    __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__ */