use criterion::{Criterion, Throughput, criterion_group, criterion_main};
use std::hint::black_box;
use std::time::{Duration, Instant};
use yo_reactor::Reactor;
use yo_resp::engine::{Cmd, ConnId, Sink, Wire, pump};
#[derive(Default)]
struct Null {
bytes: usize,
}
impl Sink for Null {
fn write(&mut self, _conn: ConnId, bytes: &[u8]) -> usize {
self.bytes += bytes.len();
bytes.len()
}
}
fn wire(args: &[&[u8]]) -> Vec<u8> {
let mut b = format!("*{}\r\n", args.len()).into_bytes();
for a in args {
b.extend_from_slice(format!("${}\r\n", a.len()).as_bytes());
b.extend_from_slice(a);
b.extend_from_slice(b"\r\n");
}
b
}
fn pipelined(args: &[&[u8]], depth: usize) -> Vec<u8> {
let one = wire(args);
let mut all = Vec::with_capacity(one.len() * depth);
for _ in 0..depth {
all.extend_from_slice(&one);
}
all
}
fn ready(sink: Null) -> (Reactor<Wire<Null>>, ConnId, Vec<Cmd>) {
let mut r = Reactor::inline(Wire::new(sink));
let conn = r.engine_mut().accept();
let warm = pipelined(&[b"SET", b"warm", b"warm"], 64);
let mut batch = Vec::new();
for _ in 0..4 {
r.engine_mut().feed(conn, &warm);
pump(&mut r, &mut batch);
}
(r, conn, batch)
}
fn bench_command(c: &mut Criterion, name: &str, args: &[&[u8]]) {
bench_over(c, name, args, |_, _| {});
}
fn bench_over(
c: &mut Criterion,
name: &str,
args: &[&[u8]],
fill: impl Fn(&mut Reactor<Wire<Null>>, ConnId),
) {
let mut g = c.benchmark_group(format!("engine/{name}"));
for depth in [1usize, 16, 64] {
let stream = pipelined(args, depth);
g.throughput(Throughput::Elements(depth as u64));
g.bench_function(format!("p{depth}"), |b| {
let (mut r, conn, mut batch) = ready(Null::default());
fill(&mut r, conn);
b.iter(|| {
r.engine_mut().feed(conn, black_box(&stream));
black_box(pump(&mut r, &mut batch))
});
});
}
g.finish();
}
const HOT: usize = 100_000;
fn fill_named(r: &mut Reactor<Wire<Null>>, conn: ConnId, key: &[u8]) {
let mut batch = Vec::new();
for chunk in 0..HOT / 64 {
let mut stream = Vec::new();
for i in 0..64 {
let m = format!("member:{:012}", chunk * 64 + i);
stream.extend_from_slice(&wire(&[b"SADD", key, m.as_bytes()]));
}
r.engine_mut().feed(conn, &stream);
pump(r, &mut batch);
}
}
fn fill_hot(r: &mut Reactor<Wire<Null>>, conn: ConnId) {
fill_named(r, conn, b"set:hot");
}
fn fill_both(r: &mut Reactor<Wire<Null>>, conn: ConnId) {
fill_named(r, conn, b"set:hot");
fill_named(r, conn, b"set:alt");
}
fn bench_set(c: &mut Criterion) {
bench_command(
c,
"set",
&[b"SET", b"key:000000000001", b"value00000000001"],
);
}
fn bench_get(c: &mut Criterion) {
bench_command(c, "get", &[b"GET", b"key:000000000001"]);
}
fn bench_incr(c: &mut Criterion) {
bench_command(c, "incr", &[b"INCR", b"hits"]);
}
fn bench_exists(c: &mut Criterion) {
bench_command(c, "exists", &[b"EXISTS", b"key:nothing"]);
}
fn bench_sadd(c: &mut Criterion) {
bench_over(
c,
"sadd",
&[b"SADD", b"set:hot", b"member:000000000001"],
fill_hot,
);
}
fn bench_sadd_alternating(c: &mut Criterion) {
let hot = wire(&[b"SADD", b"set:hot", b"member:000000000001"]);
let alt = wire(&[b"SADD", b"set:alt", b"member:000000000001"]);
let mut g = c.benchmark_group("engine/sadd-alternating");
for depth in [1usize, 16, 64] {
let mut stream = Vec::new();
for i in 0..depth {
stream.extend_from_slice(if i % 2 == 0 { &hot } else { &alt });
}
g.throughput(Throughput::Elements(depth as u64));
g.bench_function(format!("p{depth}"), |b| {
let (mut r, conn, mut batch) = ready(Null::default());
fill_both(&mut r, conn);
b.iter(|| {
r.engine_mut().feed(conn, black_box(&stream));
black_box(pump(&mut r, &mut batch))
});
});
}
g.finish();
}
fn bench_srandmember(c: &mut Criterion) {
bench_over(c, "srandmember", &[b"SRANDMEMBER", b"set:hot"], fill_hot);
}
const ENTRIES: u64 = 100_000;
const CAP: &[u8] = b"100000";
fn fill_stream(r: &mut Reactor<Wire<Null>>, conn: ConnId) {
fill_stream_of(r, conn, ENTRIES);
}
fn fill_stream_of(r: &mut Reactor<Wire<Null>>, conn: ConnId, entries: u64) {
let mut batch = Vec::new();
let mut stream = Vec::new();
for i in 1..=entries {
let id = format!("{i}-0");
stream.extend_from_slice(&wire(&[b"XADD", b"stream:hot", id.as_bytes(), b"f", b"v"]));
if i % 64 == 0 {
r.engine_mut().feed(conn, &stream);
pump(r, &mut batch);
stream.clear();
}
}
r.engine_mut().feed(conn, &stream);
pump(r, &mut batch);
r.engine_mut().feed(
conn,
&wire(&[b"XGROUP", b"CREATE", b"stream:hot", b"g", b"$"]),
);
pump(r, &mut batch);
}
fn bench_xadd(c: &mut Criterion) {
bench_over(
c,
"xadd",
&[
b"XADD",
b"stream:hot",
b"MAXLEN",
b"~",
CAP,
b"*",
b"f",
b"v",
],
fill_stream,
);
}
fn bench_xlen(c: &mut Criterion) {
bench_over(c, "xlen", &[b"XLEN", b"stream:hot"], fill_stream);
}
fn bench_xrange(c: &mut Criterion) {
bench_over(
c,
"xrange",
&[b"XRANGE", b"stream:hot", b"-", b"+", b"COUNT", b"10"],
fill_stream,
);
}
const GROUP_ENTRIES: u64 = 1024;
fn add_entries(r: &mut Reactor<Wire<Null>>, conn: ConnId, next: &mut u64, n: usize) -> Vec<u64> {
let mut batch = Vec::new();
let mut stream = Vec::new();
let mut ids = Vec::with_capacity(n);
for _ in 0..n {
*next += 1;
ids.push(*next);
let id = format!("{next}-0");
stream.extend_from_slice(&wire(&[
b"XADD",
b"stream:hot",
b"MAXLEN",
b"~",
CAP,
id.as_bytes(),
b"f",
b"v",
]));
}
r.engine_mut().feed(conn, &stream);
pump(r, &mut batch);
ids
}
fn bench_xreadgroup(c: &mut Criterion) {
let mut g = c.benchmark_group("engine/xreadgroup");
for depth in [1usize, 16, 64] {
let stream = pipelined(
&[
b"XREADGROUP",
b"GROUP",
b"g",
b"c1",
b"COUNT",
b"1",
b"STREAMS",
b"stream:hot",
b">",
],
depth,
);
g.throughput(Throughput::Elements(depth as u64));
g.bench_function(format!("p{depth}"), |b| {
b.iter_custom(|iters| {
let (mut r, conn, mut batch) = ready(Null::default());
fill_stream_of(&mut r, conn, GROUP_ENTRIES);
let mut next = GROUP_ENTRIES;
let mut total = Duration::ZERO;
for _ in 0..iters {
add_entries(&mut r, conn, &mut next, depth);
let at = Instant::now();
r.engine_mut().feed(conn, black_box(&stream));
black_box(pump(&mut r, &mut batch));
total += at.elapsed();
}
total
});
});
}
g.finish();
}
fn bench_xack(c: &mut Criterion) {
let mut g = c.benchmark_group("engine/xack");
for depth in [1usize, 16, 64] {
let read = pipelined(
&[
b"XREADGROUP",
b"GROUP",
b"g",
b"c1",
b"COUNT",
b"1",
b"STREAMS",
b"stream:hot",
b">",
],
depth,
);
g.throughput(Throughput::Elements(depth as u64));
g.bench_function(format!("p{depth}"), |b| {
b.iter_custom(|iters| {
let (mut r, conn, mut batch) = ready(Null::default());
fill_stream_of(&mut r, conn, GROUP_ENTRIES);
let mut next = GROUP_ENTRIES;
let mut total = Duration::ZERO;
for _ in 0..iters {
let ids = add_entries(&mut r, conn, &mut next, depth);
r.engine_mut().feed(conn, &read);
pump(&mut r, &mut batch);
let mut acks = Vec::new();
for id in &ids {
let id = format!("{id}-0");
acks.extend_from_slice(&wire(&[
b"XACK",
b"stream:hot",
b"g",
id.as_bytes(),
]));
}
let at = Instant::now();
r.engine_mut().feed(conn, black_box(&acks));
black_box(pump(&mut r, &mut batch));
total += at.elapsed();
}
total
});
});
}
g.finish();
}
fn bench_fanout(c: &mut Criterion) {
let one = wire(&[b"GET", b"key:000000000001"]);
let mut g = c.benchmark_group("engine/fanout");
g.throughput(Throughput::Elements(16));
g.bench_function("16x1", |b| {
let mut r = Reactor::inline(Wire::new(Null::default()));
let conns: Vec<ConnId> = (0..16).map(|_| r.engine_mut().accept()).collect();
let mut batch = Vec::new();
for _ in 0..4 {
for &conn in &conns {
r.engine_mut().feed(conn, &one);
}
pump(&mut r, &mut batch);
}
b.iter(|| {
for &conn in &conns {
r.engine_mut().feed(conn, black_box(&one));
}
black_box(pump(&mut r, &mut batch))
});
});
g.bench_function("1x16", |b| {
let stream = pipelined(&[b"GET", b"key:000000000001"], 16);
let (mut r, conn, mut batch) = ready(Null::default());
b.iter(|| {
r.engine_mut().feed(conn, black_box(&stream));
black_box(pump(&mut r, &mut batch))
});
});
g.finish();
}
criterion_group!(
benches,
bench_set,
bench_get,
bench_incr,
bench_exists,
bench_sadd,
bench_sadd_alternating,
bench_srandmember,
bench_xadd,
bench_xlen,
bench_xrange,
bench_xreadgroup,
bench_xack,
bench_fanout
);
criterion_main!(benches);