From 0b9f3adbccc4e342ab008de6f4b8079e8e1b104e Mon Sep 17 00:00:00 2001
From: Alex Ostrovski <aov@matterlabs.dev>
Date: Thu, 24 Oct 2024 13:36:38 +0300
Subject: [PATCH] Add `/debug/stale-keys` endpoint

---
 core/lib/merkle_tree/src/domain.rs            |  7 ++++++
 .../src/api_server/metrics.rs                 |  1 +
 .../metadata_calculator/src/api_server/mod.rs | 25 +++++++++++++++++++
 .../src/api_server/tests.rs                   | 22 ++++++++++++++++
 core/node/metadata_calculator/src/helpers.rs  |  6 +++++
 5 files changed, 61 insertions(+)

diff --git a/core/lib/merkle_tree/src/domain.rs b/core/lib/merkle_tree/src/domain.rs
index 82100ba7fb11..bb69bda209cc 100644
--- a/core/lib/merkle_tree/src/domain.rs
+++ b/core/lib/merkle_tree/src/domain.rs
@@ -13,6 +13,7 @@ use crate::{
         ValueHash, TREE_DEPTH,
     },
     BlockOutput, HashTree, MerkleTree, MerkleTreePruner, MerkleTreePrunerHandle, NoVersionError,
+    PruneDatabase,
 };
 
 impl TreeInstruction<StorageKey> {
@@ -460,6 +461,12 @@ impl ZkSyncTreeReader {
             .collect()
     }
 
+    /// Returns raw stale keys obsoleted in the specified version of the tree.
+    pub fn raw_stale_keys(&self, l1_batch_number: L1BatchNumber) -> Vec<NodeKey> {
+        let version = u64::from(l1_batch_number.0);
+        self.0.db.stale_keys(version)
+    }
+
     /// Verifies consistency of the tree at the specified L1 batch number.
     ///
     /// # Errors
diff --git a/core/node/metadata_calculator/src/api_server/metrics.rs b/core/node/metadata_calculator/src/api_server/metrics.rs
index a763dad3e333..92f948e09702 100644
--- a/core/node/metadata_calculator/src/api_server/metrics.rs
+++ b/core/node/metadata_calculator/src/api_server/metrics.rs
@@ -10,6 +10,7 @@ pub(super) enum MerkleTreeApiMethod {
     Info,
     GetProofs,
     GetNodes,
+    GetStaleKeys,
 }
 
 /// Metrics for Merkle tree API.
diff --git a/core/node/metadata_calculator/src/api_server/mod.rs b/core/node/metadata_calculator/src/api_server/mod.rs
index f505197cbebd..4612d859a3dd 100644
--- a/core/node/metadata_calculator/src/api_server/mod.rs
+++ b/core/node/metadata_calculator/src/api_server/mod.rs
@@ -182,6 +182,16 @@ struct TreeNodesResponse {
     nodes: HashMap<HexNodeKey, ApiRawNode>,
 }
 
+#[derive(Debug, Deserialize)]
+struct StaleKeysRequest {
+    l1_batch_number: L1BatchNumber,
+}
+
+#[derive(Debug, Serialize)]
+struct StaleKeysResponse {
+    stale_keys: Vec<HexNodeKey>,
+}
+
 /// Server-side tree API error.
 #[derive(Debug)]
 enum TreeApiServerError {
@@ -466,6 +476,17 @@ impl AsyncTreeReader {
         Json(response)
     }
 
+    async fn get_stale_keys_handler(
+        State(this): State<Self>,
+        Json(request): Json<StaleKeysRequest>,
+    ) -> Json<StaleKeysResponse> {
+        let latency = API_METRICS.latency[&MerkleTreeApiMethod::GetStaleKeys].start();
+        let stale_keys = this.clone().raw_stale_keys(request.l1_batch_number).await;
+        let stale_keys = stale_keys.into_iter().map(HexNodeKey).collect();
+        latency.observe();
+        Json(StaleKeysResponse { stale_keys })
+    }
+
     async fn create_api_server(
         self,
         bind_address: &SocketAddr,
@@ -477,6 +498,10 @@ impl AsyncTreeReader {
             .route("/", routing::get(Self::info_handler))
             .route("/proofs", routing::post(Self::get_proofs_handler))
             .route("/debug/nodes", routing::post(Self::get_nodes_handler))
+            .route(
+                "/debug/stale-keys",
+                routing::post(Self::get_stale_keys_handler),
+            )
             .with_state(self);
 
         let listener = tokio::net::TcpListener::bind(bind_address)
diff --git a/core/node/metadata_calculator/src/api_server/tests.rs b/core/node/metadata_calculator/src/api_server/tests.rs
index a9509ea2192f..815522a4cd8e 100644
--- a/core/node/metadata_calculator/src/api_server/tests.rs
+++ b/core/node/metadata_calculator/src/api_server/tests.rs
@@ -84,6 +84,18 @@ async fn merkle_tree_api() {
     let raw_nodes_response: serde_json::Value = raw_nodes_response.json().await.unwrap();
     assert_raw_nodes_response(&raw_nodes_response);
 
+    let raw_stale_keys_response = api_client
+        .inner
+        .post(format!("http://{local_addr}/debug/stale-keys"))
+        .json(&serde_json::json!({ "l1_batch_number": 1 }))
+        .send()
+        .await
+        .unwrap()
+        .error_for_status()
+        .unwrap();
+    let raw_stale_keys_response: serde_json::Value = raw_stale_keys_response.json().await.unwrap();
+    assert_raw_stale_keys_response(&raw_stale_keys_response);
+
     // Stop the calculator and the tree API server.
     stop_sender.send_replace(true);
     api_server_task.await.unwrap().unwrap();
@@ -111,6 +123,16 @@ fn assert_raw_nodes_response(response: &serde_json::Value) {
     );
 }
 
+fn assert_raw_stale_keys_response(response: &serde_json::Value) {
+    let response = response.as_object().expect("not an object");
+    let stale_keys = response["stale_keys"].as_array().expect("not an array");
+    assert!(!stale_keys.is_empty()); // At least the root is always obsoleted
+    for stale_key in stale_keys {
+        let stale_key = stale_key.as_str().expect("not a string");
+        stale_key.parse::<NodeKey>().unwrap();
+    }
+}
+
 #[tokio::test]
 async fn api_client_connection_error() {
     // Use an address that will definitely fail on a timeout.
diff --git a/core/node/metadata_calculator/src/helpers.rs b/core/node/metadata_calculator/src/helpers.rs
index 26f02ce57f74..3f370afaf77e 100644
--- a/core/node/metadata_calculator/src/helpers.rs
+++ b/core/node/metadata_calculator/src/helpers.rs
@@ -414,6 +414,12 @@ impl AsyncTreeReader {
             .await
             .unwrap()
     }
+
+    pub(crate) async fn raw_stale_keys(self, l1_batch_number: L1BatchNumber) -> Vec<NodeKey> {
+        tokio::task::spawn_blocking(move || self.inner.raw_stale_keys(l1_batch_number))
+            .await
+            .unwrap()
+    }
 }
 
 /// Version of async tree reader that holds a weak reference to RocksDB. Used in [`MerkleTreeHealthCheck`].