Skip to content

perf(memtrack): reduce tracking overhead #518

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
not-matthias wants to merge 4 commits into from
Open

perf(memtrack): reduce tracking overhead #518

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3b433ab
perf(memtrack): attach each probe site once
not-matthias Aug 21, 2026
3d73583
perf(memtrack): hand off allocator arguments in task-local storage
not-matthias Aug 21, 2026
f39317c
perf(memtrack): decide ring buffer wakeups producer-side
not-matthias Aug 21, 2026
37d61db
feat(memtrack): log with millisecond timestamps
not-matthias Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 65 additions & 114 deletions crates/memtrack/src/ebpf/c/allocator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,24 +5,23 @@
#include "utils/map_helpers.h"
#include "utils/process_tracking.h"

#define UPROBE_ARG_RET(name, arg_expr, submit_block) \
BPF_HASH_MAP(name##_arg, __u64, __u64, 10000); \
SEC(UPROBE_SEC) \
int uprobe_##name(struct pt_regs* ctx) { \
return store_param(&name##_arg, arg_expr); \
} \
SEC(URETPROBE_SEC) \
int uretprobe_##name(struct pt_regs* ctx) { \
__u64* arg_ptr = take_param(&name##_arg); \
if (!arg_ptr) { \
return 0; \
} \
__u64 ret_val = PT_REGS_RC(ctx); \
if (ret_val == 0) { \
return 0; \
} \
__u64 arg0 = *arg_ptr; \
submit_block; \
#define UPROBE_ARG_RET(name, slot, arg_expr, submit_block) \
SEC(UPROBE_SEC) \
int uprobe_##name(struct pt_regs* ctx) { \
return store_arg(slot, arg_expr); \
} \
SEC(URETPROBE_SEC) \
int uretprobe_##name(struct pt_regs* ctx) { \
struct memtrack_task_state* st = take_slot(slot); \
if (!st) { \
return 0; \
} \
__u64 ret_val = PT_REGS_RC(ctx); \
if (ret_val == 0) { \
return 0; \
} \
__u64 arg0 = st->arg0[slot]; \
submit_block; \
}

#define UPROBE_RET(name, arg_expr, submit_block) \
Expand All @@ -32,65 +31,48 @@
if (arg0 == 0) { \
return 0; \
} \
if (!tracked_state()) { \
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_SEC) \
int uprobe_##name(struct pt_regs* ctx) { \
struct task_ids ids = current_task_ids(); \
__u64 tid = ids.tid; \
\
if (!is_tracked(ids.tgid)) { \
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_SEC) \
int uretprobe_##name(struct pt_regs* ctx) { \
__u64 tid = current_tid(); \
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; \
#define UPROBE_ARGS_RET(name, slot, arg0_expr, arg1_expr, submit_block) \
SEC(UPROBE_SEC) \
int uprobe_##name(struct pt_regs* ctx) { \
return store_args(slot, arg0_expr, arg1_expr); \
} \
SEC(URETPROBE_SEC) \
int uretprobe_##name(struct pt_regs* ctx) { \
struct memtrack_task_state* st = take_slot(slot); \
if (!st) { \
return 0; \
} \
__u64 ret_val = PT_REGS_RC(ctx); \
if (ret_val == 0) { \
return 0; \
} \
__u64 arg0 = st->arg0[slot]; \
__u64 arg1 = st->arg1[slot]; \
submit_block; \
}

UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx), { return submit_alloc_event(arg0, ret_val); })
UPROBE_ARG_RET(malloc, SLOT_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),
UPROBE_ARG_RET(calloc, SLOT_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),
UPROBE_ARGS_RET(realloc, SLOT_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),
UPROBE_ARG_RET(aligned_alloc, SLOT_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); })
UPROBE_ARG_RET(memalign, SLOT_MEMALIGN, PT_REGS_PARM2(ctx),
{ return submit_aligned_alloc_event(arg0, ret_val); })

/*
* posix_memalign(void** memptr, size_t alignment, size_t size)
Expand All @@ -101,74 +83,42 @@ UPROBE_ARG_RET(memalign, PT_REGS_PARM2(ctx), { return submit_aligned_alloc_event
* ret == 0 (not a non-NULL return), and the address must be read back from
* *memptr once the call returns.
*/
struct posix_memalign_args_t {
__u64 memptr;
__u64 size;
};
BPF_HASH_MAP(posix_memalign_args, __u64, struct posix_memalign_args_t, 10000);

SEC(UPROBE_SEC)
int uprobe_posix_memalign(struct pt_regs* ctx) {
struct task_ids ids = current_task_ids();
__u64 tid = ids.tid;
if (!is_tracked(ids.tgid)) {
return 0;
}

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;
return store_args(SLOT_POSIX_MEMALIGN, PT_REGS_PARM1(ctx), PT_REGS_PARM3(ctx));
}

SEC(URETPROBE_SEC)
int uretprobe_posix_memalign(struct pt_regs* ctx) {
__u64 tid = current_tid();
struct posix_memalign_args_t* args = bpf_map_lookup_elem(&posix_memalign_args, &tid);
if (!args) {
struct memtrack_task_state* st = take_slot(SLOT_POSIX_MEMALIGN);
if (!st) {
return 0;
}

struct posix_memalign_args_t a = *args;
bpf_map_delete_elem(&posix_memalign_args, &tid);

if (PT_REGS_RC(ctx) != 0) {
return 0;
}

__u64 memptr = st->arg0[SLOT_POSIX_MEMALIGN];
__u64 size = st->arg1[SLOT_POSIX_MEMALIGN];

__u64 addr = 0;
if (bpf_probe_read_user(&addr, sizeof(addr), (void*)a.memptr) != 0 || addr == 0) {
if (bpf_probe_read_user(&addr, sizeof(addr), (void*)memptr) != 0 || addr == 0) {
return 0;
}

return submit_aligned_alloc_event(a.size, addr);
}

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) {
struct task_ids ids = current_task_ids();
__u64 tid = ids.tid;
if (is_tracked(ids.tgid)) {
struct mmap_args args = {.addr = addr, .len = len};
bpf_map_update_elem(&mmap_temp, &tid, &args, BPF_ANY);
}
return submit_aligned_alloc_event(size, addr);
}

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;
return store_args(SLOT_MMAP, ctx->args[0], ctx->args[1]);
}

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) {
struct memtrack_task_state* st = take_slot(SLOT_MMAP);
if (!st) {
return 0;
}

Expand All @@ -177,7 +127,7 @@ int tracepoint_sys_exit_mmap(struct trace_event_raw_sys_exit* ctx) {
return 0;
}

return submit_mmap_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
return submit_mmap_event((__u64)ret, st->arg1[SLOT_MMAP], EVENT_TYPE_MMAP);
}

SEC("tracepoint/syscalls/sys_enter_munmap")
Expand All @@ -189,26 +139,27 @@ int tracepoint_sys_enter_munmap(struct trace_event_raw_sys_enter* ctx) {
return 0;
}

if (!tracked_state()) {
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;
return store_arg(SLOT_BRK, ctx->args[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) {
struct memtrack_task_state* st = take_slot(SLOT_BRK);
if (!st) {
return 0;
}

__u64 new_brk = ctx->ret;
__u64 req_brk = *requested_brk;
__u64 req_brk = st->arg0[SLOT_BRK];

if (req_brk == 0 || new_brk <= 0) {
return 0;
Expand Down
Loading