use bytes::{BufMut, BytesMut};
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rustis::resp::{
BenchDecoder, BenchTape, bench_decode_to, bench_parse_only, bench_tape_footprint,
};
use std::alloc::{GlobalAlloc, Layout, System};
use std::hint::black_box;
use std::sync::atomic::{AtomicUsize, Ordering};
struct Counting;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);
static ALLOC_BYTES: AtomicUsize = AtomicUsize::new(0);
static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0);
unsafe impl GlobalAlloc for Counting {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCS.fetch_add(1, Ordering::Relaxed);
ALLOC_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
LIVE_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
unsafe { System.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed);
unsafe { System.dealloc(ptr, layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
ALLOCS.fetch_add(1, Ordering::Relaxed);
ALLOC_BYTES.fetch_add(new_size.saturating_sub(layout.size()), Ordering::Relaxed);
LIVE_BYTES.fetch_add(new_size, Ordering::Relaxed);
LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed);
unsafe { System.realloc(ptr, layout, new_size) }
}
}
fn live_bytes() -> usize {
LIVE_BYTES.load(Ordering::Relaxed)
}
#[global_allocator]
static ALLOCATOR: Counting = Counting;
fn measure_allocs<T>(f: impl FnOnce() -> T) -> (T, usize, usize) {
let allocs_before = ALLOCS.load(Ordering::Relaxed);
let bytes_before = ALLOC_BYTES.load(Ordering::Relaxed);
let out = f();
let allocs = ALLOCS.load(Ordering::Relaxed) - allocs_before;
let bytes = ALLOC_BYTES.load(Ordering::Relaxed) - bytes_before;
(out, allocs, bytes)
}
fn build_array(n: usize, elem_len: usize) -> Vec<u8> {
let elem = "x".repeat(elem_len);
let mut buf = Vec::new();
buf.extend_from_slice(format!("*{n}\r\n").as_bytes());
for _ in 0..n {
buf.extend_from_slice(format!("${elem_len}\r\n").as_bytes());
buf.extend_from_slice(elem.as_bytes());
buf.extend_from_slice(b"\r\n");
}
buf
}
fn build_nested(rows: usize, cols: usize, elem_len: usize) -> Vec<u8> {
let elem = "x".repeat(elem_len);
let mut buf = Vec::new();
buf.extend_from_slice(format!("*{rows}\r\n").as_bytes());
for _ in 0..rows {
buf.extend_from_slice(format!("*{cols}\r\n").as_bytes());
for _ in 0..cols {
buf.extend_from_slice(format!("${elem_len}\r\n").as_bytes());
buf.extend_from_slice(elem.as_bytes());
buf.extend_from_slice(b"\r\n");
}
}
buf
}
struct Shape {
label: &'static str,
reply: Vec<u8>,
expected_nodes: usize,
}
fn build_integer_array(n: usize) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(format!("*{n}\r\n").as_bytes());
for _ in 0..n {
buf.extend_from_slice(b":1\r\n");
}
buf
}
fn shapes() -> Vec<Shape> {
vec![
Shape {
label: "array 100k x :1 (4B elements: absolute worst case)",
reply: build_integer_array(100_000),
expected_nodes: 100_000 + 2,
},
Shape {
label: "array 100k x 8B (SMEMBERS of short ids: worst case)",
reply: build_array(100_000, 8),
expected_nodes: 100_000 + 2,
},
Shape {
label: "array 100k x 50B (LRANGE, the shape ARCH-01 cites)",
reply: build_array(100_000, 50),
expected_nodes: 100_000 + 2,
},
Shape {
label: "array 1k x 4KiB (large elements: tape must vanish)",
reply: build_array(1_000, 4096),
expected_nodes: 1_000 + 2,
},
Shape {
label: "nested 10k x 10 x 20B (FT.AGGREGATE)",
reply: build_nested(10_000, 10, 20),
expected_nodes: 10_000 * 12 + 2,
},
]
}
fn report_footprint() {
println!("\n=== tape_memory/footprint ===");
println!(
"{:<48} {:>12} {:>12} {:>7} {:>8} {:>12}",
"shape", "reply B", "tape B", "tape %", "allocs", "alloc B"
);
for shape in shapes() {
let ((frame_len, tape_bytes), allocs, alloc_bytes) =
measure_allocs(|| bench_tape_footprint(&shape.reply));
let expected = shape.expected_nodes * 8;
assert_eq!(
tape_bytes, expected,
"{}: tape is {tape_bytes} B, the cost model says {expected} B",
shape.label
);
let share = 100.0 * tape_bytes as f64 / (frame_len + tape_bytes) as f64;
println!(
"{:<48} {frame_len:>12} {tape_bytes:>12} {share:>6.1}% {allocs:>8} {alloc_bytes:>12}",
shape.label
);
}
println!();
}
fn report_retained() {
println!("=== tape_memory/retained ===");
println!(
" `tail cap` is what the decoder's own BytesMut reports; `live` is what is\n \
actually allocated and unfreed. They differ because a split tape leaves the\n \
decoder holding a short tail of a block it still pins entirely — so the tail\n \
capacity understates the memory held, and `live` is the figure that counts."
);
let spike = build_array(100_000, 8);
let quiet = build_array(4, 8);
let mut decoder = BenchDecoder::new();
let baseline = live_bytes();
let row = |label: String, decoder: &BenchDecoder| {
println!(
" {label:<30} tail cap {:>9} B live {:>9} B",
decoder.retained_tape_capacity(),
live_bytes().saturating_sub(baseline)
);
};
row("fresh decoder".to_string(), &decoder);
let spike_tape = decoder.feed(&spike).expect("valid spike reply");
row(format!("spike, tape {spike_tape} B"), &decoder);
for frame in 1..=20 {
decoder.feed(&quiet).expect("valid quiet reply");
row(format!("quiet frame {frame}"), &decoder);
}
println!();
}
fn drive_nodes_u64(buf: &mut BytesMut, nodes: usize) -> u64 {
buf.clear();
for i in 0..nodes {
let payload = (i as u64).wrapping_mul(13);
buf.put_u64_le(((b'$' as u64) << 56) | (payload & 0x00FF_FFFF_FFFF_FFFF));
}
let mut sum = 0u64;
for chunk in buf.chunks_exact(8) {
let word = u64::from_le_bytes(chunk.try_into().unwrap());
sum = sum.wrapping_add(word & 0x00FF_FFFF_FFFF_FFFF);
}
sum
}
fn drive_nodes_u32(buf: &mut BytesMut, nodes: usize) -> u64 {
buf.clear();
for i in 0..nodes {
let payload = (i as u64).wrapping_mul(13);
buf.put_u32_le(((b'$' as u32) << 24) | (payload as u32 & 0x00FF_FFFF));
}
let mut sum = 0u64;
for chunk in buf.chunks_exact(4) {
let word = u32::from_le_bytes(chunk.try_into().unwrap());
sum = sum.wrapping_add((word & 0x00FF_FFFF) as u64);
}
sum
}
fn bench_tape_memory(c: &mut Criterion) {
report_footprint();
report_retained();
let mut by_len = c.benchmark_group("tape_memory/throughput_by_elem_len");
for &elem_len in &[8usize, 16, 64, 256, 1024] {
let reply = build_array(50_000, elem_len);
by_len.throughput(Throughput::Elements(50_000));
by_len.bench_with_input(
BenchmarkId::new("parse_only", elem_len),
&reply,
|b, reply| {
let mut tape = BenchTape::new();
b.iter(|| bench_parse_only(black_box(reply), &mut tape))
},
);
by_len.bench_with_input(
BenchmarkId::new("decode_to", elem_len),
&reply,
|b, reply| {
b.iter(|| {
let out: Vec<String> = bench_decode_to(black_box(reply)).unwrap();
black_box(out);
})
},
);
}
by_len.finish();
let mut width = c.benchmark_group("tape_memory/node_width");
for &nodes in &[10_000usize, 100_000] {
let mut buf = BytesMut::with_capacity(nodes * 8);
width.throughput(Throughput::Elements(nodes as u64));
width.bench_with_input(BenchmarkId::new("u64_8B", nodes), &nodes, |b, &nodes| {
b.iter(|| black_box(drive_nodes_u64(&mut buf, nodes)))
});
let mut buf = BytesMut::with_capacity(nodes * 8);
width.bench_with_input(BenchmarkId::new("u32_4B", nodes), &nodes, |b, &nodes| {
b.iter(|| black_box(drive_nodes_u32(&mut buf, nodes)))
});
}
width.finish();
}
criterion_group!(benches, bench_tape_memory);
criterion_main!(benches);