Skip to content

Commit 03ba372

Browse files
committed
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.
1 parent 58d994a commit 03ba372

14 files changed

Lines changed: 880 additions & 869 deletions

File tree

crates/memtrack/AGENTS.md

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
# Repository Guidelines
2+
3+
`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`.
4+
5+
## Project Overview
6+
7+
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`.
8+
9+
## Architecture & Data Flow
10+
11+
Allocation → disk pipeline:
12+
13+
1. Allocator entry uprobe stores args in a per-tid BPF hash map (`<name>_arg` / `<name>_args`).
14+
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 16 MiB `events` ring buffer. Reserve failure atomically bumps `dropped_events`.
15+
3. `RingBufferPoller` (`src/ebpf/poller.rs`) polls every 10 ms → `parse_event` (`src/ebpf/events.rs`) casts raw bytes to the bindgen `event` struct and maps `EVENT_TYPE_*``runner_shared::MemtrackEventKind` → mpsc channel.
16+
4. `Tracker::track` (`src/ebpf/tracker.rs`) forwards poller events over a keep-alive thread.
17+
5. `main.rs` runs a 2-stage pipeline: Stage A drain thread empties the ring buffer fast → Stage B writer thread batches via `MemtrackWriter`. After the child exits, it checks `dropped_events_count()` and **bails if non-zero** (incomplete trace).
18+
19+
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.
20+
21+
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.
22+
23+
> 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.
24+
25+
## Key Directories
26+
27+
- `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`.
28+
- `src/allocators/` — allocator classification: `mod.rs`, `dynamic.rs`, `static_linked.rs`.
29+
- `tests/` — integration tests + `snapshots/` (insta).
30+
- `testdata/` — allocation fixtures: `*.c` (gcc), `alloc_cpp/` (cmkr/CMake), `alloc_rust/` + `spawn_wrapper/` (standalone Cargo workspaces).
31+
- `.agents/docs/`, `.claude/` — design notes (some current, some historical/stale).
32+
33+
## Development Commands
34+
35+
```bash
36+
cargo build # default features include `ebpf`
37+
cargo check
38+
cargo fmt
39+
cargo clippy
40+
cargo test --lib # unit tests (no root)
41+
42+
# Run the tracker (needs root); tracks a shell command's whole process tree:
43+
sudo -E cargo run --bin codspeed-memtrack -- track "<command>" --output <dir>
44+
# e.g. sudo -E cargo run --bin codspeed-memtrack -- track "ls / >/dev/null" --output .
45+
46+
# Integration tests need BPF privilege + GITHUB_ACTIONS gate + single-threaded:
47+
export GITHUB_ACTIONS=1
48+
sudo -E cargo test --test c_tests -- --test-threads 1
49+
```
50+
51+
`--test-threads 1` is **mandatory** — eBPF probes cannot self-overlap. CI runs `sudo -E cargo test --lib --test <name> -- --test-threads 1`; the main workspace `tests` job excludes memtrack (`--exclude memtrack`).
52+
53+
## Code Conventions & Common Patterns
54+
55+
- **Errors:** `anyhow` only (`Result`/`Context`/`bail`/`ensure` via `src/prelude.rs`); no `thiserror`.
56+
- **Concurrency:** no async runtime — `std::thread` + `std::sync::mpsc` + `Arc<Mutex<_>>`.
57+
- **Naming:** snake_case modules, PascalCase types. Macro-generated `try_<name>` (fallible) vs `<name>_if_found` (best-effort, trace-logs errors) attach helpers via `paste!`.
58+
- **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`.
59+
- **Module layout:** `src/lib.rs` re-exports and `#[cfg(feature = "ebpf")]` gates the whole BPF stack + binary; `prelude.rs` centralizes error/log imports.
60+
- **Cross-crate types:** `MemtrackEvent` / `MemtrackEventKind` live in `runner-shared/src/artifacts/memtrack.rs` (a breaking change there ripples here).
61+
62+
## Important Files
63+
64+
- `src/main.rs` — CLI entry (`codspeed-memtrack track`, requires `ebpf`); drops sudo privileges to `SUDO_UID`/`SUDO_GID` so the tracked child runs unprivileged.
65+
- `src/lib.rs` — crate root / re-exports.
66+
- `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.
67+
- `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.
68+
- `wrapper.h` — bindgen entry (`stdint.h` + `src/ebpf/c/event.h`).
69+
- `Cargo.toml` — lib `memtrack` + bin `codspeed-memtrack`; feature `ebpf` (default) pulls `libbpf-rs`/`libbpf-cargo`/`vmlinux`.
70+
71+
## Runtime / Tooling Prerequisites
72+
73+
- **Linux only** — consumed from workspace root under `cfg(target_os = "linux")`, targets `x86_64`/`aarch64-unknown-linux-gnu`.
74+
- **Root / BPF privilege** required at runtime (bumps `RLIMIT_MEMLOCK`, loads BPF).
75+
- **Build toolchain:** `clang` + BTF/vmlinux headers, `libbpf-dev`, `zlib1g-dev`, `pkgconf`, `build-essential`; vendored libbpf also needs `autopoint`/`bison`/`flex`.
76+
- `vmlinux.h` is pinned to a specific git rev; `libbpf-rs` uses the `vendored` feature (dist links `libbpf-rs/static`).
77+
78+
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).
79+
80+
## Testing & QA
81+
82+
- **Frameworks:** `insta` (snapshots), `rstest` (parametrized `#[case]`), `test-log`, `test-with` (`#[test_with::env(GITHUB_ACTIONS)]`), `tempfile`.
83+
- **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]`.
84+
- **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`).
85+
- **Snapshots:** `tests/snapshots/<binary>__<case>.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.

crates/memtrack/build.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ fn build_ebpf() {
1717
.expect("CARGO_CFG_TARGET_ARCH must be set in build script");
1818
let memtrack_out = PathBuf::from(env::var("OUT_DIR").unwrap()).join("memtrack.skel.rs");
1919
SkeletonBuilder::new()
20-
.source("src/ebpf/c/memtrack.bpf.c")
20+
.source("src/ebpf/c/main.bpf.c")
2121
.clang_args([
2222
"-I",
2323
&vmlinux::include_path_root().join(arch).to_string_lossy(),
Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
#ifndef __ALLOCATOR_H__
2+
#define __ALLOCATOR_H__
3+
4+
#include "utils/event_helpers.h"
5+
#include "utils/map_helpers.h"
6+
#include "utils/tracking.h"
7+
8+
#define UPROBE_ARG_RET(name, arg_expr, submit_block) \
9+
BPF_HASH_MAP(name##_arg, __u64, __u64, 10000); \
10+
SEC("uprobe") \
11+
int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_arg, arg_expr); } \
12+
SEC("uretprobe") \
13+
int uretprobe_##name(struct pt_regs* ctx) { \
14+
__u64* arg_ptr = take_param(&name##_arg); \
15+
if (!arg_ptr) { \
16+
return 0; \
17+
} \
18+
__u64 ret_val = PT_REGS_RC(ctx); \
19+
if (ret_val == 0) { \
20+
return 0; \
21+
} \
22+
__u64 arg0 = *arg_ptr; \
23+
submit_block; \
24+
}
25+
26+
#define UPROBE_RET(name, arg_expr, submit_block) \
27+
SEC("uprobe") \
28+
int uprobe_##name(struct pt_regs* ctx) { \
29+
__u64 arg0 = arg_expr; \
30+
if (arg0 == 0) { \
31+
return 0; \
32+
} \
33+
submit_block; \
34+
}
35+
36+
#define UPROBE_ARGS_RET(name, arg0_expr, arg1_expr, submit_block) \
37+
struct name##_args_t { \
38+
__u64 arg0; \
39+
__u64 arg1; \
40+
}; \
41+
BPF_HASH_MAP(name##_args, __u64, struct name##_args_t, 10000); \
42+
SEC("uprobe") \
43+
int uprobe_##name(struct pt_regs* ctx) { \
44+
__u64 tid = bpf_get_current_pid_tgid(); \
45+
__u32 pid = tid >> 32; \
46+
\
47+
if (!is_tracked(pid)) { \
48+
return 0; \
49+
} \
50+
\
51+
struct name##_args_t args = {.arg0 = arg0_expr, .arg1 = arg1_expr}; \
52+
\
53+
bpf_map_update_elem(&name##_args, &tid, &args, BPF_ANY); \
54+
return 0; \
55+
} \
56+
SEC("uretprobe") \
57+
int uretprobe_##name(struct pt_regs* ctx) { \
58+
__u64 tid = bpf_get_current_pid_tgid(); \
59+
struct name##_args_t* args = bpf_map_lookup_elem(&name##_args, &tid); \
60+
\
61+
if (!args) { \
62+
return 0; \
63+
} \
64+
\
65+
struct name##_args_t a = *args; \
66+
bpf_map_delete_elem(&name##_args, &tid); \
67+
\
68+
__u64 ret_val = PT_REGS_RC(ctx); \
69+
if (ret_val == 0) { \
70+
return 0; \
71+
} \
72+
\
73+
__u64 arg0 = a.arg0; \
74+
__u64 arg1 = a.arg1; \
75+
submit_block; \
76+
}
77+
78+
UPROBE_ARG_RET(malloc, PT_REGS_PARM1(ctx), { return submit_alloc_event(arg0, ret_val); })
79+
80+
UPROBE_RET(free, PT_REGS_PARM1(ctx), { return submit_free_event(arg0); })
81+
82+
UPROBE_ARG_RET(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx),
83+
{ return submit_calloc_event(arg0, ret_val); })
84+
85+
UPROBE_ARGS_RET(realloc, PT_REGS_PARM2(ctx), PT_REGS_PARM1(ctx),
86+
{ return submit_realloc_event(arg1, ret_val, arg0); })
87+
88+
UPROBE_ARG_RET(aligned_alloc, PT_REGS_PARM2(ctx),
89+
{ return submit_aligned_alloc_event(arg0, ret_val); })
90+
91+
UPROBE_ARG_RET(memalign, PT_REGS_PARM2(ctx), { return submit_aligned_alloc_event(arg0, ret_val); })
92+
93+
struct mmap_args {
94+
__u64 addr;
95+
__u64 len;
96+
};
97+
98+
BPF_HASH_MAP(mmap_temp, __u64, struct mmap_args, 10000);
99+
100+
static __always_inline void store_mmap_args(__u64 addr, __u64 len) {
101+
__u64 tid = bpf_get_current_pid_tgid();
102+
__u32 pid = tid >> 32;
103+
if (is_tracked(pid)) {
104+
struct mmap_args args = {.addr = addr, .len = len};
105+
bpf_map_update_elem(&mmap_temp, &tid, &args, BPF_ANY);
106+
}
107+
}
108+
109+
SEC("tracepoint/syscalls/sys_enter_mmap")
110+
int tracepoint_sys_enter_mmap(struct trace_event_raw_sys_enter* ctx) {
111+
store_mmap_args(ctx->args[0], ctx->args[1]);
112+
return 0;
113+
}
114+
115+
SEC("tracepoint/syscalls/sys_exit_mmap")
116+
int tracepoint_sys_exit_mmap(struct trace_event_raw_sys_exit* ctx) {
117+
struct mmap_args* args = (struct mmap_args*)take_param(&mmap_temp);
118+
if (!args) {
119+
return 0;
120+
}
121+
122+
__s64 ret = ctx->ret;
123+
if (ret <= 0) {
124+
return 0;
125+
}
126+
127+
return submit_mmap_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
128+
}
129+
130+
SEC("tracepoint/syscalls/sys_enter_munmap")
131+
int tracepoint_sys_enter_munmap(struct trace_event_raw_sys_enter* ctx) {
132+
__u64 addr = ctx->args[0];
133+
__u64 len = ctx->args[1];
134+
135+
if (addr == 0 || len == 0) {
136+
return 0;
137+
}
138+
139+
return submit_mmap_event(addr, len, EVENT_TYPE_MUNMAP);
140+
}
141+
142+
BPF_HASH_MAP(brk_temp, __u64, __u64, 10000);
143+
144+
SEC("tracepoint/syscalls/sys_enter_brk")
145+
int tracepoint_sys_enter_brk(struct trace_event_raw_sys_enter* ctx) {
146+
store_param(&brk_temp, ctx->args[0]);
147+
return 0;
148+
}
149+
150+
SEC("tracepoint/syscalls/sys_exit_brk")
151+
int tracepoint_sys_exit_brk(struct trace_event_raw_sys_exit* ctx) {
152+
__u64* requested_brk = take_param(&brk_temp);
153+
if (!requested_brk) {
154+
return 0;
155+
}
156+
157+
__u64 new_brk = ctx->ret;
158+
__u64 req_brk = *requested_brk;
159+
160+
if (req_brk == 0 || new_brk <= 0) {
161+
return 0;
162+
}
163+
164+
return submit_mmap_event(new_brk, 0, EVENT_TYPE_BRK);
165+
}
166+
167+
#endif /* __ALLOCATOR_H__ */
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
// clang-format off
2+
#include "vmlinux.h"
3+
// clang-format on
4+
#include <bpf/bpf_core_read.h>
5+
#include <bpf/bpf_helpers.h>
6+
#include <bpf/bpf_tracing.h>
7+
8+
#include "allocator.h"
9+
#include "event.h"
10+
#include "utils/event_helpers.h"
11+
#include "utils/map_helpers.h"
12+
#include "utils/tracking.h"
13+
14+
char LICENSE[] SEC("license") = "GPL";

0 commit comments

Comments
 (0)