use commonware_cryptography::{sha256, Sha256};
use commonware_math::algebra::Random as _;
use commonware_parallel::Sequential;
use commonware_runtime::{
benchmarks::{context, tokio},
buffer::paged::CacheRef,
tokio::{Config, Context},
BufferPooler, Supervisor as _,
};
use commonware_storage::merkle::{self, full, Bagging::ForwardFold, Family};
use commonware_utils::{test_rng, NZUsize, NZU16, NZU64};
use criterion::{criterion_group, Criterion};
use std::{
num::{NonZeroU16, NonZeroU64, NonZeroUsize},
time::{Duration, Instant},
};
type StandardHasher<H> = merkle::hasher::Standard<H>;
const ITEMS_PER_BLOB: NonZeroU64 = NZU64!(10_000_000);
const PAGE_SIZE: NonZeroU16 = NZU16!(16384);
const PAGE_CACHE_SIZE: NonZeroUsize = NZUsize!(512);
const WRITE_BUFFER_SIZE: NonZeroUsize = NZUsize!(2 * 1024 * 1024);
const REBUILD_NODE_BUDGET: usize = 2_000_000;
#[cfg(not(full_bench))]
const N_LEAVES: [usize; 2] = [10_000, 100_000];
#[cfg(full_bench)]
const N_LEAVES: [usize; 3] = [10_000, 100_000, 1_000_000];
fn merkle_cfg(ctx: &impl BufferPooler, family: &str) -> full::Config<Sequential> {
full::Config {
journal_partition: format!("journal-bench-flush-{family}"),
metadata_partition: format!("metadata-bench-flush-{family}"),
items_per_blob: ITEMS_PER_BLOB,
write_buffer: WRITE_BUFFER_SIZE,
strategy: Sequential,
page_cache: CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE),
}
}
fn bench_flush_family<F: Family>(c: &mut Criterion, family: &'static str) {
let runner = tokio::Runner::new(Config::default());
for n in N_LEAVES {
let cycles = (REBUILD_NODE_BUDGET / (2 * n)).max(1);
c.bench_function(&format!("{}/n={n} family={family}", module_path!()), |b| {
b.to_async(&runner).iter_custom(move |iters| async move {
let ctx = context::get::<Context>();
let hasher = StandardHasher::<Sha256>::new(ForwardFold);
let mut rng = test_rng();
let mut total = Duration::ZERO;
let mut remaining = iters;
while remaining > 0 {
let mut merkle = full::Merkle::<F, _, sha256::Digest, _>::init(
ctx.child(family),
&hasher,
merkle_cfg(&ctx, family),
)
.await
.unwrap();
let flushes = (cycles as u64).min(remaining);
for _ in 0..flushes {
let mut batch = merkle.new_batch();
for _ in 0..n {
batch = batch.add(&hasher, &sha256::Digest::random(&mut rng));
}
let batch = merkle.with_mem(|mem| batch.merkleize(mem, &hasher));
merkle.apply_batch(&batch).unwrap();
let start = Instant::now();
merkle.flush().await.unwrap();
total += start.elapsed();
}
merkle.destroy().await.unwrap();
remaining -= flushes;
}
total
});
});
}
}
fn bench_flush(c: &mut Criterion) {
bench_flush_family::<commonware_storage::mmr::Family>(c, "mmr");
bench_flush_family::<commonware_storage::mmb::Family>(c, "mmb");
}
criterion_group! {
name = benches;
config = Criterion::default().sample_size(10);
targets = bench_flush
}