feat(memtrack): capture allocation stacks in eBPF

.github/workflows/ci.yml

@@ -93,7 +93,7 @@ jobs:
         # 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, rss_tests]
+        test: [c_tests, cpp_tests, rust_tests, spawn_tests, dlopen_tests, rss_tests, stack_tests]
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:

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

@@ -9,6 +9,7 @@
     BPF_HASH_MAP(name##_arg, __u64, __u64, 10000);   \
     SEC(UPROBE_SEC)                                  \
     int uprobe_##name(struct pt_regs* ctx) {         \
+        stash_stack_hash(capture_stack(ctx));        \
         return store_param(&name##_arg, arg_expr);   \
     }                                                \
     SEC(URETPROBE_SEC)                               \
@@ -17,6 +18,7 @@
         if (!arg_ptr) {                              \
             return 0;                                \
         }                                            \
+        __u64 stack_hash = take_stack_hash();        \
         __u64 ret_val = PT_REGS_RC(ctx);             \
         if (ret_val == 0) {                          \
             return 0;                                \
@@ -32,6 +34,7 @@
         if (arg0 == 0) {                         \
             return 0;                            \
         }                                        \
+        __u64 stack_hash = capture_stack(ctx);   \
         submit_block;                            \
     }
 
@@ -50,6 +53,8 @@
             return 0;                                                         \
         }                                                                     \
                                                                               \
+        stash_stack_hash(capture_stack(ctx));                                 \
+                                                                              \
         struct name##_args_t args = {.arg0 = arg0_expr, .arg1 = arg1_expr};   \
                                                                               \
         bpf_map_update_elem(&name##_args, &tid, &args, BPF_ANY);              \
@@ -63,6 +68,7 @@
         if (!args) {                                                          \
             return 0;                                                         \
         }                                                                     \
+        __u64 stack_hash = take_stack_hash();                                 \
                                                                               \
         struct name##_args_t a = *args;                                       \
         bpf_map_delete_elem(&name##_args, &tid);                              \
@@ -77,20 +83,22 @@
         submit_block;                                                         \
     }
 
-UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx), { return submit_alloc_event(arg0, ret_val); })
+UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx),
+               { return submit_alloc_event(arg0, ret_val, stack_hash); })
 
-UPROBE_RET(free, PT_REGS_PARM1(ctx), { return submit_free_event(arg0); })
+UPROBE_RET(free, PT_REGS_PARM1(ctx), { return submit_free_event(arg0, stack_hash); })
 
 UPROBE_ARG_RET(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx),
-               { return submit_calloc_event(arg0, ret_val); })
+               { return submit_calloc_event(arg0, ret_val, stack_hash); })
 
 UPROBE_ARGS_RET(realloc, PT_REGS_PARM2(ctx), PT_REGS_PARM1(ctx),
-                { return submit_realloc_event(arg1, ret_val, arg0); })
+                { return submit_realloc_event(arg1, ret_val, arg0, stack_hash); })
 
 UPROBE_ARG_RET(aligned_alloc, PT_REGS_PARM2(ctx),
-               { return submit_aligned_alloc_event(arg0, ret_val); })
+               { return submit_aligned_alloc_event(arg0, ret_val, stack_hash); })
 
-UPROBE_ARG_RET(memalign, PT_REGS_PARM2(ctx), { return submit_aligned_alloc_event(arg0, ret_val); })
+UPROBE_ARG_RET(memalign, PT_REGS_PARM2(ctx),
+               { return submit_aligned_alloc_event(arg0, ret_val, stack_hash); })
 
 /*
  * posix_memalign(void** memptr, size_t alignment, size_t size)
@@ -115,6 +123,8 @@ int uprobe_posix_memalign(struct pt_regs* ctx) {
         return 0;
     }
 
+    stash_stack_hash(capture_stack(ctx));
+
     struct posix_memalign_args_t args = {.memptr = PT_REGS_PARM1(ctx), .size = PT_REGS_PARM3(ctx)};
     bpf_map_update_elem(&posix_memalign_args, &tid, &args, BPF_ANY);
     return 0;
@@ -127,6 +137,7 @@ int uretprobe_posix_memalign(struct pt_regs* ctx) {
     if (!args) {
         return 0;
     }
+    __u64 stack_hash = take_stack_hash();
 
     struct posix_memalign_args_t a = *args;
     bpf_map_delete_elem(&posix_memalign_args, &tid);
@@ -140,7 +151,7 @@ int uretprobe_posix_memalign(struct pt_regs* ctx) {
         return 0;
     }
 
-    return submit_aligned_alloc_event(a.size, addr);
+    return submit_aligned_alloc_event(a.size, addr, stack_hash);
 }
 
 struct mmap_args {

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

@@ -15,6 +15,48 @@
 #define EVENT_TYPE_RSS 12
 #define EVENT_TYPE_RMAP 13
 
+/* Largest user-stack copy one definition can carry. The scratch buffer holding
+ * header plus bytes is a per-CPU map value, capped at PCPU_MIN_UNIT_SIZE
+ * (32 KiB) by the kernel allocator. */
+#define MEMTRACK_MAX_STACK_COPY (32 * 1024 - 512)
+
+/* Registers, indexed by the capturing architecture's DWARF register number
+ * (x86_64: 0=rax .. 7=rsp, 8..15=r8-r15, 16=rip; aarch64: 0..30=x0-x30,
+ * 31=sp, 32=pc). Slots the architecture does not define stay zero. An offline
+ * DWARF unwinder needs the callee-saved ones to evaluate CFA rules, not just
+ * ip/sp/bp. */
+#define MEMTRACK_STACK_REGS 33
+
+/* Counter slots in the stack_counters array map. */
+#define MEMTRACK_STACK_COUNTER_COPY_FAILED 0
+#define MEMTRACK_STACK_COUNTER_HASH_MAP_FULL 1
+/* bpf_get_stackid() has several negative outcomes (no user callchain,
+ * hash-bucket collision, or no free bucket), so this counts only missing ids. */
+#define MEMTRACK_STACK_COUNTER_STACKID_FAILED 2
+#define MEMTRACK_STACK_COUNTER_TRUNCATED 3
+#define MEMTRACK_STACK_COUNTER_RING_FULL 4
+#define MEMTRACK_STACK_COUNTER_PREEMPTED 5
+#define MEMTRACK_STACK_COUNTER_COUNT 6
+
+struct stack_regs {
+    uint64_t reg[MEMTRACK_STACK_REGS];
+};
+
+/* Head of a stack record; `copy_len` raw stack bytes read upwards from `sp`
+ * follow it. */
+struct stack_header {
+    uint64_t hash;
+    uint64_t timestamp; /* monotonic time in nanoseconds (CLOCK_MONOTONIC) */
+    int64_t stackid;    /* bpf_get_stackid() result; negative means unavailable */
+    uint64_t sp;        /* user stack pointer the copy starts at */
+    uint32_t pid;
+    uint32_t tid;
+    uint32_t copy_len;
+    uint8_t truncated; /* the copy hit the size cap */
+    uint8_t _pad[3];
+    struct stack_regs regs;
+};
+
 /* Common header shared by all event types */
 struct event_header {
     uint8_t event_type; /* See EVENT_TYPE_* constants above */
@@ -29,20 +71,23 @@ struct event {
     union {
         /* Allocation events (malloc, calloc, aligned_alloc) */
         struct {
-            uint64_t addr; /* address returned */
-            uint64_t size; /* size requested */
+            uint64_t addr;       /* address returned */
+            uint64_t size;       /* size requested */
+            uint64_t stack_hash; /* caller stack identity; 0 = not captured */
         } alloc;
 
         /* Deallocation event (free) */
         struct {
-            uint64_t addr; /* address to free */
+            uint64_t addr;       /* address to free */
+            uint64_t stack_hash; /* caller stack identity; 0 = not captured */
         } free;
 
         /* Reallocation event - includes both old and new addresses */
         struct {
-            uint64_t old_addr; /* previous address (can be NULL) */
-            uint64_t new_addr; /* new address returned */
-            uint64_t size;     /* new size requested */
+            uint64_t old_addr;   /* previous address (can be NULL) */
+            uint64_t new_addr;   /* new address returned */
+            uint64_t size;       /* new size requested */
+            uint64_t stack_hash; /* caller stack identity; 0 = not captured */
         } realloc;
 
         /* Memory mapping events (mmap, munmap, brk) */

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

@@ -11,6 +11,7 @@
 #include "process_tracking.bpf.h"
 #include "rmap.bpf.h"
 #include "rss.bpf.h"
+#include "stack_capture.bpf.h"
 #include "utils/event_helpers.h"
 #include "utils/folio.h"
 #include "utils/map_helpers.h"