feat(metadata-calculator): Add debug endpoints for tree API

core/bin/external_node/src/node_builder.rs

@@ -11,7 +11,9 @@ use zksync_config::{
     },
     PostgresConfig,
 };
-use zksync_metadata_calculator::{MetadataCalculatorConfig, MetadataCalculatorRecoveryConfig};
+use zksync_metadata_calculator::{
+    MerkleTreeReaderConfig, MetadataCalculatorConfig, MetadataCalculatorRecoveryConfig,
+};
 use zksync_node_api_server::web3::Namespace;
 use zksync_node_framework::{
     implementations::layers::{
@@ -25,7 +27,7 @@ use zksync_node_framework::{
         logs_bloom_backfill::LogsBloomBackfillLayer,
         main_node_client::MainNodeClientLayer,
         main_node_fee_params_fetcher::MainNodeFeeParamsFetcherLayer,
-        metadata_calculator::MetadataCalculatorLayer,
+        metadata_calculator::{MetadataCalculatorLayer, TreeApiServerLayer},
         node_storage_init::{
             external_node_strategy::{ExternalNodeInitStrategyLayer, SnapshotRecoveryConfig},
             NodeStorageInitializerLayer,
@@ -385,6 +387,29 @@ impl ExternalNodeBuilder {
         Ok(self)
     }
 
+    fn add_isolated_tree_api_layer(mut self) -> anyhow::Result<Self> {
+        let reader_config = MerkleTreeReaderConfig {
+            db_path: self.config.required.merkle_tree_path.clone(),
+            max_open_files: self.config.optional.merkle_tree_max_open_files,
+            multi_get_chunk_size: self.config.optional.merkle_tree_multi_get_chunk_size,
+            block_cache_capacity: self.config.optional.merkle_tree_block_cache_size(),
+            include_indices_and_filters_in_block_cache: self
+                .config
+                .optional
+                .merkle_tree_include_indices_and_filters_in_block_cache,
+        };
+        let api_config = MerkleTreeApiConfig {
+            port: self
+                .config
+                .tree_component
+                .api_port
+                .context("should contain tree api port")?,
+        };
+        self.node
+            .add_layer(TreeApiServerLayer::new(reader_config, api_config));
+        Ok(self)
+    }
+
     fn add_tx_sender_layer(mut self) -> anyhow::Result<Self> {
         let postgres_storage_config = PostgresStorageCachesConfig {
             factory_deps_cache_size: self.config.optional.factory_deps_cache_size() as u64,
@@ -607,11 +632,11 @@ impl ExternalNodeBuilder {
                     self = self.add_metadata_calculator_layer(with_tree_api)?;
                 }
                 Component::TreeApi => {
-                    anyhow::ensure!(
-                        components.contains(&Component::Tree),
-                        "Merkle tree API cannot be started without a tree component"
-                    );
-                    // Do nothing, will be handled by the `Tree` component.
+                    if components.contains(&Component::Tree) {
+                        // Do nothing, will be handled by the `Tree` component.
+                    } else {
+                        self = self.add_isolated_tree_api_layer()?;
+                    }
                 }
                 Component::TreeFetcher => {
                     self = self.add_tree_data_fetcher_layer()?;

core/bin/external_node/src/tests/mod.rs

@@ -17,7 +17,7 @@ mod utils;
 const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
 const POLL_INTERVAL: Duration = Duration::from_millis(100);
 
-#[test_casing(3, ["all", "core", "api"])]
+#[test_casing(4, ["all", "core", "api", "core,tree_api"])]
 #[tokio::test]
 #[tracing::instrument] // Add args to the test logs
 async fn external_node_basics(components_str: &'static str) {
@@ -170,40 +170,3 @@ async fn running_tree_without_core_is_not_allowed() {
         err
     );
 }
-
-#[tokio::test]
-async fn running_tree_api_without_tree_is_not_allowed() {
-    let _guard = zksync_vlog::ObservabilityBuilder::new().try_build().ok(); // Enable logging to simplify debugging
-    let (env, _env_handles) = utils::TestEnvironment::with_genesis_block("core,tree_api").await;
-
-    let l2_client = utils::mock_l2_client(&env);
-    let eth_client = utils::mock_eth_client(env.config.diamond_proxy_address());
-
-    let node_handle = tokio::task::spawn_blocking(move || {
-        std::thread::spawn(move || {
-            let mut node = ExternalNodeBuilder::new(env.config)?;
-            inject_test_layers(
-                &mut node,
-                env.sigint_receiver,
-                env.app_health_sender,
-                eth_client,
-                l2_client,
-            );
-
-            // We're only interested in the error, so we drop the result.
-            node.build(env.components.0.into_iter().collect()).map(drop)
-        })
-        .join()
-        .unwrap()
-    });
-
-    // Check that we cannot build the node without the core component.
-    let result = node_handle.await.expect("Building the node panicked");
-    let err = result.expect_err("Building the node with tree api but without tree should fail");
-    assert!(
-        err.to_string()
-            .contains("Merkle tree API cannot be started without a tree component"),
-        "Unexpected errror: {}",
-        err
-    );
-}

core/lib/merkle_tree/src/domain.rs

@@ -9,10 +9,11 @@ use crate::{
     consistency::ConsistencyError,
     storage::{PatchSet, Patched, RocksDBWrapper},
     types::{
-        Key, Root, TreeEntry, TreeEntryWithProof, TreeInstruction, TreeLogEntry, ValueHash,
-        TREE_DEPTH,
+        Key, NodeKey, RawNode, Root, TreeEntry, TreeEntryWithProof, TreeInstruction, TreeLogEntry,
+        ValueHash, TREE_DEPTH,
     },
     BlockOutput, HashTree, MerkleTree, MerkleTreePruner, MerkleTreePrunerHandle, NoVersionError,
+    PruneDatabase,
 };
 
 impl TreeInstruction<StorageKey> {
@@ -444,6 +445,28 @@ impl ZkSyncTreeReader {
         self.0.entries_with_proofs(version, keys)
     }
 
+    /// Returns raw nodes for the specified `keys`.
+    pub fn raw_nodes(&self, keys: &[NodeKey]) -> Vec<Option<RawNode>> {
+        let raw_nodes = self.0.db.raw_nodes(keys).into_iter();
+        raw_nodes
+            .zip(keys)
+            .map(|(slice, key)| {
+                let slice = slice?;
+                Some(if key.is_empty() {
+                    RawNode::deserialize_root(&slice)
+                } else {
+                    RawNode::deserialize(&slice)
+                })
+            })
+            .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

core/lib/merkle_tree/src/errors.rs

@@ -22,6 +22,8 @@ pub enum DeserializeErrorKind {
     /// Bit mask specifying a child kind in an internal tree node is invalid.
     #[error("invalid bit mask specifying a child kind in an internal tree node")]
     InvalidChildKind,
+    #[error("data left after deserialization")]
+    Leftovers,
 
     /// Missing required tag in the tree manifest.
     #[error("missing required tag `{0}` in tree manifest")]

core/lib/merkle_tree/src/lib.rs

@@ -82,7 +82,7 @@ mod utils;
 pub mod unstable {
     pub use crate::{
         errors::DeserializeError,
-        types::{Manifest, Node, NodeKey, ProfiledTreeOperation, Root},
+        types::{Manifest, Node, NodeKey, ProfiledTreeOperation, RawNode, Root},
     };
 }
 

core/lib/merkle_tree/src/storage/rocksdb.rs

@@ -53,6 +53,23 @@ impl NamedColumnFamily for MerkleTreeColumnFamily {
 
 type LocalProfiledOperation = RefCell<Option<Arc<ProfiledOperation>>>;
 
+/// Unifies keys that can be used to load raw data from RocksDB.
+pub(crate) trait ToDbKey: Sync {
+    fn to_db_key(&self) -> Vec<u8>;
+}
+
+impl ToDbKey for NodeKey {
+    fn to_db_key(&self) -> Vec<u8> {
+        NodeKey::to_db_key(*self)
+    }
+}
+
+impl ToDbKey for (NodeKey, bool) {
+    fn to_db_key(&self) -> Vec<u8> {
+        NodeKey::to_db_key(self.0)
+    }
+}
+
 /// Main [`Database`] implementation wrapping a [`RocksDB`] reference.
 ///
 /// # Cloning
@@ -112,7 +129,7 @@ impl RocksDBWrapper {
             .expect("Failed reading from RocksDB")
     }
 
-    fn raw_nodes(&self, keys: &NodeKeys) -> Vec<Option<DBPinnableSlice<'_>>> {
+    pub(crate) fn raw_nodes<T: ToDbKey>(&self, keys: &[T]) -> Vec<Option<DBPinnableSlice<'_>>> {
         // Propagate the currently profiled operation to rayon threads used in the parallel iterator below.
         let profiled_operation = self
             .profiled_operation
@@ -126,7 +143,7 @@ impl RocksDBWrapper {
                 let _guard = profiled_operation
                     .as_ref()
                     .and_then(ProfiledOperation::start_profiling);
-                let keys = chunk.iter().map(|(key, _)| key.to_db_key());
+                let keys = chunk.iter().map(ToDbKey::to_db_key);
                 let results = self.db.multi_get_cf(MerkleTreeColumnFamily::Tree, keys);
                 results
                     .into_iter()
@@ -144,9 +161,9 @@ impl RocksDBWrapper {
         // If we didn't succeed with the patch set, or the key version is old,
         // access the underlying storage.
         let node = if is_leaf {
-            LeafNode::deserialize(raw_node).map(Node::Leaf)
+            LeafNode::deserialize(raw_node, false).map(Node::Leaf)
         } else {
-            InternalNode::deserialize(raw_node).map(Node::Internal)
+            InternalNode::deserialize(raw_node, false).map(Node::Internal)
         };
         node.map_err(|err| {
             err.with_context(if is_leaf {
@@ -187,7 +204,7 @@ impl Database for RocksDBWrapper {
         let Some(raw_root) = self.raw_node(&NodeKey::empty(version).to_db_key()) else {
             return Ok(None);
         };
-        Root::deserialize(&raw_root)
+        Root::deserialize(&raw_root, false)
             .map(Some)
             .map_err(|err| err.with_context(ErrorContext::Root(version)))
     }

core/lib/merkle_tree/src/storage/serialization.rs

@@ -5,7 +5,7 @@ use std::{collections::HashMap, str};
 use crate::{
     errors::{DeserializeError, DeserializeErrorKind, ErrorContext},
     types::{
-        ChildRef, InternalNode, Key, LeafNode, Manifest, Node, Root, TreeTags, ValueHash,
+        ChildRef, InternalNode, Key, LeafNode, Manifest, Node, RawNode, Root, TreeTags, ValueHash,
         HASH_SIZE, KEY_SIZE,
     },
 };
@@ -15,7 +15,7 @@ use crate::{
 const LEB128_SIZE_ESTIMATE: usize = 3;
 
 impl LeafNode {
-    pub(super) fn deserialize(bytes: &[u8]) -> Result<Self, DeserializeError> {
+    pub(super) fn deserialize(bytes: &[u8], strict: bool) -> Result<Self, DeserializeError> {
         if bytes.len() < KEY_SIZE + HASH_SIZE {
             return Err(DeserializeErrorKind::UnexpectedEof.into());
         }
@@ -26,6 +26,10 @@ impl LeafNode {
         let leaf_index = leb128::read::unsigned(&mut bytes).map_err(|err| {
             DeserializeErrorKind::Leb128(err).with_context(ErrorContext::LeafIndex)
         })?;
+        if strict && !bytes.is_empty() {
+            return Err(DeserializeErrorKind::Leftovers.into());
+        }
+
         Ok(Self {
             full_key,
             value_hash,
@@ -105,7 +109,7 @@ impl ChildRef {
 }
 
 impl InternalNode {
-    pub(super) fn deserialize(bytes: &[u8]) -> Result<Self, DeserializeError> {
+    pub(super) fn deserialize(bytes: &[u8], strict: bool) -> Result<Self, DeserializeError> {
         if bytes.len() < 4 {
             let err = DeserializeErrorKind::UnexpectedEof;
             return Err(err.with_context(ErrorContext::ChildrenMask));
@@ -134,6 +138,9 @@ impl InternalNode {
             }
             bitmap >>= 2;
         }
+        if strict && !bytes.is_empty() {
+            return Err(DeserializeErrorKind::Leftovers.into());
+        }
         Ok(this)
     }
 
@@ -161,8 +168,36 @@ impl InternalNode {
     }
 }
 
+impl RawNode {
+    pub(crate) fn deserialize(bytes: &[u8]) -> Self {
+        Self {
+            raw: bytes.to_vec(),
+            leaf: LeafNode::deserialize(bytes, true).ok(),
+            internal: InternalNode::deserialize(bytes, true).ok(),
+        }
+    }
+
+    pub(crate) fn deserialize_root(bytes: &[u8]) -> Self {
+        let root = Root::deserialize(bytes, true).ok();
+        let node = root.and_then(|root| match root {
+            Root::Empty => None,
+            Root::Filled { node, .. } => Some(node),
+        });
+        let (leaf, internal) = match node {
+            None => (None, None),
+            Some(Node::Leaf(leaf)) => (Some(leaf), None),
+            Some(Node::Internal(node)) => (None, Some(node)),
+        };
+        Self {
+            raw: bytes.to_vec(),
+            leaf,
+            internal,
+        }
+    }
+}
+
 impl Root {
-    pub(super) fn deserialize(mut bytes: &[u8]) -> Result<Self, DeserializeError> {
+    pub(super) fn deserialize(mut bytes: &[u8], strict: bool) -> Result<Self, DeserializeError> {
         let leaf_count = leb128::read::unsigned(&mut bytes).map_err(|err| {
             DeserializeErrorKind::Leb128(err).with_context(ErrorContext::LeafCount)
         })?;
@@ -172,11 +207,11 @@ impl Root {
                 // Try both the leaf and internal node serialization; in some cases, a single leaf
                 // may still be persisted as an internal node. Since serialization of an internal node with a single child
                 // is always shorter than that a leaf, the order (first leaf, then internal node) is chosen intentionally.
-                LeafNode::deserialize(bytes)
+                LeafNode::deserialize(bytes, strict)
                     .map(Node::Leaf)
-                    .or_else(|_| InternalNode::deserialize(bytes).map(Node::Internal))?
+                    .or_else(|_| InternalNode::deserialize(bytes, strict).map(Node::Internal))?
             }
-            _ => Node::Internal(InternalNode::deserialize(bytes)?),
+            _ => Node::Internal(InternalNode::deserialize(bytes, strict)?),
         };
         Ok(Self::new(leaf_count, node))
     }
@@ -440,7 +475,7 @@ mod tests {
         assert_eq!(buffer[64], 42); // leaf index
         assert_eq!(buffer.len(), 65);
 
-        let leaf_copy = LeafNode::deserialize(&buffer).unwrap();
+        let leaf_copy = LeafNode::deserialize(&buffer, true).unwrap();
         assert_eq!(leaf_copy, leaf);
     }
 
@@ -471,7 +506,7 @@ mod tests {
         let child_count = bitmap.count_ones();
         assert_eq!(child_count, 2);
 
-        let node_copy = InternalNode::deserialize(&buffer).unwrap();
+        let node_copy = InternalNode::deserialize(&buffer, true).unwrap();
         assert_eq!(node_copy, node);
     }
 
@@ -482,7 +517,7 @@ mod tests {
         root.serialize(&mut buffer);
         assert_eq!(buffer, [0]);
 
-        let root_copy = Root::deserialize(&buffer).unwrap();
+        let root_copy = Root::deserialize(&buffer, true).unwrap();
         assert_eq!(root_copy, root);
     }
 
@@ -494,7 +529,7 @@ mod tests {
         root.serialize(&mut buffer);
         assert_eq!(buffer[0], 1);
 
-        let root_copy = Root::deserialize(&buffer).unwrap();
+        let root_copy = Root::deserialize(&buffer, true).unwrap();
         assert_eq!(root_copy, root);
     }
 
@@ -506,7 +541,7 @@ mod tests {
         root.serialize(&mut buffer);
         assert_eq!(buffer[0], 2);
 
-        let root_copy = Root::deserialize(&buffer).unwrap();
+        let root_copy = Root::deserialize(&buffer, true).unwrap();
         assert_eq!(root_copy, root);
     }
 }