From c838383fd53610da01bfee4f7555f219f84853d0 Mon Sep 17 00:00:00 2001
From: Arya <aryasolhi@gmail.com>
Date: Thu, 1 Dec 2022 06:35:12 -0500
Subject: [PATCH] change(state): Check block and transaction Sprout anchors in
 parallel (#5742)

* parallelize anchors checks for blocks

* parallelize anchors checks for unmined_tx

* reverts par_iter in block_sapling_orchard_anchors_refer_to_final_treestates

* moves fetch_sprout_final_treestates out of rayon thread
---
 zebra-state/src/service/check/anchors.rs | 102 +++++++++++++++--------
 1 file changed, 66 insertions(+), 36 deletions(-)

diff --git a/zebra-state/src/service/check/anchors.rs b/zebra-state/src/service/check/anchors.rs
index aebc47a460f..a7b2de0f241 100644
--- a/zebra-state/src/service/check/anchors.rs
+++ b/zebra-state/src/service/check/anchors.rs
@@ -3,6 +3,8 @@
 
 use std::{collections::HashMap, sync::Arc};
 
+use rayon::prelude::*;
+
 use zebra_chain::{
     block::{Block, Height},
     sprout,
@@ -406,21 +408,35 @@ pub(crate) fn block_sprout_anchors_refer_to_treestates(
         "received sprout final treestate anchors",
     );
 
-    block
-        .transactions
-        .iter()
-        .enumerate()
-        .try_for_each(|(tx_index_in_block, transaction)| {
-            sprout_anchors_refer_to_treestates(
-                &sprout_final_treestates,
-                transaction,
-                transaction_hashes[tx_index_in_block],
-                Some(tx_index_in_block),
-                Some(height),
-            )?;
-
-            Ok(())
-        })
+    let check_tx_sprout_anchors = |(tx_index_in_block, transaction)| {
+        sprout_anchors_refer_to_treestates(
+            &sprout_final_treestates,
+            transaction,
+            transaction_hashes[tx_index_in_block],
+            Some(tx_index_in_block),
+            Some(height),
+        )?;
+
+        Ok(())
+    };
+
+    // The overhead for a parallel iterator is unwarranted if sprout_final_treestates is empty
+    // because it will either return an error for the first transaction or only check that `joinsplit_data`
+    // is `None` for each transaction.
+    if sprout_final_treestates.is_empty() {
+        // The block has no valid sprout anchors
+        block
+            .transactions
+            .iter()
+            .enumerate()
+            .try_for_each(check_tx_sprout_anchors)
+    } else {
+        block
+            .transactions
+            .par_iter()
+            .enumerate()
+            .try_for_each(check_tx_sprout_anchors)
+    }
 }
 
 /// Accepts a [`ZebraDb`], an optional [`Option<Arc<Chain>>`](Chain), and an [`UnminedTx`].
@@ -444,30 +460,44 @@ pub(crate) fn tx_anchors_refer_to_final_treestates(
         None,
     )?;
 
-    let mut sprout_final_treestates = HashMap::new();
+    // If there are no sprout transactions in the block, avoid running a rayon scope
+    if unmined_tx.transaction.has_sprout_joinsplit_data() {
+        let mut sprout_final_treestates = HashMap::new();
 
-    fetch_sprout_final_treestates(
-        &mut sprout_final_treestates,
-        finalized_state,
-        parent_chain,
-        &unmined_tx.transaction,
-        None,
-        None,
-    );
+        fetch_sprout_final_treestates(
+            &mut sprout_final_treestates,
+            finalized_state,
+            parent_chain,
+            &unmined_tx.transaction,
+            None,
+            None,
+        );
 
-    tracing::trace!(
-        sprout_final_treestate_count = ?sprout_final_treestates.len(),
-        ?sprout_final_treestates,
-        "received sprout final treestate anchors",
-    );
+        let mut sprout_anchors_result = None;
+        rayon::in_place_scope_fifo(|s| {
+            // This check is expensive, because it updates a note commitment tree for each sprout JoinSplit.
+            // Since we could be processing attacker-controlled mempool transactions, we need to run each one
+            // in its own thread, separately from tokio's blocking I/O threads. And if we are under heavy load,
+            // we want verification to finish in order, so that later transactions can't delay earlier ones.
+            s.spawn_fifo(|_s| {
+                tracing::trace!(
+                    sprout_final_treestate_count = ?sprout_final_treestates.len(),
+                    ?sprout_final_treestates,
+                    "received sprout final treestate anchors",
+                );
 
-    sprout_anchors_refer_to_treestates(
-        &sprout_final_treestates,
-        &unmined_tx.transaction,
-        unmined_tx.id.mined_id(),
-        None,
-        None,
-    )?;
+                sprout_anchors_result = Some(sprout_anchors_refer_to_treestates(
+                    &sprout_final_treestates,
+                    &unmined_tx.transaction,
+                    unmined_tx.id.mined_id(),
+                    None,
+                    None,
+                ));
+            });
+        });
+
+        sprout_anchors_result.expect("scope has finished")?;
+    }
 
     Ok(())
 }