[smt] Parallelize smt batch update with rayon::join

Closes: diem#8215
sitalkedia · Apr 16, 2021 · 1aca58d · 1aca58d
1 parent 55372e7
commit 1aca58d
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 31 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/storage/scratchpad/Cargo.toml b/storage/scratchpad/Cargo.toml
@@ -12,6 +12,7 @@ edition = "2018"
 [dependencies]
 arc-swap = "1.2.0"
 itertools = "0.10.0"
+rayon = "1.5.0"
 
 diem-crypto = { path = "../../crypto/crypto" }
 diem-infallible = { path = "../../common/infallible" }

diff --git a/storage/scratchpad/src/sparse_merkle/mod.rs b/storage/scratchpad/src/sparse_merkle/mod.rs
@@ -157,7 +157,7 @@ pub struct SparseMerkleTree<V> {
 
 impl<V> SparseMerkleTree<V>
 where
-    V: Clone + CryptoHash,
+    V: Clone + CryptoHash + Send + Sync,
 {
     /// Constructs a Sparse Merkle Tree with a root hash. This is often used when we restart and
     /// the scratch pad and the storage have identical state, so we use a single root hash to
@@ -716,8 +716,10 @@ where
                 } else if idx == n {
                     bits_on_path.push(false);
                 } else {
-                    let left_child = Self::batch_construct_subtree_from_empty(&kvs[..idx], depth);
-                    let right_child = Self::batch_construct_subtree_from_empty(&kvs[idx..], depth);
+                    let (left_child, right_child) = rayon::join(
+                        || Self::batch_construct_subtree_from_empty(&kvs[..idx], depth),
+                        || Self::batch_construct_subtree_from_empty(&kvs[idx..], depth),
+                    );
                     let fork_node = SubTree::new_internal(left_child, right_child);
                     let siblings_num = bits_on_path.len();
                     return Self::construct_subtree(
@@ -788,13 +790,19 @@ where
                     });
                     bits_on_path.push(idx == 0);
                 } else {
-                    let left_child = Self::batch_create_subtree_from_key_proof_pairs(
-                        &key_proof_pairs[..idx],
-                        depth + 1,
-                    );
-                    let right_child = Self::batch_create_subtree_from_key_proof_pairs(
-                        &key_proof_pairs[idx..],
-                        depth + 1,
+                    let (left_child, right_child) = rayon::join(
+                        || {
+                            Self::batch_create_subtree_from_key_proof_pairs(
+                                &key_proof_pairs[..idx],
+                                depth + 1,
+                            )
+                        },
+                        || {
+                            Self::batch_create_subtree_from_key_proof_pairs(
+                                &key_proof_pairs[idx..],
+                                depth + 1,
+                            )
+                        },
                     );
                     let fork_node = SubTree::new_internal(left_child, right_child);
                     return Self::construct_subtree(
@@ -862,24 +870,28 @@ where
                     }
                 } else {
                     let (left_child, right_child) = if existing_leaf_key.bit(depth) {
-                        (
-                            Self::batch_construct_subtree_from_empty(&kvs[..idx], depth + 1),
-                            Self::batch_construct_subtree_with_existing_leaf(
-                                &kvs[idx..],
-                                existing_leaf,
-                                existing_leaf_key,
-                                depth + 1,
-                            ),
+                        rayon::join(
+                            || Self::batch_construct_subtree_from_empty(&kvs[..idx], depth + 1),
+                            || {
+                                Self::batch_construct_subtree_with_existing_leaf(
+                                    &kvs[idx..],
+                                    existing_leaf,
+                                    existing_leaf_key,
+                                    depth + 1,
+                                )
+                            },
                         )
                     } else {
-                        (
-                            Self::batch_construct_subtree_with_existing_leaf(
-                                &kvs[..idx],
-                                existing_leaf,
-                                existing_leaf_key,
-                                depth + 1,
-                            ),
-                            Self::batch_construct_subtree_from_empty(&kvs[idx..], depth + 1),
+                        rayon::join(
+                            || {
+                                Self::batch_construct_subtree_with_existing_leaf(
+                                    &kvs[..idx],
+                                    existing_leaf,
+                                    existing_leaf_key,
+                                    depth + 1,
+                                )
+                            },
+                            || Self::batch_construct_subtree_from_empty(&kvs[idx..], depth + 1),
                         )
                     };
                     let fork_node = SubTree::new_internal(left_child, right_child);
@@ -908,15 +920,15 @@ where
 
 impl<V> Default for SparseMerkleTree<V>
 where
-    V: Clone + CryptoHash,
+    V: Clone + CryptoHash + Send + Sync,
 {
     fn default() -> Self {
         SparseMerkleTree::new(*SPARSE_MERKLE_PLACEHOLDER_HASH)
     }
 }
 
 /// A type that implements `ProofRead` can provide proof for keys in persistent storage.
-pub trait ProofRead<V> {
+pub trait ProofRead<V>: Sync {
     /// Gets verified proof for this key in persistent storage.
     fn get_proof(&self, key: HashValue) -> Option<&SparseMerkleProof<V>>;
 }

diff --git a/storage/scratchpad/src/sparse_merkle/sparse_merkle_test.rs b/storage/scratchpad/src/sparse_merkle/sparse_merkle_test.rs
@@ -34,19 +34,19 @@ fn hash_leaf(key: HashValue, value_hash: HashValue) -> HashValue {
 
 struct ProofReader<V>(HashMap<HashValue, SparseMerkleProof<V>>);
 
-impl<V> ProofReader<V> {
+impl<V: Sync> ProofReader<V> {
     fn new(key_with_proof: Vec<(HashValue, SparseMerkleProof<V>)>) -> Self {
         ProofReader(key_with_proof.into_iter().collect())
     }
 }
 
-impl<V> Default for ProofReader<V> {
+impl<V: Sync> Default for ProofReader<V> {
     fn default() -> Self {
         Self(HashMap::new())
     }
 }
 
-impl<V> ProofRead<V> for ProofReader<V> {
+impl<V: Sync> ProofRead<V> for ProofReader<V> {
     fn get_proof(&self, key: HashValue) -> Option<&SparseMerkleProof<V>> {
         self.0.get(&key)
     }