1. fix(perf): Run CPU-intensive state updates in parallel rayon threads

tower-batch/src/worker.rs

@@ -99,7 +99,7 @@ where
 {
     /// Creates a new batch worker.
     ///
-    /// See [`Service::new()`](crate::Service::new) for details.
+    /// See [`Batch::new()`](crate::Batch::new) for details.
     pub(crate) fn new(
         service: T,
         rx: mpsc::UnboundedReceiver<Message<Request, T::Future>>,
@@ -190,7 +190,7 @@ where
 
     /// Run loop for batch requests, which implements the batch policies.
     ///
-    /// See [`Service::new()`](crate::Service::new) for details.
+    /// See [`Batch::new()`](crate::Batch::new) for details.
     pub async fn run(mut self) {
         loop {
             // Wait on either a new message or the batch timer.

zebra-chain/Cargo.toml

@@ -66,6 +66,7 @@ serde-big-array = "0.4.1"
 # Processing
 futures = "0.3.21"
 itertools = "0.10.3"
+rayon = "1.5.3"
 
 # ZF deps
 ed25519-zebra = "3.0.0"

zebra-chain/src/block/arbitrary.rs

@@ -449,6 +449,9 @@ impl Block {
                         // would be awkward since the genesis block is handled separatedly there.
                         // This forces us to skip the genesis block here too in order to able to use
                         // this to test the finalized state.
+                        //
+                        // TODO: run note commitment tree updates in parallel rayon threads,
+                        //       using `NoteCommitmentTrees::update_trees_parallel()`
                         if generate_valid_commitments && *height != Height(0) {
                             for sapling_note_commitment in transaction.sapling_note_commitments() {
                                 sapling_tree.append(*sapling_note_commitment).unwrap();

zebra-chain/src/lib.rs

@@ -21,6 +21,7 @@ pub mod chain_tip;
 pub mod fmt;
 pub mod history_tree;
 pub mod orchard;
+pub mod parallel;
 pub mod parameters;
 pub mod primitives;
 pub mod sapling;

zebra-chain/src/orchard/tree.rs

@@ -11,7 +11,6 @@
 //! A root of a note commitment tree is associated with each treestate.
 
 #![allow(clippy::derive_hash_xor_eq)]
-#![allow(dead_code)]
 
 use std::{
     fmt,
@@ -34,6 +33,11 @@ use crate::serialization::{
     serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
 };
 
+/// The type that is used to update the note commitment tree.
+///
+/// Unfortunately, this is not the same as `orchard::NoteCommitment`.
+pub type NoteCommitmentUpdate = pallas::Base;
+
 pub(super) const MERKLE_DEPTH: usize = 32;
 
 /// MerkleCRH^Orchard Hash Function
@@ -248,8 +252,8 @@ impl<'de> serde::Deserialize<'de> for Node {
     }
 }
 
-#[allow(dead_code, missing_docs)]
-#[derive(Error, Debug, Clone, PartialEq, Eq)]
+#[derive(Error, Copy, Clone, Debug, Eq, PartialEq, Hash)]
+#[allow(missing_docs)]
 pub enum NoteCommitmentTreeError {
     #[error("The note commitment tree is full")]
     FullTree,
@@ -302,7 +306,7 @@ impl NoteCommitmentTree {
     ///
     /// Returns an error if the tree is full.
     #[allow(clippy::unwrap_in_result)]
-    pub fn append(&mut self, cm_x: pallas::Base) -> Result<(), NoteCommitmentTreeError> {
+    pub fn append(&mut self, cm_x: NoteCommitmentUpdate) -> Result<(), NoteCommitmentTreeError> {
         if self.inner.append(&cm_x.into()) {
             // Invalidate cached root
             let cached_root = self

zebra-chain/src/parallel.rs

@@ -0,0 +1,3 @@
+//! Parallel chain update methods.
+
+pub mod tree;

zebra-chain/src/parallel/tree.rs

@@ -0,0 +1,195 @@
+//! Parallel note commitment tree update methods.
+
+use std::{collections::BTreeMap, sync::Arc};
+
+use thiserror::Error;
+
+use crate::{
+    block::{Block, Height},
+    orchard, sapling, sprout,
+};
+
+/// An argument wrapper struct for note commitment trees.
+#[derive(Clone, Debug)]
+pub struct NoteCommitmentTrees {
+    /// The sprout note commitment tree.
+    pub sprout: Arc<sprout::tree::NoteCommitmentTree>,
+
+    /// The sapling note commitment tree.
+    pub sapling: Arc<sapling::tree::NoteCommitmentTree>,
+
+    /// The orchard note commitment tree.
+    pub orchard: Arc<orchard::tree::NoteCommitmentTree>,
+}
+
+/// Note commitment tree errors.
+#[derive(Error, Copy, Clone, Debug, Eq, PartialEq, Hash)]
+pub enum NoteCommitmentTreeError {
+    /// A sprout tree error
+    #[error("sprout error: {0}")]
+    Sprout(#[from] sprout::tree::NoteCommitmentTreeError),
+
+    /// A sapling tree error
+    #[error("sapling error: {0}")]
+    Sapling(#[from] sapling::tree::NoteCommitmentTreeError),
+
+    /// A orchard tree error
+    #[error("orchard error: {0}")]
+    Orchard(#[from] orchard::tree::NoteCommitmentTreeError),
+}
+
+impl NoteCommitmentTrees {
+    /// Updates the note commitment trees using the transactions in `block`,
+    /// then re-calculates the cached tree roots, using parallel `rayon` threads.
+    ///
+    /// If any of the tree updates cause an error,
+    /// it will be returned at the end of the parallel batches.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn update_trees_parallel(
+        &mut self,
+        block: &Arc<Block>,
+    ) -> Result<(), NoteCommitmentTreeError> {
+        self.update_trees_parallel_list(
+            [(
+                block
+                    .coinbase_height()
+                    .expect("height was already validated"),
+                block.clone(),
+            )]
+            .into_iter()
+            .collect(),
+        )
+    }
+
+    /// Updates the note commitment trees using the transactions in `block`,
+    /// then re-calculates the cached tree roots, using parallel `rayon` threads.
+    ///
+    /// If any of the tree updates cause an error,
+    /// it will be returned at the end of the parallel batches.
+    pub fn update_trees_parallel_list(
+        &mut self,
+        block_list: BTreeMap<Height, Arc<Block>>,
+    ) -> Result<(), NoteCommitmentTreeError> {
+        // Prepare arguments for parallel threads
+        let NoteCommitmentTrees {
+            sprout,
+            sapling,
+            orchard,
+        } = self.clone();
+
+        let sprout_note_commitments: Vec<_> = block_list
+            .values()
+            .flat_map(|block| block.transactions.iter())
+            .flat_map(|tx| tx.sprout_note_commitments())
+            .cloned()
+            .collect();
+        let sapling_note_commitments: Vec<_> = block_list
+            .values()
+            .flat_map(|block| block.transactions.iter())
+            .flat_map(|tx| tx.sapling_note_commitments())
+            .cloned()
+            .collect();
+        let orchard_note_commitments: Vec<_> = block_list
+            .values()
+            .flat_map(|block| block.transactions.iter())
+            .flat_map(|tx| tx.orchard_note_commitments())
+            .cloned()
+            .collect();
+
+        let mut sprout_result = None;
+        let mut sapling_result = None;
+        let mut orchard_result = None;
+
+        rayon::in_place_scope_fifo(|scope| {
+            if !sprout_note_commitments.is_empty() {
+                scope.spawn_fifo(|_scope| {
+                    sprout_result = Some(Self::update_sprout_note_commitment_tree(
+                        sprout,
+                        sprout_note_commitments,
+                    ));
+                });
+            }
+
+            if !sapling_note_commitments.is_empty() {
+                scope.spawn_fifo(|_scope| {
+                    sapling_result = Some(Self::update_sapling_note_commitment_tree(
+                        sapling,
+                        sapling_note_commitments,
+                    ));
+                });
+            }
+
+            if !orchard_note_commitments.is_empty() {
+                scope.spawn_fifo(|_scope| {
+                    orchard_result = Some(Self::update_orchard_note_commitment_tree(
+                        orchard,
+                        orchard_note_commitments,
+                    ));
+                });
+            }
+        });
+
+        if let Some(sprout_result) = sprout_result {
+            self.sprout = sprout_result?;
+        }
+        if let Some(sapling_result) = sapling_result {
+            self.sapling = sapling_result?;
+        }
+        if let Some(orchard_result) = orchard_result {
+            self.orchard = orchard_result?;
+        }
+
+        Ok(())
+    }
+
+    /// Update the sprout note commitment tree.
+    fn update_sprout_note_commitment_tree(
+        mut sprout: Arc<sprout::tree::NoteCommitmentTree>,
+        sprout_note_commitments: Vec<sprout::NoteCommitment>,
+    ) -> Result<Arc<sprout::tree::NoteCommitmentTree>, NoteCommitmentTreeError> {
+        let sprout_nct = Arc::make_mut(&mut sprout);
+
+        for sprout_note_commitment in sprout_note_commitments {
+            sprout_nct.append(sprout_note_commitment)?;
+        }
+
+        // Re-calculate and cache the tree root.
+        let _ = sprout_nct.root();
+
+        Ok(sprout)
+    }
+
+    /// Update the sapling note commitment tree.
+    fn update_sapling_note_commitment_tree(
+        mut sapling: Arc<sapling::tree::NoteCommitmentTree>,
+        sapling_note_commitments: Vec<sapling::tree::NoteCommitmentUpdate>,
+    ) -> Result<Arc<sapling::tree::NoteCommitmentTree>, NoteCommitmentTreeError> {
+        let sapling_nct = Arc::make_mut(&mut sapling);
+
+        for sapling_note_commitment in sapling_note_commitments {
+            sapling_nct.append(sapling_note_commitment)?;
+        }
+
+        // Re-calculate and cache the tree root.
+        let _ = sapling_nct.root();
+
+        Ok(sapling)
+    }
+
+    /// Update the orchard note commitment tree.
+    fn update_orchard_note_commitment_tree(
+        mut orchard: Arc<orchard::tree::NoteCommitmentTree>,
+        orchard_note_commitments: Vec<orchard::tree::NoteCommitmentUpdate>,
+    ) -> Result<Arc<orchard::tree::NoteCommitmentTree>, NoteCommitmentTreeError> {
+        let orchard_nct = Arc::make_mut(&mut orchard);
+
+        for orchard_note_commitment in orchard_note_commitments {
+            orchard_nct.append(orchard_note_commitment)?;
+        }
+
+        // Re-calculate and cache the tree root.
+        let _ = orchard_nct.root();
+
+        Ok(orchard)
+    }
+}

zebra-chain/src/sapling/tree.rs

@@ -10,8 +10,6 @@
 //!
 //! A root of a note commitment tree is associated with each treestate.
 
-#![allow(dead_code)]
-
 use std::{
     fmt,
     hash::{Hash, Hasher},
@@ -38,6 +36,11 @@ use crate::serialization::{
     serde_helpers, ReadZcashExt, SerializationError, ZcashDeserialize, ZcashSerialize,
 };
 
+/// The type that is used to update the note commitment tree.
+///
+/// Unfortunately, this is not the same as `sapling::NoteCommitment`.
+pub type NoteCommitmentUpdate = jubjub::Fq;
+
 pub(super) const MERKLE_DEPTH: usize = 32;
 
 /// MerkleCRH^Sapling Hash Function
@@ -252,8 +255,8 @@ impl<'de> serde::Deserialize<'de> for Node {
     }
 }
 
-#[allow(dead_code, missing_docs)]
-#[derive(Error, Debug, Clone, PartialEq, Eq)]
+#[derive(Error, Copy, Clone, Debug, Eq, PartialEq, Hash)]
+#[allow(missing_docs)]
 pub enum NoteCommitmentTreeError {
     #[error("The note commitment tree is full")]
     FullTree,
@@ -307,7 +310,7 @@ impl NoteCommitmentTree {
     ///
     /// Returns an error if the tree is full.
     #[allow(clippy::unwrap_in_result)]
-    pub fn append(&mut self, cm_u: jubjub::Fq) -> Result<(), NoteCommitmentTreeError> {
+    pub fn append(&mut self, cm_u: NoteCommitmentUpdate) -> Result<(), NoteCommitmentTreeError> {
         if self.inner.append(&cm_u.into()) {
             // Invalidate cached root
             let cached_root = self

zebra-chain/src/sprout.rs

@@ -14,6 +14,7 @@ pub mod keys;
 pub mod note;
 pub mod tree;
 
+pub use commitment::NoteCommitment;
 pub use joinsplit::JoinSplit;
 pub use joinsplit::RandomSeed;
 pub use note::{EncryptedNote, Note, Nullifier};

zebra-chain/src/sprout/tree.rs

@@ -180,8 +180,8 @@ impl<'de> serde::Deserialize<'de> for Node {
     }
 }
 
-#[allow(dead_code, missing_docs)]
-#[derive(Error, Debug, Clone, PartialEq, Eq)]
+#[derive(Error, Copy, Clone, Debug, Eq, PartialEq, Hash)]
+#[allow(missing_docs)]
 pub enum NoteCommitmentTreeError {
     #[error("the note commitment tree is full")]
     FullTree,

zebra-state/Cargo.toml

@@ -19,22 +19,24 @@ itertools = "0.10.3"
 lazy_static = "1.4.0"
 metrics = "0.18.1"
 mset = "0.1.0"
-proptest = { version = "0.10.1", optional = true }
-proptest-derive = { version = "0.3.0", optional = true }
 regex = "1.5.6"
 rlimit = "0.8.3"
 rocksdb = { version = "0.18.0", default_features = false, features = ["lz4"] }
 serde = { version = "1.0.137", features = ["serde_derive"] }
 tempfile = "3.3.0"
 thiserror = "1.0.31"
 
+rayon = "1.5.3"
 tokio = { version = "1.20.0", features = ["sync", "tracing"] }
 tower = { version = "0.4.13", features = ["buffer", "util"] }
 tracing = "0.1.31"
 
 zebra-chain = { path = "../zebra-chain" }
 zebra-test = { path = "../zebra-test/", optional = true }
 
+proptest = { version = "0.10.1", optional = true }
+proptest-derive = { version = "0.3.0", optional = true }
+
 [dev-dependencies]
 color-eyre = "0.6.1"
 once_cell = "1.12.0"