diff --git a/crates/memtrack/examples/stack_codec_ratio.rs b/crates/memtrack/examples/stack_codec_ratio.rs new file mode 100644 index 00000000..89c9a5e6 --- /dev/null +++ b/crates/memtrack/examples/stack_codec_ratio.rs @@ -0,0 +1,149 @@ +use std::fs::File; +use std::io::{BufReader, Read}; +use std::path::PathBuf; + +use clap::Parser; +use memtrack::stack_codec::{RawStack, StackDecoder, StackEncoder, fnv_stack_hash}; + +#[derive(Parser, Debug)] +#[command(name = "stack_codec_ratio")] +struct Args { + #[arg(long, default_value_t = 100_000)] + limit: usize, + + #[arg(required = true)] + dumps: Vec, +} + +struct DumpRecord { + hash: u64, + timestamp: u64, + pid: u32, + tid: u32, + stack: RawStack, +} + +// Format defined in .agents/scripts/stackdump.py: +// struct.pack(" std::io::Result> { + let mut hdr = [0u8; 8 + 8 + 4 + 4 + 8 + 4 + 1 + 1 + 2]; + match r.read_exact(&mut hdr) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None), + Err(e) => return Err(e), + } + + let expected_hash = u64::from_le_bytes(hdr[0..8].try_into().unwrap()); + let timestamp = u64::from_le_bytes(hdr[8..16].try_into().unwrap()); + let pid = u32::from_le_bytes(hdr[16..20].try_into().unwrap()); + let tid = u32::from_le_bytes(hdr[20..24].try_into().unwrap()); + let sp = u64::from_le_bytes(hdr[24..32].try_into().unwrap()); + let copy_len = u32::from_le_bytes(hdr[32..36].try_into().unwrap()) as usize; + let truncated = hdr[36] != 0; + let nregs = hdr[37] as usize; + let nfp = u16::from_le_bytes(hdr[38..40].try_into().unwrap()) as usize; + + let mut reg_bytes = vec![0u8; nregs * 8]; + r.read_exact(&mut reg_bytes)?; + let mut fp_bytes = vec![0u8; nfp * 8]; + r.read_exact(&mut fp_bytes)?; + let mut bytes = vec![0u8; copy_len]; + r.read_exact(&mut bytes)?; + + let mut regs = [0u64; 33]; + for (i, chunk) in reg_bytes.chunks_exact(8).enumerate() { + if i < 33 { + regs[i] = u64::from_le_bytes(chunk.try_into().unwrap()); + } + } + + Ok(Some(DumpRecord { + hash: expected_hash, + timestamp, + pid, + tid, + stack: RawStack { + sp, + regs, + bytes, + truncated, + }, + })) +} + +fn process_dump(path: &PathBuf, limit: usize) -> anyhow::Result<()> { + let file = File::open(path)?; + let mut reader = BufReader::with_capacity(1 << 20, file); + + let mut encoder = StackEncoder::default(); + let mut decoder = StackDecoder::default(); + + let mut record_count = 0usize; + let mut total_raw_bytes = 0u64; + let mut total_encoded_bytes = 0u64; + + while record_count < limit { + let Some(DumpRecord { + hash: expected_hash, + timestamp, + pid, + tid, + stack, + }) = read_record(&mut reader)? + else { + break; + }; + + // Validate FNV implementation against the recorded capture hash + let computed_hash = fnv_stack_hash(&stack.bytes); + assert_eq!( + computed_hash, + expected_hash, + "FNV hash mismatch in {}: record {record_count}, expected {expected_hash:#x}, got {computed_hash:#x}", + path.display() + ); + + let raw_record_size = (312 + stack.bytes.len()) as u64; // raw record = 312 B header + copy_len + total_raw_bytes += raw_record_size; + + let encoded = encoder.encode(pid, tid, timestamp, 0, &stack); + total_encoded_bytes += encoded.len() as u64; + + let (event, _) = decoder.decode(&encoded).expect("decode failed"); + if let runner_shared::artifacts::MemtrackEventKind::Stack { record } = event.kind { + assert_eq!(record.hash, expected_hash); + assert_eq!(record.bytes, stack.bytes); + assert_eq!(record.sp, stack.sp); + assert_eq!(&record.regs[..], &stack.regs[..]); + assert_eq!(record.truncated, stack.truncated); + } else { + panic!("expected Stack event"); + } + + record_count += 1; + } + + let ratio = total_raw_bytes as f64 / total_encoded_bytes as f64; + let bytes_per_record = total_encoded_bytes as f64 / record_count as f64; + + println!( + "{}: records={}, raw_bytes={}, encoded_bytes={}, ratio={:.2}x, bytes/record={:.1}", + path.file_name().unwrap_or_default().to_string_lossy(), + record_count, + total_raw_bytes, + total_encoded_bytes, + ratio, + bytes_per_record + ); + + Ok(()) +} + +fn main() -> anyhow::Result<()> { + let args = Args::parse(); + for dump in &args.dumps { + process_dump(dump, args.limit)?; + } + Ok(()) +} diff --git a/crates/memtrack/src/ebpf/c/event.h b/crates/memtrack/src/ebpf/c/event.h index b18af874..49697e43 100644 --- a/crates/memtrack/src/ebpf/c/event.h +++ b/crates/memtrack/src/ebpf/c/event.h @@ -33,26 +33,77 @@ #define MEMTRACK_STACK_COUNTER_STACKID_FAILED 2 #define MEMTRACK_STACK_COUNTER_TRUNCATED 3 #define MEMTRACK_STACK_COUNTER_RING_FULL 4 -#define MEMTRACK_STACK_COUNTER_COUNT 5 +/* Delta encoding could not get a reference slot and fell back to a raw record. */ +#define MEMTRACK_STACK_COUNTER_DELTA_FALLBACK 5 +#define MEMTRACK_STACK_COUNTER_COUNT 6 struct stack_regs { uint64_t reg[MEMTRACK_STACK_REGS]; }; -/* Fixed header followed by `copy_len` bytes read upward from `sp`. */ +/* Both stack ring record layouts start with `kind` so the consumer can + * dispatch on it; raw and delta records share one ring. */ +#define STACK_RECORD_RAW 1 +#define STACK_RECORD_DELTA 2 + +/* Raw record: fixed header followed by `copy_len` bytes read upward from `sp`. */ struct stack_header { + uint32_t kind; /* STACK_RECORD_RAW */ + uint32_t copy_len; uint64_t hash; uint64_t timestamp; /* monotonic time in nanoseconds (CLOCK_MONOTONIC) */ int64_t stackid; /* bpf_get_stackid() result; negative means unavailable */ uint64_t sp; /* user stack pointer the copy starts at */ uint32_t pid; uint32_t tid; - uint32_t copy_len; uint8_t truncated; /* the copy hit the size cap */ - uint8_t _pad[3]; + uint8_t _pad[7]; struct stack_regs regs; }; +/* Delta record. The stack is expressed as an XOR against the previous record + * emitted for the same tid (the reference), aligned by absolute address: + * word i of this copy (bytes [8i, 8i+8) above `sp`) pairs with reference word + * j = i + (sp - ref.sp) / 8, or with 0 when j is outside the reference copy. + * A keyframe (ref_hash == 0) encodes against an all-zero, empty reference, + * so the same layout carries a plain sparse copy. + * + * header + * u64 reg literal x popcount(regs_mask) (ascending register) + * for each set bit g of group_mask, ascending: + * u64 word_bitmap bit k set <=> delta word 64g+k != 0 + * u64 literal x popcount(word_bitmap) (ascending word) + * + * A group is 64 words (512 bytes). Groups whose delta is all zero are omitted + * and have their group_mask bit clear. `hash` covers the reconstructed raw + * bytes and uses the same function as the raw record, so the consumer can + * check that it decoded against the right reference. + */ +#define MEMTRACK_STACK_GROUP_WORDS 64 +#define MEMTRACK_STACK_MAX_WORDS (MEMTRACK_MAX_STACK_COPY / 8) +#define MEMTRACK_STACK_MAX_GROUPS (MEMTRACK_STACK_MAX_WORDS / MEMTRACK_STACK_GROUP_WORDS) +#define MEMTRACK_STACK_DELTA_MAX_PAYLOAD \ + (MEMTRACK_STACK_REGS * 8 + MEMTRACK_STACK_MAX_GROUPS * 8 + MEMTRACK_MAX_STACK_COPY) + +#define STACK_DELTA_FLAG_TRUNCATED 1 + +struct stack_delta_header { + uint32_t kind; /* STACK_RECORD_DELTA */ + uint32_t copy_len; /* reconstructed raw byte count, multiple of 512 */ + uint64_t hash; /* hash of the reconstructed raw bytes */ + uint64_t ref_hash; /* hash of the reference record; 0 on a keyframe */ + uint64_t timestamp; /* monotonic time in nanoseconds (CLOCK_MONOTONIC) */ + int64_t stackid; /* bpf_get_stackid() result; negative means unavailable */ + uint64_t sp; /* user stack pointer the copy starts at */ + uint32_t pid; + uint32_t tid; + uint32_t payload_len; /* bytes following this header */ + uint8_t flags; /* STACK_DELTA_FLAG_* */ + uint8_t _pad[3]; + uint64_t group_mask; /* bit g set <=> group g present in the payload */ + uint64_t regs_mask; /* bit r set <=> register r literal present */ +}; + /* Common header shared by all event types */ struct event_header { uint8_t event_type; /* See EVENT_TYPE_* constants above */ diff --git a/crates/memtrack/src/ebpf/c/stack_capture.bpf.h b/crates/memtrack/src/ebpf/c/stack_capture.bpf.h index 92097891..7454df11 100644 --- a/crates/memtrack/src/ebpf/c/stack_capture.bpf.h +++ b/crates/memtrack/src/ebpf/c/stack_capture.bpf.h @@ -89,37 +89,57 @@ static __always_inline void fill_stack_regs(struct stack_regs* out, struct pt_re #error "stack capture needs a DWARF register mapping for this architecture" #endif -static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct task_ids ids) { - void* slot = bpf_ringbuf_reserve(&stacks, sizeof(struct stack_header) + stack_copy_budget, 0); - if (!slot) { - bump_stack_counter(MEMTRACK_STACK_COUNTER_RING_FULL); - memtrack_check_ring_pressure(&stacks, ids.tgid); - return 0; - } - - /* Keep hashing scratch in the unpublished record. Large kprobe-family BPF - * stacks may use per-CPU storage, which nested uprobes can overwrite. */ - struct stack_header* header = (struct stack_header*)slot; - __u64* lanes = &header->hash; +static __always_inline void fnv64_lanes_init(__u64 lanes[4]) { lanes[0] = FNV64_OFFSET ^ 0; lanes[1] = FNV64_OFFSET ^ 1; lanes[2] = FNV64_OFFSET ^ 2; lanes[3] = FNV64_OFFSET ^ 3; - __u32 got = 0; +} + +/* Length distinguishes a full copy from the same bytes as a truncated prefix. + * Zero is reserved for allocation events without a stack. */ +static __always_inline __u64 fnv64_finish(const __u64 lanes[4], __u32 got) { + __u64 hash = + (((lanes[0] * FNV64_PRIME) ^ lanes[1]) * FNV64_PRIME ^ lanes[2]) * FNV64_PRIME ^ lanes[3]; + hash = (hash ^ got) * FNV64_PRIME; + return hash ? hash : FNV64_OFFSET; +} - /* Chunked reads stop at the first unreadable stack region. - * Loop bound is checked against stack_copy_budget (a frozen rodata constant) - * so every slot access is provably in range. */ +/* Copies the user stack into `dst` in STACK_COPY_CHUNK pieces, stopping at the + * first unreadable region, and hashes each chunk. Returns the bytes copied. + * The bound is the frozen stack_copy_budget, so `dst` must have room for the + * full budget for every access to be provably in range. */ +static __always_inline __u32 read_stack_chunks(__u8* dst, __u64 lanes[4], __u64 sp) { + __u32 got = 0; + fnv64_lanes_init(lanes); #pragma clang loop unroll(disable) for (__u32 off = 0; off + STACK_COPY_CHUNK <= stack_copy_budget; off += STACK_COPY_CHUNK) { - if (bpf_probe_read_user((__u8*)slot + sizeof(struct stack_header) + off, STACK_COPY_CHUNK, - (void*)(PT_REGS_SP(ctx) + off)) != 0) { + if (bpf_probe_read_user(dst + off, STACK_COPY_CHUNK, (void*)(sp + off)) != 0) { break; } - fnv64_hash_chunk(lanes, (const __u64*)((__u8*)slot + sizeof(struct stack_header) + off)); + fnv64_hash_chunk(lanes, (const __u64*)(dst + off)); got = off + STACK_COPY_CHUNK; } + return got; +} + +#include "stack_delta.bpf.h" + +static __always_inline __u64 capture_stack_raw(struct pt_regs* ctx, struct task_ids ids) { + void* slot = bpf_ringbuf_reserve(&stacks, sizeof(struct stack_header) + stack_copy_budget, 0); + if (!slot) { + bump_stack_counter(MEMTRACK_STACK_COUNTER_RING_FULL); + memtrack_check_ring_pressure(&stacks, ids.tgid); + return 0; + } + + /* Keep hashing scratch in the unpublished record. Large kprobe-family BPF + * stacks may use per-CPU storage, which nested uprobes can overwrite. */ + struct stack_header* header = (struct stack_header*)slot; + __u64* lanes = &header->hash; + __u32 got = + read_stack_chunks((__u8*)slot + sizeof(struct stack_header), lanes, PT_REGS_SP(ctx)); if (got == 0) { bpf_ringbuf_discard(slot, 0); @@ -133,15 +153,7 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas bump_stack_counter(MEMTRACK_STACK_COUNTER_TRUNCATED); } - __u64 hash = - (((lanes[0] * FNV64_PRIME) ^ lanes[1]) * FNV64_PRIME ^ lanes[2]) * FNV64_PRIME ^ lanes[3]; - - /* Length distinguishes a full copy from the same bytes as a truncated prefix. - * Zero is reserved for allocation events without a stack. */ - hash = (hash ^ got) * FNV64_PRIME; - if (hash == 0) { - hash = FNV64_OFFSET; - } + __u64 hash = fnv64_finish(lanes, got); header->hash = hash; long gate_result = @@ -161,17 +173,19 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas bump_stack_counter(MEMTRACK_STACK_COUNTER_STACKID_FAILED); } + header->kind = STACK_RECORD_RAW; + header->copy_len = got; header->hash = hash; header->timestamp = bpf_ktime_get_ns(); header->stackid = stackid; header->sp = PT_REGS_SP(ctx); header->pid = ids.tgid; header->tid = ids.tid; - header->copy_len = got; header->truncated = truncated; - header->_pad[0] = 0; - header->_pad[1] = 0; - header->_pad[2] = 0; +#pragma unroll + for (int i = 0; i < 7; i++) { + header->_pad[i] = 0; + } fill_stack_regs(&header->regs, ctx); bpf_ringbuf_submit(slot, 0); @@ -179,6 +193,17 @@ static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct tas return hash; } +static __always_inline __u64 capture_stack_inner(struct pt_regs* ctx, struct task_ids ids) { + if (stack_compression_enabled) { + struct stack_ref* ref = stack_ref_get(ids.tid); + if (ref) { + return capture_stack_delta(ctx, ids, ref); + } + bump_stack_counter(MEMTRACK_STACK_COUNTER_DELTA_FALLBACK); + } + return capture_stack_raw(ctx, ids); +} + static __always_inline __u64 capture_stack(struct pt_regs* ctx) { if (!capture_stacks_enabled || !is_enabled()) { return 0; diff --git a/crates/memtrack/src/ebpf/c/stack_delta.bpf.h b/crates/memtrack/src/ebpf/c/stack_delta.bpf.h new file mode 100644 index 00000000..95148d0d --- /dev/null +++ b/crates/memtrack/src/ebpf/c/stack_delta.bpf.h @@ -0,0 +1,212 @@ +#ifndef __STACK_DELTA_BPF_H__ +#define __STACK_DELTA_BPF_H__ + +/* Delta-encoded stack records (layout in event.h). Included from + * stack_capture.bpf.h after the shared read/hash helpers. + * + * Each tid keeps the last emitted copy as the reference and a staging slot + * for the capture in flight; the slots swap only after the record reached + * the ring, so the reference always matches what userspace decoded last. + * + * The encoder is branch-free per word: every literal is stored + * speculatively and the output index advances by 8 only when the delta is + * nonzero. With one verifier state per word the cost stays linear in the + * copy budget. Values that clang would otherwise turn back into + * compare-and-branch or wide expression trees go through barrier_var(). + */ + +const volatile __u8 stack_compression_enabled = 0; + +struct stack_ref_slot { + __u64 sp; + __u64 hash; + __u32 copy_len; + __u32 _pad; + __u64 regs[MEMTRACK_STACK_REGS]; + __u8 data[MEMTRACK_MAX_STACK_COPY]; +}; + +struct stack_ref { + __u32 cur; /* slot index holding the reference; 1 - cur is staging */ + struct stack_ref_slot slot[2]; + /* +8: the last speculative literal store may land past the payload. */ + __u8 out[sizeof(struct stack_delta_header) + MEMTRACK_STACK_DELTA_MAX_PAYLOAD + 8]; +}; + +/* Userspace shrinks this to one entry when compression is off. */ +BPF_LRU_HASH_MAP(stack_refs, __u32, struct stack_ref, 512); + +/* Map helpers need the insert value in map memory; the value is too large for + * the BPF stack. */ +static const struct stack_ref zero_stack_ref = {}; + +static __always_inline struct stack_ref* stack_ref_get(__u32 tid) { + struct stack_ref* ref = bpf_map_lookup_elem(&stack_refs, &tid); + if (ref) { + return ref; + } + bpf_map_update_elem(&stack_refs, &tid, &zero_stack_ref, BPF_NOEXIST); + return bpf_map_lookup_elem(&stack_refs, &tid); +} + +/* BPF has no set-on-condition instruction, so any `x != 0` or `a < b` that + * reaches the backend is lowered to a branch, and the verifier then forks a + * state per word. These helpers keep the predicates as arithmetic; the + * barriers stop clang from recognising them as comparisons. */ +static __always_inline __u64 nonzero_bit(__u64 x) { + __u64 t = x | (0 - x); + barrier_var(t); + return t >> 63; +} + +/* All ones when j < limit, treating j >= 2^63 (a wrapped negative index) as + * out of range; zero otherwise. Requires limit < 2^63. Both operands are + * opaque so clang cannot prove they are booleans and turn the AND into a + * select. */ +static __always_inline __u64 below_mask(__u64 j, __u64 limit) { + __u64 not_negative = ~j; + barrier_var(not_negative); + __u64 diff = j - limit; + barrier_var(diff); + return (__u64)((__s64)not_negative >> 63) & (__u64)((__s64)diff >> 63); +} + +static __always_inline __u64 capture_stack_delta(struct pt_regs* ctx, struct task_ids ids, + struct stack_ref* ref) { + __u32 cur = ref->cur & 1; + struct stack_ref_slot* prev = &ref->slot[cur]; + struct stack_ref_slot* stg = &ref->slot[cur ^ 1]; + __u8* out = ref->out; + + /* Hash lanes live in the output scratch; the header overwrites them later. */ + __u64* lanes = (__u64*)out; + __u64 sp = PT_REGS_SP(ctx); + __u32 got = read_stack_chunks(stg->data, lanes, sp); + if (got == 0) { + bump_stack_counter(MEMTRACK_STACK_COUNTER_COPY_FAILED); + memtrack_check_ring_pressure(&stacks, ids.tgid); + return 0; + } + + stg->hash = fnv64_finish(lanes, got); + stg->sp = sp; + stg->copy_len = got; + fill_stack_regs((struct stack_regs*)stg->regs, ctx); + __u64 hash = stg->hash; + + long gate_result = + bpf_map_update_elem(&seen_stack_hashes, &stg->hash, &seen_stack_marker, BPF_NOEXIST); + if (gate_result == -17) { /* -EEXIST */ + memtrack_check_ring_pressure(&stacks, ids.tgid); + return hash; + } + + /* Every distinct verifier state entering the encoder loop walks all of it + * again, so the loop must see one state. The read loop exits with a + * different constant `got` per chunk; reloading it from map memory makes + * it an unknown scalar the exit paths share. All other branches + * (counters, stackid, flags) run after the loop for the same reason. */ + got = *(volatile __u32*)&stg->copy_len; + + /* A keyframe encodes against the zeroed slot: ref_words is 0 and every + * reference word is masked. */ + __u64 ref_words = prev->copy_len / 8; + __u64 shift = (__u64)((__s64)(sp - prev->sp) >> 3); + + __u32 o = sizeof(struct stack_delta_header); + __u64 regs_mask = 0; +#pragma unroll + for (__u32 r = 0; r < MEMTRACK_STACK_REGS; r++) { + __u64 d = stg->regs[r] ^ prev->regs[r]; + __u64 nz = nonzero_bit(d); + *(__u64*)(out + o) = d; + o += (__u32)(nz << 3); + regs_mask |= nz << r; + barrier_var(regs_mask); + } + + const __u64* cur_words = (const __u64*)stg->data; + const __u64* ref_data = (const __u64*)prev->data; + __u64 group_mask = 0; + + /* Same frozen bound as the read loop; `got` ends it early. */ +#pragma clang loop unroll(disable) + for (__u32 off = 0; off + STACK_COPY_CHUNK <= stack_copy_budget; off += STACK_COPY_CHUNK) { + if (off >= got) { + break; + } + + __u32 base = off / 8; + __u64 bitmap = 0; + __u32 bitmap_pos = o; + /* Literal offsets are tracked relative to the group so the verifier + * keeps a tight [0, 512] range instead of a difference of two + * independently bounded indices. */ + __u32 lit = 8; +#pragma unroll + for (__u32 k = 0; k < MEMTRACK_STACK_GROUP_WORDS; k++) { + __u64 j = (__u64)(base + k) + shift; + __u64 refw = ref_data[j & (MEMTRACK_STACK_MAX_WORDS - 1)] & below_mask(j, ref_words); + __u64 d = cur_words[base + k] ^ refw; + __u64 nz = nonzero_bit(d); + *(__u64*)(out + bitmap_pos + lit) = d; + lit += (__u32)(nz << 3); + bitmap |= nz << k; + barrier_var(bitmap); + } + + *(__u64*)(out + bitmap_pos) = bitmap; + __u64 any = nonzero_bit(bitmap); + group_mask |= any << (off / STACK_COPY_CHUNK); + /* An all-zero group contributes nothing; its speculative bitmap and + * literal stores are overwritten by the next group or ignored. */ + o = bitmap_pos + (lit & (__u32)(0 - any)); + barrier_var(o); + } + + __u8 truncated = got >= stack_copy_budget; + if (truncated) { + bump_stack_counter(MEMTRACK_STACK_COUNTER_TRUNCATED); + } + if (gate_result != 0) { + bump_stack_counter(MEMTRACK_STACK_COUNTER_HASH_MAP_FULL); + } + __s64 stackid = bpf_get_stackid(ctx, &stack_traces, BPF_F_USER_STACK); + if (stackid < 0) { + bump_stack_counter(MEMTRACK_STACK_COUNTER_STACKID_FAILED); + } + + struct stack_delta_header* header = (struct stack_delta_header*)out; + header->kind = STACK_RECORD_DELTA; + header->copy_len = got; + header->hash = hash; + header->ref_hash = prev->hash; + header->timestamp = bpf_ktime_get_ns(); + header->stackid = stackid; + header->sp = sp; + header->pid = ids.tgid; + header->tid = ids.tid; + header->payload_len = o - sizeof(struct stack_delta_header); + header->flags = truncated ? STACK_DELTA_FLAG_TRUNCATED : 0; + header->_pad[0] = 0; + header->_pad[1] = 0; + header->_pad[2] = 0; + header->group_mask = group_mask; + header->regs_mask = regs_mask; + + if (bpf_ringbuf_output(&stacks, out, o, 0) != 0) { + bump_stack_counter(MEMTRACK_STACK_COUNTER_RING_FULL); + /* Let the next capture of this stack emit it again. */ + if (gate_result == 0) { + bpf_map_delete_elem(&seen_stack_hashes, &stg->hash); + } + memtrack_check_ring_pressure(&stacks, ids.tgid); + return 0; + } + + ref->cur = cur ^ 1; + memtrack_check_ring_pressure(&stacks, ids.tgid); + return hash; +} + +#endif /* __STACK_DELTA_BPF_H__ */ diff --git a/crates/memtrack/src/ebpf/events.rs b/crates/memtrack/src/ebpf/events.rs index 27159984..2d6ee636 100644 --- a/crates/memtrack/src/ebpf/events.rs +++ b/crates/memtrack/src/ebpf/events.rs @@ -104,6 +104,11 @@ pub fn parse_event(data: &[u8]) -> Option { }) } +/// Every stack ring record starts with its `STACK_RECORD_*` kind. +pub fn stack_record_kind(data: &[u8]) -> Option { + Some(u32::from_ne_bytes(data.get(..4)?.try_into().ok()?)) +} + /// Decode one stack record from the ring buffer, returning it alongside the /// `bpf_get_stackid()` result its frame-pointer chain is stored under. pub fn parse_stack(data: &[u8]) -> Option<(MemtrackEvent, i64)> { @@ -119,6 +124,11 @@ pub fn parse_stack(data: &[u8]) -> Option<(MemtrackEvent, i64)> { return None; }; + if header.kind != STACK_RECORD_RAW { + warn!("unexpected stack record kind: {}", header.kind); + return None; + } + let record_len = header_len + header.copy_len as usize; if data.len() < record_len { warn!( @@ -370,6 +380,7 @@ mod stack_tests { fn header(copy_len: u32) -> stack_header { stack_header { + kind: STACK_RECORD_RAW, hash: 0x0123_4567_89ab_cdef, timestamp: 987_654_321, stackid: -17, @@ -378,7 +389,7 @@ mod stack_tests { tid: 42, copy_len, truncated: 1, - _pad: [0; 3], + _pad: [0; 7], regs: stack_regs { reg: std::array::from_fn(|index| 0x1000 + index as u64), }, diff --git a/crates/memtrack/src/ebpf/memtrack/mod.rs b/crates/memtrack/src/ebpf/memtrack/mod.rs index dbe4680a..b204c557 100644 --- a/crates/memtrack/src/ebpf/memtrack/mod.rs +++ b/crates/memtrack/src/ebpf/memtrack/mod.rs @@ -1,8 +1,8 @@ use crate::ebpf::poller::{RingBufferPoller, ThreadedRingBufferPoller}; use crate::prelude::*; -use libbpf_rs::Link; use libbpf_rs::skel::OpenSkel; use libbpf_rs::skel::SkelBuilder; +use libbpf_rs::{Link, MapCore}; use std::collections::HashMap; use std::mem::MaybeUninit; use std::os::fd::{AsFd, AsRawFd, RawFd}; @@ -142,6 +142,7 @@ impl MemtrackBpf { }); let physical = options.physical; let capture_stacks = options.stack_capture; + let stack_compression = options.stack_compression; let stack_copy_budget = ((options.stack_budget / 512) * 512) .clamp(512, crate::ebpf::events::bindings::MEMTRACK_MAX_STACK_COPY); let page_shift = page_shift()?; @@ -173,6 +174,9 @@ impl MemtrackBpf { if capture_stacks { rodata.capture_stacks_enabled = 1; rodata.stack_copy_budget = stack_copy_budget; + if stack_compression { + rodata.stack_compression_enabled = 1; + } } } @@ -185,6 +189,10 @@ impl MemtrackBpf { open_skel.maps.pending_stack_hash.set_max_entries(1)?; } + if !capture_stacks || !stack_compression { + open_skel.maps.stack_refs.set_max_entries(1)?; + } + // Autoload is decided before load(), so missing fentry targets must be off here. macro_rules! disable_rmap_prog { ($name:ident) => { @@ -259,26 +267,61 @@ impl MemtrackBpf { poll_interval_ms: u64, tx: std::sync::mpsc::Sender>, ) -> Result { - use crate::ebpf::events; + use crate::ebpf::events::{self, bindings::*}; + use crate::ebpf::stack_codec::{StackDecodeError, StackDecoder}; use runner_shared::artifacts::MemtrackEventKind; - // The resolver owns the map handle because it outlives this skeleton borrow. - let stack_traces = with_skel!(self, skel => { - libbpf_rs::MapHandle::try_from(&skel.maps.stack_traces) - .context("Failed to create handle for stack_traces map")? + enum IncomingRecord { + Raw(runner_shared::artifacts::MemtrackEvent, i64), + Delta(Vec), + } + + // The resolver owns the map handles because it outlives this skeleton borrow. + let (stack_traces, stack_refs, seen_stack_hashes) = with_skel!(self, skel => { + let traces = libbpf_rs::MapHandle::try_from(&skel.maps.stack_traces) + .context("Failed to create handle for stack_traces map")?; + let refs = libbpf_rs::MapHandle::try_from(&skel.maps.stack_refs) + .context("Failed to create handle for stack_refs map")?; + let hashes = libbpf_rs::MapHandle::try_from(&skel.maps.seen_stack_hashes) + .context("Failed to create handle for seen_stack_hashes map")?; + (traces, refs, hashes) }); + let mut decoder = StackDecoder::default(); + let resolve = - move |(mut event, stackid): (runner_shared::artifacts::MemtrackEvent, i64)| { + move |record: IncomingRecord| -> Option { + let (mut event, stackid) = match record { + IncomingRecord::Raw(event, stackid) => (event, stackid), + IncomingRecord::Delta(bytes) => match decoder.decode(&bytes) { + Ok(res) => res, + Err(e) => { + debug!("failed to decode delta stack record: {e}"); + if let StackDecodeError::Desync { tid, hash, .. } = e { + let _ = stack_refs.delete(&tid.to_ne_bytes()); + let _ = seen_stack_hashes.delete(&hash.to_ne_bytes()); + } + return None; + } + }, + }; + if let MemtrackEventKind::Stack { record } = &mut event.kind { record.fp_chain = events::fp_chain(&stack_traces, stackid); } - event + Some(event) }; with_skel!(self, skel => ThreadedRingBufferPoller::new( &skel.maps.stacks, - events::parse_stack, + |data| match events::stack_record_kind(data)? { + STACK_RECORD_RAW => { + let (event, stackid) = events::parse_stack(data)?; + Some(IncomingRecord::Raw(event, stackid)) + } + STACK_RECORD_DELTA => Some(IncomingRecord::Delta(data.to_vec())), + _ => None, + }, resolve, tx, poll_interval_ms, diff --git a/crates/memtrack/src/ebpf/mod.rs b/crates/memtrack/src/ebpf/mod.rs index 76c72622..02ca4980 100644 --- a/crates/memtrack/src/ebpf/mod.rs +++ b/crates/memtrack/src/ebpf/mod.rs @@ -5,6 +5,7 @@ pub(crate) mod pause; pub(crate) mod poller; mod proc_fs; mod spawn; +pub mod stack_codec; mod stacks; pub mod stats; mod tracker; diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index f383e63c..c516ca06 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/crates/memtrack/src/ebpf/poller.rs @@ -91,7 +91,7 @@ pub struct RingBufferPoller { impl RingBufferPoller { pub fn new( rb_map: &M, - parse: F, + mut parse: F, tx: Sender>, poll_interval_ms: u64, on_drained: Option, @@ -99,7 +99,7 @@ impl RingBufferPoller { where M: MapCore, T: Send + 'static, - F: Fn(&[u8]) -> Option + Send + 'static, + F: FnMut(&[u8]) -> Option + Send + 'static, { // `Arc>` rather than `Rc>`: the built `RingBuffer` moves // into the poll thread, so the callback must be `Send`. @@ -206,7 +206,7 @@ impl ThreadedRingBufferPoller { pub fn new( rb_map: &M, parse: F, - resolve: R, + mut resolve: R, tx: Sender>, poll_interval_ms: u64, on_drained: Option, @@ -215,15 +215,17 @@ impl ThreadedRingBufferPoller { M: MapCore, T: Send + 'static, U: Send + 'static, - F: Fn(&[u8]) -> Option + Send + 'static, - R: Fn(T) -> U + Send + 'static, + F: FnMut(&[u8]) -> Option + Send + 'static, + R: FnMut(T) -> Option + Send + 'static, { let (parsed_tx, parsed_rx) = mpsc::channel::>(); let ring = RingBufferPoller::new(rb_map, parse, parsed_tx, poll_interval_ms, on_drained)?; let resolver = std::thread::spawn(move || { for batch in parsed_rx { - let resolved = batch.into_iter().map(&resolve).collect(); - let _ = tx.send(resolved); + let resolved: Vec = batch.into_iter().filter_map(&mut resolve).collect(); + if !resolved.is_empty() { + let _ = tx.send(resolved); + } } }); diff --git a/crates/memtrack/src/ebpf/stack_codec/decode.rs b/crates/memtrack/src/ebpf/stack_codec/decode.rs new file mode 100644 index 00000000..b717a72e --- /dev/null +++ b/crates/memtrack/src/ebpf/stack_codec/decode.rs @@ -0,0 +1,156 @@ +use std::collections::HashMap; + +use runner_shared::artifacts::{MemtrackEvent, MemtrackEventKind, StackRecord}; + +use super::super::events::bindings::{ + STACK_DELTA_FLAG_TRUNCATED, STACK_RECORD_DELTA, stack_delta_header, +}; +use super::{ + GROUP_BYTES, GROUP_WORDS, HEADER_LEN, NREGS, StackDecodeError, StackReference, fnv_stack_hash, +}; +use crate::prelude::*; + +#[derive(Default)] +pub struct StackDecoder { + refs: HashMap, + desyncs: u64, +} + +impl Drop for StackDecoder { + fn drop(&mut self) { + if self.desyncs != 0 { + warn!("{} delta stack records failed to decode", self.desyncs); + } + } +} + +struct Payload<'a> { + data: &'a [u8], +} + +impl Payload<'_> { + fn next(&mut self) -> Result { + let Some((word, rest)) = self.data.split_first_chunk::<8>() else { + return Err("payload too short"); + }; + self.data = rest; + Ok(u64::from_le_bytes(*word)) + } +} + +impl StackDecoder { + /// Decodes one ring record starting at its `stack_delta_header`. Returns the stack + /// event (without `fp_chain`) and its stackid; the reference advances only on success. + pub fn decode(&mut self, data: &[u8]) -> Result<(MemtrackEvent, i64), StackDecodeError> { + if data.len() < HEADER_LEN { + return Err(StackDecodeError::Short); + } + // SAFETY: length checked; bindgen-generated C ABI struct. + let header: stack_delta_header = unsafe { std::ptr::read_unaligned(data.as_ptr().cast()) }; + let tid = header.tid; + let hash = header.hash; + + match self.decode_inner(&header, data) { + Ok(res) => Ok(res), + Err(reason) => { + self.refs.remove(&tid); + self.desyncs += 1; + Err(StackDecodeError::Desync { tid, hash, reason }) + } + } + } + + fn decode_inner( + &mut self, + header: &stack_delta_header, + data: &[u8], + ) -> Result<(MemtrackEvent, i64), &'static str> { + if header.kind != STACK_RECORD_DELTA { + return Err("not a delta record"); + } + let copy_len = header.copy_len as usize; + if copy_len % GROUP_BYTES != 0 { + return Err("copy_len not a multiple of 512"); + } + let ngroups = copy_len / GROUP_BYTES; + if ngroups < 64 && header.group_mask >> ngroups != 0 { + return Err("group beyond copy_len"); + } + if header.regs_mask >> NREGS != 0 { + return Err("unknown register bit"); + } + let Some(payload) = data.get(HEADER_LEN..HEADER_LEN + header.payload_len as usize) else { + return Err("payload_len exceeds record"); + }; + let tid = header.tid; + + let empty = StackReference::default(); + let reference = if header.ref_hash == 0 { + &empty + } else { + let Some(reference) = self.refs.get(&tid) else { + return Err("missing reference"); + }; + if reference.hash != header.ref_hash { + return Err("reference mismatch"); + } + reference + }; + + let mut payload = Payload { data: payload }; + let mut regs = reference.regs; + for (r, reg) in regs.iter_mut().enumerate() { + if header.regs_mask & (1 << r) != 0 { + *reg ^= payload.next()?; + } + } + + let shift = reference.shift(header.sp); + let mut bytes = Vec::with_capacity(copy_len); + for g in 0..ngroups { + let present = header.group_mask & (1 << g) != 0; + let bitmap = if present { payload.next()? } else { 0 }; + if present && bitmap == 0 { + return Err("present group with empty bitmap"); + } + for k in 0..GROUP_WORDS { + let mut word = reference.word(shift, g * GROUP_WORDS + k); + if bitmap & (1 << k) != 0 { + word ^= payload.next()?; + } + bytes.extend_from_slice(&word.to_le_bytes()); + } + } + if !payload.data.is_empty() { + return Err("trailing payload bytes"); + } + + let actual = fnv_stack_hash(&bytes); + if actual != header.hash { + return Err("hash mismatch"); + } + + self.refs + .entry(tid) + .or_default() + .update(header.sp, header.hash, regs, &bytes); + + let event = MemtrackEvent { + pid: header.pid as i32, + tid: tid as i32, + timestamp: header.timestamp, + addr: 0, + kind: MemtrackEventKind::Stack { + record: Box::new(StackRecord { + hash: header.hash, + sp: header.sp, + regs: regs.to_vec(), + bytes, + fp_chain: Vec::new(), + truncated: header.flags & STACK_DELTA_FLAG_TRUNCATED as u8 != 0, + }), + }, + }; + Ok((event, header.stackid)) + } +} diff --git a/crates/memtrack/src/ebpf/stack_codec/encode.rs b/crates/memtrack/src/ebpf/stack_codec/encode.rs new file mode 100644 index 00000000..d5bba5fa --- /dev/null +++ b/crates/memtrack/src/ebpf/stack_codec/encode.rs @@ -0,0 +1,93 @@ +use std::collections::HashMap; + +use super::super::events::bindings::{ + STACK_DELTA_FLAG_TRUNCATED, STACK_RECORD_DELTA, stack_delta_header, +}; +use super::{ + GROUP_BYTES, GROUP_WORDS, HEADER_LEN, NREGS, RawStack, StackReference, fnv_stack_hash, +}; + +#[derive(Default)] +pub struct StackEncoder { + refs: HashMap, +} + +impl StackEncoder { + /// Byte-exact mirror of the kernel encoder: returns header + payload as the ring would carry it. + pub fn encode( + &mut self, + pid: u32, + tid: u32, + timestamp: u64, + stackid: i64, + stack: &RawStack, + ) -> Vec { + let hash = fnv_stack_hash(&stack.bytes); + let empty = StackReference::default(); + let reference = self.refs.get(&tid).unwrap_or(&empty); + + let mut out = vec![0u8; HEADER_LEN]; + let mut regs_mask = 0u64; + for r in 0..NREGS { + let d = stack.regs[r] ^ reference.regs[r]; + if d != 0 { + regs_mask |= 1 << r; + out.extend_from_slice(&d.to_le_bytes()); + } + } + + let shift = reference.shift(stack.sp); + let mut group_mask = 0u64; + for (g, group) in stack.bytes.chunks_exact(GROUP_BYTES).enumerate() { + let bitmap_pos = out.len(); + out.extend_from_slice(&[0; 8]); + let mut bitmap = 0u64; + for (k, chunk) in group.chunks_exact(8).enumerate() { + let w = u64::from_le_bytes(chunk.try_into().unwrap()); + let d = w ^ reference.word(shift, g * GROUP_WORDS + k); + if d != 0 { + bitmap |= 1 << k; + out.extend_from_slice(&d.to_le_bytes()); + } + } + if bitmap == 0 { + out.truncate(bitmap_pos); + continue; + } + group_mask |= 1 << g; + out[bitmap_pos..bitmap_pos + 8].copy_from_slice(&bitmap.to_le_bytes()); + } + + let mut flags = 0u8; + if stack.truncated { + flags |= STACK_DELTA_FLAG_TRUNCATED as u8; + } + // SAFETY: plain-old-data C struct; all-zero is a valid value. + let mut header: stack_delta_header = unsafe { std::mem::zeroed() }; + header.kind = STACK_RECORD_DELTA; + header.copy_len = stack.bytes.len() as u32; + header.hash = hash; + header.ref_hash = reference.hash; + header.timestamp = timestamp; + header.stackid = stackid; + header.sp = stack.sp; + header.pid = pid; + header.tid = tid; + header.payload_len = (out.len() - HEADER_LEN) as u32; + header.flags = flags; + header.group_mask = group_mask; + header.regs_mask = regs_mask; + // SAFETY: `out` holds at least HEADER_LEN bytes. + unsafe { std::ptr::write_unaligned(out.as_mut_ptr().cast(), header) }; + + self.refs + .entry(tid) + .or_default() + .update(stack.sp, hash, stack.regs, &stack.bytes); + out + } + + pub fn forget(&mut self, tid: u32) { + self.refs.remove(&tid); + } +} diff --git a/crates/memtrack/src/ebpf/stack_codec/fnv.rs b/crates/memtrack/src/ebpf/stack_codec/fnv.rs new file mode 100644 index 00000000..bf37c85f --- /dev/null +++ b/crates/memtrack/src/ebpf/stack_codec/fnv.rs @@ -0,0 +1,21 @@ +const FNV64_OFFSET: u64 = 0xcbf2_9ce4_8422_2325; +const FNV64_PRIME: u64 = 0x0000_0100_0000_01b3; + +/// Replica of the in-kernel capture hash (`stack_capture.bpf.h`): four FNV-1a-style +/// lanes over 512-byte chunks, folded, then mixed with the copy length. +pub fn fnv_stack_hash(bytes: &[u8]) -> u64 { + let mut lanes: [u64; 4] = std::array::from_fn(|lane| FNV64_OFFSET ^ lane as u64); + for chunk in bytes.chunks_exact(super::GROUP_BYTES) { + for quad in chunk.chunks_exact(32) { + for (lane, word) in lanes.iter_mut().zip(quad.chunks_exact(8)) { + let word = u64::from_le_bytes(word.try_into().unwrap()); + *lane = (*lane ^ word).wrapping_mul(FNV64_PRIME); + } + } + } + let mut hash = lanes[0].wrapping_mul(FNV64_PRIME) ^ lanes[1]; + hash = hash.wrapping_mul(FNV64_PRIME) ^ lanes[2]; + hash = hash.wrapping_mul(FNV64_PRIME) ^ lanes[3]; + hash = (hash ^ bytes.len() as u64).wrapping_mul(FNV64_PRIME); + if hash == 0 { FNV64_OFFSET } else { hash } +} diff --git a/crates/memtrack/src/ebpf/stack_codec/mod.rs b/crates/memtrack/src/ebpf/stack_codec/mod.rs new file mode 100644 index 00000000..a7acd53d --- /dev/null +++ b/crates/memtrack/src/ebpf/stack_codec/mod.rs @@ -0,0 +1,103 @@ +//! Userspace side of the delta-compressed stack records (`struct stack_delta_header` +//! in `event.h`). The decoder reverses the kernel encoder; the encoder is a +//! byte-exact mirror of it, used for offline measurement and tests. + +mod decode; +mod encode; +mod fnv; + +pub use decode::StackDecoder; +pub use encode::StackEncoder; +pub use fnv::fnv_stack_hash; + +use super::events::bindings::{MEMTRACK_STACK_REGS, stack_delta_header}; + +pub(crate) const NREGS: usize = MEMTRACK_STACK_REGS as usize; +pub(crate) const GROUP_WORDS: usize = 64; +pub(crate) const GROUP_BYTES: usize = GROUP_WORDS * 8; +pub(crate) const HEADER_LEN: usize = std::mem::size_of::(); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RawStack { + pub sp: u64, + pub regs: [u64; NREGS], + /// Length is a multiple of 512. + pub bytes: Vec, + pub truncated: bool, +} + +/// The last record emitted/decoded for a tid, which the next one is XORed against. +pub struct StackReference { + pub sp: u64, + pub hash: u64, + pub regs: [u64; NREGS], + pub bytes: Vec, +} + +/// The all-zero reference a keyframe is encoded against. +impl Default for StackReference { + fn default() -> Self { + Self { + sp: 0, + hash: 0, + regs: [0; NREGS], + bytes: Vec::new(), + } + } +} + +impl StackReference { + fn update(&mut self, sp: u64, hash: u64, regs: [u64; NREGS], bytes: &[u8]) { + self.sp = sp; + self.hash = hash; + self.regs = regs; + self.bytes.clear(); + self.bytes.extend_from_slice(bytes); + } + + /// Reference word paired with word `i` of a copy starting at `sp`. + fn word(&self, shift: i64, i: usize) -> u64 { + let Ok(j) = usize::try_from(i as i64 + shift) else { + return 0; + }; + self.bytes + .get(j * 8..j * 8 + 8) + .map_or(0, |w| u64::from_le_bytes(w.try_into().unwrap())) + } + + fn shift(&self, sp: u64) -> i64 { + (sp.wrapping_sub(self.sp) as i64) / 8 + } +} + +#[derive(Debug)] +pub enum StackDecodeError { + /// Too short to hold a header; the kernel never emits this. + Short, + /// The tid's reference was dropped; the kernel must be told to resync `tid` + /// and re-emit `hash`. + Desync { + tid: u32, + hash: u64, + reason: &'static str, + }, +} + +impl std::fmt::Display for StackDecodeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Short => write!(f, "delta stack record shorter than its header"), + Self::Desync { tid, hash, reason } => { + write!( + f, + "delta stack {hash:#x} of tid {tid} undecodable: {reason}" + ) + } + } + } +} + +impl std::error::Error for StackDecodeError {} + +#[cfg(test)] +mod tests; diff --git a/crates/memtrack/src/ebpf/stack_codec/tests.rs b/crates/memtrack/src/ebpf/stack_codec/tests.rs new file mode 100644 index 00000000..58fe0f26 --- /dev/null +++ b/crates/memtrack/src/ebpf/stack_codec/tests.rs @@ -0,0 +1,196 @@ +use runner_shared::artifacts::MemtrackEventKind; + +use super::super::events::bindings::stack_delta_header; +use super::*; + +const TID: u32 = 7; + +fn stack(sp: u64, words: impl IntoIterator) -> RawStack { + let bytes: Vec = words.into_iter().flat_map(u64::to_le_bytes).collect(); + assert_eq!(bytes.len() % GROUP_BYTES, 0); + let mut regs = [0; NREGS]; + regs[7] = sp; + regs[16] = 0x5555_0000_1234; + RawStack { + sp, + regs, + bytes, + truncated: false, + } +} + +/// Distinct words that depend on the absolute address, as a real stack mostly does. +fn stack_at(sp: u64, nwords: usize) -> RawStack { + stack( + sp, + (0..nwords as u64).map(|i| (sp + 8 * i).wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1), + ) +} + +fn header(record: &[u8]) -> stack_delta_header { + unsafe { std::ptr::read_unaligned(record.as_ptr().cast()) } +} + +fn roundtrip(enc: &mut StackEncoder, dec: &mut StackDecoder, s: &RawStack) -> Vec { + let record = enc.encode(1, TID, 99, 5, s); + assert_eq!( + record.len(), + HEADER_LEN + header(&record).payload_len as usize + ); + let (event, stackid) = dec.decode(&record).expect("decode"); + assert_eq!(stackid, 5); + let MemtrackEventKind::Stack { record: r } = event.kind else { + panic!("not a stack") + }; + assert_eq!( + (r.sp, &r.regs[..], &r.bytes, r.truncated), + (s.sp, &s.regs[..], &s.bytes, s.truncated) + ); + assert_eq!(r.hash, fnv_stack_hash(&s.bytes)); + record +} + +#[test] +fn zero_stack_hash() { + assert_eq!(fnv_stack_hash(&[0; 512]), 0xc419_5517_f716_585c); +} + +#[test] +fn keyframe_then_small_change() { + let (mut enc, mut dec) = (StackEncoder::default(), StackDecoder::default()); + let a = stack_at(0x7000_0000, 128); + let key = roundtrip(&mut enc, &mut dec, &a); + assert_eq!(header(&key).ref_hash, 0); + + let mut b = a.clone(); + b.bytes[8 * 3] ^= 0xff; + let rec = roundtrip(&mut enc, &mut dec, &b); + let h = header(&rec); + assert_eq!( + (h.flags, h.group_mask, h.regs_mask, h.ref_hash), + (0, 1, 0, fnv_stack_hash(&a.bytes)) + ); + assert_eq!(h.payload_len, 16); // one bitmap + one literal +} + +#[test] +fn shifted_sp_aligns_by_address() { + for delta_words in [-3i64, 5] { + let (mut enc, mut dec) = (StackEncoder::default(), StackDecoder::default()); + let a = stack_at(0x7000_1000, 128); + roundtrip(&mut enc, &mut dec, &a); + let sp = (a.sp as i64 + 8 * delta_words) as u64; + let b = stack_at(sp, 128); + let h = header(&roundtrip(&mut enc, &mut dec, &b)); + // Only words not covered by the reference differ (plus rsp). + let new_words = delta_words.unsigned_abs() as u32; + assert_eq!(h.regs_mask, 1 << 7); + assert_eq!(h.payload_len, 8 + 8 + 8 * new_words); + } +} + +#[test] +fn copy_len_shrinks_and_grows() { + let (mut enc, mut dec) = (StackEncoder::default(), StackDecoder::default()); + let big = stack_at(0x7000_0000, 256); + let small = stack_at(0x7000_0000, 64); + roundtrip(&mut enc, &mut dec, &big); + let h = header(&roundtrip(&mut enc, &mut dec, &small)); + assert_eq!((h.payload_len, h.group_mask), (0, 0)); + let h = header(&roundtrip(&mut enc, &mut dec, &big)); + assert_eq!(h.group_mask, 0b1110); + assert_eq!(h.payload_len as usize, 3 * (8 + GROUP_BYTES)); +} + +#[test] +fn unchanged_stack_carries_only_changed_regs() { + let (mut enc, mut dec) = (StackEncoder::default(), StackDecoder::default()); + let a = stack_at(0x7000_0000, 64); + roundtrip(&mut enc, &mut dec, &a); + let h = header(&roundtrip(&mut enc, &mut dec, &a)); + assert_eq!((h.payload_len, h.group_mask, h.regs_mask), (0, 0, 0)); + let mut b = a.clone(); + b.regs[0] = 42; + let h = header(&roundtrip(&mut enc, &mut dec, &b)); + assert_eq!((h.payload_len, h.group_mask, h.regs_mask), (8, 0, 1)); +} + +#[test] +fn all_changed_is_bounded() { + let (mut enc, mut dec) = (StackEncoder::default(), StackDecoder::default()); + let a = stack_at(0x7000_0000, 128); + roundtrip(&mut enc, &mut dec, &a); + let mut b = stack(a.sp, (0..128u64).map(|i| !i)); + b.regs = std::array::from_fn(|r| !(r as u64) ^ a.regs[r] ^ 1); + let h = header(&roundtrip(&mut enc, &mut dec, &b)); + assert_eq!(h.payload_len as usize, NREGS * 8 + 2 * 8 + b.bytes.len()); +} + +#[test] +fn max_budget_stack() { + let (mut enc, mut dec) = (StackEncoder::default(), StackDecoder::default()); + let mut a = stack_at(0x7fff_0000_0000, 4096); + a.truncated = true; + roundtrip(&mut enc, &mut dec, &a); + let mut b = a.clone(); + b.bytes[32 * 1024 - 1] ^= 1; + let h = header(&roundtrip(&mut enc, &mut dec, &b)); + assert_eq!((h.group_mask, h.payload_len), (1 << 63, 16)); +} + +#[test] +fn decode_errors() { + let a = stack_at(0x7000_0000, 64); + let mut b = a.clone(); + b.bytes[0] ^= 1; + let mut enc = StackEncoder::default(); + enc.encode(1, TID, 0, -1, &a); + let delta = enc.encode(1, TID, 0, -1, &b); + + let mut dec = StackDecoder::default(); + assert!(matches!( + dec.decode(&delta), + Err(StackDecodeError::Desync { + tid: TID, + reason: "missing reference", + .. + }) + )); + + // Reference from a different history (encoder forgot, decoder did not). + let mut enc2 = StackEncoder::default(); + dec.decode(&enc2.encode(1, TID, 0, -1, &b)).unwrap(); + enc2.forget(TID); + enc2.encode(1, TID, 0, -1, &a); + let stale = enc2.encode(1, TID, 0, -1, &b); + assert!(matches!( + dec.decode(&stale), + Err(StackDecodeError::Desync { + tid: TID, + reason: "reference mismatch", + .. + }) + )); + + let mut dec = StackDecoder::default(); + let mut enc = StackEncoder::default(); + let mut key = enc.encode(1, TID, 0, -1, &a); + *key.last_mut().unwrap() ^= 0x10; + assert!(matches!( + dec.decode(&key), + Err(StackDecodeError::Desync { + tid: TID, + reason: "hash mismatch", + .. + }) + )); + // A failed decode must not install a reference. + assert!(matches!( + dec.decode(&enc.encode(1, TID, 0, -1, &b)), + Err(StackDecodeError::Desync { + tid: TID, + reason: "missing reference", + .. + }) + )); +} diff --git a/crates/memtrack/src/ebpf/stacks.rs b/crates/memtrack/src/ebpf/stacks.rs index ff2b8a32..f825935e 100644 --- a/crates/memtrack/src/ebpf/stacks.rs +++ b/crates/memtrack/src/ebpf/stacks.rs @@ -8,6 +8,7 @@ pub struct StackCaptureFailureStats { pub stackid_failed: u64, pub truncated: u64, pub ring_full: u64, + pub delta_fallback: u64, } impl StackCaptureFailureStats { @@ -18,6 +19,7 @@ impl StackCaptureFailureStats { stackid_failed: slot(map, MEMTRACK_STACK_COUNTER_STACKID_FAILED)?, truncated: slot(map, MEMTRACK_STACK_COUNTER_TRUNCATED)?, ring_full: slot(map, MEMTRACK_STACK_COUNTER_RING_FULL)?, + delta_fallback: slot(map, MEMTRACK_STACK_COUNTER_DELTA_FALLBACK)?, }) } } diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 1280733a..dcffcefc 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -31,6 +31,9 @@ pub struct TrackerOptions { /// raw stack copying. #[builder(default = false)] pub stack_capture: bool, + /// Compress user stack captures in eBPF via delta encoding. + #[builder(default = true)] + pub stack_compression: bool, /// Maximum bytes of user stack to copy per captured call stack. #[builder(default = 8192)] pub stack_budget: u32, @@ -51,6 +54,10 @@ impl TrackerOptions { .stack_capture( std::env::var("CODSPEED_MEMTRACK_CAPTURE_STACKS").is_ok_and(|v| v == "1"), ) + .stack_compression(!matches!( + std::env::var("CODSPEED_MEMTRACK_STACK_COMPRESSION").as_deref(), + Ok("0") | Ok("false") + )) .stack_budget( std::env::var("CODSPEED_MEMTRACK_STACK_BUDGET") .ok() @@ -262,3 +269,42 @@ impl Tracker { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tracker_options_defaults() { + let opts = TrackerOptions::builder().build(); + assert!(opts.stack_compression); + } + + #[test] + fn tracker_options_from_env_compression_override() { + let key = "CODSPEED_MEMTRACK_STACK_COMPRESSION"; + let prev = std::env::var(key).ok(); + + unsafe { + std::env::set_var(key, "0"); + } + assert!(!TrackerOptions::from_env().stack_compression); + + unsafe { + std::env::set_var(key, "false"); + } + assert!(!TrackerOptions::from_env().stack_compression); + + unsafe { + std::env::set_var(key, "1"); + } + assert!(TrackerOptions::from_env().stack_compression); + + unsafe { + match prev { + Some(v) => std::env::set_var(key, v), + None => std::env::remove_var(key), + } + } + } +} diff --git a/crates/memtrack/testdata/stack_churn.c b/crates/memtrack/testdata/stack_churn.c new file mode 100644 index 00000000..53023a3a --- /dev/null +++ b/crates/memtrack/testdata/stack_churn.c @@ -0,0 +1,50 @@ +#include +#include + +/* + * Many allocations from a few call depths, with locals that change between + * iterations. Consecutive captures on the same thread therefore share most of + * their stack bytes but never hash equal, and the stack pointer moves both + * deeper and shallower between captures. + */ + +static volatile void* escaped_pointer; + +__attribute__((noinline)) static void leaf(int i) { + volatile char scratch[256]; + memset((char*)scratch, i & 0xff, sizeof(scratch)); + void* p = malloc(16 + (i % 7)); + escaped_pointer = p; + free(p); +} + +__attribute__((noinline)) static void middle(int i) { + volatile long pad[8]; + pad[i & 7] = i; + void* p = malloc(64); + escaped_pointer = p; + leaf(i); + free(p); +} + +__attribute__((noinline)) static void deep(int i) { + volatile long pad[32]; + pad[i & 31] = i; + middle(i); + void* p = malloc(128); + escaped_pointer = p; + free(p); +} + +int main() { + for (int i = 0; i < 2000; i++) { + if (i % 3 == 0) { + deep(i); + } else if (i % 3 == 1) { + middle(i); + } else { + leaf(i); + } + } + return 0; +} diff --git a/crates/memtrack/tests/shared.rs b/crates/memtrack/tests/shared.rs index 1d01c8ec..cc2b3b85 100644 --- a/crates/memtrack/tests/shared.rs +++ b/crates/memtrack/tests/shared.rs @@ -358,7 +358,7 @@ fn track_command_with_tracker(command: Command, tracker: Tracker) -> TrackResult /// Run `command` to completion under `tracker` and drain its events, handing /// back the still-live tracker so BPF state can be read before teardown. /// `checkpoint` runs while the tracked tree is still live. -fn run_tracked( +pub fn run_tracked( command: Command, tracker: Tracker, checkpoint: impl FnOnce(&Tracker, i32) -> anyhow::Result, diff --git a/crates/memtrack/tests/stack_budget_tests.rs b/crates/memtrack/tests/stack_budget_tests.rs index d5f80ba9..eecc271b 100644 --- a/crates/memtrack/tests/stack_budget_tests.rs +++ b/crates/memtrack/tests/stack_budget_tests.rs @@ -1,6 +1,7 @@ //! The stack copy budget is a frozen rodata constant, so the verifier's cost of //! the capture program scales with it. Loading at the default proves nothing -//! about the maximum; both must load. +//! about the maximum; both must load. Capture must be enabled: with it frozen +//! off the verifier prunes the capture path and checks nothing. use memtrack::{BpfVariant, MemtrackBpf, TrackerOptions}; use rstest::rstest; @@ -9,14 +10,18 @@ use rstest::rstest; #[case(8192)] #[case(u32::MAX)] #[test_log::test] -fn skeleton_loads_at_stack_budget(#[case] budget: u32) { +fn skeleton_loads_at_stack_budget(#[case] budget: u32, #[values(false, true)] compression: bool) { for variant in [BpfVariant::Legacy, BpfVariant::Token] { let options = TrackerOptions::builder() .variant(Some(variant)) + .stack_capture(true) + .stack_compression(compression) .stack_budget(budget) .build(); MemtrackBpf::new(&options).unwrap_or_else(|e| { - panic!("{variant:?} skeleton failed to load at budget {budget}: {e:#}") + panic!( + "{variant:?} skeleton failed to load at budget {budget} (compression {compression}): {e:#}" + ) }); } } diff --git a/crates/memtrack/tests/stack_compression_tests.rs b/crates/memtrack/tests/stack_compression_tests.rs new file mode 100644 index 00000000..93822b78 --- /dev/null +++ b/crates/memtrack/tests/stack_compression_tests.rs @@ -0,0 +1,88 @@ +//! Delta-compressed stack records must decode to exactly the bytes the kernel +//! hashed. The decoder drops any record whose reconstruction fails the hash +//! check, so a run with many records and every allocation's stack hash +//! resolvable proves the kernel encoder and userspace decoder agree. +#[macro_use] +mod shared; + +use memtrack::stack_codec::fnv_stack_hash; +use memtrack::{Tracker, TrackerOptions}; +use rstest::rstest; +use runner_shared::artifacts::MemtrackEventKind; +use std::collections::HashSet; +use std::process::Command; +use tempfile::TempDir; + +fn alloc_stack_hash(kind: &MemtrackEventKind) -> Option { + match kind { + MemtrackEventKind::Malloc { stack_hash, .. } + | MemtrackEventKind::Calloc { stack_hash, .. } + | MemtrackEventKind::AlignedAlloc { stack_hash, .. } + | MemtrackEventKind::Realloc { stack_hash, .. } + | MemtrackEventKind::Free { stack_hash } => Some(*stack_hash), + _ => None, + } +} + +#[test_with::env(GITHUB_ACTIONS)] +#[rstest] +#[case(8192)] +#[case(32768)] +#[test_log::test] +fn delta_records_decode_to_hashed_bytes( + #[case] budget: u32, +) -> Result<(), Box> { + let temp_dir = TempDir::new()?; + let binary = shared::compile_c_source( + include_str!("../testdata/stack_churn.c"), + "stack_churn", + temp_dir.path(), + )?; + let options = TrackerOptions::builder() + .stack_capture(true) + .stack_budget(budget) + .build(); + let tracker = Tracker::with_options(options)?; + let (tracker, events, ()) = shared::run_tracked(Command::new(binary), tracker, |_, _| Ok(()))?; + + let capture_stats = tracker.stack_capture_stats()?; + let teardown = std::thread::spawn(move || drop(tracker)); + + let mut emitted = HashSet::new(); + let mut stack_records = 0usize; + for event in &events { + let MemtrackEventKind::Stack { record } = &event.kind else { + continue; + }; + stack_records += 1; + assert_eq!(record.bytes.len() % 512, 0, "copy length is chunked"); + assert!(record.bytes.len() as u32 <= budget); + assert_eq!(fnv_stack_hash(&record.bytes), record.hash); + emitted.insert(record.hash); + } + + let referenced: Vec = events + .iter() + .filter_map(|e| alloc_stack_hash(&e.kind)) + .filter(|&hash| hash != 0) + .collect(); + let unresolved = referenced + .iter() + .filter(|hash| !emitted.contains(hash)) + .count(); + + eprintln!( + "budget {budget}: {stack_records} stack records, {} referencing allocations, {capture_stats:?}", + referenced.len() + ); + assert!(stack_records >= 1000, "expected many distinct stacks"); + assert_eq!(capture_stats.ring_full, 0); + assert_eq!(capture_stats.delta_fallback, 0); + assert_eq!( + unresolved, 0, + "every allocation stack hash has a decoded record" + ); + + teardown.join().expect("tracker teardown thread panicked"); + Ok(()) +}