feat(engine): allow to override has_enough_parallelism (#14608)

crates/engine/tree/src/tree/config.rs

@@ -1,5 +1,6 @@
 //! Engine tree configuration.
 
+use crate::tree::root::has_enough_parallelism;
 use alloy_eips::merge::EPOCH_SLOTS;
 
 /// The largest gap for which the tree will be used for sync. See docs for `pipeline_run_threshold`
@@ -53,6 +54,8 @@ pub struct TreeConfig {
     use_caching_and_prewarming: bool,
     /// Cross-block cache size in bytes.
     cross_block_cache_size: u64,
+    /// Wether the host has enough parallelism to run state root task.
+    has_enough_parallelism: bool,
 }
 
 impl Default for TreeConfig {
@@ -67,6 +70,7 @@ impl Default for TreeConfig {
             always_compare_trie_updates: false,
             use_caching_and_prewarming: false,
             cross_block_cache_size: DEFAULT_CROSS_BLOCK_CACHE_SIZE,
+            has_enough_parallelism: has_enough_parallelism(),
         }
     }
 }
@@ -84,6 +88,7 @@ impl TreeConfig {
         always_compare_trie_updates: bool,
         use_caching_and_prewarming: bool,
         cross_block_cache_size: u64,
+        has_enough_parallelism: bool,
     ) -> Self {
         Self {
             persistence_threshold,
@@ -95,6 +100,7 @@ impl TreeConfig {
             always_compare_trie_updates,
             use_caching_and_prewarming,
             cross_block_cache_size,
+            has_enough_parallelism,
         }
     }
 
@@ -211,4 +217,15 @@ impl TreeConfig {
         self.cross_block_cache_size = cross_block_cache_size;
         self
     }
+
+    /// Setter for has enough parallelism.
+    pub const fn with_has_enough_parallelism(mut self, has_enough_parallelism: bool) -> Self {
+        self.has_enough_parallelism = has_enough_parallelism;
+        self
+    }
+
+    /// Wether or not to use state root task
+    pub(crate) fn use_state_root_task(&self) -> bool {
+        self.has_enough_parallelism && !self.legacy_state_root
+    }
 }

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

@@ -2408,11 +2408,8 @@ where
         // Atomic bool for letting the prewarm tasks know when to stop
         let cancel_execution = ManualCancel::default();
 
-        let use_legacy_state_root =
-            self.config.legacy_state_root() || !root::has_enough_parallelism();
-
         let (state_root_handle, state_root_task_config, state_root_sender, state_hook) =
-            if is_descendant_of_persisting_blocks && !use_legacy_state_root {
+            if is_descendant_of_persisting_blocks && self.config.use_state_root_task() {
                 let consistent_view = ConsistentDbView::new_with_latest_tip(self.provider.clone())?;
 
                 // Compute trie input
@@ -2537,7 +2534,21 @@ where
         // a different database transaction per thread and it might end up with a
         // different view of the database.
         let (state_root, trie_output, root_elapsed) = if is_descendant_of_persisting_blocks {
-            if use_legacy_state_root {
+            if self.config.use_state_root_task() {
+                let state_root_handle = state_root_handle
+                    .expect("state root handle must exist if legacy_state_root is false");
+                let state_root_config = state_root_task_config.expect("task config is present");
+
+                // Handle state root result from task using handle
+                self.handle_state_root_result(
+                    state_root_handle,
+                    state_root_config,
+                    block.sealed_block(),
+                    &hashed_state,
+                    &state_provider,
+                    root_time,
+                )?
+            } else {
                 match self.compute_state_root_parallel(block.header().parent_hash(), &hashed_state)
                 {
                     Ok(result) => {
@@ -2557,20 +2568,6 @@ where
                     }
                     Err(error) => return Err(InsertBlockErrorKind::Other(Box::new(error))),
                 }
-            } else {
-                let state_root_handle = state_root_handle
-                    .expect("state root handle must exist if legacy_state_root is false");
-                let state_root_config = state_root_task_config.expect("task config is present");
-
-                // Handle state root result from task using handle
-                self.handle_state_root_result(
-                    state_root_handle,
-                    state_root_config,
-                    block.sealed_block(),
-                    &hashed_state,
-                    &state_provider,
-                    root_time,
-                )?
             }
         } else {
             debug!(target: "engine::tree", block=?block_num_hash, ?is_descendant_of_persisting_blocks, "Failed to compute state root in parallel");
@@ -3361,7 +3358,10 @@ mod tests {
                 PersistenceState::default(),
                 payload_builder,
                 // TODO: fix tests for state root task https://github.com/paradigmxyz/reth/issues/14376
-                TreeConfig::default().with_legacy_state_root(true),
+                // always assume enough parallelism for tests
+                TreeConfig::default()
+                    .with_legacy_state_root(true)
+                    .with_has_enough_parallelism(true),
                 EngineApiKind::Ethereum,
                 evm_config,
             );