Improve benchmarks, small improvements

3 files changed, 118 insertions, 20 deletions

diff --git a/benches/benchmark.rs b/benches/benchmark.rs
index b9ecbb3..c999811 100644
--- a/benches/benchmark.rs
+++ b/benches/benchmark.rs
@@ -12,8 +12,8 @@ enum Field {
     Data,
 }
 
-pub fn upsert_benchmark(c: &mut Criterion) {
-    let mut group = c.benchmark_group("db.upsert");
+pub fn upsert_various_initial_sizes(c: &mut Criterion) {
+    let mut group = c.benchmark_group("upsert_various_initial_sizes");
 
     for size in [0, 10, 100, 1000, 10000] {
         let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir");
@@ -31,21 +31,55 @@ pub fn upsert_benchmark(c: &mut Criterion) {
             .primary_key(Field::Id)
             .initialize()
             .expect("Failed to initialize DB");
-        prefill_db_with_n_records(&mut db, size).expect("Failed to prefill DB");
+        prefill_db(&mut db, size).expect("Failed to prefill DB");
 
         group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| {
             b.iter(|| {
-                let record = random_record();
+                let record = random_record(0, size as i64);
                 let _ = db.upsert(black_box(&record));
             });
         });
     }
 }
 
-pub fn get_benchmark(c: &mut Criterion) {
-    let mut group = c.benchmark_group("db.get");
+pub fn upsert_write_durability(c: &mut Criterion) {
+    let mut group = c.benchmark_group("upsert_write_durability");
 
-    for size in [0, 10, 100, 1000, 10000] {
+    for mode in [
+        WriteDurability::Async,
+        WriteDurability::Flush,
+        WriteDurability::FlushSync,
+    ] {
+        group.bench_with_input(BenchmarkId::from_parameter(&mode), &mode, |b, _mode| {
+            let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir");
+            let data_dir = &data_dir_obj
+                .path()
+                .to_str()
+                .expect("Failed to convert tmpdir path to str");
+            let mut db = DB::configure()
+                .data_dir(&data_dir)
+                .fields(&vec![
+                    (Field::Id, RecordFieldType::Int),
+                    (Field::Name, RecordFieldType::String),
+                    (Field::Data, RecordFieldType::Bytes),
+                ])
+                .write_durability(mode.clone())
+                .primary_key(Field::Id)
+                .initialize()
+                .expect("Failed to initialize DB");
+
+            b.iter(|| {
+                let record = random_record(0, 1000);
+                let _ = db.upsert(black_box(&record));
+            });
+        });
+    }
+}
+
+pub fn get_from_disk_various_initial_sizes(c: &mut Criterion) {
+    let mut group = c.benchmark_group("get_from_disk_various_initial_sizes");
+
+    for size in [0, 10, 100, 1000, 3300, 6700, 10000, 50000, 100_000] {
         let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir");
         let data_dir = &data_dir_obj
             .path()
@@ -53,6 +87,7 @@ pub fn get_benchmark(c: &mut Criterion) {
             .expect("Failed to convert tmpdir path to str");
         let mut db = DB::configure()
             .data_dir(&data_dir)
+            .memtable_capacity(0)
             .fields(&vec![
                 (Field::Id, RecordFieldType::Int),
                 (Field::Name, RecordFieldType::String),
@@ -61,11 +96,60 @@ pub fn get_benchmark(c: &mut Criterion) {
             .primary_key(Field::Id)
             .initialize()
             .expect("Failed to initialize DB");
-        prefill_db_with_n_records(&mut db, size).expect("Failed to prefill DB");
+        prefill_db(&mut db, size).expect("Failed to prefill DB");
 
         group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| {
             b.iter(|| {
-                let id = random_int();
+                let id = random_int(0, size as i64 + 1);
+                let _ = db.get(black_box(&RecordValue::Int(id)));
+            });
+        });
+    }
+}
+
+pub fn get_various_memtable_capacities(c: &mut Criterion) {
+    let mut group = c.benchmark_group("get_various_memtable_capacities");
+
+    const PREFILL_N: usize = 10000;
+    let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir");
+    let data_dir = &data_dir_obj
+        .path()
+        .to_str()
+        .expect("Failed to convert tmpdir path to str");
+
+    // Create a db instance for prefilling
+    let mut db = DB::configure()
+        .data_dir(&data_dir)
+        .fields(&vec![
+            (Field::Id, RecordFieldType::Int),
+            (Field::Name, RecordFieldType::String),
+            (Field::Data, RecordFieldType::Bytes),
+        ])
+        .primary_key(Field::Id)
+        .initialize()
+        .expect("Failed to initialize DB");
+
+    prefill_db(&mut db, PREFILL_N).expect("Failed to prefill DB");
+    drop(db);
+
+    // prefill_db generates IDs between 0..1000, so having memtable_capacity = 1000
+    // effectively indexes the whole DB
+    for size in (0..).map(|x| x * 100).take_while(|&x| x <= 1000) {
+        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| {
+            let mut db = DB::configure()
+                .data_dir(&data_dir)
+                .fields(&vec![
+                    (Field::Id, RecordFieldType::Int),
+                    (Field::Name, RecordFieldType::String),
+                    (Field::Data, RecordFieldType::Bytes),
+                ])
+                .memtable_capacity(size)
+                .primary_key(Field::Id)
+                .initialize()
+                .expect("Failed to initialize DB");
+
+            b.iter(|| {
+                let id = random_int(0, 1000 + 1);
                 let _ = db.get(black_box(&RecordValue::Int(id)));
             });
         });
@@ -73,5 +157,11 @@ pub fn get_benchmark(c: &mut Criterion) {
 }
 
 // Register the benchmark group
-criterion_group!(benches, upsert_benchmark, get_benchmark);
+criterion_group!(
+    benches,
+    upsert_various_initial_sizes,
+    upsert_write_durability,
+    get_from_disk_various_initial_sizes,
+    get_various_memtable_capacities,
+);
 criterion_main!(benches);
diff --git a/benches/utils.rs b/benches/utils.rs
index 9f20e3c..63b3d84 100644
--- a/benches/utils.rs
+++ b/benches/utils.rs
@@ -5,9 +5,9 @@ use std::fmt::Debug;
 use std::io;
 
 // Function to generate a random integer
-pub fn random_int() -> i64 {
+pub fn random_int(from: i64, to: i64) -> i64 {
     let mut rng = rand::thread_rng();
-    rng.gen_range(0..1000) // Random number between 0 and 1000
+    rng.gen_range(from..to)
 }
 
 // Function to generate a random string
@@ -23,22 +23,22 @@ pub fn random_bytes(len: usize) -> Vec<u8> {
 }
 
 // Function to generate a random record
-pub fn random_record() -> Record {
+pub fn random_record(from_id: i64, to_id: i64) -> Record {
     Record {
         values: vec![
-            RecordValue::Int(random_int()),        // Random int value
-            RecordValue::String(random_string(5)), // Random string of length 5
-            RecordValue::Bytes(random_bytes(10)),  // Random bytes of length 10
+            RecordValue::Int(random_int(from_id, to_id)), // Random int value between 0..1000
+            RecordValue::String(random_string(5)),        // Random string of length 5
+            RecordValue::Bytes(random_bytes(10)),         // Random bytes of length 10
         ],
     }
 }
 
-pub fn prefill_db_with_n_records<T: Eq + Clone + Debug>(
+pub fn prefill_db<T: Eq + Clone + Debug>(
     db: &mut DB<T>,
-    n: usize,
+    n_records: usize,
 ) -> Result<(), io::Error> {
-    for _ in 0..n {
-        let record = random_record();
+    for _ in 0..n_records {
+        let record = random_record(0, n_records as i64);
         db.upsert(&record)?;
     }
 
diff --git a/src/common.rs b/src/common.rs
index 3d31263..2f0f942 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -1,3 +1,4 @@
+use std::fmt::Display;
 use std::fs::{metadata, File};
 use std::io::{self};
 use std::path::PathBuf;
@@ -84,6 +85,13 @@ pub enum WriteDurability {
     FlushSync,
 }
 
+impl Display for WriteDurability {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
+        write!(f, "{:?}", self)?;
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
 pub enum IndexableValue {
     Int(i64),