diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bbbb05a93..15c3467ce 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -141,6 +141,32 @@ jobs: mode: ${{ matrix.mode }} run: cargo codspeed run -p runner-shared + memtrack-benchmarks: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + submodules: true + + - uses: ./.github/actions/install-rust + - uses: ./.github/actions/install-bpf-deps + + - name: Install memtrack + run: | + cargo install --path crates/memtrack --locked + + - name: Grant memtrack file capabilities + run: cargo r -- setup --mode memory + + - name: Build the codspeed CLI + run: cargo build --release + + - name: Prepare memtrack output directory + run: mkdir -p /tmp/codspeed-memtrack-bench + + - name: Run memtrack walltime benchmarks + run: ./target/release/codspeed --config crates/memtrack/codspeed.yml run -m walltime + check: runs-on: ubuntu-latest if: always() @@ -152,6 +178,7 @@ jobs: - macos-basic-run-test - bpf-tests - benchmarks + - memtrack-benchmarks steps: - uses: re-actors/alls-green@release/v1 with: diff --git a/crates/memtrack/codspeed.yml b/crates/memtrack/codspeed.yml new file mode 100644 index 000000000..9b57b8972 --- /dev/null +++ b/crates/memtrack/codspeed.yml @@ -0,0 +1,27 @@ +$schema: https://raw.githubusercontent.com/CodSpeedHQ/codspeed/refs/heads/main/schemas/codspeed.schema.json + +# Walltime benchmarks measuring codspeed-memtrack's own overhead (eBPF probe +# attach + tracking) across a few representative workloads, not the memory +# usage of the tracked command. +# +# The warmup/max times are generous because a single tracked run already pays a +# fixed BPF load + uprobe attach cost, which is far above the defaults tuned +# for near-instant commands. +options: + warmup-time: "5s" + max-time: "60s" + +benchmarks: + # Read-only, low-allocation baseline. The tracked command string is run + # through `bash -c`, so output can be redirected away: otherwise every round + # dumps the whole listing into the runner log. + - name: "memtrack track ls" + exec: codspeed-memtrack track "ls -la /usr/lib/x86_64-linux-gnu > /dev/null" --output /tmp/codspeed-memtrack-bench + + # Allocation- and I/O-heavy: many small file reads. + - name: "memtrack track tar" + exec: codspeed-memtrack track "tar -cf /tmp/memtrack-bench.tar /usr/lib/x86_64-linux-gnu" --output /tmp/codspeed-memtrack-bench + + # I/O-heavy with minimal allocation. + - name: "memtrack track dd" + exec: codspeed-memtrack track "dd if=/dev/zero of=/tmp/memtrack-bench-dd.bin bs=1M count=64 2> /dev/null" --output /tmp/codspeed-memtrack-bench diff --git a/crates/memtrack/src/ebpf/c/allocator.h b/crates/memtrack/src/ebpf/c/allocator.h index bf7f7e850..606342b36 100644 --- a/crates/memtrack/src/ebpf/c/allocator.h +++ b/crates/memtrack/src/ebpf/c/allocator.h @@ -5,22 +5,21 @@ #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) \ @@ -30,65 +29,46 @@ 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) @@ -99,74 +79,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; } @@ -175,7 +123,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") @@ -187,26 +135,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; diff --git a/crates/memtrack/src/ebpf/c/utils/event_helpers.h b/crates/memtrack/src/ebpf/c/utils/event_helpers.h index 89c8be14f..49a94bc49 100644 --- a/crates/memtrack/src/ebpf/c/utils/event_helpers.h +++ b/crates/memtrack/src/ebpf/c/utils/event_helpers.h @@ -8,43 +8,136 @@ 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. */ +/* Wake the consumer once this much data has been submitted. 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 a 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; +/* Bytes submitted per CPU since the last forced wakeup. + * + * Counting what this CPU produced, rather than asking the ring buffer how much + * is unconsumed, keeps the decision on the producer side: bpf_ringbuf_query() + * reads the consumer position, a cache line the polling thread on another CPU + * writes continuously, so querying it per event costs a cross-CPU miss on every + * event. */ +BPF_PERCPU_ARRAY_MAP(submitted_bytes, __u64, 1); + +static __always_inline long wake_flags(__u64 event_size) { + __u32 zero = 0; + __u64* pending = bpf_map_lookup_elem(&submitted_bytes, &zero); + if (!pending) { + /* Can't track the watermark, so don't risk a stalled consumer. */ + return BPF_RB_FORCE_WAKEUP; + } + + *pending += event_size; + if (*pending < WAKEUP_DATA_SIZE) { + return BPF_RB_NO_WAKEUP; + } + + *pending = 0; + return BPF_RB_FORCE_WAKEUP; +} + +/* Per-thread scratch for the allocator entry/exit hand-off. + * + * One slot per instrumented entry point rather than a single shared slot: an + * allocator may call another (glibc realloc() reaches malloc()), and nested + * calls on one thread must not clobber each other's saved arguments. + * + * `valid` marks which slots hold a value, so a zero argument is still + * distinguishable from an absent one, and a return probe that fires without a + * matching entry probe (attach raced with a call already in flight) is ignored. + */ +enum arg_slot { + SLOT_MALLOC, + SLOT_CALLOC, + SLOT_REALLOC, + SLOT_ALIGNED_ALLOC, + SLOT_MEMALIGN, + SLOT_POSIX_MEMALIGN, + SLOT_MMAP, + SLOT_BRK, + SLOT__COUNT, +}; + +struct memtrack_task_state { + __u64 arg0[SLOT__COUNT]; + __u64 arg1[SLOT__COUNT]; + __u32 valid; + /* Memoized positive result of is_tracked(). Tracking is monotonic: pids are + * only ever added to tracked_pids (from userspace or on fork), never + * removed, so a task that is tracked stays tracked and the answer can be + * cached. A negative result is never cached, since the tracker may register + * this task later. */ + __u8 tracked; +}; + +BPF_TASK_STORAGE(task_state, struct memtrack_task_state); + +/* Task state for the current task if it is tracked, else NULL. + * + * Hot path is a single task-storage lookup; the hashed is_tracked() walk runs + * once per task, on the first hook that observes it. */ +static __always_inline struct memtrack_task_state* tracked_state(void) { + struct task_struct* task = (struct task_struct*)bpf_get_current_task_btf(); + struct memtrack_task_state* st = bpf_task_storage_get(&task_state, task, NULL, 0); + if (st && st->tracked) { + return st; + } + + if (!is_tracked(current_tgid())) { + return NULL; + } + + if (!st) { + st = bpf_task_storage_get(&task_state, task, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE); + if (!st) { + return NULL; + } + } + st->tracked = 1; + return st; } -static __always_inline int store_param(void* map, __u64 value) { - /* Key by the tid: unique per thread, so it survives the entry/exit pair even - * when several threads are inside the same allocator call. */ - struct task_ids ids = current_task_ids(); - __u64 tid = ids.tid; - if (is_tracked(ids.tgid)) { - bpf_map_update_elem(map, &tid, &value, BPF_ANY); +static __always_inline int store_arg(enum arg_slot slot, __u64 value) { + struct memtrack_task_state* st = tracked_state(); + if (!st) { + return 0; } + st->arg0[slot] = value; + st->valid |= (1u << slot); return 0; } -static __always_inline __u64* take_param(void* map) { - __u64 tid = current_tid(); - __u64* value = bpf_map_lookup_elem(map, &tid); - if (value) { - bpf_map_delete_elem(map, &tid); +static __always_inline int store_args(enum arg_slot slot, __u64 arg0, __u64 arg1) { + struct memtrack_task_state* st = tracked_state(); + if (!st) { + return 0; } - return value; + st->arg0[slot] = arg0; + st->arg1[slot] = arg1; + st->valid |= (1u << slot); + return 0; } +/* Consume a slot: returns the state with the slot cleared, or NULL if the entry + * probe never ran for this call. */ +static __always_inline struct memtrack_task_state* take_slot(enum arg_slot slot) { + struct task_struct* task = (struct task_struct*)bpf_get_current_task_btf(); + struct memtrack_task_state* st = bpf_task_storage_get(&task_state, task, NULL, 0); + if (!st || !(st->valid & (1u << slot))) { + return NULL; + } + st->valid &= ~(1u << slot); + return st; +} + +/* Emit an event for a task already known to be tracked. */ #define SUBMIT_EVENT(evt_type, fill_data) \ { \ - struct task_ids ids = current_task_ids(); \ - \ - if (!is_tracked(ids.tgid) || !is_enabled()) { \ + if (!is_enabled()) { \ return 0; \ } \ \ @@ -58,14 +151,15 @@ static __always_inline __u64* take_param(void* map) { return 0; \ } \ \ + __u64 pid_tgid = memtrack_current_pid_tgid(); \ e->header.timestamp = bpf_ktime_get_ns(); \ - e->header.pid = ids.tgid; \ - e->header.tid = ids.tid; \ + e->header.pid = pid_tgid >> 32; \ + e->header.tid = (__u32)pid_tgid; \ e->header.event_type = evt_type; \ \ fill_data; \ \ - bpf_ringbuf_submit(e, wake_flags()); \ + bpf_ringbuf_submit(e, wake_flags(sizeof(*e))); \ return 0; \ } diff --git a/crates/memtrack/src/ebpf/c/utils/map_helpers.h b/crates/memtrack/src/ebpf/c/utils/map_helpers.h index 484fe9703..ef0cbe6b6 100644 --- a/crates/memtrack/src/ebpf/c/utils/map_helpers.h +++ b/crates/memtrack/src/ebpf/c/utils/map_helpers.h @@ -17,6 +17,25 @@ __type(value, value_type); \ } name SEC(".maps") +/* Task-local storage: one value per task_struct, reached by pointer chase off + * the task rather than a hashed lookup, and freed with the task. NO_PREALLOC is + * mandatory for this map type. */ +#define BPF_TASK_STORAGE(name, value_type) \ + struct { \ + __uint(type, BPF_MAP_TYPE_TASK_STORAGE); \ + __uint(map_flags, BPF_F_NO_PREALLOC); \ + __type(key, int); \ + __type(value, value_type); \ + } name SEC(".maps") + +#define BPF_PERCPU_ARRAY_MAP(name, value_type, max_ents) \ + struct { \ + __uint(type, BPF_MAP_TYPE_PERCPU_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); \ diff --git a/crates/memtrack/src/ebpf/memtrack/macros.rs b/crates/memtrack/src/ebpf/memtrack/macros.rs index d7b7847cf..debd7dabe 100644 --- a/crates/memtrack/src/ebpf/memtrack/macros.rs +++ b/crates/memtrack/src/ebpf/memtrack/macros.rs @@ -57,21 +57,25 @@ 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 = attach_one!(self, $prog_entry, lib_path, offset, false) - .context(format!( - "Failed to attach uprobe at offset {:#x} in {}", - offset, - lib_path.display() - ))?; - self.probes.push(link); + if self.claim_site(stringify!($prog_entry), lib_path, offset, false) { + let link = attach_one!(self, $prog_entry, lib_path, offset, false) + .context(format!( + "Failed to attach uprobe at offset {:#x} in {}", + offset, + lib_path.display() + ))?; + self.probes.push(link); + } - let link = attach_one!(self, $prog_return, lib_path, offset, true) - .context(format!( - "Failed to attach uretprobe at offset {:#x} in {}", - offset, - lib_path.display() - ))?; - self.probes.push(link); + if self.claim_site(stringify!($prog_return), lib_path, offset, true) { + let link = attach_one!(self, $prog_return, lib_path, offset, true) + .context(format!( + "Failed to attach uretprobe at offset {:#x} in {}", + offset, + lib_path.display() + ))?; + self.probes.push(link); + } Ok(()) } @@ -102,13 +106,15 @@ macro_rules! attach_uprobe { ($name:ident, $prog:ident) => { paste! { fn [](&mut self, lib_path: &Path, offset: usize) -> Result<()> { - let link = attach_one!(self, $prog, lib_path, offset, false) - .context(format!( - "Failed to attach uprobe at offset {:#x} in {}", - offset, - lib_path.display() - ))?; - self.probes.push(link); + if self.claim_site(stringify!($prog), lib_path, offset, false) { + let link = attach_one!(self, $prog, lib_path, offset, false) + .context(format!( + "Failed to attach uprobe at offset {:#x} in {}", + offset, + lib_path.display() + ))?; + self.probes.push(link); + } Ok(()) } diff --git a/crates/memtrack/src/ebpf/memtrack/mod.rs b/crates/memtrack/src/ebpf/memtrack/mod.rs index f0b80b797..7bd98f553 100644 --- a/crates/memtrack/src/ebpf/memtrack/mod.rs +++ b/crates/memtrack/src/ebpf/memtrack/mod.rs @@ -4,7 +4,7 @@ use libbpf_rs::skel::OpenSkel; use libbpf_rs::skel::SkelBuilder; use std::collections::HashMap; use std::mem::MaybeUninit; -use std::path::Path; +use std::path::{Path, PathBuf}; use crate::ebpf::poller::RingBufferPoller; @@ -108,6 +108,25 @@ impl ResolvedSymbols { pub struct MemtrackBpf { pub(super) skel: Skel, pub(super) probes: Vec, + /// Attach sites already claimed, as (program, library, offset, retprobe). + /// Allocator entry points share addresses through aliases (`free`, + /// `cfree` and `__libc_free` are one symbol in glibc), and attaching the + /// same program twice at one address makes it run twice per call. + pub(super) attached_sites: std::collections::HashSet<(&'static str, PathBuf, usize, bool)>, +} + +impl MemtrackBpf { + /// Reserve an attach site, returning false if it is already instrumented. + pub(super) fn claim_site( + &mut self, + prog: &'static str, + lib_path: &Path, + offset: usize, + retprobe: bool, + ) -> bool { + self.attached_sites + .insert((prog, lib_path.to_path_buf(), offset, retprobe)) + } } impl MemtrackBpf { @@ -160,6 +179,7 @@ impl MemtrackBpf { Ok(Self { skel, probes: Vec::new(), + attached_sites: std::collections::HashSet::new(), }) }