refactor(dursto): operations queue for benchmark

Refactors the database benchmark to use a queue of operations in lieu of
hardcoded operations.

Author: kurtisc

durable-storage/Cargo.toml

@@ -15,11 +15,11 @@ cfg-if.workspace = true
 log.workspace = true
 hex.workspace = true
 octez-riscv-data.workspace = true
+rand.workspace = true
 rocksdb.workspace = true
 tempfile.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
-rand.workspace = true
 
 [dev-dependencies]
 blake3.workspace = true
@@ -29,6 +29,7 @@ rand.workspace = true
 
 [features]
 bench = []
+bench_cloning = []
 
 [[bench]]
 name = "avl_tree"

durable-storage/benches/database.rs

@@ -6,7 +6,9 @@ use criterion::Criterion;
 use criterion::criterion_group;
 use criterion::criterion_main;
 use octez_riscv_durable_storage::database::Database;
+use octez_riscv_durable_storage::database::DatabaseError;
 use octez_riscv_durable_storage::merkle_layer::Key;
+use octez_riscv_durable_storage::persistence_layer::PersistenceLayerError;
 use octez_riscv_durable_storage::persistence_layer::utils::TestableTmpdir;
 use octez_riscv_durable_storage::random::generate_keys;
 use octez_riscv_durable_storage::random::generate_random_bytes;
@@ -15,48 +17,63 @@ use rand::prelude::*;
 use rand::rng;
 use tokio::runtime::Handle;
 
-struct BenchmarkState {
+#[cfg_attr(
+    not(feature = "bench_cloning"),
+    expect(
+        dead_code,
+        reason = "The handle and repo are only used for cloning the database"
+    )
+)]
+struct BenchmarkState<'a> {
     database: Database,
-    large_read_buffer: Vec<u8>,
-    large_read_choices: Vec<usize>,
-    prepopulated_large_node_keys: Vec<Key>,
+    operations: Vec<Operation>,
     prepopulated_node_keys: Vec<Key>,
     prepopulated_values: Vec<Bytes>,
-    small_read_buffer: Vec<u8>,
-    small_read_choices: Vec<usize>,
-    small_write_choices: Vec<(usize, usize)>,
+    read_buffer: Vec<u8>,
+    handle: &'a Handle,
+    repo: &'a DirectoryManager,
 }
 
-impl BenchmarkState {
-    fn clone_with(&self, handle: &Handle, repo: &DirectoryManager) -> Self {
+impl<'a> BenchmarkState<'a> {
+    fn clone_with(&self, handle: &'a Handle, repo: &'a DirectoryManager) -> Self {
         Self {
             database: self
                 .database
                 .try_clone_with(handle, repo)
                 .expect("Cloning the database should work"),
-            large_read_buffer: vec![0u8; LARGE_READ_MIN_SIZE],
-            large_read_choices: self.large_read_choices.clone(),
-            prepopulated_large_node_keys: self.prepopulated_large_node_keys.clone(),
+            operations: self.operations.clone(),
+            read_buffer: vec![0u8; LARGE_READ_MAX_SIZE],
             prepopulated_node_keys: self.prepopulated_node_keys.clone(),
             prepopulated_values: self.prepopulated_values.clone(),
-            small_read_buffer: vec![0u8; SMALL_READ_SIZE],
-            small_read_choices: self.small_read_choices.clone(),
-            small_write_choices: self.small_write_choices.clone(),
+            handle,
+            repo,
         }
     }
 }
 
+#[derive(Clone)]
+#[expect(dead_code, reason = "Some operations aren't benchmarked yet")]
+enum Operation {
+    Clone,
+    Delete { key_index: usize },
+    Exists { key_index: usize },
+    Hash,
+    Read { key_index: usize, size: usize },
+    ValueLength { key_index: usize },
+    Write { key_index: usize, data_index: usize },
+}
+
 const BLOCK_FREQUENCY: usize = 5_000;
 const ERC_20_TRANSACTIONS: usize = 10_000;
-const LARGE_READ_MIN_SIZE: usize = 48000;
+const LARGE_READ_MIN_SIZE: usize = 96000;
 const LARGE_READ_MAX_SIZE: usize = LARGE_READ_MIN_SIZE * 2;
 const PREPOPULATED_LARGE_NODE_KEYS_COUNT: usize = 100;
 const PREPOPULATED_NODE_KEYS_COUNT: usize = 10_000_000 - PREPOPULATED_LARGE_NODE_KEYS_COUNT;
 const SMALL_READS: usize = 13;
 const SMALL_READ_SIZE: usize = 32;
 const SMALL_WRITES: usize = 7;
 
-fn setup_benchmark_state(handle: &Handle, repo: &DirectoryManager) -> BenchmarkState {
+fn setup_benchmark_state<'a>(handle: &'a Handle, repo: &'a DirectoryManager) -> BenchmarkState<'a> {
     let mut database = Database::try_new(handle, repo).expect("Creating a database should succeed");
     let mut rng = rng();
 
@@ -92,68 +109,135 @@ fn setup_benchmark_state(handle: &Handle, repo: &DirectoryManager) -> BenchmarkS
         prepopulated_large_node_keys.push(key);
     }
 
-    let large_read_choices = (0..ERC_20_TRANSACTIONS)
-        .map(|_| rng.random_range(0..prepopulated_large_node_keys.len()))
-        .collect();
+    let mut operations =
+        Vec::<Operation>::with_capacity(ERC_20_TRANSACTIONS * (1 + SMALL_READS + SMALL_WRITES));
+
+    for i in 0..ERC_20_TRANSACTIONS {
+        let key_index = rng.random_range(0..prepopulated_large_node_keys.len());
+        let size = rng.random_range(LARGE_READ_MIN_SIZE..LARGE_READ_MAX_SIZE);
+        let operation = Operation::Read { key_index, size };
 
-    let small_read_choices = (0..ERC_20_TRANSACTIONS * SMALL_READS)
-        .map(|_| rng.random_range(0..prepopulated_node_keys.len()))
-        .collect();
+        let mut vec = vec![0u8; size];
+        rng.fill(vec.as_mut_slice());
 
-    let small_write_choices = (0..ERC_20_TRANSACTIONS * SMALL_WRITES)
-        .map(|_| {
-            (
-                rng.random_range(0..prepopulated_node_keys.len()),
-                rng.random_range(0..prepopulated_values.len()),
+        // Write to the node so there's something meaningful and large to read
+        database
+            .write(
+                prepopulated_large_node_keys[key_index].clone(),
+                0,
+                Bytes::from(vec),
             )
-        })
-        .collect();
+            .expect("The write should succeed");
+        operations.push(operation);
+
+        for _ in 0..SMALL_READS {
+            let key_index = rng.random_range(0..prepopulated_large_node_keys.len());
+            let size = SMALL_READ_SIZE;
+            operations.push(Operation::Read { key_index, size });
+
+            let mut vec = vec![0u8; size];
+            rng.fill(vec.as_mut_slice());
+
+            // Write to the node so there's something meaningful to read
+            database
+                .write(
+                    prepopulated_node_keys[key_index].clone(),
+                    0,
+                    Bytes::from(vec),
+                )
+                .expect("The write should succeed");
+        }
+
+        for _ in 0..SMALL_WRITES {
+            operations.push(Operation::Write {
+                key_index: rng.random_range(0..prepopulated_node_keys.len()),
+                data_index: rng.random_range(0..prepopulated_values.len()),
+            });
+        }
+
+        if i % BLOCK_FREQUENCY == 0 {
+            operations.push(Operation::Hash);
+        }
+    }
 
     BenchmarkState {
         database,
-        large_read_buffer: vec![0u8; LARGE_READ_MAX_SIZE],
-        large_read_choices,
-        prepopulated_large_node_keys,
+        operations,
         prepopulated_node_keys,
         prepopulated_values,
-        small_read_buffer: vec![0u8; SMALL_READ_SIZE],
-        small_read_choices,
-        small_write_choices,
+        read_buffer: vec![0u8; LARGE_READ_MAX_SIZE],
+        handle,
+        repo,
     }
 }
 
 #[inline(never)]
 fn bench_run(mut state: BenchmarkState) {
-    let large_node_keys = &state.prepopulated_large_node_keys;
-    let node_keys = &state.prepopulated_node_keys;
-    let values = &state.prepopulated_values;
-
-    for i in 0..ERC_20_TRANSACTIONS {
-        let key = &large_node_keys[state.large_read_choices[i]];
-        state
-            .database
-            .read(key, 0, &mut state.large_read_buffer)
-            .expect("The read should succeed");
-
-        for _ in 0..SMALL_READS {
-            let key = &node_keys[state.small_read_choices[i]];
-            state
+    // `KeyNotFound` errors are allowed, since the (conceptual, not std::hint) black box that is the transaction inputs may delete keys
+    // which are later used.
+    for operation in state.operations {
+        match operation {
+            Operation::Clone => {
+                // TODO
+                // Remove this feature after resolving RV-855 (Persistence layer cloning is slow).
+                #[cfg(feature = "bench_cloning")]
+                black_box(
+                    state
+                        .database
+                        .try_clone_with(state.handle, state.repo)
+                        .expect("The clone should succeed"),
+                );
+            }
+            Operation::Delete { key_index } => match state
                 .database
-                .read(key, 0, &mut state.small_read_buffer)
-                .expect("The read should succeed");
-        }
-
-        for _ in 0..SMALL_WRITES {
-            let key = node_keys[state.small_write_choices[i].0].clone();
-            let value = values[state.small_write_choices[i].1].clone();
-            state
+                .delete(state.prepopulated_node_keys[key_index].clone())
+            {
+                Ok(_)
+                | Err(DatabaseError::PersistenceLayer(PersistenceLayerError::KeyNotFound)) => {}
+                Err(e) => panic!("The deletion should succeed: {e:?}"),
+            },
+            Operation::Exists { key_index } => {
+                state
+                    .database
+                    .exists(&state.prepopulated_node_keys[key_index])
+                    .expect("The existence check should succeed");
+            }
+            Operation::Hash => {
+                black_box(state.database.hash());
+            }
+            Operation::Read { key_index, size } => {
+                match state.database.read(
+                    &state.prepopulated_node_keys[key_index],
+                    0,
+                    &mut state.read_buffer[0..size],
+                ) {
+                    Ok(_)
+                    | Err(DatabaseError::PersistenceLayer(PersistenceLayerError::KeyNotFound)) => {}
+                    Err(e) => panic!("The read should succeed: {e:?}"),
+                }
+            }
+            Operation::ValueLength { key_index } => match state
                 .database
-                .write(key, 0, value)
-                .expect("The read should succeed");
-        }
-
-        if i % BLOCK_FREQUENCY == 0 {
-            black_box(state.database.hash());
+                .value_length(&state.prepopulated_node_keys[key_index])
+            {
+                Ok(_)
+                | Err(DatabaseError::PersistenceLayer(PersistenceLayerError::KeyNotFound)) => {}
+                Err(e) => panic!("The value length calculation should succeed: {e:?}"),
+            },
+            Operation::Write {
+                key_index,
+                data_index,
+            } => {
+                match state.database.write(
+                    state.prepopulated_node_keys[key_index].clone(),
+                    0,
+                    state.prepopulated_values[data_index].clone(),
+                ) {
+                    Ok(_)
+                    | Err(DatabaseError::PersistenceLayer(PersistenceLayerError::KeyNotFound)) => {}
+                    Err(e) => panic!("The write should succeed: {e:?}"),
+                }
+            }
         }
     }
 }

durable-storage/src/database.rs

@@ -46,17 +46,15 @@ pub enum DatabaseError {
 }
 
 impl Database {
-    #[cfg_attr(not(test), expect(dead_code, reason = "Added in RV-827"))]
     /// Remove a key from the database.
-    pub(crate) fn delete(&mut self, key: Key) -> Result<(), DatabaseError> {
+    pub fn delete(&mut self, key: Key) -> Result<(), DatabaseError> {
         self.persistent.delete(key.as_ref())?;
         self.merkle.delete(key);
         Ok(())
     }
 
-    #[cfg_attr(not(test), expect(dead_code, reason = "Implemented in RV-827"))]
     /// Returns true if the provided key exists in the database, false if it does not.
-    pub(crate) fn exists(&self, key: &Key) -> Result<bool, DatabaseError> {
+    pub fn exists(&self, key: &Key) -> Result<bool, DatabaseError> {
         match self.persistent.get(key.as_ref()) {
             Ok(_) => Ok(true),
             Err(PersistenceLayerError::KeyNotFound) => Ok(false),
@@ -135,12 +133,11 @@ impl Database {
         })
     }
 
-    #[cfg_attr(not(test), expect(dead_code, reason = "Implemented in RV-827"))]
     /// Retrieve the length of the value associated with the provided key.
     ///
     /// Fails if:
     ///  - The key does not exist in the database.
-    pub(crate) fn value_length(&self, key: &Key) -> Result<usize, DatabaseError> {
+    pub fn value_length(&self, key: &Key) -> Result<usize, DatabaseError> {
         let value = self.persistent.get(key.as_ref())?;
         Ok(value.as_ref().len())
     }