From 408393c7d8ceee0ae95cbc1f2b24a3375e345e97 Mon Sep 17 00:00:00 2001
From: Alex Ostrovski <slowli@users.noreply.github.com>
Date: Tue, 25 Jun 2024 12:37:24 +0300
Subject: [PATCH] fix(merkle-tree): Change `LazyAsyncTreeReader::wait()`
 signature (#2314)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## What ❔

Makes `LazyAsyncTreeReader::wait()` return an `Option` to make it
clearer that not initializing is not necessarily an error (this decision
is up to the caller).

## Why ❔

Recent errors on ENs were caused by unwrapping
`LazyAsyncTreeReader::wait()`. They partially masked the real error
cause (not related to this task).

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Documentation comments have been added / updated.
- [x] Code has been formatted via `zk fmt` and `zk lint`.
---
 core/bin/external_node/src/main.rs            | 19 ++++++++++++-------
 core/node/metadata_calculator/src/helpers.rs  | 19 +++++++++----------
 .../layers/metadata_calculator.rs             | 17 ++++++++++-------
 3 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/core/bin/external_node/src/main.rs b/core/bin/external_node/src/main.rs
index 5d5de22cabf9..25b6f81a6b5a 100644
--- a/core/bin/external_node/src/main.rs
+++ b/core/bin/external_node/src/main.rs
@@ -185,14 +185,19 @@ async fn run_tree(
     if let Some(api_config) = api_config {
         let address = (Ipv4Addr::UNSPECIFIED, api_config.port).into();
         let tree_reader = metadata_calculator.tree_reader();
-        let stop_receiver = stop_receiver.clone();
+        let mut stop_receiver = stop_receiver.clone();
         task_futures.push(tokio::spawn(async move {
-            tree_reader
-                .wait()
-                .await
-                .context("Cannot initialize tree reader")?
-                .run_api_server(address, stop_receiver)
-                .await
+            if let Some(reader) = tree_reader.wait().await {
+                reader.run_api_server(address, stop_receiver).await
+            } else {
+                // Tree is dropped before initialized, e.g. because the node is getting shut down.
+                // We don't want to treat this as an error since it could mask the real shutdown cause in logs etc.
+                tracing::warn!(
+                    "Tree is dropped before initialized, not starting the tree API server"
+                );
+                stop_receiver.changed().await?;
+                Ok(())
+            }
         }));
     }
 
diff --git a/core/node/metadata_calculator/src/helpers.rs b/core/node/metadata_calculator/src/helpers.rs
index d6918b7a5e87..c71c0ecf9255 100644
--- a/core/node/metadata_calculator/src/helpers.rs
+++ b/core/node/metadata_calculator/src/helpers.rs
@@ -91,9 +91,10 @@ impl MerkleTreeHealthCheck {
         let weak_reader = Arc::<OnceCell<WeakAsyncTreeReader>>::default();
         let weak_reader_for_task = weak_reader.clone();
         tokio::spawn(async move {
-            weak_reader_for_task
-                .set(reader.wait().await.unwrap().downgrade())
-                .ok();
+            if let Some(reader) = reader.wait().await {
+                weak_reader_for_task.set(reader.downgrade()).ok();
+            }
+            // Otherwise, the tree is dropped before getting initialized; this is not an error in this context.
         });
 
         Self {
@@ -393,16 +394,14 @@ impl LazyAsyncTreeReader {
         self.0.borrow().clone()
     }
 
-    /// Waits until the tree is initialized and returns a reader for it.
-    pub async fn wait(mut self) -> anyhow::Result<AsyncTreeReader> {
+    /// Waits until the tree is initialized and returns a reader for it. If the tree is dropped before
+    /// getting initialized, returns `None`.
+    pub async fn wait(mut self) -> Option<AsyncTreeReader> {
         loop {
             if let Some(reader) = self.0.borrow().clone() {
-                break Ok(reader);
+                break Some(reader);
             }
-            self.0
-                .changed()
-                .await
-                .context("Tree dropped without getting ready; not resolving tree reader")?;
+            self.0.changed().await.ok()?;
         }
     }
 }
diff --git a/core/node/node_framework/src/implementations/layers/metadata_calculator.rs b/core/node/node_framework/src/implementations/layers/metadata_calculator.rs
index 1d0164c5f518..9185aeea553c 100644
--- a/core/node/node_framework/src/implementations/layers/metadata_calculator.rs
+++ b/core/node/node_framework/src/implementations/layers/metadata_calculator.rs
@@ -162,13 +162,16 @@ impl Task for TreeApiTask {
         "tree_api".into()
     }
 
-    async fn run(self: Box<Self>, stop_receiver: StopReceiver) -> anyhow::Result<()> {
-        self.tree_reader
-            .wait()
-            .await
-            .context("Cannot initialize tree reader")?
-            .run_api_server(self.bind_addr, stop_receiver.0)
-            .await
+    async fn run(self: Box<Self>, mut stop_receiver: StopReceiver) -> anyhow::Result<()> {
+        if let Some(reader) = self.tree_reader.wait().await {
+            reader.run_api_server(self.bind_addr, stop_receiver.0).await
+        } else {
+            // Tree is dropped before initialized, e.g. because the node is getting shut down.
+            // We don't want to treat this as an error since it could mask the real shutdown cause in logs etc.
+            tracing::warn!("Tree is dropped before initialized, not starting the tree API server");
+            stop_receiver.0.changed().await?;
+            Ok(())
+        }
     }
 }