feat(memtrack): delta-encode stack captures in eBPF

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<PathBuf>,
+}
+
+struct DumpRecord {
+    hash: u64,
+    timestamp: u64,
+    pid: u32,
+    tid: u32,
+    stack: RawStack,
+}
+
+// Format defined in .agents/scripts/stackdump.py:
+// struct.pack("<QQIIQIBBH", hash, timestamp, pid, tid, sp, copy_len, truncated, nregs, nfp)
+// followed by nregs * u64, nfp * u64, and copy_len * u8.
+fn read_record(r: &mut impl Read) -> std::io::Result<Option<DumpRecord>> {
+    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(())
+}

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 */

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,24 +173,37 @@ 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);
     memtrack_check_ring_pressure(&stacks, ids.tgid);
     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;