From b2b1a3c68ccd8abf6f7327af5447e86be2ca09dc Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Wed, 30 Apr 2025 23:23:19 +0100
Subject: [PATCH] perf(engine): drop sparse trie after task returned result
 (#16000)

---
 .../tree/src/tree/payload_processor/mod.rs    |  2 +-
 .../src/tree/payload_processor/sparse_trie.rs | 28 +++++++++++++------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs
index af9a2548ce..6ecd529ddc 100644
--- a/crates/engine/tree/src/tree/payload_processor/mod.rs
+++ b/crates/engine/tree/src/tree/payload_processor/mod.rs
@@ -171,7 +171,7 @@ where
             multi_proof_task.run();
         });
 
-        let sparse_trie_task = SparseTrieTask::new(
+        let mut sparse_trie_task = SparseTrieTask::new(
             self.executor.clone(),
             sparse_trie_rx,
             proof_task.handle(),
diff --git a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs
index db3730d68a..7119c519c1 100644
--- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs
+++ b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs
@@ -24,13 +24,21 @@ use tracing::{debug, trace, trace_span};
 const SPARSE_TRIE_INCREMENTAL_LEVEL: usize = 2;
 
 /// A task responsible for populating the sparse trie.
-pub(super) struct SparseTrieTask<BPF> {
+pub(super) struct SparseTrieTask<BPF>
+where
+    BPF: BlindedProviderFactory + Send + Sync,
+    BPF::AccountNodeProvider: BlindedProvider + Send + Sync,
+    BPF::StorageNodeProvider: BlindedProvider + Send + Sync,
+{
     /// Executor used to spawn subtasks.
     #[expect(unused)] // TODO use this for spawning trie tasks
     pub(super) executor: WorkloadExecutor,
     /// Receives updates from the state root task.
     pub(super) updates: mpsc::Receiver<SparseTrieUpdate>,
-    pub(super) blinded_provider_factory: BPF,
+    /// Sparse Trie initialized with the blinded provider factory.
+    ///
+    /// It's kept as a field on the struct to prevent blocking on de-allocation in [`Self::run`].
+    pub(super) trie: SparseStateTrie<BPF>,
     pub(super) metrics: MultiProofTaskMetrics,
 }
 
@@ -41,13 +49,18 @@ where
     BPF::StorageNodeProvider: BlindedProvider + Send + Sync,
 {
     /// Creates a new sparse trie task.
-    pub(super) const fn new(
+    pub(super) fn new(
         executor: WorkloadExecutor,
         updates: mpsc::Receiver<SparseTrieUpdate>,
         blinded_provider_factory: BPF,
         metrics: MultiProofTaskMetrics,
     ) -> Self {
-        Self { executor, updates, blinded_provider_factory, metrics }
+        Self {
+            executor,
+            updates,
+            metrics,
+            trie: SparseStateTrie::new(blinded_provider_factory).with_updates(true),
+        }
     }
 
     /// Runs the sparse trie task to completion.
@@ -58,11 +71,10 @@ where
     ///
     /// NOTE: This function does not take `self` by value to prevent blocking on [`SparseStateTrie`]
     /// drop.
-    pub(super) fn run(&self) -> Result<StateRootComputeOutcome, ParallelStateRootError> {
+    pub(super) fn run(&mut self) -> Result<StateRootComputeOutcome, ParallelStateRootError> {
         let now = Instant::now();
 
         let mut num_iterations = 0;
-        let mut trie = SparseStateTrie::new(&self.blinded_provider_factory).with_updates(true);
 
         while let Ok(mut update) = self.updates.recv() {
             num_iterations += 1;
@@ -80,7 +92,7 @@ where
                 "Updating sparse trie"
             );
 
-            let elapsed = update_sparse_trie(&mut trie, update).map_err(|e| {
+            let elapsed = update_sparse_trie(&mut self.trie, update).map_err(|e| {
                 ParallelStateRootError::Other(format!("could not calculate state root: {e:?}"))
             })?;
             self.metrics.sparse_trie_update_duration_histogram.record(elapsed);
@@ -90,7 +102,7 @@ where
         debug!(target: "engine::root", num_iterations, "All proofs processed, ending calculation");
 
         let start = Instant::now();
-        let (state_root, trie_updates) = trie.root_with_updates().map_err(|e| {
+        let (state_root, trie_updates) = self.trie.root_with_updates().map_err(|e| {
             ParallelStateRootError::Other(format!("could not calculate state root: {e:?}"))
         })?;