Merge pull request #883 from zcash/put_blocks-parallel-tree-building

zcash_client_sqlite: Build subtrees from new commitments in a threadpool

Cargo.toml

@@ -19,6 +19,6 @@ panic = 'abort'
 codegen-units = 1
 
 [patch.crates-io]
-incrementalmerkletree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "bae25ad89c0c192bee625252d2d419bd56638e48" }
-shardtree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "bae25ad89c0c192bee625252d2d419bd56638e48" }
+incrementalmerkletree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "2a667f500958d517c6b2ea608477140afd907fd8" }
+shardtree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "2a667f500958d517c6b2ea608477140afd907fd8" }
 orchard = { git = "https://github.com/zcash/orchard.git", rev = "6ef89d5f154de2cf7b7dd87edb8d8c49158beebb" }

zcash_client_sqlite/CHANGELOG.md

@@ -9,6 +9,9 @@ and this library adheres to Rust's notion of
 ### Added
 - `zcash_client_sqlite::serialization` Serialization formats for data stored
   as SQLite BLOBs in the wallet database.
+- A new default-enabled feature flag `multicore`. This allows users to disable
+  multicore support by setting `default_features = false` on their
+  `zcash_primitives`, `zcash_proofs`, and `zcash_client_sqlite` dependencies.
 
 ### Changed
 - MSRV is now 1.65.0.

zcash_client_sqlite/Cargo.toml

@@ -48,6 +48,7 @@ uuid = "1.1"
 
 # Dependencies used internally:
 # (Breaking upgrades to these are usually backwards-compatible, but check MSRVs.)
+maybe-rayon = {version = "0.1.0", default-features = false}
 
 [dev-dependencies]
 assert_matches = "1.5"
@@ -63,6 +64,8 @@ zcash_primitives = { version = "0.12", path = "../zcash_primitives", features =
 zcash_address = { version = "0.3", path = "../components/zcash_address", features = ["test-dependencies"] }
 
 [features]
+default = ["multicore"]
+multicore = ["maybe-rayon/threads", "zcash_primitives/multicore"]
 mainnet = []
 test-dependencies = [
     "incrementalmerkletree/test-dependencies",

zcash_client_sqlite/src/lib.rs

@@ -33,6 +33,10 @@
 #![deny(rustdoc::broken_intra_doc_links)]
 
 use either::Either;
+use maybe_rayon::{
+    prelude::{IndexedParallelIterator, ParallelIterator},
+    slice::ParallelSliceMut,
+};
 use rusqlite::{self, Connection};
 use secrecy::{ExposeSecret, SecretVec};
 use std::{borrow::Borrow, collections::HashMap, convert::AsRef, fmt, io, ops::Range, path::Path};
@@ -474,7 +478,7 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
                 }));
 
                 last_scanned_height = Some(block.height());
-                sapling_commitments.extend(block.into_sapling_commitments().into_iter());
+                sapling_commitments.extend(block.into_sapling_commitments().into_iter().map(Some));
             }
 
             // Prune the nullifier map of entries we no longer need.
@@ -490,10 +494,31 @@ impl<P: consensus::Parameters> WalletWrite for WalletDb<rusqlite::Connection, P>
             if let Some(((start_height, start_position), last_scanned_height)) =
                 start_positions.zip(last_scanned_height)
             {
+                // Create subtrees from the note commitments in parallel.
+                const CHUNK_SIZE: usize = 1024;
+                let subtrees = sapling_commitments
+                    .par_chunks_mut(CHUNK_SIZE)
+                    .enumerate()
+                    .filter_map(|(i, chunk)| {
+                        let start = start_position + (i * CHUNK_SIZE) as u64;
+                        let end = start + chunk.len() as u64;
+
+                        shardtree::LocatedTree::from_iter(
+                            start..end,
+                            SAPLING_SHARD_HEIGHT.into(),
+                            chunk.iter_mut().map(|n| n.take().expect("always Some")),
+                        )
+                    })
+                    .map(|res| (res.subtree, res.checkpoints))
+                    .collect::<Vec<_>>();
+
                 // Update the Sapling note commitment tree with all newly read note commitments
-                let mut sapling_commitments = sapling_commitments.into_iter();
+                let mut subtrees = subtrees.into_iter();
                 wdb.with_sapling_tree_mut::<_, _, SqliteClientError>(move |sapling_tree| {
-                    sapling_tree.batch_insert(start_position, &mut sapling_commitments)?;
+                    for (tree, checkpoints) in &mut subtrees {
+                        sapling_tree.insert_tree(tree, checkpoints)?;
+                    }
+
                     Ok(())
                 })?;