Refactors to avoid broken benchmarks

durable-storage/Cargo.toml

@@ -28,15 +28,10 @@ octez-riscv-test-utils.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 
-[features]
-bench = []
-
 [[bench]]
 name = "avl_tree"
 harness = false
-required-features = ["bench"]
 
 [[bench]]
 name = "database"
 harness = false
-required-features = ["bench"]

durable-storage/benches/avl_tree.rs

@@ -1,26 +1,27 @@
 // SPDX-FileCopyrightText: 2025 TriliTech <contact@trili.tech>
 // SPDX-License-Identifier: MIT
 
+mod random;
+
 use std::collections::BTreeMap;
 use std::time::Duration;
 
 use bytes::Bytes;
 use criterion::Criterion;
 use criterion::criterion_group;
 use criterion::criterion_main;
-use octez_riscv_durable_storage::avl::Tree;
 use octez_riscv_durable_storage::avl::resolver::ArcResolver;
+use octez_riscv_durable_storage::avl::tree::Tree;
 use octez_riscv_durable_storage::key::Key;
-use octez_riscv_durable_storage::random::generate_keys;
-use octez_riscv_durable_storage::random::generate_random_bytes_in_range;
 use rand::prelude::*;
+use random::generate_keys;
+use random::generate_random_bytes_in_range;
 
 const KEY_COUNT: usize = 10_000_000;
 const OPERATIONS_PER_SAMPLE: usize = 10_000;
 
 #[derive(Debug, Clone)]
-pub enum Operation {
-    Get(Key),
+enum Operation {
     Upsert(Key, Bytes),
     Delete(Key),
 }
@@ -32,9 +33,8 @@ fn get_operations_batch(mut rng: &mut impl Rng, keys: &[Key], batch_size: usize)
                 .choose(rng)
                 .expect("The keys array is not empty")
                 .clone();
-            match rng.random_range(0..3) {
-                0 => Operation::Get(key),
-                1 => Operation::Upsert(key, generate_random_bytes_in_range(&mut rng, 1..20).into()),
+            match rng.random_range(0..2) {
+                0 => Operation::Upsert(key, generate_random_bytes_in_range(&mut rng, 1..20).into()),
                 _ => Operation::Delete(key),
             }
         })
@@ -43,7 +43,7 @@ fn get_operations_batch(mut rng: &mut impl Rng, keys: &[Key], batch_size: usize)
 
 /// This bench inserts half of the [`KEY_COUNT`]
 /// generated keys into an AVL tree and samples from
-/// all of them for the get, set and delete operations
+/// all of them for the set and delete operations
 /// on the tree.
 fn bench_avl_tree_operations(c: &mut Criterion) {
     let mut rng = rand::rng();
@@ -63,9 +63,6 @@ fn bench_avl_tree_operations(c: &mut Criterion) {
             |operations| {
                 for operation in operations {
                     match operation {
-                        Operation::Get(key) => {
-                            tree.get(&key, &resolver);
-                        }
                         Operation::Upsert(key, value) => {
                             tree.set(&key, &value, &mut resolver);
                         }
@@ -102,9 +99,6 @@ fn reference(c: &mut Criterion) {
             |operations| {
                 for operation in operations {
                     match operation {
-                        Operation::Get(key) => {
-                            tree.get(&key);
-                        }
                         Operation::Upsert(key, value) => {
                             tree.insert(key, value);
                         }

durable-storage/benches/database.rs

@@ -3,6 +3,8 @@
 //
 // SPDX-License-Identifier: MIT
 
+mod random;
+
 use std::hint::black_box;
 
 use bytes::Bytes;
@@ -12,14 +14,14 @@ 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::database::DirectoryManager;
+use octez_riscv_durable_storage::database::PersistenceLayerError;
 use octez_riscv_durable_storage::key::Key;
-use octez_riscv_durable_storage::persistence_layer::PersistenceLayerError;
-use octez_riscv_durable_storage::random::generate_keys;
-use octez_riscv_durable_storage::random::generate_random_bytes;
-use octez_riscv_durable_storage::random::generate_random_bytes_in_range;
-use octez_riscv_durable_storage::repo::DirectoryManager;
 use octez_riscv_test_utils::TestableTmpdir;
 use rand::rng;
+use random::generate_keys;
+use random::generate_random_bytes;
+use random::generate_random_bytes_in_range;
 use serde::Deserialize;
 use tokio::runtime::Handle;
 

durable-storage/src/random.rs → durable-storage/benches/random/mod.rs

@@ -2,16 +2,13 @@
 //
 // SPDX-License-Identifier: MIT
 
-#![cfg(feature = "bench")]
-
 use std::collections::HashSet;
 use std::ops::Range;
 
+use octez_riscv_durable_storage::key::KEY_MAX_SIZE;
+use octez_riscv_durable_storage::key::Key;
 use rand::prelude::*;
 
-use crate::key::KEY_MAX_SIZE;
-use crate::key::Key;
-
 pub fn generate_keys(rng: &mut impl Rng, length: usize) -> Vec<Key> {
     let mut tmp: HashSet<Key> = HashSet::with_capacity(length);
     while tmp.len() < length {

durable-storage/src/avl.rs

@@ -4,22 +4,6 @@
 
 //! Implementation of a Merkleisable AVL tree.
 
-pub(crate) mod node;
-
-cfg_if::cfg_if! {
-    if #[cfg(feature = "bench")] {
-        pub mod resolver;
-        pub mod tree;
-    } else {
-        pub(crate) mod resolver;
-        pub(crate) mod tree;
-    }
-}
-
-// Re-exports for tests and benchmarks
-#[cfg(test)]
-pub(crate) use node::Node;
-#[cfg(test)]
-pub(crate) use node::hash;
-#[cfg(any(test, feature = "bench"))]
-pub use tree::Tree;
+pub mod node;
+pub mod resolver;
+pub mod tree;

durable-storage/src/avl/node.rs

@@ -52,12 +52,6 @@ struct NodeHashRepresentation<'a, Value> {
 }
 
 impl Node {
-    #[inline]
-    /// The data stored in the [`Node`].
-    pub(crate) fn data(&self) -> &Value {
-        &self.data
-    }
-
     /// Create a new leaf [`Node`] from the given key and data.
     pub(crate) fn new(key: Key, data: impl Into<Value>) -> Self {
         Node {
@@ -97,52 +91,6 @@ impl Node {
         &mut self.balance_factor
     }
 
-    /// The data stored in a [`Node`] within the subtree of this [`Node`] with a given [`Key`] .
-    pub(super) fn get<'a>(
-        mut node: &'a Arc<Node>,
-        key: &Key,
-        resolver: &impl Resolver<Arc<Node>, Node>,
-    ) -> Option<&'a Value> {
-        loop {
-            let resolved_node = resolver.resolve(node);
-            match resolved_node.key().cmp(key) {
-                Ordering::Equal => return Some(resolved_node.data()),
-                Ordering::Greater => node = resolved_node.left_ref().root()?,
-                Ordering::Less => node = resolved_node.right_ref().root()?,
-            }
-        }
-    }
-
-    /// A mutable reference to the data stored in a [`Node`] within the subtree of this [`Node`]
-    /// with a given [`Key`] .
-    pub(super) fn get_mut<'a>(
-        mut node: &'a mut Arc<Node>,
-        key: &Key,
-        resolver: &mut impl Resolver<Arc<Node>, Node>,
-    ) -> Option<&'a mut Value> {
-        loop {
-            let ordering = {
-                let node_ref = resolver.resolve(node);
-                node_ref.key().cmp(key)
-            };
-            match ordering {
-                Ordering::Equal => {
-                    let node_mut = resolver.resolve_mut(node);
-                    node_mut.invalidate_hash();
-                    return Some(&mut node_mut.data);
-                }
-                Ordering::Greater => {
-                    let node_mut = resolver.resolve_mut(node);
-                    node = node_mut.left_mut().root_mut()?;
-                }
-                Ordering::Less => {
-                    let node_mut = resolver.resolve_mut(node);
-                    node = node_mut.right_mut().root_mut()?;
-                }
-            }
-        }
-    }
-
     #[inline]
     /// The [`Key`] used for determining the [`Node`].
     pub(super) fn key(&self) -> &Key {
@@ -610,6 +558,28 @@ pub(crate) fn hash<'a>(node: &'a Arc<Node>, resolver: &impl Resolver<Arc<Node>,
 
 #[cfg(test)]
 impl Node {
+    #[inline]
+    /// The data stored in the [`Node`].
+    pub(crate) fn data(&self) -> &Value {
+        &self.data
+    }
+
+    /// The data stored in a [`Node`] within the subtree of this [`Node`] with a given [`Key`] .
+    pub(super) fn get<'a>(
+        mut node: &'a Arc<Node>,
+        key: &Key,
+        resolver: &impl Resolver<Arc<Node>, Node>,
+    ) -> Option<&'a Value> {
+        loop {
+            let resolved_node = resolver.resolve(node);
+            match resolved_node.key().cmp(key) {
+                Ordering::Equal => return Some(resolved_node.data()),
+                Ordering::Greater => node = resolved_node.left_ref().root()?,
+                Ordering::Less => node = resolved_node.right_ref().root()?,
+            }
+        }
+    }
+
     /// Returns true if the balance factors stored in the [`Node`]'s subtree are correct.
     pub(super) fn has_correct_balance_factors(
         &self,

durable-storage/src/avl/tree.rs

@@ -73,13 +73,6 @@ impl Tree {
         }
     }
 
-    #[inline]
-    /// The data stored in a [`Node`] in the [`Tree`] with a given [`Key`].
-    pub fn get(&self, key: &Key, resolver: &impl Resolver<Arc<Node>, Node>) -> Option<&Value> {
-        let node = self.root()?;
-        Node::get(node, key, resolver)
-    }
-
     #[inline]
     /// Set the value of the [`Node`] with a given key.
     ///
@@ -93,17 +86,6 @@ impl Tree {
         self.upsert(key, 0, |old_data| old_data.set(data), resolver)
     }
 
-    #[inline]
-    /// A mutable reference to the data stored in a [`Node`] in the [`Tree`] with a given [`Key`].
-    pub(crate) fn get_mut(
-        &mut self,
-        key: &Key,
-        resolver: &mut impl Resolver<Arc<Node>, Node>,
-    ) -> Option<&mut Value> {
-        let node = self.root_mut()?;
-        Node::get_mut(node, key, resolver)
-    }
-
     /// Returns the node [`struct@Hash`], potentially re-hashing uncached [`Node`]s.
     ///
     /// If the [`struct@Hash`] has been cached, the memo is returned. Otherwise, the
@@ -292,6 +274,13 @@ impl<'a> Iterator for TreeIterator<'a> {
 
 #[cfg(test)]
 impl Tree {
+    #[inline]
+    /// The data stored in a [`Node`] in the [`Tree`] with a given [`Key`].
+    pub fn get(&self, key: &Key, resolver: &impl Resolver<Arc<Node>, Node>) -> Option<&Value> {
+        let node = self.root()?;
+        Node::get(node, key, resolver)
+    }
+
     /// Asserts that the [`Tree`] is a valid AVL tree
     pub(crate) fn check(&self, resolver: &impl Resolver<Arc<Node>, Node>) {
         let inorder = self.is_inorder(resolver);

durable-storage/src/database.rs

@@ -19,8 +19,8 @@ use crate::key::Key;
 use crate::merkle_worker::MerkleWorker;
 use crate::merkle_worker::MerkleWorkerError;
 use crate::persistence_layer::PersistenceLayer;
-use crate::persistence_layer::PersistenceLayerError;
-use crate::repo::DirectoryManager;
+pub use crate::persistence_layer::PersistenceLayerError;
+pub use crate::repo::DirectoryManager;
 
 /// An isolated key-space, independent from other [`Database`]s, on which database operations can
 /// be performed, e.g. read, write, delete.

durable-storage/src/lib.rs

@@ -28,24 +28,13 @@
 //! - **Persistence layer**: Responsible for actually persisting the basic get, set, delete
 //!   operations on disk.
 
-cfg_if::cfg_if! {
-    if #[cfg(feature = "bench")] {
-        pub mod avl;
-        pub mod merkle_layer;
-        pub mod merkle_worker;
-        pub mod persistence_layer;
-        pub mod random;
-    } else {
-        mod avl;
-        mod merkle_layer;
-        mod merkle_worker;
-        #[cfg_attr(not(test), expect(dead_code, reason = "Incomplete"))]
-        pub(crate) mod persistence_layer;
-    }
-}
-
+pub mod avl;
 pub mod commit;
 pub mod database;
 pub mod key;
+mod merkle_layer;
+mod merkle_worker;
+#[cfg_attr(not(test), expect(dead_code, reason = "Incomplete"))]
+pub(crate) mod persistence_layer;
 pub mod registry;
 pub mod repo;