Skip to content

Commit abe57c5

Browse files
not-matthiasnot-matthias
committed
refactor(memtrack): restructure event handling with a common header and improve event parsing
This commit also saves the old address for `realloc` events. We previously only saved the new address, which lead to inaccuracies when building the timeline.
1 parent 93448ad commit abe57c5

8 files changed

Lines changed: 357 additions & 115 deletions

File tree

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

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,42 @@
1010
#define EVENT_TYPE_MUNMAP 7
1111
#define EVENT_TYPE_BRK 8
1212

13-
/* Event structure - shared between BPF and userspace */
14-
struct event {
13+
/* Common header shared by all event types */
14+
struct event_header {
1515
uint8_t event_type; /* See EVENT_TYPE_* constants above */
1616
uint64_t timestamp; /* monotonic time in nanoseconds (CLOCK_MONOTONIC) */
1717
uint32_t pid;
1818
uint32_t tid;
19-
uint64_t addr; /* address returned/freed */
20-
uint64_t size; /* size requested */
19+
};
20+
21+
/* Tagged union event structure */
22+
struct event {
23+
struct event_header header;
24+
union {
25+
/* Allocation events (malloc, calloc, aligned_alloc) */
26+
struct {
27+
uint64_t addr; /* address returned */
28+
uint64_t size; /* size requested */
29+
} alloc;
30+
31+
/* Deallocation event (free) */
32+
struct {
33+
uint64_t addr; /* address to free */
34+
} free;
35+
36+
/* Reallocation event - includes both old and new addresses */
37+
struct {
38+
uint64_t old_addr; /* previous address (can be NULL) */
39+
uint64_t new_addr; /* new address returned */
40+
uint64_t size; /* new size requested */
41+
} realloc;
42+
43+
/* Memory mapping events (mmap, munmap, brk) */
44+
struct {
45+
uint64_t addr; /* address of mapping */
46+
uint64_t size; /* size of mapping */
47+
} mmap;
48+
} data;
2149
};
2250

2351
#endif /* __EVENT_H__ */

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

Lines changed: 141 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -123,77 +123,171 @@ static __always_inline __u64* take_param(void* map) {
123123
return value;
124124
}
125125

126-
/* Helper to submit an event to the ring buffer */
127-
static __always_inline int submit_event(__u64 addr, __u64 size, __u8 event_type) {
128-
__u64 tid = bpf_get_current_pid_tgid();
129-
__u32 pid = tid >> 32;
130-
131-
if (!is_tracked(pid) || !is_enabled()) {
132-
return 0;
126+
/* Macro to handle common event submission boilerplate
127+
* Usage: SUBMIT_EVENT(event_type, { e->data.foo = bar; })
128+
*/
129+
#define SUBMIT_EVENT(evt_type, fill_data) \
130+
{ \
131+
__u64 tid = bpf_get_current_pid_tgid(); \
132+
__u32 pid = tid >> 32; \
133+
\
134+
if (!is_tracked(pid) || !is_enabled()) { \
135+
return 0; \
136+
} \
137+
\
138+
struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0); \
139+
if (!e) { \
140+
return 0; \
141+
} \
142+
\
143+
e->header.timestamp = bpf_ktime_get_ns(); \
144+
e->header.pid = pid; \
145+
e->header.tid = tid & 0xFFFFFFFF; \
146+
e->header.event_type = evt_type; \
147+
\
148+
fill_data; \
149+
\
150+
bpf_ringbuf_submit(e, 0); \
151+
return 0; \
133152
}
134153

135-
struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
136-
if (!e) {
137-
return 0;
138-
}
154+
/* Helper to submit an allocation event (malloc, calloc, aligned_alloc) */
155+
static __always_inline int submit_alloc_event(__u64 size, __u64 addr) {
156+
SUBMIT_EVENT(EVENT_TYPE_MALLOC, {
157+
e->data.alloc.addr = addr;
158+
e->data.alloc.size = size;
159+
});
160+
}
139161

140-
e->timestamp = bpf_ktime_get_ns();
141-
e->pid = pid;
142-
e->tid = tid & 0xFFFFFFFF;
143-
e->event_type = event_type;
144-
e->addr = addr;
145-
e->size = size;
146-
bpf_ringbuf_submit(e, 0);
162+
/* Helper to submit a free event */
163+
static __always_inline int submit_free_event(__u64 addr) {
164+
SUBMIT_EVENT(EVENT_TYPE_FREE, {
165+
e->data.free.addr = addr;
166+
});
167+
}
147168

148-
return 0;
169+
/* Helper to submit a realloc event with both old and new addresses */
170+
static __always_inline int submit_realloc_event(__u64 old_addr, __u64 new_addr, __u64 size) {
171+
SUBMIT_EVENT(EVENT_TYPE_REALLOC, {
172+
e->data.realloc.old_addr = old_addr;
173+
e->data.realloc.new_addr = new_addr;
174+
e->data.realloc.size = size;
175+
});
149176
}
150177

151-
/* Macro to generate uprobe/uretprobe pairs for allocation functions */
152-
#define UPROBE_WITH_ARGS(name, size_expr, addr_expr, event_type) \
153-
BPF_HASH_MAP(name##_size, __u64, __u64, 10000); \
178+
/* Helper to submit a memory mapping event (mmap, munmap, brk) */
179+
static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_type) {
180+
SUBMIT_EVENT(event_type, {
181+
e->data.mmap.addr = addr;
182+
e->data.mmap.size = size;
183+
});
184+
}
185+
186+
/* Macro to generate uprobe/uretprobe pairs for allocation functions with 1 argument */
187+
#define UPROBE_ARG_RET(name, arg_expr, submit_block) \
188+
BPF_HASH_MAP(name##_arg, __u64, __u64, 10000); \
154189
SEC("uprobe") \
155-
int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_size, size_expr); } \
190+
int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_arg, arg_expr); } \
156191
SEC("uretprobe") \
157192
int uretprobe_##name(struct pt_regs* ctx) { \
158-
__u64* size_ptr = take_param(&name##_size); \
159-
if (!size_ptr) { \
193+
__u64* arg_ptr = take_param(&name##_arg); \
194+
if (!arg_ptr) { \
160195
return 0; \
161196
} \
162-
__u64 addr = addr_expr; \
163-
if (addr == 0) { \
197+
__u64 ret_val = PT_REGS_RC(ctx); \
198+
if (ret_val == 0) { \
164199
return 0; \
165200
} \
166-
return submit_event(addr, *size_ptr, event_type); \
201+
__u64 arg0 = *arg_ptr; \
202+
submit_block; \
203+
}
204+
205+
/* Macro for simple return value only functions like free */
206+
#define UPROBE_RET(name, arg_expr, submit_block) \
207+
SEC("uprobe") \
208+
int uprobe_##name(struct pt_regs* ctx) { \
209+
__u64 arg0 = arg_expr; \
210+
if (arg0 == 0) { \
211+
return 0; \
212+
} \
213+
submit_block; \
167214
}
168215

169-
/* Macro for simple address-only functions like free */
170-
#define UPROBE_ADDR_ONLY(name, addr_expr, event_type) \
171-
SEC("uprobe") \
172-
int uprobe_##name(struct pt_regs* ctx) { \
173-
__u64 addr = addr_expr; \
174-
if (addr == 0) { \
175-
return 0; \
176-
} \
177-
return submit_event(addr, 0, event_type); \
216+
/* Macro to generate uprobe/uretprobe pairs for functions with 2 arguments */
217+
#define UPROBE_ARGS_RET(name, arg0_expr, arg1_expr, submit_block) \
218+
struct name##_args_t { \
219+
__u64 arg0; \
220+
__u64 arg1; \
221+
}; \
222+
BPF_HASH_MAP(name##_args, __u64, struct name##_args_t, 10000); \
223+
SEC("uprobe") \
224+
int uprobe_##name(struct pt_regs* ctx) { \
225+
__u64 tid = bpf_get_current_pid_tgid(); \
226+
__u32 pid = tid >> 32; \
227+
\
228+
if (!is_tracked(pid)) { \
229+
return 0; \
230+
} \
231+
\
232+
struct name##_args_t args = { \
233+
.arg0 = arg0_expr, \
234+
.arg1 = arg1_expr \
235+
}; \
236+
\
237+
bpf_map_update_elem(&name##_args, &tid, &args, BPF_ANY); \
238+
return 0; \
239+
} \
240+
SEC("uretprobe") \
241+
int uretprobe_##name(struct pt_regs* ctx) { \
242+
__u64 tid = bpf_get_current_pid_tgid(); \
243+
struct name##_args_t* args = bpf_map_lookup_elem(&name##_args, &tid); \
244+
\
245+
if (!args) { \
246+
return 0; \
247+
} \
248+
\
249+
struct name##_args_t a = *args; \
250+
bpf_map_delete_elem(&name##_args, &tid); \
251+
\
252+
__u64 ret_val = PT_REGS_RC(ctx); \
253+
if (ret_val == 0) { \
254+
return 0; \
255+
} \
256+
\
257+
__u64 arg0 = a.arg0; \
258+
__u64 arg1 = a.arg1; \
259+
submit_block; \
178260
}
179261

180262
/* malloc: allocates with size parameter */
181-
UPROBE_WITH_ARGS(malloc, PT_REGS_PARM1(ctx), PT_REGS_RC(ctx), EVENT_TYPE_MALLOC)
263+
UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx), {
264+
return submit_alloc_event(arg0, ret_val);
265+
})
182266

183267
/* free: deallocates by address */
184-
UPROBE_ADDR_ONLY(free, PT_REGS_PARM1(ctx), EVENT_TYPE_FREE)
268+
UPROBE_RET(free, PT_REGS_PARM1(ctx), {
269+
return submit_free_event(arg0);
270+
})
185271

186272
/* calloc: allocates with nmemb * size */
187-
UPROBE_WITH_ARGS(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_CALLOC)
273+
UPROBE_ARG_RET(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx), {
274+
return submit_alloc_event(arg0, ret_val);
275+
})
188276

189-
/* realloc: reallocates with new size */
190-
UPROBE_WITH_ARGS(realloc, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_REALLOC)
277+
/* realloc: reallocates with old_addr and new size */
278+
UPROBE_ARGS_RET(realloc, PT_REGS_PARM2(ctx), PT_REGS_PARM1(ctx), {
279+
return submit_realloc_event(arg1, ret_val, arg0);
280+
})
191281

192282
/* aligned_alloc: allocates with alignment and size */
193-
UPROBE_WITH_ARGS(aligned_alloc, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_ALIGNED_ALLOC)
283+
UPROBE_ARG_RET(aligned_alloc, PT_REGS_PARM2(ctx), {
284+
return submit_alloc_event(arg0, ret_val);
285+
})
194286

195287
/* memalign: allocates with alignment and size (legacy interface) */
196-
UPROBE_WITH_ARGS(memalign, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_ALIGNED_ALLOC)
288+
UPROBE_ARG_RET(memalign, PT_REGS_PARM2(ctx), {
289+
return submit_alloc_event(arg0, ret_val);
290+
})
197291

198292
/* Map to store mmap parameters between entry and return */
199293
struct mmap_args {
@@ -234,7 +328,7 @@ int tracepoint_sys_exit_mmap(struct trace_event_raw_sys_exit* ctx) {
234328
return 0;
235329
}
236330

237-
return submit_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
331+
return submit_mmap_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
238332
}
239333

240334
/* munmap tracking */
@@ -247,7 +341,7 @@ int tracepoint_sys_enter_munmap(struct trace_event_raw_sys_enter* ctx) {
247341
return 0;
248342
}
249343

250-
return submit_event(addr, len, EVENT_TYPE_MUNMAP);
344+
return submit_mmap_event(addr, len, EVENT_TYPE_MUNMAP);
251345
}
252346

253347
/* brk tracking - adjusts the program break (heap boundary) */
@@ -281,5 +375,5 @@ int tracepoint_sys_exit_brk(struct trace_event_raw_sys_exit* ctx) {
281375
return 0;
282376
}
283377

284-
return submit_event(new_brk, 0, EVENT_TYPE_BRK);
378+
return submit_mmap_event(new_brk, 0, EVENT_TYPE_BRK);
285379
}

0 commit comments

Comments
 (0)