From 0fa79c6c65dfbd3273cf111db802bd52d11832be Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Mon, 27 Jan 2025 15:31:44 +0000
Subject: [PATCH] fix(root): drop sparse trie tx sender when state updates
 finished (#14004)

---
 crates/engine/tree/src/tree/root.rs | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
index 3b0128025..dd96692a7 100644
--- a/crates/engine/tree/src/tree/root.rs
+++ b/crates/engine/tree/src/tree/root.rs
@@ -630,6 +630,19 @@ where
                     StateRootMessage::FinishedStateUpdates => {
                         trace!(target: "engine::root", "processing StateRootMessage::FinishedStateUpdates");
                         updates_finished = true;
+
+                        let all_proofs_received = proofs_processed >= updates_received;
+                        let no_pending = !self.proof_sequencer.has_pending();
+                        if all_proofs_received && no_pending {
+                            // drop the sender
+                            sparse_trie_tx.take();
+                            debug!(
+                                target: "engine::root",
+                                total_updates = updates_received,
+                                total_proofs = proofs_processed,
+                                "State updates finished and all proofs processed, ending calculation"
+                            );
+                        }
                     }
                     StateRootMessage::ProofCalculated(proof_calculated) => {
                         trace!(target: "engine::root", "processing StateRootMessage::ProofCalculated");
@@ -660,7 +673,12 @@ where
                         if all_proofs_received && no_pending && updates_finished {
                             // drop the sender
                             sparse_trie_tx.take();
-                            debug!(target: "engine::root", total_updates = updates_received, total_proofs = proofs_processed, "All proofs processed, ending calculation");
+                            debug!(
+                                target: "engine::root",
+                                total_updates = updates_received,
+                                total_proofs = proofs_processed,
+                                "All proofs processed, ending calculation"
+                            );
                         }
                     }
                     StateRootMessage::RootCalculated { state_root, trie_updates, iterations } => {