From 29d2e2ee1bd844106c1417814e8dc80e01e4cb49 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 19:54:11 +0200 Subject: [PATCH 1/2] refactor(memtrack): split monolithic eBPF C and Rust into domain files Split the single-file eBPF source (memtrack.bpf.c) and its Rust attach code (memtrack.rs) into small, single-responsibility files. C side: - main.bpf.c: includes + LICENSE only, no programs or maps - utils/map_helpers.h: BPF_HASH_MAP/BPF_ARRAY_MAP/BPF_RINGBUF macros - utils/tracking.h: tracked_pids/pids_ppid/tracking_enabled maps + is_tracked/is_enabled/track_child helpers + sched_fork program - utils/event_helpers.h: events ringbuf, dropped_events, SUBMIT_EVENT, submit_*_event, store_param/take_param - allocator.h: UPROBE_* macros + all allocator uprobes + mmap/munmap/brk tracepoints Encapsulation improvements: - track_child() helper extracted from sched_fork body - store_mmap_args() helper extracted from mmap enter handler - brk enter now uses store_param() instead of raw bpf_map_update_elem Rust side: - memtrack/mod.rs: skel include, MemtrackBpf struct, new(), Drop - memtrack/macros.rs: attach_uprobe_uretprobe!/attach_uprobe!/ attach_tracepoint! macros + ensure_symbol_exists - memtrack/maps.rs: add_tracked_pid, enable/disable_tracking, dropped_events_count - memtrack/allocator.rs: all attach_* probe methods + per-allocator attach methods - memtrack/tracking.rs: attach_sched_fork / attach_tracepoints Skeleton struct names change from MemtrackSkel to MainSkel (derived from main.bpf.c filename). Output file memtrack.skel.rs unchanged. --- crates/memtrack/build.rs | 2 +- crates/memtrack/src/ebpf/c/allocator.h | 167 +++++ crates/memtrack/src/ebpf/c/main.bpf.c | 14 + crates/memtrack/src/ebpf/c/memtrack.bpf.c | 403 ----------- .../memtrack/src/ebpf/c/utils/event_helpers.h | 111 +++ .../memtrack/src/ebpf/c/utils/map_helpers.h | 26 + crates/memtrack/src/ebpf/c/utils/tracking.h | 57 ++ crates/memtrack/src/ebpf/memtrack.rs | 648 ------------------ .../memtrack/src/ebpf/memtrack/allocator.rs | 274 ++++++++ crates/memtrack/src/ebpf/memtrack/macros.rs | 136 ++++ crates/memtrack/src/ebpf/memtrack/maps.rs | 66 ++ crates/memtrack/src/ebpf/memtrack/mod.rs | 168 +++++ crates/memtrack/src/ebpf/memtrack/tracking.rs | 12 + 13 files changed, 1032 insertions(+), 1052 deletions(-) create mode 100644 crates/memtrack/src/ebpf/c/allocator.h create mode 100644 crates/memtrack/src/ebpf/c/main.bpf.c delete mode 100644 crates/memtrack/src/ebpf/c/memtrack.bpf.c create mode 100644 crates/memtrack/src/ebpf/c/utils/event_helpers.h create mode 100644 crates/memtrack/src/ebpf/c/utils/map_helpers.h create mode 100644 crates/memtrack/src/ebpf/c/utils/tracking.h delete mode 100644 crates/memtrack/src/ebpf/memtrack.rs create mode 100644 crates/memtrack/src/ebpf/memtrack/allocator.rs create mode 100644 crates/memtrack/src/ebpf/memtrack/macros.rs create mode 100644 crates/memtrack/src/ebpf/memtrack/maps.rs create mode 100644 crates/memtrack/src/ebpf/memtrack/mod.rs create mode 100644 crates/memtrack/src/ebpf/memtrack/tracking.rs diff --git a/crates/memtrack/build.rs b/crates/memtrack/build.rs index a9227d13..1c337ebd 100644 --- a/crates/memtrack/build.rs +++ b/crates/memtrack/build.rs @@ -17,7 +17,7 @@ fn build_ebpf() { .expect("CARGO_CFG_TARGET_ARCH must be set in build script"); let memtrack_out = PathBuf::from(env::var("OUT_DIR").unwrap()).join("memtrack.skel.rs"); SkeletonBuilder::new() - .source("src/ebpf/c/memtrack.bpf.c") + .source("src/ebpf/c/main.bpf.c") .clang_args([ "-I", &vmlinux::include_path_root().join(arch).to_string_lossy(), diff --git a/crates/memtrack/src/ebpf/c/allocator.h b/crates/memtrack/src/ebpf/c/allocator.h new file mode 100644 index 00000000..a48571a4 --- /dev/null +++ b/crates/memtrack/src/ebpf/c/allocator.h @@ -0,0 +1,167 @@ +#ifndef __ALLOCATOR_H__ +#define __ALLOCATOR_H__ + +#include "utils/event_helpers.h" +#include "utils/map_helpers.h" +#include "utils/tracking.h" + +#define UPROBE_ARG_RET(name, arg_expr, submit_block) \ + BPF_HASH_MAP(name##_arg, __u64, __u64, 10000); \ + SEC("uprobe") \ + int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_arg, arg_expr); } \ + SEC("uretprobe") \ + 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) \ + SEC("uprobe") \ + int uprobe_##name(struct pt_regs* ctx) { \ + __u64 arg0 = arg_expr; \ + if (arg0 == 0) { \ + return 0; \ + } \ + submit_block; \ + } + +#define UPROBE_ARGS_RET(name, arg0_expr, arg1_expr, submit_block) \ + struct name##_args_t { \ + __u64 arg0; \ + __u64 arg1; \ + }; \ + BPF_HASH_MAP(name##_args, __u64, struct name##_args_t, 10000); \ + SEC("uprobe") \ + int uprobe_##name(struct pt_regs* ctx) { \ + __u64 tid = bpf_get_current_pid_tgid(); \ + __u32 pid = tid >> 32; \ + \ + if (!is_tracked(pid)) { \ + return 0; \ + } \ + \ + struct name##_args_t args = {.arg0 = arg0_expr, .arg1 = arg1_expr}; \ + \ + bpf_map_update_elem(&name##_args, &tid, &args, BPF_ANY); \ + return 0; \ + } \ + SEC("uretprobe") \ + int uretprobe_##name(struct pt_regs* ctx) { \ + __u64 tid = bpf_get_current_pid_tgid(); \ + struct name##_args_t* args = bpf_map_lookup_elem(&name##_args, &tid); \ + \ + if (!args) { \ + return 0; \ + } \ + \ + struct name##_args_t a = *args; \ + bpf_map_delete_elem(&name##_args, &tid); \ + \ + __u64 ret_val = PT_REGS_RC(ctx); \ + if (ret_val == 0) { \ + return 0; \ + } \ + \ + __u64 arg0 = a.arg0; \ + __u64 arg1 = a.arg1; \ + submit_block; \ + } + +UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx), { return submit_alloc_event(arg0, ret_val); }) + +UPROBE_RET(free, PT_REGS_PARM1(ctx), { return submit_free_event(arg0); }) + +UPROBE_ARG_RET(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx), + { return submit_calloc_event(arg0, ret_val); }) + +UPROBE_ARGS_RET(realloc, PT_REGS_PARM2(ctx), PT_REGS_PARM1(ctx), + { return submit_realloc_event(arg1, ret_val, arg0); }) + +UPROBE_ARG_RET(aligned_alloc, 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); }) + +struct mmap_args { + __u64 addr; + __u64 len; +}; + +BPF_HASH_MAP(mmap_temp, __u64, struct mmap_args, 10000); + +static __always_inline void store_mmap_args(__u64 addr, __u64 len) { + __u64 tid = bpf_get_current_pid_tgid(); + __u32 pid = tid >> 32; + if (is_tracked(pid)) { + struct mmap_args args = {.addr = addr, .len = len}; + bpf_map_update_elem(&mmap_temp, &tid, &args, BPF_ANY); + } +} + +SEC("tracepoint/syscalls/sys_enter_mmap") +int tracepoint_sys_enter_mmap(struct trace_event_raw_sys_enter* ctx) { + store_mmap_args(ctx->args[0], ctx->args[1]); + return 0; +} + +SEC("tracepoint/syscalls/sys_exit_mmap") +int tracepoint_sys_exit_mmap(struct trace_event_raw_sys_exit* ctx) { + struct mmap_args* args = (struct mmap_args*)take_param(&mmap_temp); + if (!args) { + return 0; + } + + __s64 ret = ctx->ret; + if (ret <= 0) { + return 0; + } + + return submit_mmap_event((__u64)ret, args->len, EVENT_TYPE_MMAP); +} + +SEC("tracepoint/syscalls/sys_enter_munmap") +int tracepoint_sys_enter_munmap(struct trace_event_raw_sys_enter* ctx) { + __u64 addr = ctx->args[0]; + __u64 len = ctx->args[1]; + + if (addr == 0 || len == 0) { + return 0; + } + + return submit_mmap_event(addr, len, EVENT_TYPE_MUNMAP); +} + +BPF_HASH_MAP(brk_temp, __u64, __u64, 10000); + +SEC("tracepoint/syscalls/sys_enter_brk") +int tracepoint_sys_enter_brk(struct trace_event_raw_sys_enter* ctx) { + store_param(&brk_temp, ctx->args[0]); + return 0; +} + +SEC("tracepoint/syscalls/sys_exit_brk") +int tracepoint_sys_exit_brk(struct trace_event_raw_sys_exit* ctx) { + __u64* requested_brk = take_param(&brk_temp); + if (!requested_brk) { + return 0; + } + + __u64 new_brk = ctx->ret; + __u64 req_brk = *requested_brk; + + if (req_brk == 0 || new_brk <= 0) { + return 0; + } + + return submit_mmap_event(new_brk, 0, EVENT_TYPE_BRK); +} + +#endif /* __ALLOCATOR_H__ */ diff --git a/crates/memtrack/src/ebpf/c/main.bpf.c b/crates/memtrack/src/ebpf/c/main.bpf.c new file mode 100644 index 00000000..10442c21 --- /dev/null +++ b/crates/memtrack/src/ebpf/c/main.bpf.c @@ -0,0 +1,14 @@ +// clang-format off +#include "vmlinux.h" +// clang-format on +#include +#include +#include + +#include "allocator.h" +#include "event.h" +#include "utils/event_helpers.h" +#include "utils/map_helpers.h" +#include "utils/tracking.h" + +char LICENSE[] SEC("license") = "GPL"; diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c deleted file mode 100644 index ad64d73b..00000000 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ /dev/null @@ -1,403 +0,0 @@ -// clang-format off -// Prevent clang-format from reformatting the include statement, which is -// needed for the bpf headers below. -#include "vmlinux.h" -// clang-format on - -#include -#include -#include - -#include "event.h" - -char LICENSE[] SEC("license") = "GPL"; - -/* Macros for common map definitions */ -#define BPF_HASH_MAP(name, key_type, value_type, max_ents) \ - struct { \ - __uint(type, BPF_MAP_TYPE_HASH); \ - __uint(max_entries, max_ents); \ - __type(key, key_type); \ - __type(value, value_type); \ - } name SEC(".maps") - -#define BPF_ARRAY_MAP(name, value_type, max_ents) \ - struct { \ - __uint(type, BPF_MAP_TYPE_ARRAY); \ - __uint(max_entries, max_ents); \ - __type(key, __u32); \ - __type(value, value_type); \ - } name SEC(".maps") - -#define BPF_RINGBUF(name, size) \ - struct { \ - __uint(type, BPF_MAP_TYPE_RINGBUF); \ - __uint(max_entries, size); \ - } name SEC(".maps") - -/* Map to store PIDs we're tracking */ -BPF_HASH_MAP(tracked_pids, __u32, __u8, 10000); -/* Map to store parent-child relationships to detect hierarchy */ -BPF_HASH_MAP(pids_ppid, __u32, __u32, 10000); -/* Ring buffer for sending events to userspace */ -BPF_RINGBUF(events, 256 * 1024 * 1024); -/* Map to control whether tracking is enabled (0 = disabled, 1 = enabled) */ -BPF_ARRAY_MAP(tracking_enabled, __u8, 1); -/* Counter for events that couldn't be added to the ring buffer */ -BPF_ARRAY_MAP(dropped_events, __u64, 1); - -/* == Code that tracks process forks and execs == */ - -/* Helper to check if a PID or any of its ancestors should be tracked */ -static __always_inline int is_tracked(__u32 pid) { - /* Direct check */ - if (bpf_map_lookup_elem(&tracked_pids, &pid)) { - return 1; - } - -/* Check parent recursively (up to 5 levels) */ -#pragma unroll - for (int i = 0; i < 5; i++) { - __u32* ppid = bpf_map_lookup_elem(&pids_ppid, &pid); - if (!ppid) { - break; - } - pid = *ppid; - if (bpf_map_lookup_elem(&tracked_pids, &pid)) { - return 1; - } - } - - return 0; -} - -SEC("tracepoint/sched/sched_process_fork") -int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) { - __u32 parent_pid = ctx->parent_pid; - __u32 child_pid = ctx->child_pid; - - /* Print process fork with PIDs */ - // bpf_printk("sched_fork: parent_pid=%u child_pid=%u", parent_pid, child_pid); - - /* Check if parent is being tracked */ - if (is_tracked(parent_pid)) { - /* Auto-track this child */ - __u8 marker = 1; - bpf_map_update_elem(&tracked_pids, &child_pid, &marker, BPF_ANY); - bpf_map_update_elem(&pids_ppid, &child_pid, &parent_pid, BPF_ANY); - - // bpf_printk("auto-tracking child process: child_pid=%u", child_pid); - } - - return 0; -} - -/* == Helper functions for the allocation tracking == */ - -/* Wake the consumer only once this much unconsumed data has accumulated. - * Per-event wakeups dominate submission cost at high event rates; batching - * them behind a data watermark amortizes the wakeup to ~1 per thousand - * events. The userspace poller's poll timeout flushes the tail that never - * reaches the watermark. */ -#define WAKEUP_DATA_SIZE (64 * 1024) - -static __always_inline long wake_flags(void) { - long avail = bpf_ringbuf_query(&events, BPF_RB_AVAIL_DATA); - return avail >= WAKEUP_DATA_SIZE ? BPF_RB_FORCE_WAKEUP : BPF_RB_NO_WAKEUP; -} - -/* Helper to check if tracking is currently enabled */ -static __always_inline int is_enabled(void) { - __u32 key = 0; - __u8* enabled = bpf_map_lookup_elem(&tracking_enabled, &key); - - /* Default to enabled if map not initialized */ - if (!enabled) { - return 1; - } - - return *enabled; -} - -/* Helper to store parameter value in map for tracking between entry and return - */ -static __always_inline int store_param(void* map, __u64 value) { - __u64 tid = bpf_get_current_pid_tgid(); - __u32 pid = tid >> 32; - if (is_tracked(pid)) { - bpf_map_update_elem(map, &tid, &value, BPF_ANY); - } - return 0; -} - -/* Helper to take parameter value from map (lookup and delete) */ -static __always_inline __u64* take_param(void* map) { - __u64 tid = bpf_get_current_pid_tgid(); - __u64* value = bpf_map_lookup_elem(map, &tid); - if (value) { - bpf_map_delete_elem(map, &tid); - } - return value; -} - -/* Macro to handle common event submission boilerplate - * Usage: SUBMIT_EVENT(event_type, { e->data.foo = bar; }) - */ -#define SUBMIT_EVENT(evt_type, fill_data) \ - { \ - __u64 tid = bpf_get_current_pid_tgid(); \ - __u32 pid = tid >> 32; \ - \ - if (!is_tracked(pid) || !is_enabled()) { \ - return 0; \ - } \ - \ - struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0); \ - if (!e) { \ - __u32 zero = 0; \ - __u64* drops = bpf_map_lookup_elem(&dropped_events, &zero); \ - if (drops) { \ - __sync_fetch_and_add(drops, 1); \ - } \ - return 0; \ - } \ - \ - e->header.timestamp = bpf_ktime_get_ns(); \ - e->header.pid = pid; \ - e->header.tid = tid & 0xFFFFFFFF; \ - e->header.event_type = evt_type; \ - \ - fill_data; \ - \ - bpf_ringbuf_submit(e, wake_flags()); \ - return 0; \ - } - -/* Helper to submit an allocation event (malloc, calloc) */ -static __always_inline int submit_alloc_event(__u64 size, __u64 addr) { - SUBMIT_EVENT(EVENT_TYPE_MALLOC, { - e->data.alloc.addr = addr; - e->data.alloc.size = size; - }); -} - -/* Helper to submit an aligned allocation event (aligned_alloc, memalign) */ -static __always_inline int submit_aligned_alloc_event(__u64 size, __u64 addr) { - SUBMIT_EVENT(EVENT_TYPE_ALIGNED_ALLOC, { - e->data.alloc.addr = addr; - e->data.alloc.size = size; - }); -} - -/* Helper to submit a calloc event */ -static __always_inline int submit_calloc_event(__u64 size, __u64 addr) { - SUBMIT_EVENT(EVENT_TYPE_CALLOC, { - e->data.alloc.addr = addr; - e->data.alloc.size = size; - }); -} - -/* Helper to submit a free event */ -static __always_inline int submit_free_event(__u64 addr) { - SUBMIT_EVENT(EVENT_TYPE_FREE, { e->data.free.addr = addr; }); -} - -/* Helper to submit a realloc event with both old and new addresses */ -static __always_inline int submit_realloc_event(__u64 old_addr, __u64 new_addr, __u64 size) { - SUBMIT_EVENT(EVENT_TYPE_REALLOC, { - e->data.realloc.old_addr = old_addr; - e->data.realloc.new_addr = new_addr; - e->data.realloc.size = size; - }); -} - -/* Helper to submit a memory mapping event (mmap, munmap, brk) */ -static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_type) { - SUBMIT_EVENT(event_type, { - e->data.mmap.addr = addr; - e->data.mmap.size = size; - }); -} - -/* Macro to generate uprobe/uretprobe pairs for allocation functions with 1 argument */ -#define UPROBE_ARG_RET(name, arg_expr, submit_block) \ - BPF_HASH_MAP(name##_arg, __u64, __u64, 10000); \ - SEC("uprobe") \ - int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_arg, arg_expr); } \ - SEC("uretprobe") \ - 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; \ - } - -/* Macro for simple return value only functions like free */ -#define UPROBE_RET(name, arg_expr, submit_block) \ - SEC("uprobe") \ - int uprobe_##name(struct pt_regs* ctx) { \ - __u64 arg0 = arg_expr; \ - if (arg0 == 0) { \ - return 0; \ - } \ - submit_block; \ - } - -/* Macro to generate uprobe/uretprobe pairs for functions with 2 arguments */ -#define UPROBE_ARGS_RET(name, arg0_expr, arg1_expr, submit_block) \ - struct name##_args_t { \ - __u64 arg0; \ - __u64 arg1; \ - }; \ - BPF_HASH_MAP(name##_args, __u64, struct name##_args_t, 10000); \ - SEC("uprobe") \ - int uprobe_##name(struct pt_regs* ctx) { \ - __u64 tid = bpf_get_current_pid_tgid(); \ - __u32 pid = tid >> 32; \ - \ - if (!is_tracked(pid)) { \ - return 0; \ - } \ - \ - struct name##_args_t args = {.arg0 = arg0_expr, .arg1 = arg1_expr}; \ - \ - bpf_map_update_elem(&name##_args, &tid, &args, BPF_ANY); \ - return 0; \ - } \ - SEC("uretprobe") \ - int uretprobe_##name(struct pt_regs* ctx) { \ - __u64 tid = bpf_get_current_pid_tgid(); \ - struct name##_args_t* args = bpf_map_lookup_elem(&name##_args, &tid); \ - \ - if (!args) { \ - return 0; \ - } \ - \ - struct name##_args_t a = *args; \ - bpf_map_delete_elem(&name##_args, &tid); \ - \ - __u64 ret_val = PT_REGS_RC(ctx); \ - if (ret_val == 0) { \ - return 0; \ - } \ - \ - __u64 arg0 = a.arg0; \ - __u64 arg1 = a.arg1; \ - submit_block; \ - } - -/* malloc: allocates with size parameter */ -UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx), { return submit_alloc_event(arg0, ret_val); }) - -/* free: deallocates by address */ -UPROBE_RET(free, PT_REGS_PARM1(ctx), { return submit_free_event(arg0); }) - -/* calloc: allocates with nmemb * size */ -UPROBE_ARG_RET(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx), - { return submit_calloc_event(arg0, ret_val); }) - -/* realloc: reallocates with old_addr and new size */ -UPROBE_ARGS_RET(realloc, PT_REGS_PARM2(ctx), PT_REGS_PARM1(ctx), - { return submit_realloc_event(arg1, ret_val, arg0); }) - -/* aligned_alloc: allocates with alignment and size */ -UPROBE_ARG_RET(aligned_alloc, PT_REGS_PARM2(ctx), - { return submit_aligned_alloc_event(arg0, ret_val); }) - -/* memalign: allocates with alignment and size (legacy interface) */ -UPROBE_ARG_RET(memalign, PT_REGS_PARM2(ctx), { return submit_aligned_alloc_event(arg0, ret_val); }) - -/* Map to store mmap parameters between entry and return */ -struct mmap_args { - __u64 addr; - __u64 len; -}; - -BPF_HASH_MAP(mmap_temp, __u64, struct mmap_args, 10000); - -SEC("tracepoint/syscalls/sys_enter_mmap") -int tracepoint_sys_enter_mmap(struct trace_event_raw_sys_enter* ctx) { - __u64 tid = bpf_get_current_pid_tgid(); - __u32 pid = tid >> 32; - - if (is_tracked(pid)) { - struct mmap_args args = {0}; - - /* mmap(addr, len, prot, flags, fd, offset) - * We care about addr (can be 0 for kernel choice) and len */ - args.addr = ctx->args[0]; - args.len = ctx->args[1]; - - bpf_map_update_elem(&mmap_temp, &tid, &args, BPF_ANY); - } - - return 0; -} - -SEC("tracepoint/syscalls/sys_exit_mmap") -int tracepoint_sys_exit_mmap(struct trace_event_raw_sys_exit* ctx) { - struct mmap_args* args = (struct mmap_args*)take_param(&mmap_temp); - if (!args) { - return 0; - } - - __s64 ret = ctx->ret; - if (ret <= 0) { - return 0; - } - - return submit_mmap_event((__u64)ret, args->len, EVENT_TYPE_MMAP); -} - -/* munmap tracking */ -SEC("tracepoint/syscalls/sys_enter_munmap") -int tracepoint_sys_enter_munmap(struct trace_event_raw_sys_enter* ctx) { - __u64 addr = ctx->args[0]; - __u64 len = ctx->args[1]; - - if (addr == 0 || len == 0) { - return 0; - } - - return submit_mmap_event(addr, len, EVENT_TYPE_MUNMAP); -} - -/* brk tracking - adjusts the program break (heap boundary) */ -BPF_HASH_MAP(brk_temp, __u64, __u64, 10000); - -SEC("tracepoint/syscalls/sys_enter_brk") -int tracepoint_sys_enter_brk(struct trace_event_raw_sys_enter* ctx) { - __u64 tid = bpf_get_current_pid_tgid(); - __u32 pid = tid >> 32; - - if (is_tracked(pid)) { - /* brk(addr) - if addr is 0, just queries current break */ - __u64 requested_brk = ctx->args[0]; - bpf_map_update_elem(&brk_temp, &tid, &requested_brk, BPF_ANY); - } - - return 0; -} - -SEC("tracepoint/syscalls/sys_exit_brk") -int tracepoint_sys_exit_brk(struct trace_event_raw_sys_exit* ctx) { - __u64* requested_brk = take_param(&brk_temp); - if (!requested_brk) { - return 0; - } - - __u64 new_brk = ctx->ret; - __u64 req_brk = *requested_brk; - - if (req_brk == 0 || new_brk <= 0) { - return 0; - } - - return submit_mmap_event(new_brk, 0, EVENT_TYPE_BRK); -} diff --git a/crates/memtrack/src/ebpf/c/utils/event_helpers.h b/crates/memtrack/src/ebpf/c/utils/event_helpers.h new file mode 100644 index 00000000..ea8575f7 --- /dev/null +++ b/crates/memtrack/src/ebpf/c/utils/event_helpers.h @@ -0,0 +1,111 @@ +#ifndef __EVENT_HELPERS_H__ +#define __EVENT_HELPERS_H__ + +#include "../event.h" +#include "map_helpers.h" +#include "tracking.h" + +BPF_RINGBUF(events, 256 * 1024 * 1024); +BPF_ARRAY_MAP(dropped_events, __u64, 1); + +/* Wake the consumer only once this much unconsumed data has accumulated. + * Per-event wakeups dominate submission cost at high event rates; batching + * them behind a data watermark amortizes the wakeup to ~1 per thousand + * events. The userspace poller's poll timeout flushes the tail that never + * reaches the watermark. */ +#define WAKEUP_DATA_SIZE (64 * 1024) + +static __always_inline long wake_flags(void) { + long avail = bpf_ringbuf_query(&events, BPF_RB_AVAIL_DATA); + return avail >= WAKEUP_DATA_SIZE ? BPF_RB_FORCE_WAKEUP : BPF_RB_NO_WAKEUP; +} + +static __always_inline int store_param(void* map, __u64 value) { + __u64 tid = bpf_get_current_pid_tgid(); + __u32 pid = tid >> 32; + if (is_tracked(pid)) { + bpf_map_update_elem(map, &tid, &value, BPF_ANY); + } + return 0; +} + +static __always_inline __u64* take_param(void* map) { + __u64 tid = bpf_get_current_pid_tgid(); + __u64* value = bpf_map_lookup_elem(map, &tid); + if (value) { + bpf_map_delete_elem(map, &tid); + } + return value; +} + +#define SUBMIT_EVENT(evt_type, fill_data) \ + { \ + __u64 tid = bpf_get_current_pid_tgid(); \ + __u32 pid = tid >> 32; \ + \ + if (!is_tracked(pid) || !is_enabled()) { \ + return 0; \ + } \ + \ + struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0); \ + if (!e) { \ + __u32 zero = 0; \ + __u64* drops = bpf_map_lookup_elem(&dropped_events, &zero); \ + if (drops) { \ + __sync_fetch_and_add(drops, 1); \ + } \ + return 0; \ + } \ + \ + e->header.timestamp = bpf_ktime_get_ns(); \ + e->header.pid = pid; \ + e->header.tid = tid & 0xFFFFFFFF; \ + e->header.event_type = evt_type; \ + \ + fill_data; \ + \ + bpf_ringbuf_submit(e, wake_flags()); \ + return 0; \ + } + +static __always_inline int submit_alloc_event(__u64 size, __u64 addr) { + SUBMIT_EVENT(EVENT_TYPE_MALLOC, { + e->data.alloc.addr = addr; + e->data.alloc.size = size; + }); +} + +static __always_inline int submit_aligned_alloc_event(__u64 size, __u64 addr) { + SUBMIT_EVENT(EVENT_TYPE_ALIGNED_ALLOC, { + e->data.alloc.addr = addr; + e->data.alloc.size = size; + }); +} + +static __always_inline int submit_calloc_event(__u64 size, __u64 addr) { + SUBMIT_EVENT(EVENT_TYPE_CALLOC, { + e->data.alloc.addr = addr; + e->data.alloc.size = size; + }); +} + +static __always_inline int submit_free_event(__u64 addr) { + SUBMIT_EVENT(EVENT_TYPE_FREE, { e->data.free.addr = addr; }); +} + +static __always_inline int submit_realloc_event(__u64 old_addr, __u64 new_addr, __u64 size) { + SUBMIT_EVENT(EVENT_TYPE_REALLOC, { + e->data.realloc.old_addr = old_addr; + e->data.realloc.new_addr = new_addr; + e->data.realloc.size = size; + }); +} + +static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_type) { + SUBMIT_EVENT(event_type, { + e->data.mmap.addr = addr; + e->data.mmap.size = size; + }); +} + +#endif /* __EVENT_HELPERS_H__ */ diff --git a/crates/memtrack/src/ebpf/c/utils/map_helpers.h b/crates/memtrack/src/ebpf/c/utils/map_helpers.h new file mode 100644 index 00000000..484fe970 --- /dev/null +++ b/crates/memtrack/src/ebpf/c/utils/map_helpers.h @@ -0,0 +1,26 @@ +#ifndef __MAP_HELPERS_H__ +#define __MAP_HELPERS_H__ + +#define BPF_HASH_MAP(name, key_type, value_type, max_ents) \ + struct { \ + __uint(type, BPF_MAP_TYPE_HASH); \ + __uint(max_entries, max_ents); \ + __type(key, key_type); \ + __type(value, value_type); \ + } name SEC(".maps") + +#define BPF_ARRAY_MAP(name, value_type, max_ents) \ + struct { \ + __uint(type, BPF_MAP_TYPE_ARRAY); \ + __uint(max_entries, max_ents); \ + __type(key, __u32); \ + __type(value, value_type); \ + } name SEC(".maps") + +#define BPF_RINGBUF(name, size) \ + struct { \ + __uint(type, BPF_MAP_TYPE_RINGBUF); \ + __uint(max_entries, size); \ + } name SEC(".maps") + +#endif /* __MAP_HELPERS_H__ */ diff --git a/crates/memtrack/src/ebpf/c/utils/tracking.h b/crates/memtrack/src/ebpf/c/utils/tracking.h new file mode 100644 index 00000000..0d9e937b --- /dev/null +++ b/crates/memtrack/src/ebpf/c/utils/tracking.h @@ -0,0 +1,57 @@ +#ifndef __TRACKING_H__ +#define __TRACKING_H__ + +#include "map_helpers.h" + +BPF_HASH_MAP(tracked_pids, __u32, __u8, 10000); +BPF_HASH_MAP(pids_ppid, __u32, __u32, 10000); +BPF_ARRAY_MAP(tracking_enabled, __u8, 1); + +static __always_inline int is_tracked(__u32 pid) { + if (bpf_map_lookup_elem(&tracked_pids, &pid)) { + return 1; + } + +#pragma unroll + for (int i = 0; i < 5; i++) { + __u32* ppid = bpf_map_lookup_elem(&pids_ppid, &pid); + if (!ppid) { + break; + } + pid = *ppid; + if (bpf_map_lookup_elem(&tracked_pids, &pid)) { + return 1; + } + } + + return 0; +} + +static __always_inline int is_enabled(void) { + __u32 key = 0; + __u8* enabled = bpf_map_lookup_elem(&tracking_enabled, &key); + if (!enabled) { + return 1; + } + return *enabled; +} + +static __always_inline void track_child(__u32 child_pid, __u32 parent_pid) { + __u8 marker = 1; + bpf_map_update_elem(&tracked_pids, &child_pid, &marker, BPF_ANY); + bpf_map_update_elem(&pids_ppid, &child_pid, &parent_pid, BPF_ANY); +} + +SEC("tracepoint/sched/sched_process_fork") +int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) { + __u32 parent_pid = ctx->parent_pid; + __u32 child_pid = ctx->child_pid; + + if (is_tracked(parent_pid)) { + track_child(child_pid, parent_pid); + } + + return 0; +} + +#endif /* __TRACKING_H__ */ diff --git a/crates/memtrack/src/ebpf/memtrack.rs b/crates/memtrack/src/ebpf/memtrack.rs deleted file mode 100644 index 649d1797..00000000 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ /dev/null @@ -1,648 +0,0 @@ -use crate::prelude::*; -use libbpf_rs::Link; -use libbpf_rs::skel::OpenSkel; -use libbpf_rs::skel::SkelBuilder; -use libbpf_rs::{MapCore, UprobeOpts}; -use paste::paste; -use std::collections::HashMap; -use std::mem::MaybeUninit; -use std::path::Path; - -use crate::allocators::{AllocatorKind, AllocatorLib}; -use crate::ebpf::poller::RingBufferPoller; - -pub mod memtrack_skel { - include!(concat!(env!("OUT_DIR"), "/memtrack.skel.rs")); -} -pub use memtrack_skel::*; - -/// Resolve libbpf attach targets for every defined symbol in `lib_path`. -pub fn resolve_symbol_offsets(lib_path: &Path) -> Result { - use object::{Object, ObjectSymbol}; - - let data = std::fs::read(lib_path)?; - let file = object::File::parse(&*data)?; - let mut offsets = HashMap::new(); - - for symbol in file.symbols().chain(file.dynamic_symbols()) { - if !symbol.is_definition() { - continue; - } - - let Ok(name) = symbol.name() else { - continue; - }; - - if let Some(file_offset) = symbol_file_offset(&file, &symbol) { - offsets.insert(name.to_owned(), file_offset); - } - } - - Ok(ResolvedSymbols { offsets }) -} - -/// The libbpf file offset for `symbol`, or `None` when it has no address in a -/// file-backed section (absolute, `SHT_NOBITS`, ...). -fn symbol_file_offset<'a>( - file: &object::File, - symbol: &impl object::ObjectSymbol<'a>, -) -> Option { - use object::{Object, ObjectSection}; - - let address = symbol.address(); - if address == 0 { - return None; - } - - let section = file.section_by_index(symbol.section_index()?).ok()?; - let (sh_offset, _) = section.file_range()?; - Some((address - section.address() + sh_offset) as usize) -} - -/// Attach targets resolved from a library's symbol tables. -pub struct ResolvedSymbols { - offsets: HashMap, -} - -impl ResolvedSymbols { - fn offset(&self, symbol: &str) -> Option { - self.offsets.get(symbol).copied() - } -} - -/// Macro to attach a function with both entry and return probes at a resolved -/// file offset. Also generates an `attach_*_if_found` variant that skips -/// symbols absent from the offset table (returning whether it attached) and -/// propagates attach failures. -macro_rules! attach_uprobe_uretprobe { - ($name:ident, $prog_entry:ident, $prog_return:ident) => { - paste! { - fn [](&mut self, lib_path: &Path, offset: usize) -> Result<()> { - let link = self - .skel - .progs - .$prog_entry - .attach_uprobe_with_opts( - -1, - lib_path, - offset, - UprobeOpts { - retprobe: false, - ..Default::default() - }, - ) - .context(format!( - "Failed to attach uprobe at offset {:#x} in {}", - offset, - lib_path.display() - ))?; - self.probes.push(link); - - let link = self - .skel - .progs - .$prog_return - .attach_uprobe_with_opts( - -1, - lib_path, - offset, - UprobeOpts { - retprobe: true, - ..Default::default() - }, - ) - .context(format!( - "Failed to attach uretprobe at offset {:#x} in {}", - offset, - lib_path.display() - ))?; - self.probes.push(link); - - Ok(()) - } - - fn [<$name _if_found>]( - &mut self, - lib_path: &Path, - symbol: &str, - symbols: &ResolvedSymbols, - ) -> Result { - let Some(offset) = symbols.offset(symbol) else { - return Ok(false); - }; - self.[](lib_path, offset) - .with_context(|| format!("Failed to attach {symbol}"))?; - log::trace!("Attached {} at {:#x}", symbol, offset); - Ok(true) - } - } - }; -} - -/// Macro to attach a function with only an entry probe (no return probe) at a -/// resolved file offset. Also generates an `attach_*_if_found` variant that -/// skips symbols absent from the offset table (returning whether it attached) -/// and propagates attach failures. -macro_rules! attach_uprobe { - ($name:ident, $prog:ident) => { - paste! { - fn [](&mut self, lib_path: &Path, offset: usize) -> Result<()> { - let link = self - .skel - .progs - .$prog - .attach_uprobe_with_opts( - -1, - lib_path, - offset, - UprobeOpts { - retprobe: false, - ..Default::default() - }, - ) - .context(format!( - "Failed to attach uprobe at offset {:#x} in {}", - offset, - lib_path.display() - ))?; - self.probes.push(link); - Ok(()) - } - - fn [<$name _if_found>]( - &mut self, - lib_path: &Path, - symbol: &str, - symbols: &ResolvedSymbols, - ) -> Result { - let Some(offset) = symbols.offset(symbol) else { - return Ok(false); - }; - self.[](lib_path, offset) - .with_context(|| format!("Failed to attach {symbol}"))?; - log::trace!("Attached {} at {:#x}", symbol, offset); - Ok(true) - } - } - }; -} - -macro_rules! attach_tracepoint { - ($func:ident, $prog:ident) => { - fn $func(&mut self) -> Result<()> { - let link = self - .skel - .progs - .$prog - .attach() - .context(format!("Failed to attach {} tracepoint", stringify!($prog)))?; - self.probes.push(link); - Ok(()) - } - }; - ($name:ident) => { - paste! { - attach_tracepoint!([], []); - } - }; -} - -pub struct MemtrackBpf { - skel: Box>, - probes: Vec, -} - -impl MemtrackBpf { - pub fn new() -> Result { - // Build and open the syscalls BPF program - let builder = MemtrackSkelBuilder::default(); - let open_object = Box::leak(Box::new(MaybeUninit::uninit())); - let open_skel = builder - .open(open_object) - .context("Failed to open syscalls BPF skeleton")?; - - let skel = Box::new( - open_skel - .load() - .context("Failed to load syscalls BPF skeleton")?, - ); - - Ok(Self { - skel, - probes: Vec::new(), - }) - } - - pub fn add_tracked_pid(&mut self, pid: i32) -> Result<()> { - self.skel - .maps - .tracked_pids - .update( - &pid.to_le_bytes(), - &1u8.to_le_bytes(), - libbpf_rs::MapFlags::ANY, - ) - .context("Failed to add PID to uprobes tracked set")?; - - Ok(()) - } - - /// Enable event tracking - pub fn enable_tracking(&mut self) -> Result<()> { - let key = 0u32; - let value = true as u8; - self.skel - .maps - .tracking_enabled - .update( - &key.to_le_bytes(), - &value.to_le_bytes(), - libbpf_rs::MapFlags::ANY, - ) - .context("Failed to enable tracking")?; - Ok(()) - } - - /// Read the count of events dropped because the ring buffer was full. - pub fn dropped_events_count(&self) -> Result { - let key = 0u32; - let value = self - .skel - .maps - .dropped_events - .lookup(&key.to_le_bytes(), libbpf_rs::MapFlags::ANY) - .context("Failed to read dropped_events counter")? - .ok_or_else(|| anyhow!("dropped_events slot 0 missing"))?; - - let bytes: [u8; 8] = value - .as_slice() - .try_into() - .map_err(|_| anyhow!("dropped_events value has unexpected size"))?; - Ok(u64::from_le_bytes(bytes)) - } - - /// Disable event tracking - pub fn disable_tracking(&mut self) -> Result<()> { - let key = 0u32; - let value = false as u8; - self.skel - .maps - .tracking_enabled - .update( - &key.to_le_bytes(), - &value.to_le_bytes(), - libbpf_rs::MapFlags::ANY, - ) - .context("Failed to disable tracking")?; - Ok(()) - } - - // ========================================================================= - // Allocation probe functions (symbol passed at call time) - // ========================================================================= - attach_uprobe_uretprobe!(attach_malloc, uprobe_malloc, uretprobe_malloc); - attach_uprobe_uretprobe!(attach_calloc, uprobe_calloc, uretprobe_calloc); - attach_uprobe_uretprobe!(attach_realloc, uprobe_realloc, uretprobe_realloc); - attach_uprobe_uretprobe!( - attach_aligned_alloc, - uprobe_aligned_alloc, - uretprobe_aligned_alloc - ); - attach_uprobe_uretprobe!(attach_memalign, uprobe_memalign, uretprobe_memalign); - attach_uprobe!(attach_free, uprobe_free); - - // ========================================================================= - // Attach methods grouped by allocator - // ========================================================================= - - /// Attach probes for every discovered allocator library. - pub fn attach_allocators(&mut self, libs: &[AllocatorLib]) -> Result<()> { - use rayon::prelude::*; - - let resolved = libs - .par_iter() - .map(|lib| resolve_symbol_offsets(&lib.path).map(|offsets| (lib, offsets))) - .collect::>>()?; - - for (lib, offsets) in resolved { - let before = self.probes.len(); - self.attach_allocator_probes_with_offsets(lib.kind, &lib.path, &offsets)?; - debug!( - "Attached {} links to {} ({} resolved symbols)", - self.probes.len() - before, - lib.path.display(), - offsets.offsets.len() - ); - } - - Ok(()) - } - - /// Attach probes for a specific allocator kind. - /// This attaches both standard probes (if the allocator exports them) and - /// allocator-specific prefixed probes. - pub fn attach_allocator_probes(&mut self, kind: AllocatorKind, lib_path: &Path) -> Result<()> { - let offsets = resolve_symbol_offsets(lib_path)?; - self.attach_allocator_probes_with_offsets(kind, lib_path, &offsets) - } - - fn attach_allocator_probes_with_offsets( - &mut self, - kind: AllocatorKind, - lib_path: &Path, - offsets: &ResolvedSymbols, - ) -> Result<()> { - debug!( - "Attaching {} probes to: {}", - kind.name(), - lib_path.display() - ); - - match kind { - AllocatorKind::Libc => { - // Libc only has standard probes, and they must succeed: a libc - // with an uninstrumented core entry point would silently - // produce an incomplete trace. - for symbol in ["malloc", "calloc", "realloc", "free"] { - ensure!( - offsets.offset(symbol).is_some(), - "Required allocator symbol {symbol} has no resolvable file offset in {}", - lib_path.display() - ); - } - self.attach_libc_probes(lib_path, offsets) - } - AllocatorKind::LibCpp => { - // libc++ exports C++ operator new/delete symbols - self.attach_libcpp_probes(lib_path, offsets) - } - AllocatorKind::Jemalloc => { - // Jemalloc exposes libc/libcpp compatible allocator functions: - self.attach_compat_probes(lib_path, offsets); - self.attach_jemalloc_probes(lib_path, offsets) - } - AllocatorKind::Mimalloc => { - // Mimalloc exposes libc/libcpp compatible allocator functions: - self.attach_compat_probes(lib_path, offsets); - self.attach_mimalloc_probes(lib_path, offsets) - } - AllocatorKind::Tcmalloc => { - // Tcmalloc exposes libc/libcpp compatible allocator functions: - self.attach_compat_probes(lib_path, offsets); - self.attach_tcmalloc_probes(lib_path, offsets) - } - } - } - - /// Best-effort attach of the libc/libcpp compatible symbols that - /// non-libc allocators may also export. - fn attach_compat_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) { - if let Err(e) = self.attach_libc_probes(lib_path, offsets) { - warn!("libc-compatible probes for {}: {e:#}", lib_path.display()); - } - if let Err(e) = self.attach_libcpp_probes(lib_path, offsets) { - warn!("libcpp-compatible probes for {}: {e:#}", lib_path.display()); - } - } - - fn attach_standard_probes( - &mut self, - lib_path: &Path, - prefixes: &[&str], - suffixes: &[&str], - offsets: &ResolvedSymbols, - ) -> Result<()> { - // Always include "" to capture the basic case - let prefixes_with_base: Vec<&str> = std::iter::once("") - .chain(prefixes.iter().copied()) - .unique() - .collect(); - - let suffixes_with_base: Vec<&str> = std::iter::once("") - .chain(suffixes.iter().copied()) - .unique() - .collect(); - - for prefix in &prefixes_with_base { - for suffix in &suffixes_with_base { - self.attach_malloc_if_found(lib_path, &format!("{prefix}malloc{suffix}"), offsets)?; - self.attach_malloc_if_found(lib_path, &format!("{prefix}valloc{suffix}"), offsets)?; - self.attach_malloc_if_found( - lib_path, - &format!("{prefix}pvalloc{suffix}"), - offsets, - )?; - self.attach_calloc_if_found(lib_path, &format!("{prefix}calloc{suffix}"), offsets)?; - self.attach_realloc_if_found( - lib_path, - &format!("{prefix}realloc{suffix}"), - offsets, - )?; - self.attach_aligned_alloc_if_found( - lib_path, - &format!("{prefix}aligned_alloc{suffix}"), - offsets, - )?; - self.attach_memalign_if_found( - lib_path, - &format!("{prefix}memalign{suffix}"), - offsets, - )?; - self.attach_memalign_if_found( - lib_path, - &format!("{prefix}posix_memalign{suffix}"), - offsets, - )?; - self.attach_free_if_found(lib_path, &format!("{prefix}free{suffix}"), offsets)?; - self.attach_free_if_found( - lib_path, - &format!("{prefix}free_sized{suffix}"), - offsets, - )?; - self.attach_free_if_found( - lib_path, - &format!("{prefix}free_aligned_sized{suffix}"), - offsets, - )?; - self.attach_free_if_found(lib_path, &format!("{prefix}cfree{suffix}"), offsets)?; - } - } - - Ok(()) - } - - /// Attach standard library allocation probes (libc-style: malloc, free, calloc, etc.) - /// This works for libc and allocators that export standard symbol names. - /// For non-libc allocators, standard names are optional - just try them silently. - fn attach_libc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { - self.attach_standard_probes(lib_path, &[], &[], offsets) - } - - /// Attach C++ operator new/delete probes. - /// These are mangled C++ symbols that wrap the underlying allocator. - /// C++ operators have identical signatures to malloc/free, so we reuse those handlers. - fn attach_libcpp_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { - self.attach_malloc_if_found(lib_path, "_Znwm", offsets)?; // operator new(size_t) - self.attach_malloc_if_found(lib_path, "_Znam", offsets)?; // operator new[](size_t) - self.attach_malloc_if_found(lib_path, "_ZnwmSt11align_val_t", offsets)?; // operator new(size_t, std::align_val_t) - self.attach_malloc_if_found(lib_path, "_ZnamSt11align_val_t", offsets)?; // operator new[](size_t, std::align_val_t) - self.attach_free_if_found(lib_path, "_ZdlPv", offsets)?; // operator delete(void*) - self.attach_free_if_found(lib_path, "_ZdaPv", offsets)?; // operator delete[](void*) - self.attach_free_if_found(lib_path, "_ZdlPvm", offsets)?; // operator delete(void*, size_t) - C++14 sized delete - self.attach_free_if_found(lib_path, "_ZdaPvm", offsets)?; // operator delete[](void*, size_t) - C++14 sized delete - self.attach_free_if_found(lib_path, "_ZdlPvSt11align_val_t", offsets)?; // operator delete(void*, std::align_val_t) - self.attach_free_if_found(lib_path, "_ZdaPvSt11align_val_t", offsets)?; // operator delete[](void*, std::align_val_t) - self.attach_free_if_found(lib_path, "_ZdlPvmSt11align_val_t", offsets)?; // operator delete(void*, size_t, std::align_val_t) - self.attach_free_if_found(lib_path, "_ZdaPvmSt11align_val_t", offsets)?; // operator delete[](void*, size_t, std::align_val_t) - - Ok(()) - } - - /// Attach jemalloc-specific probes (prefixed and extended API). - fn attach_jemalloc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { - // The following functions are used in Rust when setting a global allocator: - // - rust_alloc: _rjem_malloc and _rjem_mallocx - // - rust_alloc_zeroed: _rjem_mallocx / _rjem_calloc - // - rust_dealloc: _rjem_sdallocx - // - rust_realloc: _rjem_realloc / _rjem_rallocx - - // je_* API (internal jemalloc functions, static linking) - // _rjem_* API (Rust jemalloc crate, dynamic linking) - let prefixes = ["je_", "_rjem_"]; - let suffixes = ["", "_default"]; - - self.attach_standard_probes(lib_path, &prefixes, &suffixes, offsets)?; - - // Non-standard API that has an additional flag parameter - // See: https://jemalloc.net/jemalloc.3.html - for prefix in prefixes { - for suffix in suffixes { - self.attach_malloc_if_found( - lib_path, - &format!("{prefix}mallocx{suffix}"), - offsets, - )?; - self.attach_realloc_if_found( - lib_path, - &format!("{prefix}rallocx{suffix}"), - offsets, - )?; - self.attach_free_if_found(lib_path, &format!("{prefix}dallocx{suffix}"), offsets)?; - self.attach_free_if_found(lib_path, &format!("{prefix}sdallocx{suffix}"), offsets)?; - } - } - - Ok(()) - } - - /// Attach mimalloc-specific probes (mi_* API). - fn attach_mimalloc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { - // The following functions are used in Rust when setting a global allocator: - // - mi_malloc_aligned - // - mi_free - // - mi_realloc_aligned - // - mi_zalloc_aligned - - self.attach_standard_probes(lib_path, &["mi_"], &[], offsets)?; - - // Zero-initialized and aligned variants - self.attach_malloc_if_found(lib_path, "mi_malloc_aligned", offsets)?; - self.attach_calloc_if_found(lib_path, "mi_zalloc", offsets)?; - self.attach_calloc_if_found(lib_path, "mi_zalloc_aligned", offsets)?; - self.attach_realloc_if_found(lib_path, "mi_realloc_aligned", offsets)?; - - Ok(()) - } - - /// Attach TCMalloc probes ( tc_* API). - /// - /// See: - /// - https://github.com/google/tcmalloc/blob/master/docs/reference.md - /// - https://github.com/gperftools/gperftools/blob/a47243150ec41097602730ff8779fafcc172d1fb/src/tcmalloc.cc#L178-L190 - fn attach_tcmalloc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { - self.attach_standard_probes(lib_path, &["tc_"], &[], offsets)?; - - self.attach_free_if_found(lib_path, "free_sized", offsets)?; - self.attach_free_if_found(lib_path, "free_aligned_sized", offsets)?; - self.attach_free_if_found(lib_path, "sdallocx", offsets)?; - - Ok(()) - } - attach_tracepoint!(sched_fork); - - pub fn attach_tracepoints(&mut self) -> Result<()> { - self.attach_sched_fork()?; - Ok(()) - } - - /// Start polling, forwarding each event over the returned channel. - pub fn start_polling_with_channel( - &self, - poll_timeout_ms: u64, - ) -> Result<( - RingBufferPoller, - crossbeam_channel::Receiver, - )> { - // Use the syscalls skeleton's ring buffer (both programs share the same one) - RingBufferPoller::with_channel(&self.skel.maps.events, poll_timeout_ms) - } -} - -impl MemtrackBpf { - /// Detach all BPF links in parallel. Closing a uprobe link blocks on two - /// RCU grace periods in the kernel, but concurrent waiters share grace - /// periods, so closing from many threads scales near-linearly. - pub fn detach_probes(&mut self) { - const DETACH_THREADS: usize = 32; - - let mut probes = std::mem::take(&mut self.probes); - if probes.is_empty() { - return; - } - - debug!("Detaching {} BPF links", probes.len()); - let start = std::time::Instant::now(); - let chunk_size = probes.len().div_ceil(DETACH_THREADS); - std::thread::scope(|scope| { - while !probes.is_empty() { - let split_at = probes.len().saturating_sub(chunk_size); - let chunk = probes.split_off(split_at); - scope.spawn(move || drop(chunk)); - } - }); - debug!("Detached BPF links in {:?}", start.elapsed()); - } -} - -impl Drop for MemtrackBpf { - fn drop(&mut self) { - self.detach_probes(); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Allocator entry points must resolve to file offsets; a symbol that - /// silently fails to resolve attaches nothing and loses all events. - #[test] - fn libc_allocator_symbols_resolve_to_offsets() { - let maps = std::fs::read_to_string("/proc/self/maps").unwrap(); - let libc_path = maps - .lines() - .find_map(|line| { - let path = line.split_whitespace().last()?; - path.contains("libc.so.6").then(|| path.to_owned()) - }) - .expect("test process has no mapped libc.so.6"); - - let symbols = resolve_symbol_offsets(Path::new(&libc_path)).unwrap(); - for symbol in ["malloc", "calloc", "realloc", "free"] { - assert!( - symbols.offset(symbol).is_some(), - "{symbol} in {libc_path} did not resolve to a file offset" - ); - } - } -} diff --git a/crates/memtrack/src/ebpf/memtrack/allocator.rs b/crates/memtrack/src/ebpf/memtrack/allocator.rs new file mode 100644 index 00000000..410914c5 --- /dev/null +++ b/crates/memtrack/src/ebpf/memtrack/allocator.rs @@ -0,0 +1,274 @@ +use crate::prelude::*; +use libbpf_rs::UprobeOpts; +use paste::paste; +use std::path::Path; + +use super::{MemtrackBpf, ResolvedSymbols, resolve_symbol_offsets}; +use crate::allocators::{AllocatorKind, AllocatorLib}; + +impl MemtrackBpf { + attach_uprobe_uretprobe!(attach_malloc, uprobe_malloc, uretprobe_malloc); + attach_uprobe_uretprobe!(attach_calloc, uprobe_calloc, uretprobe_calloc); + attach_uprobe_uretprobe!(attach_realloc, uprobe_realloc, uretprobe_realloc); + attach_uprobe_uretprobe!( + attach_aligned_alloc, + uprobe_aligned_alloc, + uretprobe_aligned_alloc + ); + attach_uprobe_uretprobe!(attach_memalign, uprobe_memalign, uretprobe_memalign); + attach_uprobe!(attach_free, uprobe_free); + + /// Attach probes for every discovered allocator library. + pub fn attach_allocators(&mut self, libs: &[AllocatorLib]) -> Result<()> { + use rayon::prelude::*; + + let resolved = libs + .par_iter() + .map(|lib| resolve_symbol_offsets(&lib.path).map(|offsets| (lib, offsets))) + .collect::>>()?; + + for (lib, offsets) in resolved { + let before = self.probes.len(); + self.attach_allocator_probes_with_offsets(lib.kind, &lib.path, &offsets)?; + debug!( + "Attached {} links to {} ({} resolved symbols)", + self.probes.len() - before, + lib.path.display(), + offsets.offsets.len() + ); + } + + Ok(()) + } + + /// Attach probes for a specific allocator kind. + /// This attaches both standard probes (if the allocator exports them) and + /// allocator-specific prefixed probes. + pub fn attach_allocator_probes(&mut self, kind: AllocatorKind, lib_path: &Path) -> Result<()> { + let offsets = resolve_symbol_offsets(lib_path)?; + self.attach_allocator_probes_with_offsets(kind, lib_path, &offsets) + } + + fn attach_allocator_probes_with_offsets( + &mut self, + kind: AllocatorKind, + lib_path: &Path, + offsets: &ResolvedSymbols, + ) -> Result<()> { + debug!( + "Attaching {} probes to: {}", + kind.name(), + lib_path.display() + ); + + match kind { + AllocatorKind::Libc => { + // Libc only has standard probes, and they must succeed: a libc + // with an uninstrumented core entry point would silently + // produce an incomplete trace. + for symbol in ["malloc", "calloc", "realloc", "free"] { + ensure!( + offsets.offset(symbol).is_some(), + "Required allocator symbol {symbol} has no resolvable file offset in {}", + lib_path.display() + ); + } + self.attach_libc_probes(lib_path, offsets) + } + AllocatorKind::LibCpp => { + // libc++ exports C++ operator new/delete symbols + self.attach_libcpp_probes(lib_path, offsets) + } + AllocatorKind::Jemalloc => { + // Jemalloc exposes libc/libcpp compatible allocator functions: + self.attach_compat_probes(lib_path, offsets); + self.attach_jemalloc_probes(lib_path, offsets) + } + AllocatorKind::Mimalloc => { + // Mimalloc exposes libc/libcpp compatible allocator functions: + self.attach_compat_probes(lib_path, offsets); + self.attach_mimalloc_probes(lib_path, offsets) + } + AllocatorKind::Tcmalloc => { + // Tcmalloc exposes libc/libcpp compatible allocator functions: + self.attach_compat_probes(lib_path, offsets); + self.attach_tcmalloc_probes(lib_path, offsets) + } + } + } + + /// Best-effort attach of the libc/libcpp compatible symbols that + /// non-libc allocators may also export. + fn attach_compat_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) { + if let Err(e) = self.attach_libc_probes(lib_path, offsets) { + warn!("libc-compatible probes for {}: {e:#}", lib_path.display()); + } + if let Err(e) = self.attach_libcpp_probes(lib_path, offsets) { + warn!("libcpp-compatible probes for {}: {e:#}", lib_path.display()); + } + } + + fn attach_standard_probes( + &mut self, + lib_path: &Path, + prefixes: &[&str], + suffixes: &[&str], + offsets: &ResolvedSymbols, + ) -> Result<()> { + // Always include "" to capture the basic case + let prefixes_with_base: Vec<&str> = std::iter::once("") + .chain(prefixes.iter().copied()) + .unique() + .collect(); + + let suffixes_with_base: Vec<&str> = std::iter::once("") + .chain(suffixes.iter().copied()) + .unique() + .collect(); + + for prefix in &prefixes_with_base { + for suffix in &suffixes_with_base { + self.attach_malloc_if_found(lib_path, &format!("{prefix}malloc{suffix}"), offsets)?; + self.attach_malloc_if_found(lib_path, &format!("{prefix}valloc{suffix}"), offsets)?; + self.attach_malloc_if_found( + lib_path, + &format!("{prefix}pvalloc{suffix}"), + offsets, + )?; + self.attach_calloc_if_found(lib_path, &format!("{prefix}calloc{suffix}"), offsets)?; + self.attach_realloc_if_found( + lib_path, + &format!("{prefix}realloc{suffix}"), + offsets, + )?; + self.attach_aligned_alloc_if_found( + lib_path, + &format!("{prefix}aligned_alloc{suffix}"), + offsets, + )?; + self.attach_memalign_if_found( + lib_path, + &format!("{prefix}memalign{suffix}"), + offsets, + )?; + self.attach_memalign_if_found( + lib_path, + &format!("{prefix}posix_memalign{suffix}"), + offsets, + )?; + self.attach_free_if_found(lib_path, &format!("{prefix}free{suffix}"), offsets)?; + self.attach_free_if_found( + lib_path, + &format!("{prefix}free_sized{suffix}"), + offsets, + )?; + self.attach_free_if_found( + lib_path, + &format!("{prefix}free_aligned_sized{suffix}"), + offsets, + )?; + self.attach_free_if_found(lib_path, &format!("{prefix}cfree{suffix}"), offsets)?; + } + } + + Ok(()) + } + + /// Attach standard library allocation probes (libc-style: malloc, free, calloc, etc.) + /// This works for libc and allocators that export standard symbol names. + /// For non-libc allocators, standard names are optional - just try them silently. + fn attach_libc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { + self.attach_standard_probes(lib_path, &[], &[], offsets) + } + + /// Attach C++ operator new/delete probes. + /// These are mangled C++ symbols that wrap the underlying allocator. + /// C++ operators have identical signatures to malloc/free, so we reuse those handlers. + fn attach_libcpp_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { + self.attach_malloc_if_found(lib_path, "_Znwm", offsets)?; // operator new(size_t) + self.attach_malloc_if_found(lib_path, "_Znam", offsets)?; // operator new[](size_t) + self.attach_malloc_if_found(lib_path, "_ZnwmSt11align_val_t", offsets)?; // operator new(size_t, std::align_val_t) + self.attach_malloc_if_found(lib_path, "_ZnamSt11align_val_t", offsets)?; // operator new[](size_t, std::align_val_t) + self.attach_free_if_found(lib_path, "_ZdlPv", offsets)?; // operator delete(void*) + self.attach_free_if_found(lib_path, "_ZdaPv", offsets)?; // operator delete[](void*) + self.attach_free_if_found(lib_path, "_ZdlPvm", offsets)?; // operator delete(void*, size_t) - C++14 sized delete + self.attach_free_if_found(lib_path, "_ZdaPvm", offsets)?; // operator delete[](void*, size_t) - C++14 sized delete + self.attach_free_if_found(lib_path, "_ZdlPvSt11align_val_t", offsets)?; // operator delete(void*, std::align_val_t) + self.attach_free_if_found(lib_path, "_ZdaPvSt11align_val_t", offsets)?; // operator delete[](void*, std::align_val_t) + self.attach_free_if_found(lib_path, "_ZdlPvmSt11align_val_t", offsets)?; // operator delete(void*, size_t, std::align_val_t) + self.attach_free_if_found(lib_path, "_ZdaPvmSt11align_val_t", offsets)?; // operator delete[](void*, size_t, std::align_val_t) + + Ok(()) + } + + /// Attach jemalloc-specific probes (prefixed and extended API). + fn attach_jemalloc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { + // The following functions are used in Rust when setting a global allocator: + // - rust_alloc: _rjem_malloc and _rjem_mallocx + // - rust_alloc_zeroed: _rjem_mallocx / _rjem_calloc + // - rust_dealloc: _rjem_sdallocx + // - rust_realloc: _rjem_realloc / _rjem_rallocx + + // je_* API (internal jemalloc functions, static linking) + // _rjem_* API (Rust jemalloc crate, dynamic linking) + let prefixes = ["je_", "_rjem_"]; + let suffixes = ["", "_default"]; + + self.attach_standard_probes(lib_path, &prefixes, &suffixes, offsets)?; + + // Non-standard API that has an additional flag parameter + // See: https://jemalloc.net/jemalloc.3.html + for prefix in prefixes { + for suffix in suffixes { + self.attach_malloc_if_found( + lib_path, + &format!("{prefix}mallocx{suffix}"), + offsets, + )?; + self.attach_realloc_if_found( + lib_path, + &format!("{prefix}rallocx{suffix}"), + offsets, + )?; + self.attach_free_if_found(lib_path, &format!("{prefix}dallocx{suffix}"), offsets)?; + self.attach_free_if_found(lib_path, &format!("{prefix}sdallocx{suffix}"), offsets)?; + } + } + + Ok(()) + } + + /// Attach mimalloc-specific probes (mi_* API). + fn attach_mimalloc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { + // The following functions are used in Rust when setting a global allocator: + // - mi_malloc_aligned + // - mi_free + // - mi_realloc_aligned + // - mi_zalloc_aligned + + self.attach_standard_probes(lib_path, &["mi_"], &[], offsets)?; + + // Zero-initialized and aligned variants + self.attach_malloc_if_found(lib_path, "mi_malloc_aligned", offsets)?; + self.attach_calloc_if_found(lib_path, "mi_zalloc", offsets)?; + self.attach_calloc_if_found(lib_path, "mi_zalloc_aligned", offsets)?; + self.attach_realloc_if_found(lib_path, "mi_realloc_aligned", offsets)?; + + Ok(()) + } + + /// Attach TCMalloc probes ( tc_* API). + /// + /// See: + /// - https://github.com/google/tcmalloc/blob/master/docs/reference.md + /// - https://github.com/gperftools/gperftools/blob/a47243150ec41097602730ff8779fafcc172d1fb/src/tcmalloc.cc#L178-L190 + fn attach_tcmalloc_probes(&mut self, lib_path: &Path, offsets: &ResolvedSymbols) -> Result<()> { + self.attach_standard_probes(lib_path, &["tc_"], &[], offsets)?; + + self.attach_free_if_found(lib_path, "free_sized", offsets)?; + self.attach_free_if_found(lib_path, "free_aligned_sized", offsets)?; + self.attach_free_if_found(lib_path, "sdallocx", offsets)?; + + Ok(()) + } +} diff --git a/crates/memtrack/src/ebpf/memtrack/macros.rs b/crates/memtrack/src/ebpf/memtrack/macros.rs new file mode 100644 index 00000000..88da64a6 --- /dev/null +++ b/crates/memtrack/src/ebpf/memtrack/macros.rs @@ -0,0 +1,136 @@ +/// Macro to attach a function with both entry and return probes at a resolved +/// file offset. Also generates an `attach_*_if_found` variant that skips +/// symbols absent from the offset table (returning whether it attached) and +/// propagates attach failures. +macro_rules! attach_uprobe_uretprobe { + ($name:ident, $prog_entry:ident, $prog_return:ident) => { + paste! { + fn [](&mut self, lib_path: &Path, offset: usize) -> Result<()> { + let link = self + .skel + .progs + .$prog_entry + .attach_uprobe_with_opts( + -1, + lib_path, + offset, + UprobeOpts { + retprobe: false, + ..Default::default() + }, + ) + .context(format!( + "Failed to attach uprobe at offset {:#x} in {}", + offset, + lib_path.display() + ))?; + self.probes.push(link); + + let link = self + .skel + .progs + .$prog_return + .attach_uprobe_with_opts( + -1, + lib_path, + offset, + UprobeOpts { + retprobe: true, + ..Default::default() + }, + ) + .context(format!( + "Failed to attach uretprobe at offset {:#x} in {}", + offset, + lib_path.display() + ))?; + self.probes.push(link); + + Ok(()) + } + + fn [<$name _if_found>]( + &mut self, + lib_path: &Path, + symbol: &str, + symbols: &ResolvedSymbols, + ) -> Result { + let Some(offset) = symbols.offset(symbol) else { + return Ok(false); + }; + self.[](lib_path, offset) + .with_context(|| format!("Failed to attach {symbol}"))?; + log::trace!("Attached {} at {:#x}", symbol, offset); + Ok(true) + } + } + }; +} + +/// Macro to attach a function with only an entry probe (no return probe) at a +/// resolved file offset. Also generates an `attach_*_if_found` variant that +/// skips symbols absent from the offset table (returning whether it attached) +/// and propagates attach failures. +macro_rules! attach_uprobe { + ($name:ident, $prog:ident) => { + paste! { + fn [](&mut self, lib_path: &Path, offset: usize) -> Result<()> { + let link = self + .skel + .progs + .$prog + .attach_uprobe_with_opts( + -1, + lib_path, + offset, + UprobeOpts { + retprobe: false, + ..Default::default() + }, + ) + .context(format!( + "Failed to attach uprobe at offset {:#x} in {}", + offset, + lib_path.display() + ))?; + self.probes.push(link); + Ok(()) + } + + fn [<$name _if_found>]( + &mut self, + lib_path: &Path, + symbol: &str, + symbols: &ResolvedSymbols, + ) -> Result { + let Some(offset) = symbols.offset(symbol) else { + return Ok(false); + }; + self.[](lib_path, offset) + .with_context(|| format!("Failed to attach {symbol}"))?; + log::trace!("Attached {} at {:#x}", symbol, offset); + Ok(true) + } + } + }; +} + +macro_rules! attach_tracepoint { + ($func:ident, $prog:ident) => { + fn $func(&mut self) -> Result<()> { + let link = self + .skel + .progs + .$prog + .attach() + .context(format!("Failed to attach {} tracepoint", stringify!($prog)))?; + self.probes.push(link); + Ok(()) + } + }; + ($name:ident) => { + paste! { + attach_tracepoint!([], []); + } + }; +} diff --git a/crates/memtrack/src/ebpf/memtrack/maps.rs b/crates/memtrack/src/ebpf/memtrack/maps.rs new file mode 100644 index 00000000..560ca46a --- /dev/null +++ b/crates/memtrack/src/ebpf/memtrack/maps.rs @@ -0,0 +1,66 @@ +use super::MemtrackBpf; +use crate::prelude::*; +use libbpf_rs::MapCore; + +impl MemtrackBpf { + pub fn add_tracked_pid(&mut self, pid: i32) -> Result<()> { + self.skel + .maps + .tracked_pids + .update( + &pid.to_le_bytes(), + &1u8.to_le_bytes(), + libbpf_rs::MapFlags::ANY, + ) + .context("Failed to add PID to uprobes tracked set")?; + + Ok(()) + } + + pub fn enable_tracking(&mut self) -> Result<()> { + let key = 0u32; + let value = true as u8; + self.skel + .maps + .tracking_enabled + .update( + &key.to_le_bytes(), + &value.to_le_bytes(), + libbpf_rs::MapFlags::ANY, + ) + .context("Failed to enable tracking")?; + Ok(()) + } + + pub fn disable_tracking(&mut self) -> Result<()> { + let key = 0u32; + let value = false as u8; + self.skel + .maps + .tracking_enabled + .update( + &key.to_le_bytes(), + &value.to_le_bytes(), + libbpf_rs::MapFlags::ANY, + ) + .context("Failed to disable tracking")?; + Ok(()) + } + + pub fn dropped_events_count(&self) -> Result { + let key = 0u32; + let value = self + .skel + .maps + .dropped_events + .lookup(&key.to_le_bytes(), libbpf_rs::MapFlags::ANY) + .context("Failed to read dropped_events counter")? + .ok_or_else(|| anyhow!("dropped_events slot 0 missing"))?; + + let bytes: [u8; 8] = value + .as_slice() + .try_into() + .map_err(|_| anyhow!("dropped_events value has unexpected size"))?; + Ok(u64::from_le_bytes(bytes)) + } +} diff --git a/crates/memtrack/src/ebpf/memtrack/mod.rs b/crates/memtrack/src/ebpf/memtrack/mod.rs new file mode 100644 index 00000000..6030a29f --- /dev/null +++ b/crates/memtrack/src/ebpf/memtrack/mod.rs @@ -0,0 +1,168 @@ +use crate::prelude::*; +use libbpf_rs::Link; +use libbpf_rs::skel::OpenSkel; +use libbpf_rs::skel::SkelBuilder; +use std::collections::HashMap; +use std::mem::MaybeUninit; +use std::path::Path; + +use crate::ebpf::poller::RingBufferPoller; + +pub mod memtrack_skel { + include!(concat!(env!("OUT_DIR"), "/memtrack.skel.rs")); +} +pub use memtrack_skel::*; + +#[macro_use] +mod macros; +mod allocator; +mod maps; +mod tracking; + +/// Resolve libbpf attach targets for every defined symbol in `lib_path`. +pub fn resolve_symbol_offsets(lib_path: &Path) -> Result { + use object::{Object, ObjectSymbol}; + + let data = std::fs::read(lib_path)?; + let file = object::File::parse(&*data)?; + let mut offsets = HashMap::new(); + + for symbol in file.symbols().chain(file.dynamic_symbols()) { + if !symbol.is_definition() { + continue; + } + + let Ok(name) = symbol.name() else { + continue; + }; + + if let Some(file_offset) = symbol_file_offset(&file, &symbol) { + offsets.insert(name.to_owned(), file_offset); + } + } + + Ok(ResolvedSymbols { offsets }) +} + +/// The libbpf file offset for `symbol`, or `None` when it has no address in a +/// file-backed section (absolute, `SHT_NOBITS`, ...). +fn symbol_file_offset<'a>( + file: &object::File, + symbol: &impl object::ObjectSymbol<'a>, +) -> Option { + use object::{Object, ObjectSection}; + + let address = symbol.address(); + if address == 0 { + return None; + } + + let section = file.section_by_index(symbol.section_index()?).ok()?; + let (sh_offset, _) = section.file_range()?; + Some((address - section.address() + sh_offset) as usize) +} + +/// Attach targets resolved from a library's symbol tables. +pub struct ResolvedSymbols { + offsets: HashMap, +} + +impl ResolvedSymbols { + fn offset(&self, symbol: &str) -> Option { + self.offsets.get(symbol).copied() + } +} + +pub struct MemtrackBpf { + skel: Box>, + probes: Vec, +} + +impl MemtrackBpf { + pub fn new() -> Result { + let builder = MainSkelBuilder::default(); + let open_object = Box::leak(Box::new(MaybeUninit::uninit())); + let open_skel = builder + .open(open_object) + .context("Failed to open syscalls BPF skeleton")?; + + let skel = Box::new( + open_skel + .load() + .context("Failed to load syscalls BPF skeleton")?, + ); + + Ok(Self { + skel, + probes: Vec::new(), + }) + } + + /// Start polling, forwarding each event over the returned channel. + pub fn start_polling_with_channel( + &self, + poll_timeout_ms: u64, + ) -> Result<( + RingBufferPoller, + crossbeam_channel::Receiver, + )> { + RingBufferPoller::with_channel(&self.skel.maps.events, poll_timeout_ms) + } + + /// Detach all BPF links in parallel. Closing a uprobe link blocks on two + /// RCU grace periods in the kernel, but concurrent waiters share grace + /// periods, so closing from many threads scales near-linearly. + pub fn detach_probes(&mut self) { + const DETACH_THREADS: usize = 32; + + let mut probes = std::mem::take(&mut self.probes); + if probes.is_empty() { + return; + } + + debug!("Detaching {} BPF links", probes.len()); + let start = std::time::Instant::now(); + let chunk_size = probes.len().div_ceil(DETACH_THREADS); + std::thread::scope(|scope| { + while !probes.is_empty() { + let split_at = probes.len().saturating_sub(chunk_size); + let chunk = probes.split_off(split_at); + scope.spawn(move || drop(chunk)); + } + }); + debug!("Detached BPF links in {:?}", start.elapsed()); + } +} + +impl Drop for MemtrackBpf { + fn drop(&mut self) { + self.detach_probes(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Allocator entry points must resolve to file offsets; a symbol that + /// silently fails to resolve attaches nothing and loses all events. + #[test] + fn libc_allocator_symbols_resolve_to_offsets() { + let maps = std::fs::read_to_string("/proc/self/maps").unwrap(); + let libc_path = maps + .lines() + .find_map(|line| { + let path = line.split_whitespace().last()?; + path.contains("libc.so.6").then(|| path.to_owned()) + }) + .expect("test process has no mapped libc.so.6"); + + let symbols = resolve_symbol_offsets(Path::new(&libc_path)).unwrap(); + for symbol in ["malloc", "calloc", "realloc", "free"] { + assert!( + symbols.offset(symbol).is_some(), + "{symbol} in {libc_path} did not resolve to a file offset" + ); + } + } +} diff --git a/crates/memtrack/src/ebpf/memtrack/tracking.rs b/crates/memtrack/src/ebpf/memtrack/tracking.rs new file mode 100644 index 00000000..7661c7e5 --- /dev/null +++ b/crates/memtrack/src/ebpf/memtrack/tracking.rs @@ -0,0 +1,12 @@ +use super::MemtrackBpf; +use crate::prelude::*; +use paste::paste; + +impl MemtrackBpf { + attach_tracepoint!(sched_fork); + + pub fn attach_tracepoints(&mut self) -> Result<()> { + self.attach_sched_fork()?; + Ok(()) + } +} From 181f056fde2fccfbba331ad0ad9557740584aec7 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Thu, 16 Jul 2026 19:12:49 +0200 Subject: [PATCH 2/2] docs(memtrack): add crate AGENTS.md guide --- crates/memtrack/AGENTS.md | 85 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 crates/memtrack/AGENTS.md diff --git a/crates/memtrack/AGENTS.md b/crates/memtrack/AGENTS.md new file mode 100644 index 00000000..0c8d86d4 --- /dev/null +++ b/crates/memtrack/AGENTS.md @@ -0,0 +1,85 @@ +# Repository Guidelines + +`memtrack` is the eBPF-based memory-allocation tracker of the CodSpeed runner (workspace member, Linux-only). This guide covers the crate at `crates/memtrack/`; the workspace root has its own `AGENTS.md`. + +## Project Overview + +Attaches uprobes/uretprobes to allocator functions (`malloc`/`free`/`calloc`/`realloc`/`aligned_alloc`/`memalign`) and tracepoints to `mmap`/`munmap`/`brk` + `sched_process_fork` in a target process tree, streams allocation events through a BPF ring buffer to userspace, and writes them to a `MemtrackArtifact` file. Ships a CLI binary `codspeed-memtrack track`. + +## Architecture & Data Flow + +Allocation → disk pipeline: + +1. Allocator entry uprobe stores args in a per-tid BPF hash map (`_arg` / `_args`). +2. Uretprobe reads the stored arg + `PT_REGS_RC`. The `SUBMIT_EVENT` macro gates on `is_tracked(pid)` (`tracked_pids` map, ancestor walk ≤5 levels via `pids_ppid`, plus `sched_fork` auto-tracking of children) **and** `is_enabled()` (`tracking_enabled` map), then `bpf_ringbuf_reserve`/`submit` into the 256 MiB `events` ring buffer. Reserve failure atomically bumps `dropped_events`. `bpf_ringbuf_submit` uses `wake_flags()`: it wakes the consumer only once ≥64 KiB (`WAKEUP_DATA_SIZE`) of unconsumed data has accumulated (`BPF_RB_FORCE_WAKEUP`), otherwise `BPF_RB_NO_WAKEUP` — amortizing per-event wakeups to ~1 per thousand events. +3. `RingBufferPoller` (`src/ebpf/poller.rs`) polls every 10 ms; each `poll()` is followed by a `consume()` that drains the buffer to empty in one call (no extra `epoll_wait`), and the 10 ms timeout flushes the sub-watermark tail. `parse_event` (`src/ebpf/events.rs`) casts raw bytes to the bindgen `event` struct and maps `EVENT_TYPE_*` → `runner_shared::MemtrackEventKind` → mpsc channel. +4. `Tracker::track` (`src/ebpf/tracker.rs`) forwards poller events over a keep-alive thread. +5. `main.rs` spawns one pipeline thread running `encode_events` (`runner-shared`), which fans out across `available_parallelism() - 2` rayon workers (two cores reserved for the poll thread + tracked command so encode workers don't starve the poller and overflow the ring buffer) and batches through `MemtrackWriter`. After the child exits, it checks `dropped_events_count()` and **bails if non-zero** (incomplete trace). + +Control plane: `src/ipc.rs` exposes an out-of-band `ipc-channel` protocol (`Enable`/`Disable`/`Ping`) so the runner toggles the `tracking_enabled` map at runtime. Without `--ipc-server`, tracking is enabled up front. + +Allocator discovery (`src/allocators/`): `AllocatorLib::find_all()` = dynamic (glob shared libs incl. `/nix/store/*` hints) + static-linked (scan build-dir ELF symbols) + env (`CODSPEED_MEMTRACK_BINARIES`). Each `AllocatorKind` (`Libc`/`LibCpp`/`Jemalloc`/`Mimalloc`/`Tcmalloc`) maps to best-effort attach helpers; only libc must succeed. + +> Note: the "on-demand attach" design in `.agents/docs/` (AttachWorker, `CODSPEED_MEMTRACK_ONDEMAND`, SIGSTOP/SIGCONT) is a **plan, not yet in source**. Current behavior is upfront attach + `sched_fork` auto-tracking. + +## Key Directories + +- `src/ebpf/` — BPF stack (feature-gated `ebpf`): `tracker.rs` (facade), `memtrack/` (libbpf-rs wrapper + generated skeleton, split into `mod.rs`/`macros.rs`/`maps.rs`/`allocator.rs`/`tracking.rs`), `poller.rs`, `events.rs`, `c/main.bpf.c` + `c/event.h` + `c/utils/*.h` + `c/allocator.h`. +- `src/allocators/` — allocator classification: `mod.rs`, `dynamic.rs`, `static_linked.rs`. +- `tests/` — integration tests + `snapshots/` (insta). +- `testdata/` — allocation fixtures: `*.c` (gcc), `alloc_cpp/` (cmkr/CMake), `alloc_rust/` + `spawn_wrapper/` (standalone Cargo workspaces). +- `.agents/docs/`, `.claude/` — design notes (some current, some historical/stale). + +## Development Commands + +```bash +cargo build # default features include `ebpf` +cargo check +cargo fmt +cargo clippy +cargo test --lib # unit tests (no root) + +# Run the tracker (needs root); tracks a shell command's whole process tree: +sudo -E cargo run --bin codspeed-memtrack -- track "" --output +# e.g. sudo -E cargo run --bin codspeed-memtrack -- track "ls / >/dev/null" --output . + +# Integration tests need BPF privilege + GITHUB_ACTIONS gate + single-threaded: +export GITHUB_ACTIONS=1 +sudo -E cargo test --test c_tests -- --test-threads 1 +``` + +`--test-threads 1` is **mandatory** — eBPF probes cannot self-overlap. CI runs `sudo -E cargo test --lib --test -- --test-threads 1`; the main workspace `tests` job excludes memtrack (`--exclude memtrack`). + +## Code Conventions & Common Patterns + +- **Errors:** `anyhow` only (`Result`/`Context`/`bail`/`ensure` via `src/prelude.rs`); no `thiserror`. +- **Concurrency:** no async runtime — `std::thread` + `std::sync::mpsc` + `Arc>`. +- **Naming:** snake_case modules, PascalCase types. Macro-generated `try_` (fallible) vs `_if_found` (best-effort, trace-logs errors) attach helpers via `paste!`. +- **BPF C:** macro-heavy (`UPROBE_ARG_RET`/`UPROBE_RET`/`UPROBE_ARGS_RET`/`SUBMIT_EVENT`); `.clang-format` is Google-based, 4-space indent, 100-col, left pointers; vmlinux.h include wrapped in `// clang-format off/on`. +- **Module layout:** `src/lib.rs` re-exports and `#[cfg(feature = "ebpf")]` gates the whole BPF stack + binary; `prelude.rs` centralizes error/log imports. +- **Cross-crate types:** `MemtrackEvent` / `MemtrackEventKind` live in `runner-shared/src/artifacts/memtrack.rs` (a breaking change there ripples here). + +## Important Files + +- `src/main.rs` — CLI entry (`codspeed-memtrack track`, requires `ebpf`); drops sudo privileges to `SUDO_UID`/`SUDO_GID` so the tracked child runs unprivileged. +- `src/lib.rs` — crate root / re-exports. +- `src/ebpf/c/main.bpf.c` + `c/event.h` + `c/utils/*.h` + `c/allocator.h` — BPF program (includes only) and event struct definitions; maps/programs live in headers. +- `build.rs` — feature-gated: `libbpf_cargo::SkeletonBuilder` compiles `main.bpf.c` → `OUT_DIR/memtrack.skel.rs`; `bindgen` on `wrapper.h` → `OUT_DIR/event.rs`. Reruns on `src/ebpf/c` changes and on `GITHUB_ACTIONS` env change. +- `wrapper.h` — bindgen entry (`stdint.h` + `src/ebpf/c/event.h`). +- `Cargo.toml` — lib `memtrack` + bin `codspeed-memtrack`; feature `ebpf` (default) pulls `libbpf-rs`/`libbpf-cargo`/`vmlinux`. + +## Runtime / Tooling Prerequisites + +- **Linux only** — consumed from workspace root under `cfg(target_os = "linux")`, targets `x86_64`/`aarch64-unknown-linux-gnu`. +- **Root / BPF privilege** required at runtime (bumps `RLIMIT_MEMLOCK`, loads BPF). +- **Build toolchain:** `clang` + BTF/vmlinux headers, `libbpf-dev`, `zlib1g-dev`, `pkgconf`, `build-essential`; vendored libbpf also needs `autopoint`/`bison`/`flex`. +- `vmlinux.h` is pinned to a specific git rev; `libbpf-rs` uses the `vendored` feature (dist links `libbpf-rs/static`). + +Env vars actually wired: `CODSPEED_MEMTRACK_BINARIES` (extra static-allocator binaries), `CODSPEED_LOG` (log filter, default `info`), `SUDO_UID`/`SUDO_GID` (privilege drop), `GITHUB_ACTIONS` (build rebuild trigger + test gate). + +## Testing & QA + +- **Frameworks:** `insta` (snapshots), `rstest` (parametrized `#[case]`), `test-log`, `test-with` (`#[test_with::env(GITHUB_ACTIONS)]`), `tempfile`. +- **Harness:** `tests/shared.rs` — `assert_events_snapshot!` / `assert_events_with_marker!` (dedup by addr+discriminant, `Realloc` strips `old_addr`, `0xC0D59EED` marker windowing), `compile_rust_binary`, `track_command`/`track_binary[_with_opts]`. +- **Suites:** `c_tests` (8 gcc-compiled C fixtures, no marker), `cpp_tests` (7 cmkr targets: system + jemalloc/mimalloc/tcmalloc × static/dynamic), `rust_tests` (system/jemalloc/mimalloc via features), `spawn_tests` (static-allocator discovery across `exec`). +- **Snapshots:** `tests/snapshots/__.snap`, address-stripped and deduped. Every integration test is `GITHUB_ACTIONS`-gated and needs BPF privilege; unit tests (`src/ebpf/events.rs`) run without root.