From 465af6e0f37b1f8c7c86a1b377da6fa016184fff Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Wed, 19 Feb 2025 15:37:19 +0000
Subject: [PATCH] feat(tree): increase state root task thread pool size
 (#14455)

---
 crates/engine/tree/src/tree/mod.rs  |  2 +-
 crates/engine/tree/src/tree/root.rs | 23 +++++++++++++++--------
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index ca5b262d9..f78f3cd2e 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -671,7 +671,7 @@ where
     ) -> Self {
         let (incoming_tx, incoming) = std::sync::mpsc::channel();
 
-        let num_threads = root::thread_pool_size();
+        let num_threads = root::rayon_thread_pool_size();
 
         let thread_pool = Arc::new(
             rayon::ThreadPoolBuilder::new()
diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
index 463c41927..206992f55 100644
--- a/crates/engine/tree/src/tree/root.rs
+++ b/crates/engine/tree/src/tree/root.rs
@@ -42,14 +42,21 @@ use tracing::{debug, error, trace, trace_span};
 /// The level below which the sparse trie hashes are calculated in [`update_sparse_trie`].
 const SPARSE_TRIE_INCREMENTAL_LEVEL: usize = 2;
 
-/// Determines the size of the thread pool to be used in [`StateRootTask`].
-/// It should be at least three, one for multiproof calculations  plus two to be
-/// used internally in [`StateRootTask`].
+/// Determines the size of the rayon thread pool to be used in [`StateRootTask`].
+///
+/// The value is determined as `max(NUM_THREADS - 2, 3)`:
+/// - It should leave at least 2 threads to the rest of the system to be used in:
+///     - Engine
+///     - State Root Task spawned in [`StateRootTask::spawn`]
+/// - It should heave at least 3 threads to be used in:
+///     - Sparse Trie spawned in [`run_sparse_trie`]
+///     - Multiproof computation spawned in [`MultiproofManager::spawn_multiproof`]
+///     - Storage root computation spawned in [`ParallelProof::multiproof`]
 ///
 /// NOTE: this value can be greater than the available cores in the host, it
 /// represents the maximum number of threads that can be handled by the pool.
-pub(crate) fn thread_pool_size() -> usize {
-    std::thread::available_parallelism().map_or(3, |num| (num.get() / 2).max(3))
+pub(crate) fn rayon_thread_pool_size() -> usize {
+    std::thread::available_parallelism().map_or(3, |num| (num.get().saturating_sub(2).max(3)))
 }
 
 /// Determines if the host has enough parallelism to run the state root task.
@@ -535,7 +542,7 @@ where
             fetched_proof_targets: Default::default(),
             proof_sequencer: ProofSequencer::new(),
             thread_pool: thread_pool.clone(),
-            multiproof_manager: MultiproofManager::new(thread_pool, thread_pool_size()),
+            multiproof_manager: MultiproofManager::new(thread_pool, rayon_thread_pool_size()),
             metrics: StateRootTaskMetrics::default(),
         }
     }
@@ -1286,7 +1293,7 @@ mod tests {
             + Clone
             + 'static,
     {
-        let num_threads = thread_pool_size();
+        let num_threads = rayon_thread_pool_size();
 
         let thread_pool = rayon::ThreadPoolBuilder::new()
             .num_threads(num_threads)
@@ -1361,7 +1368,7 @@ mod tests {
             prefix_sets: Arc::new(input.prefix_sets),
         };
 
-        let num_threads = thread_pool_size();
+        let num_threads = rayon_thread_pool_size();
 
         let state_root_task_pool = rayon::ThreadPoolBuilder::new()
             .num_threads(num_threads)