From 70b3a8aea33820d5bf932b608c9e68ecc2915d4c Mon Sep 17 00:00:00 2001
From: Alex Ostrovski <slowli@users.noreply.github.com>
Date: Fri, 28 Jun 2024 16:27:21 +0300
Subject: [PATCH] fix(merkle-tree): Fix chunk recovery reporting during tree
 recovery (#2348)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## What ❔

Fixes logging / observing metrics in the case a termination signal was
received during tree recovery.

## Why ❔

Logging / observing metrics in this case is misleading.

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Code has been formatted via `zk fmt` and `zk lint`.
---
 .../metadata_calculator/src/recovery/mod.rs     | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/core/node/metadata_calculator/src/recovery/mod.rs b/core/node/metadata_calculator/src/recovery/mod.rs
index 4aee14c0c79a..dcbc0a68af92 100644
--- a/core/node/metadata_calculator/src/recovery/mod.rs
+++ b/core/node/metadata_calculator/src/recovery/mod.rs
@@ -261,8 +261,10 @@ impl AsyncTreeRecovery {
                 .acquire()
                 .await
                 .context("semaphore is never closed")?;
-            Self::recover_key_chunk(&tree, snapshot.l2_block, chunk, pool, stop_receiver).await?;
-            options.events.chunk_recovered();
+            if Self::recover_key_chunk(&tree, snapshot.l2_block, chunk, pool, stop_receiver).await?
+            {
+                options.events.chunk_recovered();
+            }
             anyhow::Ok(())
         });
         future::try_join_all(chunk_tasks).await?;
@@ -338,20 +340,21 @@ impl AsyncTreeRecovery {
         Ok(output)
     }
 
+    /// Returns `Ok(true)` if the chunk was recovered, `Ok(false)` if the recovery process was interrupted.
     async fn recover_key_chunk(
         tree: &Mutex<AsyncTreeRecovery>,
         snapshot_l2_block: L2BlockNumber,
         key_chunk: ops::RangeInclusive<H256>,
         pool: &ConnectionPool<Core>,
         stop_receiver: &watch::Receiver<bool>,
-    ) -> anyhow::Result<()> {
+    ) -> anyhow::Result<bool> {
         let acquire_connection_latency =
             RECOVERY_METRICS.chunk_latency[&ChunkRecoveryStage::AcquireConnection].start();
         let mut storage = pool.connection_tagged("metadata_calculator").await?;
         acquire_connection_latency.observe();
 
         if *stop_receiver.borrow() {
-            return Ok(());
+            return Ok(false);
         }
 
         let entries_latency =
@@ -368,7 +371,7 @@ impl AsyncTreeRecovery {
         );
 
         if *stop_receiver.borrow() {
-            return Ok(());
+            return Ok(false);
         }
 
         // Sanity check: all entry keys must be distinct. Otherwise, we may end up writing non-final values
@@ -398,7 +401,7 @@ impl AsyncTreeRecovery {
         lock_tree_latency.observe();
 
         if *stop_receiver.borrow() {
-            return Ok(());
+            return Ok(false);
         }
 
         let extend_tree_latency =
@@ -408,7 +411,7 @@ impl AsyncTreeRecovery {
         tracing::debug!(
             "Extended Merkle tree with entries for chunk {key_chunk:?} in {extend_tree_latency:?}"
         );
-        Ok(())
+        Ok(true)
     }
 }