refactor(stage): perform trie computation as separate task (#373)

Spawn a separate task on the cpu-bound blocking task for performing the actual state trie computation to avoid blocking the async executor.

Author: kariy

crates/node/src/full/mod.rs

@@ -127,7 +127,7 @@ impl Node {
         let block_downloader = BatchBlockDownloader::new_gateway(gateway_client.clone(), 20);
         pipeline.add_stage(Blocks::new(storage_provider.clone(), block_downloader));
         pipeline.add_stage(Classes::new(storage_provider.clone(), gateway_client.clone(), 20));
-        pipeline.add_stage(StateTrie::new(storage_provider.clone()));
+        pipeline.add_stage(StateTrie::new(storage_provider.clone(), task_spawner.clone()));
 
         // -- build chain tip watcher using gateway client
 

crates/sync/stage/src/trie.rs

@@ -5,6 +5,7 @@ use katana_provider::api::block::HeaderProvider;
 use katana_provider::api::state_update::StateUpdateProvider;
 use katana_provider::api::trie::TrieWriter;
 use katana_provider::{MutableProvider, ProviderFactory};
+use katana_tasks::TaskSpawner;
 use starknet::macros::short_string;
 use starknet_types_core::hash::{Poseidon, StarkHash};
 use tracing::{debug, debug_span, error};
@@ -22,19 +23,20 @@ use crate::{Stage, StageExecutionInput, StageExecutionOutput, StageResult};
 #[derive(Debug)]
 pub struct StateTrie<P> {
     storage_provider: P,
+    task_spawner: TaskSpawner,
 }
 
 impl<P> StateTrie<P> {
     /// Create a new [`StateTrie`] stage.
-    pub fn new(storage_provider: P) -> Self {
-        Self { storage_provider }
+    pub fn new(storage_provider: P, task_spawner: TaskSpawner) -> Self {
+        Self { storage_provider, task_spawner }
     }
 }
 
 impl<P> Stage for StateTrie<P>
 where
     P: ProviderFactory,
-    <P as ProviderFactory>::ProviderMut: StateUpdateProvider + HeaderProvider + TrieWriter,
+    <P as ProviderFactory>::ProviderMut: StateUpdateProvider + HeaderProvider + TrieWriter + Clone,
 {
     fn id(&self) -> &'static str {
         "StateTrie"
@@ -58,21 +60,37 @@ where
                     .state_update(block_number.into())?
                     .ok_or(Error::MissingStateUpdate(block_number))?;
 
-                let computed_contract_trie_root =
-                    provider_mut.trie_insert_contract_updates(block_number, &state_update)?;
-
-                debug!(
-                    contract_trie_root = format!("{computed_contract_trie_root:#x}"),
-                    "Computed contract trie root."
-                );
-
-                let computed_class_trie_root = provider_mut
-                    .trie_insert_declared_classes(block_number, &state_update.declared_classes)?;
-
-                debug!(
-                    classes_tri_root = format!("{computed_class_trie_root:#x}"),
-                    "Computed classes trie root."
-                );
+                let provider_mut_clone = provider_mut.clone();
+                let (computed_contract_trie_root, computed_class_trie_root) = self
+                    .task_spawner
+                    .cpu_bound()
+                    .spawn(move || {
+                        let computed_contract_trie_root = provider_mut_clone
+                            .trie_insert_contract_updates(block_number, &state_update)?;
+
+                        debug!(
+                            contract_trie_root = format!("{computed_contract_trie_root:#x}"),
+                            "Computed contract trie root."
+                        );
+
+                        let computed_class_trie_root = provider_mut_clone
+                            .trie_insert_declared_classes(
+                                block_number,
+                                &state_update.declared_classes,
+                            )?;
+
+                        debug!(
+                            classes_tri_root = format!("{computed_class_trie_root:#x}"),
+                            "Computed classes trie root."
+                        );
+
+                        Result::<(Felt, Felt), crate::Error>::Ok((
+                            computed_contract_trie_root,
+                            computed_class_trie_root,
+                        ))
+                    })
+                    .await
+                    .map_err(Error::StateComputationTaskJoinError)??;
 
                 let computed_state_root = if computed_class_trie_root == Felt::ZERO {
                     computed_contract_trie_root
@@ -121,6 +139,9 @@ pub enum Error {
     #[error("Missing state update for block {0}")]
     MissingStateUpdate(BlockNumber),
 
+    #[error("State computation task join error: {0}")]
+    StateComputationTaskJoinError(katana_tasks::JoinError),
+
     #[error(
         "State root mismatch at block {block_number}: expected (from header) {expected:#x}, \
          computed {computed:#x}"

crates/sync/stage/tests/trie.rs

@@ -12,6 +12,7 @@ use katana_provider::api::trie::TrieWriter;
 use katana_provider::{ProviderFactory, ProviderResult};
 use katana_stage::trie::StateTrie;
 use katana_stage::{Stage, StageExecutionInput};
+use katana_tasks::TaskManager;
 use rstest::rstest;
 use starknet::macros::short_string;
 use starknet_types_core::hash::{Poseidon, StarkHash};
@@ -220,6 +221,7 @@ async fn verify_state_roots_success(
     #[case] to_block: BlockNumber,
     #[case] expected_blocks: Vec<BlockNumber>,
 ) {
+    let task_manager = TaskManager::current();
     let mut provider = MockProvider::new();
 
     // Configure blocks with correct state roots
@@ -230,8 +232,10 @@ async fn verify_state_roots_success(
         provider = provider.with_header(num, header).with_state_update(num, state_update);
     }
 
+    let mut stage = StateTrie::new(provider.clone(), task_manager.task_spawner());
+
     let input = StageExecutionInput::new(from_block, to_block);
-    let result = StateTrie::new(provider.clone()).execute(&input).await;
+    let result = stage.execute(&input).await;
     assert!(result.is_ok(), "Stage execution should succeed");
 
     // Verify that trie inserts were called for each block (twice per block: classes + contracts)
@@ -251,7 +255,8 @@ async fn state_root_mismatch_returns_error() {
         .with_header(block_number, header)
         .with_state_update(block_number, state_update);
 
-    let mut stage = StateTrie::new(provider);
+    let task_manager = TaskManager::current();
+    let mut stage = StateTrie::new(provider, task_manager.task_spawner());
     let input = StageExecutionInput::new(block_number, block_number);
 
     let result = stage.execute(&input).await;