perf(tree): worker pooling for account proofs (#18901)

Co-authored-by: Brian Picciano <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Alexey Shekhirin <[email protected]>

Author: 4 people

crates/engine/primitives/src/config.rs

@@ -21,6 +21,14 @@ fn default_storage_worker_count() -> usize {
     }
 }
 
+/// Returns the default number of account worker threads.
+///
+/// Account workers coordinate storage proof collection and account trie traversal.
+/// They are set to the same count as storage workers for simplicity.
+fn default_account_worker_count() -> usize {
+    default_storage_worker_count()
+}
+
 /// The size of proof targets chunk to spawn in one multiproof calculation.
 pub const DEFAULT_MULTIPROOF_TASK_CHUNK_SIZE: usize = 10;
 
@@ -123,6 +131,8 @@ pub struct TreeConfig {
     allow_unwind_canonical_header: bool,
     /// Number of storage proof worker threads.
     storage_worker_count: usize,
+    /// Number of account proof worker threads.
+    account_worker_count: usize,
 }
 
 impl Default for TreeConfig {
@@ -150,6 +160,7 @@ impl Default for TreeConfig {
             prewarm_max_concurrency: DEFAULT_PREWARM_MAX_CONCURRENCY,
             allow_unwind_canonical_header: false,
             storage_worker_count: default_storage_worker_count(),
+            account_worker_count: default_account_worker_count(),
         }
     }
 }
@@ -180,6 +191,7 @@ impl TreeConfig {
         prewarm_max_concurrency: usize,
         allow_unwind_canonical_header: bool,
         storage_worker_count: usize,
+        account_worker_count: usize,
     ) -> Self {
         assert!(max_proof_task_concurrency > 0, "max_proof_task_concurrency must be at least 1");
         Self {
@@ -205,6 +217,7 @@ impl TreeConfig {
             prewarm_max_concurrency,
             allow_unwind_canonical_header,
             storage_worker_count,
+            account_worker_count,
         }
     }
 
@@ -482,4 +495,15 @@ impl TreeConfig {
         self.storage_worker_count = storage_worker_count;
         self
     }
+
+    /// Return the number of account proof worker threads.
+    pub const fn account_worker_count(&self) -> usize {
+        self.account_worker_count
+    }
+
+    /// Setter for the number of account proof worker threads.
+    pub const fn with_account_worker_count(mut self, account_worker_count: usize) -> Self {
+        self.account_worker_count = account_worker_count;
+        self
+    }
 }

crates/engine/tree/src/tree/payload_processor/mod.rs

@@ -192,8 +192,7 @@ where
     {
         let (to_sparse_trie, sparse_trie_rx) = channel();
         // spawn multiproof task, save the trie input
-        let (trie_input, state_root_config) =
-            MultiProofConfig::new_from_input(consistent_view, trie_input);
+        let (trie_input, state_root_config) = MultiProofConfig::from_input(trie_input);
         self.trie_input = Some(trie_input);
 
         // Create and spawn the storage proof task
@@ -202,14 +201,15 @@ where
             state_root_config.state_sorted.clone(),
             state_root_config.prefix_sets.clone(),
         );
-        let max_proof_task_concurrency = config.max_proof_task_concurrency() as usize;
         let storage_worker_count = config.storage_worker_count();
+        let account_worker_count = config.account_worker_count();
+        let max_proof_task_concurrency = config.max_proof_task_concurrency() as usize;
         let proof_task = match ProofTaskManager::new(
             self.executor.handle().clone(),
-            state_root_config.consistent_view.clone(),
+            consistent_view,
             task_ctx,
-            max_proof_task_concurrency,
             storage_worker_count,
+            account_worker_count,
         ) {
             Ok(task) => task,
             Err(error) => {