use mmtkvdb::{self as kv, traits::*}; use criterion::{black_box, BatchSize, Criterion, Throughput}; use rand::{rngs::SmallRng, seq::SliceRandom, Rng, SeedableRng}; use tempfile::tempdir; fn main() { let rng = &mut SmallRng::from_entropy(); let mut bmgr = Criterion::default() .noise_threshold(0.05) .configure_from_args(); const KEY_LEN: usize = 12; const VAL_LEN: usize = 24; const MAX_DBS: usize = 64; const TOTAL: u64 = 100_000; let env_builder = kv::EnvBuilder::new().max_dbs(MAX_DBS); { const READ_ITERATIONS: u64 = 1000; let mut used_keys: Vec<[u8; KEY_LEN]> = (0..TOTAL).map(|_| rng.gen()).collect(); let unused_keys: Vec<[u8; KEY_LEN]> = (0..TOTAL).map(|_| rng.gen()).collect(); let unused_values: Vec<[u8; VAL_LEN]> = (0..TOTAL).map(|_| rng.gen()).collect(); let location = tempdir().unwrap(); let db_opts = kv::DbBuilder::new().name("benchmark"); let env_builder = env_builder.clone().dir(location.path()); let mut env = unsafe { env_builder.open_rw() }.unwrap(); let db = unsafe { env.create_db(&db_opts) }.unwrap(); let mut txn = env.txn_rw().unwrap(); for key in used_keys.iter() { let value = rng.gen::<[u8; VAL_LEN]>(); txn.put(&db, &key[..], &value).unwrap(); } for _ in 0..TOTAL { let key = rng.gen::<[u8; KEY_LEN]>(); let value = rng.gen::<[u8; VAL_LEN]>(); txn.put(&db, &key, &value).unwrap(); } txn.commit().unwrap(); used_keys.shuffle(rng); { let mut bgrp = bmgr.benchmark_group("read_one"); bgrp.throughput(Throughput::Elements(1)); let txn = env.txn_ro().unwrap(); bgrp.bench_function("not_found", |b| { b.iter_batched( || &unused_keys.choose(rng).unwrap()[..], |key| { black_box(txn.get(&db, key).unwrap()); }, BatchSize::NumIterations(READ_ITERATIONS), ) }); bgrp.bench_function("found", |b| { b.iter_batched( || &used_keys.choose(rng).unwrap()[..], |key| { black_box(txn.get(&db, key).unwrap()); }, BatchSize::NumIterations(READ_ITERATIONS), ) }); drop(txn); bgrp.finish(); } { let mut bgrp = bmgr.benchmark_group("txn_update_one"); bgrp.throughput(Throughput::Elements(1)); bgrp.bench_function("update_commit", |b| { b.iter(|| { let mut txn = env.txn_rw().unwrap(); let key = &used_keys.choose(rng).unwrap()[..]; let value = &unused_values.choose(rng).unwrap()[..]; black_box(txn.put(&db, key, value).unwrap()); txn.commit().unwrap(); }); }); bgrp.finish(); } for count in [100, 10000] { let mut bgrp = bmgr.benchmark_group(format!("txn_update_{count}")); bgrp.sample_size(30); bgrp.throughput(Throughput::Elements(count)); bgrp.bench_function("update_commit", |b| { b.iter(|| { let mut txn = env.txn_rw().unwrap(); for _ in 0..count { let key = &used_keys.choose(rng).unwrap()[..]; let value = &unused_values.choose(rng).unwrap()[..]; black_box(txn.put(&db, key, value).unwrap()); } txn.commit().unwrap(); }); }); bgrp.finish(); } location.close().unwrap(); } { let mut bgrp = bmgr.benchmark_group("full"); bgrp.sample_size(50); bgrp.throughput(Throughput::Elements(TOTAL)); //bgrp.measurement_time(std::time::Duration::from_secs(10)); bgrp.bench_function(format!("create_insert_{TOTAL}"), |b| { b.iter(|| { let location = tempdir().unwrap(); let db_opts = kv::DbBuilder::new().name("benchmark"); let env_builder = env_builder.clone().dir(location.path()); let mut env = unsafe { env_builder.open_rw() }.unwrap(); let db = unsafe { env.create_db(&db_opts) }.unwrap(); let mut txn = env.txn_rw().unwrap(); for _ in 0..TOTAL { let key = rng.gen::<[u8; KEY_LEN]>(); let value = rng.gen::<[u8; VAL_LEN]>(); txn.put(&db, &key, &value).unwrap(); } txn.commit().unwrap(); location.close().unwrap(); }) }); bgrp.finish(); } }