use std::hint::black_box;
use std::time::Duration;
use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use mnemo_engine::{CacheFlushPolicy, CachePutOpts, Mnemo, SemanticCachePutOpts};
use tempfile::TempDir;
#[path = "corpus.rs"]
mod corpus;
use corpus::{fast_kdf_cfg, gen_clustered_vectors, DIMENSIONS, PASSPHRASE};
const NS: &str = "bench";
const EXACT_PREPOP: usize = 100;
const SEMANTIC_PREPOP: usize = 500;
const SEMANTIC_MODEL: &str = "bench-model";
fn make_db(td: &TempDir) -> Mnemo {
let path = td.path().join("cache.mnemo");
Mnemo::create(&path, PASSPHRASE, fast_kdf_cfg(DIMENSIONS)).expect("create")
}
fn bench_cache_get_exact_hit(c: &mut Criterion) {
let td = TempDir::new().expect("tempdir");
let mut db = make_db(&td);
for i in 0..EXACT_PREPOP {
db.cache_put(
NS,
&format!("k{i}"),
format!("v{i}").as_bytes(),
CachePutOpts::default(),
)
.expect("prepop cache_put");
}
db.flush().expect("flush prepop");
let mut idx = 0usize;
let mut group = c.benchmark_group("cache");
group.throughput(criterion::Throughput::Elements(1));
group.bench_function("cache_get_exact_hit", |b| {
b.iter(|| {
let key = format!("k{}", idx % EXACT_PREPOP);
idx += 1;
let hit = db.cache_get(NS, &key).expect("cache_get");
black_box(hit);
});
});
group.finish();
}
fn bench_cache_put_exact_strict(c: &mut Criterion) {
let td = TempDir::new().expect("tempdir");
let mut db = make_db(&td);
db.set_cache_flush_policy(CacheFlushPolicy::Strict);
let mut idx = 0usize;
let mut group = c.benchmark_group("cache");
group.sample_size(20);
group.throughput(criterion::Throughput::Elements(1));
group.bench_function("cache_put_exact_strict", |b| {
b.iter(|| {
let key = format!("strict-{idx}");
idx += 1;
db.cache_put(NS, &key, b"v", CachePutOpts::default())
.expect("cache_put strict");
});
});
group.finish();
}
fn bench_cache_put_exact_batched(c: &mut Criterion) {
let td = TempDir::new().expect("tempdir");
let mut db = make_db(&td);
db.set_cache_flush_policy(CacheFlushPolicy::Batched {
max_dirty: 1_000,
max_age: Duration::from_secs(60),
});
let mut idx = 0usize;
let mut group = c.benchmark_group("cache");
group.throughput(criterion::Throughput::Elements(1));
group.bench_function("cache_put_exact_batched", |b| {
b.iter(|| {
let key = format!("batched-{idx}");
idx += 1;
db.cache_put(NS, &key, b"v", CachePutOpts::default())
.expect("cache_put batched");
});
});
group.finish();
db.flush().expect("final flush");
}
fn prepop_semantic(db: &mut Mnemo) -> Vec<Vec<f32>> {
let vectors = gen_clustered_vectors(SEMANTIC_PREPOP, DIMENSIONS, 32, 0xCACE_5EED);
for (i, v) in vectors.iter().enumerate() {
db.cache_put_semantic(
NS,
&format!("sem-k{i}"),
v.clone(),
format!("sv{i}").as_bytes(),
SemanticCachePutOpts::new(SEMANTIC_MODEL),
)
.expect("cache_put_semantic prepop");
}
db.flush().expect("flush semantic prepop");
vectors
}
fn bench_cache_get_semantic_hit(c: &mut Criterion) {
let td = TempDir::new().expect("tempdir");
let mut db = make_db(&td);
let prepop_vectors = prepop_semantic(&mut db);
let queries: Vec<Vec<f32>> = gen_clustered_vectors(20, DIMENSIONS, 32, 0xCACE_B055);
let mut idx = 0usize;
let mut group = c.benchmark_group("cache");
group.throughput(criterion::Throughput::Elements(1));
group.bench_function("cache_get_semantic_hit", |b| {
b.iter(|| {
let q = &queries[idx % queries.len()];
idx += 1;
let hit = db
.cache_get_semantic(NS, q, 0.0, SEMANTIC_MODEL)
.expect("cache_get_semantic");
black_box(hit);
});
});
group.finish();
black_box(prepop_vectors);
}
fn bench_cache_put_semantic_strict(c: &mut Criterion) {
let td = TempDir::new().expect("tempdir");
let mut db = make_db(&td);
db.set_cache_flush_policy(CacheFlushPolicy::Strict);
let pool: Vec<Vec<f32>> = gen_clustered_vectors(50, DIMENSIONS, 8, 0xCACE_D04E);
let mut idx = 0usize;
let mut group = c.benchmark_group("cache");
group.sample_size(20);
group.throughput(criterion::Throughput::Elements(1));
group.bench_function("cache_put_semantic_strict", |b| {
b.iter_batched(
|| {
let i = idx;
idx += 1;
let key = format!("sem-strict-{i}");
let v = pool[i % pool.len()].clone();
(key, v)
},
|(key, v)| {
db.cache_put_semantic(NS, &key, v, b"v", SemanticCachePutOpts::new(SEMANTIC_MODEL))
.expect("cache_put_semantic strict");
},
BatchSize::SmallInput,
);
});
group.finish();
}
criterion_group!(
cache_benches,
bench_cache_get_exact_hit,
bench_cache_put_exact_strict,
bench_cache_put_exact_batched,
bench_cache_get_semantic_hit,
bench_cache_put_semantic_strict,
);
criterion_main!(cache_benches);