Skip to content

Commit 47dbac0

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 f33f588 commit 47dbac0

8 files changed

Lines changed: 367 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: 157 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -123,77 +123,187 @@ 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) */
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+
}
161+
162+
/* Helper to submit an aligned allocation event (aligned_alloc, memalign) */
163+
static __always_inline int submit_aligned_alloc_event(__u64 size, __u64 addr) {
164+
SUBMIT_EVENT(EVENT_TYPE_ALIGNED_ALLOC, {
165+
e->data.alloc.addr = addr;
166+
e->data.alloc.size = size;
167+
});
168+
}
139169

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);
170+
/* Helper to submit a calloc event */
171+
static __always_inline int submit_calloc_event(__u64 size, __u64 addr) {
172+
SUBMIT_EVENT(EVENT_TYPE_CALLOC, {
173+
e->data.alloc.addr = addr;
174+
e->data.alloc.size = size;
175+
});
176+
}
147177

148-
return 0;
178+
/* Helper to submit a free event */
179+
static __always_inline int submit_free_event(__u64 addr) {
180+
SUBMIT_EVENT(EVENT_TYPE_FREE, {
181+
e->data.free.addr = addr;
182+
});
183+
}
184+
185+
/* Helper to submit a realloc event with both old and new addresses */
186+
static __always_inline int submit_realloc_event(__u64 old_addr, __u64 new_addr, __u64 size) {
187+
SUBMIT_EVENT(EVENT_TYPE_REALLOC, {
188+
e->data.realloc.old_addr = old_addr;
189+
e->data.realloc.new_addr = new_addr;
190+
e->data.realloc.size = size;
191+
});
149192
}
150193

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); \
194+
/* Helper to submit a memory mapping event (mmap, munmap, brk) */
195+
static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_type) {
196+
SUBMIT_EVENT(event_type, {
197+
e->data.mmap.addr = addr;
198+
e->data.mmap.size = size;
199+
});
200+
}
201+
202+
/* Macro to generate uprobe/uretprobe pairs for allocation functions with 1 argument */
203+
#define UPROBE_ARG_RET(name, arg_expr, submit_block) \
204+
BPF_HASH_MAP(name##_arg, __u64, __u64, 10000); \
154205
SEC("uprobe") \
155-
int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_size, size_expr); } \
206+
int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_arg, arg_expr); } \
156207
SEC("uretprobe") \
157208
int uretprobe_##name(struct pt_regs* ctx) { \
158-
__u64* size_ptr = take_param(&name##_size); \
159-
if (!size_ptr) { \
209+
__u64* arg_ptr = take_param(&name##_arg); \
210+
if (!arg_ptr) { \
160211
return 0; \
161212
} \
162-
__u64 addr = addr_expr; \
163-
if (addr == 0) { \
213+
__u64 ret_val = PT_REGS_RC(ctx); \
214+
if (ret_val == 0) { \
164215
return 0; \
165216
} \
166-
return submit_event(addr, *size_ptr, event_type); \
217+
__u64 arg0 = *arg_ptr; \
218+
submit_block; \
219+
}
220+
221+
/* Macro for simple return value only functions like free */
222+
#define UPROBE_RET(name, arg_expr, submit_block) \
223+
SEC("uprobe") \
224+
int uprobe_##name(struct pt_regs* ctx) { \
225+
__u64 arg0 = arg_expr; \
226+
if (arg0 == 0) { \
227+
return 0; \
228+
} \
229+
submit_block; \
167230
}
168231

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); \
232+
/* Macro to generate uprobe/uretprobe pairs for functions with 2 arguments */
233+
#define UPROBE_ARGS_RET(name, arg0_expr, arg1_expr, submit_block) \
234+
struct name##_args_t { \
235+
__u64 arg0; \
236+
__u64 arg1; \
237+
}; \
238+
BPF_HASH_MAP(name##_args, __u64, struct name##_args_t, 10000); \
239+
SEC("uprobe") \
240+
int uprobe_##name(struct pt_regs* ctx) { \
241+
__u64 tid = bpf_get_current_pid_tgid(); \
242+
__u32 pid = tid >> 32; \
243+
\
244+
if (!is_tracked(pid)) { \
245+
return 0; \
246+
} \
247+
\
248+
struct name##_args_t args = { \
249+
.arg0 = arg0_expr, \
250+
.arg1 = arg1_expr \
251+
}; \
252+
\
253+
bpf_map_update_elem(&name##_args, &tid, &args, BPF_ANY); \
254+
return 0; \
255+
} \
256+
SEC("uretprobe") \
257+
int uretprobe_##name(struct pt_regs* ctx) { \
258+
__u64 tid = bpf_get_current_pid_tgid(); \
259+
struct name##_args_t* args = bpf_map_lookup_elem(&name##_args, &tid); \
260+
\
261+
if (!args) { \
262+
return 0; \
263+
} \
264+
\
265+
struct name##_args_t a = *args; \
266+
bpf_map_delete_elem(&name##_args, &tid); \
267+
\
268+
__u64 ret_val = PT_REGS_RC(ctx); \
269+
if (ret_val == 0) { \
270+
return 0; \
271+
} \
272+
\
273+
__u64 arg0 = a.arg0; \
274+
__u64 arg1 = a.arg1; \
275+
submit_block; \
178276
}
179277

180278
/* malloc: allocates with size parameter */
181-
UPROBE_WITH_ARGS(malloc, PT_REGS_PARM1(ctx), PT_REGS_RC(ctx), EVENT_TYPE_MALLOC)
279+
UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx), {
280+
return submit_alloc_event(arg0, ret_val);
281+
})
182282

183283
/* free: deallocates by address */
184-
UPROBE_ADDR_ONLY(free, PT_REGS_PARM1(ctx), EVENT_TYPE_FREE)
284+
UPROBE_RET(free, PT_REGS_PARM1(ctx), {
285+
return submit_free_event(arg0);
286+
})
185287

186288
/* 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)
289+
UPROBE_ARG_RET(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx), {
290+
return submit_calloc_event(arg0, ret_val);
291+
})
188292

189-
/* realloc: reallocates with new size */
190-
UPROBE_WITH_ARGS(realloc, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_REALLOC)
293+
/* realloc: reallocates with old_addr and new size */
294+
UPROBE_ARGS_RET(realloc, PT_REGS_PARM2(ctx), PT_REGS_PARM1(ctx), {
295+
return submit_realloc_event(arg1, ret_val, arg0);
296+
})
191297

192298
/* 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)
299+
UPROBE_ARG_RET(aligned_alloc, PT_REGS_PARM2(ctx), {
300+
return submit_aligned_alloc_event(arg0, ret_val);
301+
})
194302

195303
/* 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)
304+
UPROBE_ARG_RET(memalign, PT_REGS_PARM2(ctx), {
305+
return submit_aligned_alloc_event(arg0, ret_val);
306+
})
197307

198308
/* Map to store mmap parameters between entry and return */
199309
struct mmap_args {
@@ -234,7 +344,7 @@ int tracepoint_sys_exit_mmap(struct trace_event_raw_sys_exit* ctx) {
234344
return 0;
235345
}
236346

237-
return submit_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
347+
return submit_mmap_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
238348
}
239349

240350
/* munmap tracking */
@@ -247,7 +357,7 @@ int tracepoint_sys_enter_munmap(struct trace_event_raw_sys_enter* ctx) {
247357
return 0;
248358
}
249359

250-
return submit_event(addr, len, EVENT_TYPE_MUNMAP);
360+
return submit_mmap_event(addr, len, EVENT_TYPE_MUNMAP);
251361
}
252362

253363
/* brk tracking - adjusts the program break (heap boundary) */
@@ -281,5 +391,5 @@ int tracepoint_sys_exit_brk(struct trace_event_raw_sys_exit* ctx) {
281391
return 0;
282392
}
283393

284-
return submit_event(new_brk, 0, EVENT_TYPE_BRK);
394+
return submit_mmap_event(new_brk, 0, EVENT_TYPE_BRK);
285395
}

0 commit comments

Comments
 (0)