Benchmark firewood with Criterion

firewood/Cargo.toml

@@ -46,9 +46,9 @@ assert_cmd = "2.0.7"
 predicates = "3.0.1"
 serial_test = "2.0.0"
 clap = { version = "4.3.1", features = ['derive'] }
-bencher = "0.1.5"
 tempdir = "0.3.7"
 test-case = "3.1.0"
+pprof = { version = "0.12.1", features = ["flamegraph"] }
 
 [features]
 # proof API

firewood/benches/hashops.rs

@@ -1,70 +1,124 @@
+use criterion::{criterion_group, criterion_main, profiler::Profiler, BatchSize, Criterion};
 // hash benchmarks; run with 'cargo bench'
-use std::ops::Deref;
-
-use bencher::{benchmark_group, benchmark_main, Bencher};
 use firewood::merkle::{Merkle, TrieHash, TRIE_HASH_LEN};
 use firewood_shale::{
-    cached::PlainMem, disk_address::DiskAddress, CachedStore, Storable, StoredView,
+    cached::PlainMem, compact::CompactHeader, disk_address::DiskAddress, CachedStore, ObjCache,
+    Storable, StoredView,
 };
-use rand::{distributions::Alphanumeric, Rng, SeedableRng};
+use pprof::ProfilerGuard;
+use rand::{distributions::Alphanumeric, rngs::StdRng, Rng, SeedableRng};
+use std::{fs::File, iter::repeat_with, ops::Deref, os::raw::c_int, path::Path};
 
 const ZERO_HASH: TrieHash = TrieHash([0u8; TRIE_HASH_LEN]);
 
-fn bench_dehydrate(b: &mut Bencher) {
+// To enable flamegraph output
+// cargo bench --bench shale-bench -- --profile-time=N
+enum FlamegraphProfiler {
+    Init(c_int),
+    Active(ProfilerGuard<'static>),
+}
+
+impl Profiler for FlamegraphProfiler {
+    fn start_profiling(&mut self, _benchmark_id: &str, _benchmark_dir: &Path) {
+        if let Self::Init(frequency) = self {
+            let guard = ProfilerGuard::new(*frequency).unwrap();
+            *self = Self::Active(guard);
+        }
+    }
+
+    fn stop_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
+        std::fs::create_dir_all(benchmark_dir).unwrap();
+        let filename = "firewood-flamegraph.svg";
+        let flamegraph_path = benchmark_dir.join(filename);
+        let flamegraph_file = File::create(flamegraph_path.as_path())
+            .unwrap_or_else(|_| panic!("Error creating file `{}`", flamegraph_path.display()));
+
+        if let Self::Active(profiler) = self {
+            profiler
+                .report()
+                .build()
+                .unwrap()
+                .flamegraph(flamegraph_file)
+                .expect("Error writing flamegraph");
+        }
+    }
+}
+
+fn bench_dehydrate(criterion: &mut Criterion) {
     let mut to = [1u8; TRIE_HASH_LEN];
-    b.iter(|| {
-        ZERO_HASH.dehydrate(&mut to).unwrap();
+
+    criterion.bench_function("dehydrate", |b| {
+        b.iter(|| ZERO_HASH.dehydrate(&mut to).unwrap());
     });
 }
 
-fn bench_hydrate(b: &mut Bencher) {
+fn bench_hydrate(criterion: &mut Criterion) {
     let mut store = PlainMem::new(TRIE_HASH_LEN as u64, 0u8);
     store.write(0, ZERO_HASH.deref());
 
-    b.iter(|| {
-        TrieHash::hydrate(0, &store).unwrap();
+    criterion.bench_function("hydrate", |b| {
+        b.iter(|| TrieHash::hydrate(0, &store).unwrap());
     });
 }
 
-fn bench_insert(b: &mut Bencher) {
+fn bench_insert<const N: usize>(criterion: &mut Criterion) {
     const TEST_MEM_SIZE: u64 = 20_000_000;
-    let merkle_payload_header = DiskAddress::null();
-
-    let merkle_payload_header_ref = StoredView::ptr_to_obj(
-        &PlainMem::new(2 * firewood_shale::compact::CompactHeader::MSIZE, 9),
-        merkle_payload_header,
-        firewood_shale::compact::CompactHeader::MSIZE,
-    )
-    .unwrap();
-
-    let store = firewood_shale::compact::CompactSpace::new(
-        PlainMem::new(TEST_MEM_SIZE, 0).into(),
-        PlainMem::new(TEST_MEM_SIZE, 1).into(),
-        merkle_payload_header_ref,
-        firewood_shale::ObjCache::new(1 << 20),
-        4096,
-        4096,
-    )
-    .unwrap();
-    let mut merkle = Merkle::new(Box::new(store));
-    let root = merkle.init_root().unwrap();
-
-    let mut rng = rand::rngs::StdRng::seed_from_u64(1234);
     const KEY_LEN: usize = 4;
-    b.iter(|| {
-        // generate a random key
-        let k = (&mut rng)
-            .sample_iter(&Alphanumeric)
-            .take(KEY_LEN)
-            .collect::<Vec<u8>>();
-        merkle.insert(k, vec![b'v'], root).unwrap();
-    });
-    #[cfg(trace)]
-    {
-        merkle.dump(root, &mut io::std::stdout().lock()).unwrap();
-        println!("done\n---\n\n");
-    }
+    let mut rng = StdRng::seed_from_u64(1234);
+
+    criterion
+        .benchmark_group("firewood")
+        .sample_size(30)
+        .bench_function("insert", |b| {
+            b.iter_batched(
+                || {
+                    let merkle_payload_header = DiskAddress::from(0);
+
+                    let merkle_payload_header_ref = StoredView::ptr_to_obj(
+                        &PlainMem::new(2 * CompactHeader::MSIZE, 9),
+                        merkle_payload_header,
+                        CompactHeader::MSIZE,
+                    )
+                    .unwrap();
+
+                    let store = firewood_shale::compact::CompactSpace::new(
+                        PlainMem::new(TEST_MEM_SIZE, 0).into(),
+                        PlainMem::new(TEST_MEM_SIZE, 1).into(),
+                        merkle_payload_header_ref,
+                        ObjCache::new(1 << 20),
+                        4096,
+                        4096,
+                    )
+                    .unwrap();
+
+                    let merkle = Merkle::new(Box::new(store));
+                    let root = merkle.init_root().unwrap();
+
+                    // generate a random key
+                    let keys: Vec<Vec<u8>> = repeat_with(|| {
+                        (&mut rng)
+                            .sample_iter(&Alphanumeric)
+                            .take(KEY_LEN)
+                            .collect()
+                    })
+                    .take(N)
+                    .collect();
+
+                    (merkle, root, keys)
+                },
+                |(mut merkle, root, keys)| {
+                    keys.into_iter()
+                        .for_each(|key| merkle.insert(key, vec![b'v'], root).unwrap())
+                },
+                BatchSize::SmallInput,
+            );
+        });
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().with_profiler(FlamegraphProfiler::Init(100));
+    targets = bench_dehydrate, bench_hydrate, bench_insert::<1>
 }
 
-benchmark_group!(benches, bench_dehydrate, bench_hydrate, bench_insert);
-benchmark_main!(benches);
+criterion_main!(benches);

firewood/examples/benchmark.rs

@@ -2,7 +2,7 @@
 // See the file LICENSE.md for licensing terms.
 
 use clap::Parser;
-use criterion::Criterion;
+use criterion::{Criterion, SamplingMode, Throughput};
 use firewood::db::{BatchOp, Db, DbConfig, WalConfig};
 use rand::{rngs::StdRng, Rng, SeedableRng};
 
@@ -35,12 +35,10 @@ fn main() {
 
     println!("workload prepared");
 
-    group
-        .sampling_mode(criterion::SamplingMode::Flat)
-        .sample_size(10);
+    group.sampling_mode(SamplingMode::Flat).sample_size(10);
 
     let total = (args.nbatch * args.batch_size) as u64;
-    group.throughput(criterion::Throughput::Elements(total));
+    group.throughput(Throughput::Elements(total));
 
     group.bench_with_input(
         format!("nbatch={} batch_size={}", args.nbatch, args.batch_size),

shale/benches/shale-bench.rs

@@ -1,48 +1,41 @@
-extern crate firewood_shale as shale;
-
 use criterion::{
     black_box, criterion_group, criterion_main, profiler::Profiler, Bencher, Criterion,
 };
-use pprof::ProfilerGuard;
-use rand::Rng;
-use shale::{
+use firewood_shale::{
     cached::{DynamicMem, PlainMem},
-    compact::CompactSpaceHeader,
+    compact::{CompactHeader, CompactSpaceHeader},
     disk_address::DiskAddress,
     CachedStore, Obj, StoredView,
 };
+use pprof::ProfilerGuard;
+use rand::Rng;
 use std::{fs::File, os::raw::c_int, path::Path};
 
 const BENCH_MEM_SIZE: usize = 2_000_000;
 
 // To enable flamegraph output
 // cargo bench --bench shale-bench -- --profile-time=N
-pub struct FlamegraphProfiler<'a> {
-    frequency: c_int,
-    active_profiler: Option<ProfilerGuard<'a>>,
-}
-
-impl<'a> FlamegraphProfiler<'a> {
-    #[allow(dead_code)]
-    pub fn new(frequency: c_int) -> Self {
-        FlamegraphProfiler {
-            frequency,
-            active_profiler: None,
-        }
-    }
+enum FlamegraphProfiler {
+    Init(c_int),
+    Active(ProfilerGuard<'static>),
 }
 
-impl<'a> Profiler for FlamegraphProfiler<'a> {
+impl<'a> Profiler for FlamegraphProfiler {
     fn start_profiling(&mut self, _benchmark_id: &str, _benchmark_dir: &Path) {
-        self.active_profiler = Some(ProfilerGuard::new(self.frequency).unwrap());
+        if let Self::Init(frequency) = self {
+            let guard = ProfilerGuard::new(*frequency).unwrap();
+            *self = Self::Active(guard);
+        }
     }
 
     fn stop_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
         std::fs::create_dir_all(benchmark_dir).unwrap();
-        let flamegraph_path = benchmark_dir.join("flamegraph.svg");
-        let flamegraph_file =
-            File::create(flamegraph_path).expect("File system error while creating flamegraph.svg");
-        if let Some(profiler) = self.active_profiler.take() {
+        let filename = "shale-flamegraph.svg";
+        let flamegraph_path = benchmark_dir.join(filename);
+        let flamegraph_file = File::create(flamegraph_path)
+            .expect(&format!("File system error while creating {filename}"));
+
+        if let Self::Active(profiler) = self {
             profiler
                 .report()
                 .build()
@@ -75,8 +68,7 @@ fn get_view<C: CachedStore>(b: &mut Bencher, mut cached: C) {
 fn serialize<T: CachedStore>(m: &T) {
     let compact_header_obj: DiskAddress = DiskAddress::from(0x0);
     let _: Obj<CompactSpaceHeader> =
-        StoredView::ptr_to_obj(m, compact_header_obj, shale::compact::CompactHeader::MSIZE)
-            .unwrap();
+        StoredView::ptr_to_obj(m, compact_header_obj, CompactHeader::MSIZE).unwrap();
 }
 
 fn bench_cursors(c: &mut Criterion) {
@@ -93,7 +85,7 @@ fn bench_cursors(c: &mut Criterion) {
 
 criterion_group! {
     name = benches;
-    config = Criterion::default().with_profiler(FlamegraphProfiler::new(100));
+    config = Criterion::default().with_profiler(FlamegraphProfiler::Init(100));
     targets = bench_cursors
 }