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

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_RSS 9
+#define EVENT_TYPE_RMAP_ANON 10
+#define EVENT_TYPE_FORK 11
+#define EVENT_TYPE_EXEC 12
+#define EVENT_TYPE_EXIT 13
 
 /* Common header shared by all event types */
 struct event_header {
@@ -45,6 +50,21 @@ struct event {
             uint64_t addr; /* address of mapping */
             uint64_t size; /* size of mapping */
         } mmap;
+
+        struct {
+            int32_t member;
+            uint64_t size;
+        } rss;
+
+        struct {
+            int64_t delta;
+            uint64_t addr;
+        } rmap;
+
+        /* Process lifecycle events (fork carries the parent; exec/exit have no payload) */
+        struct {
+            uint32_t parent_pid;
+        } fork;
     } data;
 };
 

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

@@ -131,12 +131,11 @@ static __always_inline __u64* take_param(void* map) {
 /* Macro to handle common event submission boilerplate
  * Usage: SUBMIT_EVENT(event_type, { e->data.foo = bar; })
  */
-#define SUBMIT_EVENT(evt_type, fill_data)                               \
+#define SUBMIT_EVENT_AS(owner_pid, evt_type, fill_data)                 \
     {                                                                   \
         __u64 tid = bpf_get_current_pid_tgid();                         \
-        __u32 pid = tid >> 32;                                          \
                                                                         \
-        if (!is_tracked(pid) || !is_enabled()) {                        \
+        if (!is_enabled()) {                                            \
             return 0;                                                   \
         }                                                               \
                                                                         \
@@ -151,7 +150,7 @@ static __always_inline __u64* take_param(void* map) {
         }                                                               \
                                                                         \
         e->header.timestamp = bpf_ktime_get_ns();                       \
-        e->header.pid = pid;                                            \
+        e->header.pid = owner_pid;                                      \
         e->header.tid = tid & 0xFFFFFFFF;                               \
         e->header.event_type = evt_type;                                \
                                                                         \
@@ -161,6 +160,18 @@ static __always_inline __u64* take_param(void* map) {
         return 0;                                                       \
     }
 
+#define SUBMIT_EVENT(evt_type, fill_data)       \
+    {                                           \
+        __u64 tid = bpf_get_current_pid_tgid(); \
+        __u32 pid = tid >> 32;                  \
+                                                \
+        if (!is_tracked(pid)) {                 \
+            return 0;                           \
+        }                                       \
+                                                \
+        SUBMIT_EVENT_AS(pid, evt_type, fill_data); \
+    }
+
 /* Helper to submit an allocation event (malloc, calloc) */
 static __always_inline int submit_alloc_event(__u64 size, __u64 addr) {
     SUBMIT_EVENT(EVENT_TYPE_MALLOC, {
@@ -389,3 +400,48 @@ int tracepoint_sys_exit_brk(struct trace_event_raw_sys_exit* ctx) {
 
     return submit_mmap_event(new_brk, 0, EVENT_TYPE_BRK);
 }
+
+/* == Process lifecycle events ==
+ *
+ * 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("tracepoint/task/task_newtask")
+int tracepoint_task_newtask(struct trace_event_raw_task_newtask* ctx) {
+    if (ctx->clone_flags & CLONE_THREAD) {
+        return 0;
+    }
+
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 parent_pid = tid >> 32;
+    if (!is_tracked(parent_pid)) {
+        return 0;
+    }
+
+    __u32 child_pid = ctx->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) {
+    SUBMIT_EVENT(EVENT_TYPE_EXEC, {});
+}
+
+SEC("tracepoint/sched/sched_process_exit")
+int tracepoint_sched_process_exit(void* ctx) {
+    __u64 tid = bpf_get_current_pid_tgid();
+    if ((tid >> 32) != (tid & 0xFFFFFFFF)) {
+        return 0;
+    }
+
+    SUBMIT_EVENT(EVENT_TYPE_EXIT, {});
+}
+
+#include "rss.bpf.h"

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

@@ -0,0 +1,126 @@
+#ifndef __RSS_BPF_H__
+#define __RSS_BPF_H__
+
+#define FOLIO_MAPPING_ANON 0x1UL
+
+const volatile __u32 page_shift = 12;
+
+static __always_inline int submit_rss_event(__u32 owner_pid, __s32 member, __u64 size) {
+    SUBMIT_EVENT_AS(owner_pid, EVENT_TYPE_RSS, {
+        e->data.rss.member = member;
+        e->data.rss.size = size;
+    });
+}
+
+SEC("tracepoint/kmem/rss_stat")
+int tracepoint_rss_stat(struct trace_event_raw_rss_stat* ctx) {
+    if (!ctx->curr) {
+        return 0;
+    }
+
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 pid = tid >> 32;
+    if (!is_tracked(pid)) {
+        return 0;
+    }
+
+    if (ctx->member == MM_SWAPENTS) {
+        return 0;
+    }
+
+    return submit_rss_event(pid, ctx->member, ctx->size);
+}
+
+static __always_inline __u64 folio_nr_pages_est(struct folio* folio) {
+    unsigned long flags = BPF_CORE_READ(folio, flags).f;
+    if (!(flags & (1UL << PG_head))) {
+        return 1;
+    }
+    unsigned long order = BPF_CORE_READ(folio, _flags_1) & 0xff;
+    return 1UL << order;
+}
+
+static __always_inline int folio_is_anon(struct folio* folio) {
+    unsigned long mapping = (unsigned long)BPF_CORE_READ(folio, mapping);
+    return (mapping & FOLIO_MAPPING_ANON) != 0;
+}
+
+static __always_inline __u64 folio_page_address(struct folio* folio, struct page* page,
+                                                struct vm_area_struct* vma) {
+    __u64 page_idx = ((__u64)page - (__u64)folio) / bpf_core_type_size(struct page);
+    __u64 pgoff = BPF_CORE_READ(folio, index) + page_idx;
+    __u64 vm_pgoff = BPF_CORE_READ(vma, vm_pgoff);
+    __u64 vm_start = BPF_CORE_READ(vma, vm_start);
+    if (pgoff < vm_pgoff) {
+        return vm_start;
+    }
+    return vm_start + ((pgoff - vm_pgoff) << page_shift);
+}
+
+static __always_inline int submit_rmap_anon(struct vm_area_struct* vma, __s64 delta, __u64 addr) {
+    __u64 mm = (__u64)BPF_CORE_READ(vma, vm_mm);
+    struct task_struct* task = bpf_get_current_task_btf();
+    if ((__u64)BPF_CORE_READ(task, mm) != mm) {
+        return 0;
+    }
+
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 pid = tid >> 32;
+    if (!is_tracked(pid)) {
+        return 0;
+    }
+
+    SUBMIT_EVENT_AS(pid, EVENT_TYPE_RMAP_ANON, {
+        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_anon(vma, (__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_anon(vma, (__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_anon(vma, (__s64)folio_nr_pages_est(folio), address);
+}
+
+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) {
+    if (!folio_is_anon(folio)) {
+        return 0;
+    }
+    return submit_rmap_anon(vma, -(__s64)nr_pages, folio_page_address(folio, page, vma));
+}
+
+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) {
+    if (!folio_is_anon(folio)) {
+        return 0;
+    }
+    return submit_rmap_anon(vma, -(__s64)folio_nr_pages_est(folio),
+                            folio_page_address(folio, page, vma));
+}
+
+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) {
+    if (!folio_is_anon(folio)) {
+        return 0;
+    }
+    return submit_rmap_anon(vma, -(__s64)folio_nr_pages_est(folio),
+                            folio_page_address(folio, page, vma));
+}
+
+#endif /* __RSS_BPF_H__ */

crates/memtrack/src/ebpf/events.rs

@@ -74,6 +74,27 @@ pub fn parse_event(data: &[u8]) -> Option<MemtrackEvent> {
                     size: event.data.mmap.size,
                 },
             ),
+            EVENT_TYPE_RSS => (
+                0,
+                MemtrackEventKind::Rss {
+                    member: event.data.rss.member,
+                    size: event.data.rss.size,
+                },
+            ),
+            EVENT_TYPE_RMAP_ANON => (
+                event.data.rmap.addr,
+                MemtrackEventKind::RmapAnon {
+                    delta: event.data.rmap.delta,
+                },
+            ),
+            EVENT_TYPE_FORK => (
+                0,
+                MemtrackEventKind::Fork {
+                    parent_pid: event.data.fork.parent_pid as i32,
+                },
+            ),
+            EVENT_TYPE_EXEC => (0, MemtrackEventKind::Exec),
+            EVENT_TYPE_EXIT => (0, MemtrackEventKind::Exit),
             unknown => {
                 panic!("Unknown event type: {unknown}");
             }
@@ -160,4 +181,90 @@ mod tests {
             _ => panic!("Expected Malloc event kind"),
         }
     }
+
+    #[test]
+    fn test_parse_rss_event() {
+        let mut event: bindings::event = unsafe { std::mem::zeroed() };
+        event.header.event_type = bindings::EVENT_TYPE_RSS as u8;
+        event.header.timestamp = 12345678;
+        event.header.pid = 1000;
+        event.header.tid = 2000;
+        event.data.rss.member = 1;
+        event.data.rss.size = 4096 * 10;
+
+        let bytes = unsafe {
+            std::slice::from_raw_parts(
+                &event as *const _ as *const u8,
+                std::mem::size_of_val(&event),
+            )
+        };
+
+        let parsed = parse_event(bytes).unwrap();
+        assert_eq!(parsed.pid, 1000);
+        assert_eq!(parsed.addr, 0);
+
+        match parsed.kind {
+            MemtrackEventKind::Rss { member, size } => {
+                assert_eq!(member, 1);
+                assert_eq!(size, 4096 * 10);
+            }
+            _ => panic!("Expected Rss event kind"),
+        }
+    }
+
+    #[test]
+    fn test_parse_rmap_anon_event() {
+        let mut event: bindings::event = unsafe { std::mem::zeroed() };
+        event.header.event_type = bindings::EVENT_TYPE_RMAP_ANON as u8;
+        event.header.timestamp = 12345678;
+        event.header.pid = 1000;
+        event.header.tid = 2000;
+        event.data.rmap.delta = 8;
+        event.data.rmap.addr = 0x7f00;
+
+        let bytes = unsafe {
+            std::slice::from_raw_parts(
+                &event as *const _ as *const u8,
+                std::mem::size_of_val(&event),
+            )
+        };
+
+        let parsed = parse_event(bytes).unwrap();
+        assert_eq!(parsed.pid, 1000);
+        assert_eq!(parsed.addr, 0x7f00);
+
+        match parsed.kind {
+            MemtrackEventKind::RmapAnon { delta } => {
+                assert_eq!(delta, 8);
+            }
+            _ => panic!("Expected RmapAnon event kind"),
+        }
+    }
+
+    #[test]
+    fn test_parse_fork_event() {
+        let mut event: bindings::event = unsafe { std::mem::zeroed() };
+        event.header.event_type = bindings::EVENT_TYPE_FORK as u8;
+        event.header.timestamp = 12345678;
+        event.header.pid = 1001;
+        event.header.tid = 2000;
+        event.data.fork.parent_pid = 1000;
+
+        let bytes = unsafe {
+            std::slice::from_raw_parts(
+                &event as *const _ as *const u8,
+                std::mem::size_of_val(&event),
+            )
+        };
+
+        let parsed = parse_event(bytes).unwrap();
+        assert_eq!(parsed.pid, 1001);
+
+        match parsed.kind {
+            MemtrackEventKind::Fork { parent_pid } => {
+                assert_eq!(parent_pid, 1000);
+            }
+            _ => panic!("Expected Fork event kind"),
+        }
+    }
 }