[storage/qmdb/bench] chained-growth benchmark (#3633)

Author: roberto-bayardo

storage/src/qmdb/benches/bench.rs

@@ -2,9 +2,15 @@
 
 use criterion::criterion_main;
 
+mod chained_growth;
 mod common;
 mod generate;
 mod init;
 mod merkleize;
 
-criterion_main!(generate::benches, init::benches, merkleize::benches);
+criterion_main!(
+    chained_growth::benches,
+    generate::benches,
+    init::benches,
+    merkleize::benches
+);

storage/src/qmdb/benches/chained_growth.rs

@@ -0,0 +1,227 @@
+//! Benchmark for chained-growth merkleization against Current QMDB variants.
+//!
+//! Setup (untimed): seed `NUM_KEYS` keys, then grow a chain of `PREBUILT_CHAIN` batches applying
+//! each parent while the child is still alive.
+//!
+//! Timed: do `batches` more merkleize + apply iterations on top of the pre-built chain, with a
+//! single random update per batch so each overlay covers a tiny fraction of chunks.
+
+use crate::common::{seed_db, write_random_updates, Digest, WRITE_BUFFER_SIZE};
+use commonware_cryptography::Sha256;
+use commonware_runtime::{
+    benchmarks::{context, tokio},
+    buffer::paged::CacheRef,
+    tokio::{Config, Context},
+    BufferPooler, ThreadPooler,
+};
+use commonware_storage::{
+    journal::contiguous::fixed::Config as FConfig,
+    merkle::{self, journaled, mmb::Family as Mmb},
+    qmdb::{
+        any::traits::{DbAny, MerkleizedBatch as _, UnmerkleizedBatch as _},
+        current::{ordered::fixed::Db as OCFixed, unordered::fixed::Db as UCFixed},
+    },
+    translator::EightCap,
+};
+use commonware_utils::{NZUsize, NZU16, NZU64};
+use criterion::{criterion_group, Criterion};
+use rand::{rngs::StdRng, SeedableRng};
+use std::{
+    hint::black_box,
+    num::{NonZeroU16, NonZeroU64, NonZeroUsize},
+    time::{Duration, Instant},
+};
+
+// -- Config (mirrors merkleize bench) --
+
+const ITEMS_PER_BLOB: NonZeroU64 = NZU64!(10_000_000);
+const THREADS: NonZeroUsize = NZUsize!(8);
+const PAGE_SIZE: NonZeroU16 = NZU16!(4096);
+const LARGE_PAGE_CACHE_SIZE: NonZeroUsize = NZUsize!(131_072);
+const PARTITION: &str = "bench-chained-growth";
+
+const SMALL_CHUNK_SIZE: usize = 32;
+const LARGE_CHUNK_SIZE: usize = 256;
+
+type CurUFix32Mmb = UCFixed<Mmb, Context, Digest, Digest, Sha256, EightCap, SMALL_CHUNK_SIZE>;
+type CurOFix32Mmb = OCFixed<Mmb, Context, Digest, Digest, Sha256, EightCap, SMALL_CHUNK_SIZE>;
+type CurUFix256Mmb = UCFixed<Mmb, Context, Digest, Digest, Sha256, EightCap, LARGE_CHUNK_SIZE>;
+type CurOFix256Mmb = OCFixed<Mmb, Context, Digest, Digest, Sha256, EightCap, LARGE_CHUNK_SIZE>;
+
+fn merkle_cfg(ctx: &(impl BufferPooler + ThreadPooler), pc: CacheRef) -> journaled::Config {
+    journaled::Config {
+        journal_partition: format!("journal-{PARTITION}"),
+        metadata_partition: format!("metadata-{PARTITION}"),
+        items_per_blob: ITEMS_PER_BLOB,
+        write_buffer: WRITE_BUFFER_SIZE,
+        thread_pool: Some(ctx.create_thread_pool(THREADS).unwrap()),
+        page_cache: pc,
+    }
+}
+
+fn fix_log_cfg(pc: CacheRef) -> FConfig {
+    FConfig {
+        partition: format!("log-journal-{PARTITION}"),
+        items_per_blob: ITEMS_PER_BLOB,
+        page_cache: pc,
+        write_buffer: WRITE_BUFFER_SIZE,
+    }
+}
+
+fn pc(ctx: &impl BufferPooler) -> CacheRef {
+    CacheRef::from_pooler(ctx, PAGE_SIZE, LARGE_PAGE_CACHE_SIZE)
+}
+
+fn cur_fix_cfg(
+    ctx: &(impl BufferPooler + ThreadPooler),
+) -> commonware_storage::qmdb::current::FixedConfig<EightCap> {
+    let pc = pc(ctx);
+    commonware_storage::qmdb::current::FixedConfig {
+        merkle_config: merkle_cfg(ctx, pc.clone()),
+        journal_config: fix_log_cfg(pc),
+        grafted_metadata_partition: format!("grafted-metadata-{PARTITION}"),
+        translator: EightCap,
+    }
+}
+
+/// Number of pre-populated keys in the seeded database.
+const NUM_KEYS: u64 = 1_000_000;
+
+/// Random updates per batch. One update means each batch's chunk overlay covers ~1 / num_chunks
+/// of the bitmap, forcing chain reads to walk deep before finding a matching layer.
+const UPDATES_PER_BATCH: u64 = 1;
+
+/// Number of batches grown during the untimed seed phase, producing a Db::status chain of this
+/// depth that subsequent reads must walk through.
+const PREBUILT_CHAIN: u64 = 10_000;
+
+/// Number of additional batches to grow during the timed region.
+const GROW_COUNTS: [u64; 1] = [100];
+
+#[derive(Debug, Clone, Copy)]
+enum CurrentVariant {
+    UnorderedFixed32,
+    OrderedFixed32,
+    UnorderedFixed256,
+    OrderedFixed256,
+}
+
+impl CurrentVariant {
+    const fn name(self) -> &'static str {
+        match self {
+            Self::UnorderedFixed32 => "current::unordered::fixed::mmb chunk=32",
+            Self::OrderedFixed32 => "current::ordered::fixed::mmb chunk=32",
+            Self::UnorderedFixed256 => "current::unordered::fixed::mmb chunk=256",
+            Self::OrderedFixed256 => "current::ordered::fixed::mmb chunk=256",
+        }
+    }
+}
+
+const CURRENT_VARIANTS: [CurrentVariant; 4] = [
+    CurrentVariant::UnorderedFixed32,
+    CurrentVariant::OrderedFixed32,
+    CurrentVariant::UnorderedFixed256,
+    CurrentVariant::OrderedFixed256,
+];
+
+/// Construct a Current database for `$variant`, bind it as `$db`, and execute `$body`.
+macro_rules! with_current_db {
+    ($ctx:expr, $variant:expr, |mut $db:ident| $body:expr) => {{
+        macro_rules! init_db {
+            ($DbType:ty) => {{
+                #[allow(unused_mut)]
+                let mut $db = <$DbType>::init($ctx.clone(), cur_fix_cfg(&$ctx))
+                    .await
+                    .unwrap();
+                $body
+            }};
+        }
+        match $variant {
+            CurrentVariant::UnorderedFixed32 => init_db!(CurUFix32Mmb),
+            CurrentVariant::OrderedFixed32 => init_db!(CurOFix32Mmb),
+            CurrentVariant::UnorderedFixed256 => init_db!(CurUFix256Mmb),
+            CurrentVariant::OrderedFixed256 => init_db!(CurOFix256Mmb),
+        }
+    }};
+}
+
+/// Run a chained-growth sequence with a pre-built deep chain.
+///
+/// `fork_child` bridges the generic trait and the concrete `new_batch` method on a merkleized
+/// batch.
+async fn run_chained_growth<
+    F: merkle::Family,
+    C: DbAny<F, Key = Digest, Value = Digest>,
+    Fork: Fn(&C::Merkleized) -> C::Batch,
+>(
+    mut db: C,
+    grow: u64,
+    fork_child: Fork,
+) -> Duration {
+    seed_db(&mut db, NUM_KEYS).await;
+    let mut rng = StdRng::seed_from_u64(99);
+
+    // Pre-build a deep chain (untimed).
+    let initial = write_random_updates(db.new_batch(), UPDATES_PER_BATCH, NUM_KEYS, &mut rng);
+    let mut parent = initial.merkleize(&db, None).await.unwrap();
+    for _ in 0..PREBUILT_CHAIN {
+        let child_batch =
+            write_random_updates(fork_child(&parent), UPDATES_PER_BATCH, NUM_KEYS, &mut rng);
+        let child = child_batch.merkleize(&db, None).await.unwrap();
+        db.apply_batch(parent).await.unwrap();
+        parent = child;
+    }
+
+    // Flush buffered data so the timed region doesn't inherit setup fsync cost.
+    db.commit().await.unwrap();
+    db.sync().await.unwrap();
+
+    // Timed: grow more batches on top of the pre-built chain.
+    let start = Instant::now();
+    for _ in 0..grow {
+        let child_batch =
+            write_random_updates(fork_child(&parent), UPDATES_PER_BATCH, NUM_KEYS, &mut rng);
+        let child = child_batch.merkleize(&db, None).await.unwrap();
+        black_box(child.root());
+        db.apply_batch(parent).await.unwrap();
+        parent = child;
+    }
+    db.apply_batch(parent).await.unwrap();
+    let total = start.elapsed();
+
+    db.destroy().await.unwrap();
+    total
+}
+
+fn bench_chained_growth(c: &mut Criterion) {
+    let runner = tokio::Runner::new(Config::default());
+    for batches in GROW_COUNTS {
+        for &variant in &CURRENT_VARIANTS {
+            c.bench_function(
+                &format!(
+                    "{}/variant={} batches={batches}",
+                    module_path!(),
+                    variant.name()
+                ),
+                |b| {
+                    b.to_async(&runner).iter_custom(|iters| async move {
+                        let ctx = context::get::<Context>();
+                        let mut total = Duration::ZERO;
+                        for _ in 0..iters {
+                            with_current_db!(ctx.clone(), variant, |mut db| {
+                                total += run_chained_growth(db, batches, |p| p.new_batch()).await;
+                            });
+                        }
+                        total
+                    });
+                },
+            );
+        }
+    }
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = bench_chained_growth
+}

storage/src/qmdb/benches/common.rs

@@ -5,11 +5,14 @@ use commonware_cryptography::{Hasher, Sha256};
 use commonware_runtime::{buffer::paged::CacheRef, tokio::Context, BufferPooler, ThreadPooler};
 use commonware_storage::{
     journal::contiguous::{fixed::Config as FConfig, variable::Config as VConfig},
-    merkle::mmr::{journaled::Config as MmrConfig, Family},
+    merkle::{
+        self,
+        mmr::{journaled::Config as MmrConfig, Family},
+    },
     qmdb::{
         any::{
             ordered::{fixed::Db as OFixed, variable::Db as OVariable},
-            traits::{DbAny, UnmerkleizedBatch as _},
+            traits::{DbAny, UnmerkleizedBatch},
             unordered::{fixed::Db as UFixed, variable::Db as UVariable},
             FixedConfig as AnyFixedConfig, VariableConfig as AnyVariableConfig,
         },
@@ -543,6 +546,43 @@ pub fn make_fixed_value(rng: &mut StdRng) -> Digest {
     Sha256::hash(&rng.next_u32().to_be_bytes())
 }
 
+/// Pre-populate the database with `num_keys` unique keys, then commit and sync so that
+/// seed-phase buffered writes are flushed before the caller starts timing.
+pub async fn seed_db<F: merkle::Family, C: DbAny<F, Key = Digest, Value = Digest>>(
+    db: &mut C,
+    num_keys: u64,
+) {
+    let mut rng = StdRng::seed_from_u64(42);
+    let mut batch = db.new_batch();
+    for i in 0u64..num_keys {
+        let k = Sha256::hash(&i.to_be_bytes());
+        batch = batch.write(k, Some(make_fixed_value(&mut rng)));
+    }
+    let merkleized = batch.merkleize(db, None).await.unwrap();
+    db.apply_batch(merkleized).await.unwrap();
+    db.commit().await.unwrap();
+    db.sync().await.unwrap();
+}
+
+/// Write `num_updates` random key updates into a batch.
+pub fn write_random_updates<B, Db>(
+    mut batch: B,
+    num_updates: u64,
+    num_keys: u64,
+    rng: &mut StdRng,
+) -> B
+where
+    B: UnmerkleizedBatch<Db, K = Digest, V = Digest>,
+    Db: ?Sized,
+{
+    for _ in 0..num_updates {
+        let idx = rng.next_u64() % num_keys;
+        let k = Sha256::hash(&idx.to_be_bytes());
+        batch = batch.write(k, Some(make_fixed_value(rng)));
+    }
+    batch
+}
+
 /// Generate a variable-size `Vec<u8>` value (1-256 bytes).
 pub fn make_var_value(rng: &mut StdRng) -> Vec<u8> {
     let len = (rng.next_u32() as usize) % VARIABLE_VALUE_MAX_LEN + 1;

storage/src/qmdb/benches/merkleize.rs

@@ -6,8 +6,8 @@
 //! initialization time is not included in the benchmark. The page cache is large enough to hold the
 //! entire active key set to eliminate disk access delays from affecting the results.
 
-use crate::common::{make_fixed_value, Digest, CHUNK_SIZE, WRITE_BUFFER_SIZE};
-use commonware_cryptography::{Hasher, Sha256};
+use crate::common::{seed_db, write_random_updates, Digest, CHUNK_SIZE, WRITE_BUFFER_SIZE};
+use commonware_cryptography::Sha256;
 use commonware_runtime::{
     benchmarks::{context, tokio},
     buffer::paged::CacheRef,
@@ -22,7 +22,7 @@ use commonware_storage::{
 };
 use commonware_utils::{NZUsize, NZU16, NZU64};
 use criterion::{criterion_group, Criterion};
-use rand::{rngs::StdRng, RngCore, SeedableRng};
+use rand::{rngs::StdRng, SeedableRng};
 use std::{
     hint::black_box,
     num::{NonZeroU16, NonZeroU64, NonZeroUsize},
@@ -340,41 +340,6 @@ fn cur_var_cfg(
 
 // -- Benchmark helpers --
 
-/// Pre-populate the database with `num_keys` unique keys, commit, and sync.
-async fn seed_db<F: merkle::Family, C: DbAny<F, Key = Digest, Value = Digest>>(
-    db: &mut C,
-    num_keys: u64,
-) {
-    let mut rng = StdRng::seed_from_u64(42);
-    let mut batch = db.new_batch();
-    for i in 0u64..num_keys {
-        let k = Sha256::hash(&i.to_be_bytes());
-        batch = batch.write(k, Some(make_fixed_value(&mut rng)));
-    }
-    let merkleized = batch.merkleize(db, None).await.unwrap();
-    db.apply_batch(merkleized).await.unwrap();
-    db.commit().await.unwrap();
-    db.sync().await.unwrap();
-}
-
-/// Write `num_updates` random key updates into a batch.
-fn write_random_updates<
-    B: commonware_storage::qmdb::any::traits::UnmerkleizedBatch<Db, K = Digest, V = Digest>,
-    Db: ?Sized,
->(
-    mut batch: B,
-    num_updates: u64,
-    num_keys: u64,
-    rng: &mut StdRng,
-) -> B {
-    for _ in 0..num_updates {
-        let idx = rng.next_u64() % num_keys;
-        let k = Sha256::hash(&idx.to_be_bytes());
-        batch = batch.write(k, Some(make_fixed_value(rng)));
-    }
-    batch
-}
-
 /// Single-batch benchmark: create batch, write updates, merkleize, read root.
 async fn run_bench<F: merkle::Family, C: DbAny<F, Key = Digest, Value = Digest>>(
     mut db: C,