feat: share sparse trie pipeline with payload builder

Author: constbh

src/main.rs

@@ -458,6 +458,13 @@ fn main() -> eyre::Result<()> {
             reth_bsc::shared::set_engine_api_tx(node.engine_api_tx.clone().unwrap()).unwrap();
             tracing::debug!("set engine api tx successfully");
 
+            // Publish CanonicalInMemoryState so the sparse-trie spawner can resolve
+            // in-memory parents back to the on-disk anchor. Concrete `BlockchainProvider`
+            // is only reachable here (post-launch); the generic builder context isn't.
+            let _ = reth_bsc::shared::set_canonical_in_memory_state(
+                node.provider.canonical_in_memory_state(),
+            );
+
             // Set the IPC client
             reth_bsc::shared::set_ipc_client(ipc_path).await.unwrap();
 

src/node/engine.rs

@@ -111,6 +111,134 @@ where
                 MiningConfig::from_env()
             };
 
+        // Register the sparse-trie state-root spawner, if enabled.
+        //
+        // We construct a long-lived `PayloadProcessor` keyed to a fresh `BscEvmConfig`
+        // (built from chain_spec — same source as the rest of the BSC pipeline). For
+        // each build job, the registered closure constructs a one-shot
+        // `OverlayStateProviderFactory` anchored at the parent block hash and calls
+        // `spawn_state_root` to get a `StateRootHandle`.
+        //
+        // `TreeConfig::default()` is used here as a workaround — the engine-launch
+        // `TreeConfig` (which honors `--engine.*` CLI flags) is not reachable from
+        // `BscPayloadServiceBuilder`. Default values are reasonable for the sparse-trie
+        // background task; revisit if performance testing shows we need CLI-tunable
+        // worker counts or cache sizes.
+        if mining_config.use_sparse_trie_state_root
+            && !rust_eth_triedb::triedb_manager::is_triedb_active()
+        {
+            use alloy_consensus::BlockHeader;
+            use reth_chain_state::LazyOverlay;
+            use reth_engine_tree::tree::{
+                payload_processor::PayloadProcessor, precompile_cache::PrecompileCacheMap,
+                TreeConfig,
+            };
+            use reth_provider::providers::{OverlayBuilder, OverlayStateProviderFactory};
+            use reth_tasks::{Runtime, RuntimeBuilder, RuntimeConfig, TokioConfig};
+            use reth_trie_db::ChangesetCache;
+
+            let chain_spec = Arc::new(ctx.config().chain.clone().as_ref().clone());
+            let bsc_evm_config = crate::node::evm::config::BscEvmConfig::new(chain_spec);
+            let tree_config = Arc::new(TreeConfig::default());
+            let provider = ctx.provider().clone();
+
+            // Build a dedicated `reth_tasks::Runtime` for the sparse-trie pools,
+            // attached to the existing tokio handle so we don't spin up a second tokio
+            // executor. Rayon pools default to `available_parallelism()`, matching what
+            // the engine uses for its own PayloadProcessor.
+            //
+            // TODO: in a follow-up, share the engine's Runtime instead of constructing
+            // a parallel one — this currently means two sets of rayon pools competing
+            // for CPU. Acceptable for a first cut; revisit if perf testing shows
+            // contention.
+            let tokio_handle = ctx.task_executor().handle().clone();
+            let runtime = RuntimeBuilder::new(
+                RuntimeConfig::default().with_tokio(TokioConfig::existing_handle(tokio_handle)),
+            )
+            .build()
+            .map_err(|e| eyre::eyre!("failed to build sparse-trie Runtime: {e}"))?;
+            let _: &Runtime = &runtime; // type-check anchor
+
+            let payload_processor = std::sync::Arc::new(PayloadProcessor::new(
+                runtime,
+                bsc_evm_config,
+                tree_config.as_ref(),
+                PrecompileCacheMap::default(),
+            ));
+
+            let tree_config_for_closure = tree_config.clone();
+            let spawn_fn: crate::shared::SparseTrieSpawnFn = std::sync::Arc::new(
+                move |parent_hash: alloy_primitives::B256,
+                      parent_state_root: alloy_primitives::B256| {
+                    // Walk the in-memory canonical chain to find the on-disk anchor.
+                    // Without this, proof workers fail with `BlockHashNotFound` whenever
+                    // the parent hasn't been persisted yet (the common case during fast
+                    // block production with last_persisted_number lagging head).
+                    //
+                    // Mirrors `payload_validator::get_parent_lazy_overlay` from upstream
+                    // reth, which the engine's own PayloadProcessor uses.
+                    //
+                    // `canonical_in_memory_state` is published from main.rs after launch
+                    // (it's only reachable on the concrete BlockchainProvider). If for
+                    // some reason it's not yet set (race during early startup), fall back
+                    // to anchoring at parent_hash directly — proof workers will fail and
+                    // builder will use the synchronous state_root_with_updates path.
+                    let (anchor_hash, lazy_overlay) = if let Some(cim) =
+                        crate::shared::get_canonical_in_memory_state()
+                    {
+                        match cim.state_by_hash(parent_hash) {
+                            Some(state) => {
+                                // chain() yields newest-to-oldest including self, exactly
+                                // the order LazyOverlay::new requires.
+                                let blocks: Vec<ExecutedBlock<crate::BscPrimitives>> =
+                                    state.chain().map(|bs| bs.block()).collect();
+                                // Anchor = parent of the oldest in-memory block (= on-disk tip).
+                                let anchor = blocks
+                                    .last()
+                                    .map(|b| b.recovered_block().parent_hash())
+                                    .unwrap_or(parent_hash);
+                                (anchor, Some(LazyOverlay::new(blocks)))
+                            }
+                            None => {
+                                // Parent already persisted — anchor directly, no overlay.
+                                (parent_hash, None)
+                            }
+                        }
+                    } else {
+                        (parent_hash, None)
+                    };
+
+                    let overlay_builder = OverlayBuilder::<crate::BscPrimitives>::new(
+                        anchor_hash,
+                        ChangesetCache::default(),
+                    )
+                    .with_lazy_overlay(lazy_overlay);
+
+                    let overlay_factory = OverlayStateProviderFactory::new(
+                        provider.clone(),
+                        overlay_builder,
+                    );
+                    Some(payload_processor.spawn_state_root(
+                        overlay_factory,
+                        parent_state_root,
+                        false, // halve_workers
+                        tree_config_for_closure.as_ref(),
+                    ))
+                },
+            );
+
+            if crate::shared::set_sparse_trie_spawn_fn(spawn_fn).is_err() {
+                tracing::warn!(
+                    "Sparse-trie spawner already registered, keeping existing one"
+                );
+            } else {
+                info!(
+                    "Sparse-trie state-root spawner registered \
+                     (use_sparse_trie_state_root=true, triedb=inactive)"
+                );
+            }
+        }
+
         // Skip mining setup if disabled
         if !mining_config.is_mining_enabled() {
             info!("Mining is disabled in configuration");

src/node/evm/builder.rs

@@ -175,9 +175,59 @@ where
         state: impl StateProvider,
     ) -> Result<BlockBuilderOutcomeWithDiffLayer<BscPrimitives>, BlockExecutionError> {
         let finish_start = std::time::Instant::now();
+        // `executor.finish()` runs BSC's post-execution system txs (slash spoiled
+        // validator, distribute fees / finality rewards, breathe-block validator-set
+        // updates). Any `state_hook` previously installed on the executor — including
+        // the sparse-trie hook — captures those state changes here. Consuming the
+        // executor on the next line drops the hook closure, which triggers
+        // `StateHookSender::drop` → `FinishedStateUpdates` → sparse-trie task can
+        // safely return from `state_root()` below.
         let (evm, result) = self.executor.finish()?;
         let (db, evm_env) = evm.finish();
 
+        // Sparse-trie state-root collection: now that the executor (and therefore the
+        // state_hook) has been dropped, the background task has all updates and is
+        // finalizing. Take the handle from ctx, block on `state_root()`, and stash the
+        // result in the sink — the MDBX branch below reads it.
+        //
+        // Failures fall through silently to the legacy `state_root_with_updates` path,
+        // logged at WARN. A failure here is non-fatal but indicates the task panicked
+        // or its channel was dropped; investigate via the trace target below.
+        if let Some(handle_slot) = self.ctx.trie_handle.take() {
+            if let Some(mut handle) = handle_slot.lock().unwrap().take() {
+                let wait_start = std::time::Instant::now();
+                match handle.state_root() {
+                    Ok(outcome) => {
+                        let updates = std::sync::Arc::try_unwrap(outcome.trie_updates)
+                            .unwrap_or_else(|arc| (*arc).clone());
+                        if let Some(sink) = self.ctx.state_root_precomputed_sink.as_ref() {
+                            *sink.lock().unwrap() = Some((outcome.state_root, updates));
+                        } else {
+                            // No sink registered — write into the field (which the
+                            // MDBX branch also reads as fallback).
+                            self.precomputed_state_root = Some((outcome.state_root, updates));
+                        }
+                        tracing::debug!(
+                            target: "bsc::builder",
+                            parent_hash = %self.parent.hash(),
+                            block_number = %(self.parent.number + 1),
+                            wait_ms = wait_start.elapsed().as_millis(),
+                            "Sparse-trie state-root delivered post-finish()"
+                        );
+                    }
+                    Err(err) => {
+                        tracing::warn!(
+                            target: "bsc::builder",
+                            parent_hash = %self.parent.hash(),
+                            block_number = %(self.parent.number + 1),
+                            %err,
+                            "Sparse-trie task failed post-finish(); falling back to state_root_with_updates"
+                        );
+                    }
+                }
+            }
+        }
+
         let assembled_system_txs = {
             let mut inner = self.shared_ctx.inner.borrow_mut();
             std::mem::take(&mut inner.assembled_system_txs)
@@ -228,11 +278,21 @@ where
                 "Calculated state root using triedb"
             );
             (new_root, TrieUpdates::default(), Some(new_difflayer))
-        } else if let Some((root, updates)) = self.precomputed_state_root.take() {
+        } else if let Some((root, updates)) = self
+            .ctx
+            .state_root_precomputed_sink
+            .as_ref()
+            .and_then(|sink| sink.lock().unwrap().take())
+            .or_else(|| self.precomputed_state_root.take())
+        {
             // Fast path: sparse-trie background task already computed the root concurrently
             // with execution. See `crate::shared::spawn_sparse_trie_state_root` and reth 2.0
             // `--engine.share-sparse-trie-with-payload-builder` semantics for the upstream
             // mechanism we mirror.
+            //
+            // Preferred source is the sink on `self.ctx` (filled by the payload layer
+            // post-exec); the field on `Self` is a fallback that `fn finish` populates
+            // when callers route through the trait method.
             tracing::debug!(
                 target: "bsc::builder",
                 parent_hash = %self.parent.hash(),

src/node/evm/config.rs

@@ -69,6 +69,27 @@ pub struct BscNextBlockEnvAttributes {
     /// Sink for transporting `turn_length` from builder to payload layer without writing to
     /// TURN_LENGTH_CACHE prematurely.
     pub turn_length_sink: Option<Arc<Mutex<Option<u8>>>>,
+    /// Sink for precomputed `(state_root, trie_updates)` from a sparse-trie background
+    /// task. Filled by payload layer between exec and `finish_with_difflayer` so the
+    /// builder's MDBX branch can skip the blocking `state_root_with_updates` call. See
+    /// [`BscBlockExecutionCtx::state_root_precomputed_sink`] for full semantics.
+    pub state_root_precomputed_sink: Option<
+        Arc<Mutex<Option<(alloy_primitives::B256, reth_trie_common::updates::TrieUpdates)>>>,
+    >,
+    /// Sparse-trie state-root handle, threaded through to `finish_with_difflayer`.
+    ///
+    /// Stored here (in `Arc<Mutex<Option<_>>>` so `Clone` works for the type-erased
+    /// builder path) so that `state_root()` can be called **after** `executor.finish()`
+    /// runs BSC's post-execution system transactions (slash, fee distribution,
+    /// validator-set updates). Those system txs change state via the same executor
+    /// that has the `state_hook` installed; the hook is dropped naturally when the
+    /// executor is consumed by `finish()`, which sends `FinishedStateUpdates` to the
+    /// background task. Only after that drop is it safe to await `state_root()`.
+    ///
+    /// `None` when sparse-trie is disabled or in TrieDB mode.
+    pub trie_handle: Option<
+        Arc<Mutex<Option<reth_engine_tree::tree::multiproof::StateRootHandle>>>,
+    >,
 }
 
 impl<H: BlockHeader> BuildPendingEnv<H> for BscNextBlockEnvAttributes {
@@ -79,6 +100,8 @@ impl<H: BlockHeader> BuildPendingEnv<H> for BscNextBlockEnvAttributes {
             triedb_prefetcher: None,
             validator_cache_sink: None,
             turn_length_sink: None,
+            state_root_precomputed_sink: None,
+            trie_handle: None,
         }
     }
 }
@@ -135,6 +158,29 @@ pub struct BscBlockExecutionCtx<'a> {
     pub validator_cache_sink: Option<ValidatorCacheSink>,
     /// Sink for `turn_length` — same lifecycle as `validator_cache_sink`.
     pub turn_length_sink: Option<Arc<Mutex<Option<u8>>>>,
+    /// Sink for a precomputed `(state_root, trie_updates)` from a sparse-trie background
+    /// task (reth 2.0 mechanism).
+    ///
+    /// Write direction is **reversed** vs the other sinks: the payload layer fills this
+    /// **before** calling `finish_with_difflayer`, and the builder's MDBX branch reads
+    /// it to skip the synchronous `state_root_with_updates` call. `None` in the bid
+    /// simulator path and when the `--mining.use-sparse-trie-state-root` flag is off,
+    /// triggering the legacy state-root path.
+    pub state_root_precomputed_sink: Option<
+        Arc<Mutex<Option<(alloy_primitives::B256, reth_trie_common::updates::TrieUpdates)>>>,
+    >,
+    /// Sparse-trie state-root handle. The builder consumes this **after**
+    /// `executor.finish()` runs BSC's post-execution system transactions (slash,
+    /// fee distribution, validator-set updates), so those state changes are
+    /// captured by the executor's `state_hook` before the hook is dropped (which
+    /// signals the sparse-trie task to finalize). Calling `state_root()` before
+    /// the executor is dropped would deadlock the task on `FinishedStateUpdates`.
+    ///
+    /// `Arc<Mutex<Option<_>>>` because `StateRootHandle` is `!Clone` (single-use
+    /// receiver) and `BscBlockExecutionCtx` derives `Clone`.
+    pub trie_handle: Option<
+        Arc<Mutex<Option<reth_engine_tree::tree::multiproof::StateRootHandle>>>,
+    >,
 }
 
 impl<'a> BscBlockExecutionCtx<'a> {
@@ -440,6 +486,8 @@ where
             triedb_prefetcher: None,
             validator_cache_sink: None,
             turn_length_sink: None,
+            state_root_precomputed_sink: None,
+            trie_handle: None,
         })
     }
 
@@ -466,6 +514,8 @@ where
             triedb_prefetcher: attributes.triedb_prefetcher,
             validator_cache_sink: attributes.validator_cache_sink,
             turn_length_sink: attributes.turn_length_sink,
+            state_root_precomputed_sink: attributes.state_root_precomputed_sink,
+            trie_handle: attributes.trie_handle,
         })
     }
 
@@ -541,6 +591,8 @@ where
             triedb_prefetcher: None,
             validator_cache_sink: None,
             turn_length_sink: None,
+            state_root_precomputed_sink: None,
+            trie_handle: None,
         })
     }
 

src/node/miner/bid_simulator.rs

@@ -459,6 +459,10 @@ where
                     triedb_prefetcher,
                     validator_cache_sink: Some(bid_validator_cache_sink.clone()),
                     turn_length_sink: Some(bid_turn_length_sink.clone()),
+                    // Bid simulation does not run alongside a sparse-trie task —
+                    // builder will fall through to state_root_with_updates.
+                    state_root_precomputed_sink: None,
+                    trie_handle: None,
                 },
             )
             .map_err(PayloadBuilderError::other)

src/node/miner/bsc_miner.rs

@@ -733,6 +733,7 @@ where
             // Filled in by BscPayloadJob::start when sparse-trie state-root is enabled
             // and the engine has registered a spawner. Falls back to legacy path when None.
             state_root_precomputed: std::sync::Arc::new(std::sync::Mutex::new(None)),
+            trie_handle: std::sync::Arc::new(std::sync::Mutex::new(None)),
         };
 
         let parent_hash = mining_ctx.parent_header.hash();