feat: Parallel updates to Merkle trie (v2)

firewood/Cargo.toml

@@ -35,6 +35,7 @@ metrics.workspace = true
 thiserror.workspace = true
 # Regular dependencies
 typed-builder = "0.22.0"
+rayon = "1.11.0"
 
 [features]
 default = []

firewood/src/db.rs

@@ -27,6 +27,8 @@ use std::sync::Arc;
 use thiserror::Error;
 use typed_builder::TypedBuilder;
 
+use crate::merkle::parallel::ParallelMerkle;
+
 #[derive(Error, Debug)]
 #[non_exhaustive]
 /// Represents the different types of errors that can occur in the database.
@@ -200,6 +202,26 @@ impl Db {
         self.manager.view(root_hash).map_err(Into::into)
     }
 
+    /// Propose a new batch that is processed in parallel.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the revision manager cannot create a thread pool.
+    pub fn propose_parallel(
+        &self,
+        batch: impl IntoIterator<IntoIter: KeyValuePairIter>,
+    ) -> Result<Proposal<'_>, api::Error> {
+        let parent = self.manager.current_revision();
+        let mut parallel_merkle = ParallelMerkle::default();
+        let immutable =
+            parallel_merkle.create_proposal(&parent, batch, self.manager.threadpool())?;
+        self.manager.add_proposal(immutable.clone());
+        Ok(Proposal {
+            nodestore: immutable,
+            db: self,
+        })
+    }
+
     /// Dump the Trie of the latest revision.
     pub fn dump(&self, w: &mut dyn Write) -> Result<(), std::io::Error> {
         let latest_rev_nodestore = self.manager.current_revision();
@@ -555,6 +577,164 @@ mod test {
         assert_eq!(&*value, b"proposal_value");
     }
 
+    #[test]
+    fn test_propose_parallel_reopen() {
+        fn insert_commit(db: &TestDb, kv: u8) {
+            let keys: Vec<[u8; 1]> = vec![[kv; 1]];
+            let vals: Vec<Box<[u8]>> = vec![Box::new([kv; 1])];
+            let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+            let proposal = db.propose_parallel(kviter).unwrap();
+            proposal.commit().unwrap();
+        }
+
+        // Create, insert, close, open, insert
+        let db = TestDb::new();
+        insert_commit(&db, 1);
+        let db = db.reopen();
+        insert_commit(&db, 2);
+        // Check that the keys are still there after the commits
+        let committed = db.revision(db.root_hash().unwrap().unwrap()).unwrap();
+        let keys: Vec<[u8; 1]> = vec![[1; 1], [2; 1]];
+        let vals: Vec<Box<[u8]>> = vec![Box::new([1; 1]), Box::new([2; 1])];
+        let kviter = keys.iter().zip(vals.iter());
+        for (k, v) in kviter {
+            assert_eq!(&committed.val(k).unwrap().unwrap(), v);
+        }
+        drop(db);
+
+        // Open-db1, insert, open-db2, insert
+        let db1 = TestDb::new();
+        insert_commit(&db1, 1);
+        let db2 = TestDb::new();
+        insert_commit(&db2, 2);
+        let committed1 = db1.revision(db1.root_hash().unwrap().unwrap()).unwrap();
+        let committed2 = db2.revision(db2.root_hash().unwrap().unwrap()).unwrap();
+        let keys: Vec<[u8; 1]> = vec![[1; 1], [2; 1]];
+        let vals: Vec<Box<[u8]>> = vec![Box::new([1; 1]), Box::new([2; 1])];
+        let mut kviter = keys.iter().zip(vals.iter());
+        let (k, v) = kviter.next().unwrap();
+        assert_eq!(&committed1.val(k).unwrap().unwrap(), v);
+        let (k, v) = kviter.next().unwrap();
+        assert_eq!(&committed2.val(k).unwrap().unwrap(), v);
+    }
+
+    #[test]
+    fn test_propose_parallel() {
+        const N: usize = 100;
+        let db = TestDb::new();
+
+        // Test an empty proposal
+        let keys: Vec<[u8; 0]> = Vec::new();
+        let vals: Vec<Box<[u8]>> = Vec::new();
+
+        let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+        let proposal = db.propose_parallel(kviter).unwrap();
+        proposal.commit().unwrap();
+
+        // Create a proposal consisting of a single entry and an empty key.
+        let keys: Vec<[u8; 0]> = vec![[0; 0]];
+
+        // Note that if the value is [], then it is interpreted as a DeleteRange.
+        // Instead, set value to [0]
+        let vals: Vec<Box<[u8]>> = vec![Box::new([0; 1])];
+
+        let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+        let proposal = db.propose_parallel(kviter).unwrap();
+
+        let kviter = keys.iter().zip(vals.iter());
+        for (k, v) in kviter {
+            assert_eq!(&proposal.val(k).unwrap().unwrap(), v);
+        }
+        proposal.commit().unwrap();
+
+        // Check that the key is still there after the commit
+        let committed = db.revision(db.root_hash().unwrap().unwrap()).unwrap();
+        let kviter = keys.iter().zip(vals.iter());
+        for (k, v) in kviter {
+            assert_eq!(&committed.val(k).unwrap().unwrap(), v);
+        }
+
+        // Create a proposal that deletes the previous entry
+        let vals: Vec<Box<[u8]>> = vec![Box::new([0; 0])];
+        let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+        let proposal = db.propose_parallel(kviter).unwrap();
+
+        let kviter = keys.iter().zip(vals.iter());
+        for (k, _v) in kviter {
+            assert_eq!(proposal.val(k).unwrap(), None);
+        }
+        proposal.commit().unwrap();
+
+        // Create a proposal that inserts 0 to 999
+        let (keys, vals): (Vec<_>, Vec<_>) = (0..1000)
+            .map(|i| {
+                (
+                    format!("key{i}").into_bytes(),
+                    Box::from(format!("value{i}").as_bytes()),
+                )
+            })
+            .unzip();
+
+        let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+        let proposal = db.propose_parallel(kviter).unwrap();
+        let kviter = keys.iter().zip(vals.iter());
+        for (k, v) in kviter {
+            assert_eq!(&proposal.val(k).unwrap().unwrap(), v);
+        }
+        proposal.commit().unwrap();
+
+        // Create a proposal that deletes all of the even entries
+        let (keys, vals): (Vec<_>, Vec<_>) = (0..1000)
+            .filter_map(|i| {
+                if i % 2 != 0 {
+                    Some::<(Vec<u8>, Box<[u8]>)>((format!("key{i}").into_bytes(), Box::new([])))
+                } else {
+                    None
+                }
+            })
+            .unzip();
+
+        let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+        let proposal = db.propose_parallel(kviter).unwrap();
+        let kviter = keys.iter().zip(vals.iter());
+        for (k, _v) in kviter {
+            assert_eq!(proposal.val(k).unwrap(), None);
+        }
+        proposal.commit().unwrap();
+
+        // Create a proposal that deletes using empty prefix
+        let keys: Vec<[u8; 0]> = vec![[0; 0]];
+        let vals: Vec<Box<[u8]>> = vec![Box::new([0; 0])];
+        let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+        let proposal = db.propose_parallel(kviter).unwrap();
+        proposal.commit().unwrap();
+
+        // Create N keys and values like (key0, value0)..(keyN, valueN)
+        let rng = firewood_storage::SeededRng::from_env_or_random();
+        let (keys, vals): (Vec<_>, Vec<_>) = (0..N)
+            .map(|i| {
+                (
+                    rng.random::<[u8; 32]>(),
+                    Box::from(format!("value{i}").as_bytes()),
+                )
+            })
+            .unzip();
+
+        // Looping twice to test that we are reusing the thread pool.
+        for _ in 0..2 {
+            let kviter = keys.iter().zip(vals.iter()).map_into_batch();
+            let proposal = db.propose_parallel(kviter).unwrap();
+
+            // iterate over the keys and values again, checking that the values are in the correct proposal
+            let kviter = keys.iter().zip(vals.iter());
+
+            for (k, v) in kviter {
+                assert_eq!(&proposal.val(k).unwrap().unwrap(), v);
+            }
+            proposal.commit().unwrap();
+        }
+    }
+
     /// Test that proposing on a proposal works as expected
     ///
     /// Test creates two batches and proposes them, and verifies that the values are in the correct proposal.

firewood/src/manager.rs

@@ -13,18 +13,20 @@
 use std::collections::{HashMap, VecDeque};
 use std::num::NonZero;
 use std::path::PathBuf;
-use std::sync::{Arc, Mutex, RwLock};
+use std::sync::{Arc, Mutex, OnceLock, RwLock};
 
 use firewood_storage::logger::{trace, warn};
 use metrics::gauge;
+use rayon::{ThreadPool, ThreadPoolBuilder};
 use typed_builder::TypedBuilder;
 
 use crate::merkle::Merkle;
 use crate::v2::api::{ArcDynDbView, HashKey, OptionalHashKeyExt};
 
 pub use firewood_storage::CacheReadStrategy;
 use firewood_storage::{
-    Committed, FileBacked, FileIoError, HashedNodeReader, ImmutableProposal, NodeStore, TrieHash,
+    BranchNode, Committed, FileBacked, FileIoError, HashedNodeReader, ImmutableProposal, NodeStore,
+    TrieHash,
 };
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, TypedBuilder)]
@@ -75,6 +77,7 @@ pub(crate) struct RevisionManager {
     proposals: Mutex<Vec<ProposedRevision>>,
     // committing_proposals: VecDeque<Arc<ProposedImmutable>>,
     by_hash: RwLock<HashMap<TrieHash, CommittedRevision>>,
+    threadpool: OnceLock<ThreadPool>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -115,6 +118,7 @@ impl RevisionManager {
             by_hash: RwLock::new(Default::default()),
             proposals: Mutex::new(Default::default()),
             // committing_proposals: Default::default(),
+            threadpool: OnceLock::new(),
         };
 
         if let Some(hash) = nodestore.root_hash().or_default_root_hash() {
@@ -309,6 +313,23 @@ impl RevisionManager {
             .expect("there is always one revision")
             .clone()
     }
+
+    /// Gets or creates a threadpool associated with the revision manager.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the it cannot create a thread pool.
+    pub fn threadpool(&self) -> &ThreadPool {
+        // Note that OnceLock currently doesn't support get_or_try_init (it is available in a
+        // nightly release). The get_or_init should be replaced with get_or_try_init once it
+        // is available to allow the error to be passed back to the caller.
+        self.threadpool.get_or_init(|| {
+            ThreadPoolBuilder::new()
+                .num_threads(BranchNode::MAX_CHILDREN)
+                .build()
+                .expect("Error in creating threadpool")
+        })
+    }
 }
 
 #[cfg(test)]

firewood/src/merkle.rs → firewood/src/merkle/mod.rs

@@ -4,6 +4,9 @@
 #[cfg(test)]
 pub(crate) mod tests;
 
+/// Parallel merkle
+pub mod parallel;
+
 use crate::iter::{MerkleKeyValueIter, PathIterator, TryExtend};
 use crate::proof::{Proof, ProofCollection, ProofError, ProofNode};
 use crate::range_proof::RangeProof;
@@ -593,10 +596,17 @@ impl<S: ReadableStorage> Merkle<NodeStore<MutableProposal, S>> {
     /// Map `key` to `value` in the trie.
     /// Each element of key is 2 nibbles.
     pub fn insert(&mut self, key: &[u8], value: Value) -> Result<(), FileIoError> {
-        let key = Path::from_nibbles_iterator(NibblesIterator::new(key));
+        self.insert_from_iter(NibblesIterator::new(key), value)
+    }
 
+    /// Map `key` to `value` in the trie when `key` is a `NibblesIterator`
+    pub fn insert_from_iter(
+        &mut self,
+        key: NibblesIterator<'_>,
+        value: Value,
+    ) -> Result<(), FileIoError> {
+        let key = Path::from_nibbles_iterator(key);
         let root = self.nodestore.root_mut();
-
         let Some(root_node) = std::mem::take(root) else {
             // The trie is empty. Create a new leaf node with `value` and set
             // it as the root.
@@ -751,8 +761,18 @@ impl<S: ReadableStorage> Merkle<NodeStore<MutableProposal, S>> {
     /// Otherwise returns `None`.
     /// Each element of `key` is 2 nibbles.
     pub fn remove(&mut self, key: &[u8]) -> Result<Option<Value>, FileIoError> {
-        let key = Path::from_nibbles_iterator(NibblesIterator::new(key));
+        self.remove_from_iter(NibblesIterator::new(key))
+    }
 
+    /// Removes the value associated with the given `key` where `key` is a `NibblesIterator`
+    /// Returns the value that was removed, if any.
+    /// Otherwise returns `None`.
+    /// Each element of `key` is 2 nibbles.
+    pub fn remove_from_iter(
+        &mut self,
+        key: NibblesIterator<'_>,
+    ) -> Result<Option<Value>, FileIoError> {
+        let key = Path::from_nibbles_iterator(key);
         let root = self.nodestore.root_mut();
         let Some(root_node) = std::mem::take(root) else {
             // The trie is empty. There is nothing to remove.
@@ -842,8 +862,16 @@ impl<S: ReadableStorage> Merkle<NodeStore<MutableProposal, S>> {
     /// Removes any key-value pairs with keys that have the given `prefix`.
     /// Returns the number of key-value pairs removed.
     pub fn remove_prefix(&mut self, prefix: &[u8]) -> Result<usize, FileIoError> {
-        let prefix = Path::from_nibbles_iterator(NibblesIterator::new(prefix));
+        self.remove_prefix_from_iter(NibblesIterator::new(prefix))
+    }
 
+    /// Removes any key-value pairs with keys that have the given `prefix` where `prefix` is a `NibblesIterator`
+    /// Returns the number of key-value pairs removed.
+    pub fn remove_prefix_from_iter(
+        &mut self,
+        prefix: NibblesIterator<'_>,
+    ) -> Result<usize, FileIoError> {
+        let prefix = Path::from_nibbles_iterator(prefix);
         let root = self.nodestore.root_mut();
         let Some(root_node) = std::mem::take(root) else {
             // The trie is empty. There is nothing to remove.