feat(tree): Improved tree pruning (#1532)

Changes in tree pruner to allow syncing pruning state with db pruning
(pruning tree up to the same l1 batch as db)

---------

Signed-off-by: tomg10 <[email protected]>
Co-authored-by: Alex Ostrovski <[email protected]>
Co-authored-by: Alex Ostrovski <[email protected]>

core/bin/external_node/src/main.rs

@@ -160,10 +160,10 @@ async fn run_tree(
     .await
     .context("failed creating DB pool for Merkle tree recovery")?;
 
-    let metadata_calculator =
-        MetadataCalculator::new(metadata_calculator_config, None, tree_pool, recovery_pool)
-            .await
-            .context("failed initializing metadata calculator")?;
+    let metadata_calculator = MetadataCalculator::new(metadata_calculator_config, None, tree_pool)
+        .await
+        .context("failed initializing metadata calculator")?
+        .with_recovery_pool(recovery_pool);
 
     let tree_reader = Arc::new(metadata_calculator.tree_reader());
     app_health.insert_component(metadata_calculator.tree_health_check());

core/lib/merkle_tree/examples/loadtest/main.rs

@@ -107,7 +107,7 @@ impl Cli {
             }
 
             if self.prune {
-                let (mut pruner, pruner_handle) = MerkleTreePruner::new(rocksdb.clone(), 0);
+                let (mut pruner, pruner_handle) = MerkleTreePruner::new(rocksdb.clone());
                 pruner.set_poll_interval(Duration::from_secs(10));
                 let pruner_thread = thread::spawn(|| pruner.run());
                 pruner_handles = Some((pruner_handle, pruner_thread));
@@ -160,6 +160,16 @@ impl Cli {
                 let output = tree.extend(kvs.collect());
                 output.root_hash
             };
+
+            if let Some((pruner_handle, _)) = &pruner_handles {
+                if pruner_handle.set_target_retained_version(version).is_err() {
+                    tracing::error!("Pruner unexpectedly stopped");
+                    let (_, pruner_thread) = pruner_handles.unwrap();
+                    pruner_thread.join().expect("Pruner panicked");
+                    return; // unreachable
+                }
+            }
+
             let elapsed = start.elapsed();
             tracing::info!("Processed block #{version} in {elapsed:?}, root hash = {root_hash:?}");
         }
@@ -172,8 +182,8 @@ impl Cli {
         tracing::info!("Verified tree consistency in {elapsed:?}");
 
         if let Some((pruner_handle, pruner_thread)) = pruner_handles {
-            pruner_handle.abort();
-            pruner_thread.join().unwrap();
+            drop(pruner_handle);
+            pruner_thread.join().expect("Pruner panicked");
         }
     }
 

core/lib/merkle_tree/src/domain.rs

@@ -6,12 +6,13 @@ use zksync_prover_interface::inputs::{PrepareBasicCircuitsJob, StorageLogMetadat
 use zksync_types::{L1BatchNumber, StorageKey};
 
 use crate::{
+    consistency::ConsistencyError,
     storage::{PatchSet, Patched, RocksDBWrapper},
     types::{
         Key, Root, TreeEntry, TreeEntryWithProof, TreeInstruction, TreeLogEntry, ValueHash,
         TREE_DEPTH,
     },
-    BlockOutput, HashTree, MerkleTree, NoVersionError,
+    BlockOutput, HashTree, MerkleTree, MerkleTreePruner, MerkleTreePrunerHandle, NoVersionError,
 };
 
 /// Metadata for the current tree state.
@@ -42,6 +43,7 @@ pub struct ZkSyncTree {
     tree: MerkleTree<Patched<RocksDBWrapper>>,
     thread_pool: Option<ThreadPool>,
     mode: TreeMode,
+    pruning_enabled: bool,
 }
 
 impl ZkSyncTree {
@@ -93,9 +95,26 @@ impl ZkSyncTree {
             tree: MerkleTree::new(Patched::new(db)),
             thread_pool: None,
             mode,
+            pruning_enabled: false,
         }
     }
 
+    /// Returns tree pruner and a handle to stop it.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this method was already called for the tree instance; it's logically unsound to run
+    /// multiple pruners for the same tree concurrently.
+    pub fn pruner(&mut self) -> (MerkleTreePruner<RocksDBWrapper>, MerkleTreePrunerHandle) {
+        assert!(
+            !self.pruning_enabled,
+            "pruner was already obtained for the tree"
+        );
+        self.pruning_enabled = true;
+        let db = self.tree.db.inner().clone();
+        MerkleTreePruner::new(db)
+    }
+
     /// Returns a readonly handle to the tree. The handle **does not** see uncommitted changes to the tree,
     /// only ones flushed to RocksDB.
     pub fn reader(&self) -> ZkSyncTreeReader {
@@ -360,6 +379,14 @@ impl ZkSyncTreeReader {
         L1BatchNumber(number)
     }
 
+    /// Returns the minimum L1 batch number retained by the tree.
+    #[allow(clippy::missing_panics_doc)]
+    pub fn min_l1_batch_number(&self) -> Option<L1BatchNumber> {
+        self.0.first_retained_version().map(|version| {
+            L1BatchNumber(u32::try_from(version).expect("integer overflow for L1 batch number"))
+        })
+    }
+
     /// Returns the number of leaves in the tree.
     pub fn leaf_count(&self) -> u64 {
         self.0.latest_root().leaf_count()
@@ -379,4 +406,17 @@ impl ZkSyncTreeReader {
         let version = u64::from(l1_batch_number.0);
         self.0.entries_with_proofs(version, keys)
     }
+
+    /// Verifies consistency of the tree at the specified L1 batch number.
+    ///
+    /// # Errors
+    ///
+    /// Returns the first encountered verification error, should one occur.
+    pub fn verify_consistency(
+        &self,
+        l1_batch_number: L1BatchNumber,
+    ) -> Result<(), ConsistencyError> {
+        let version = l1_batch_number.0.into();
+        self.0.verify_consistency(version, true)
+    }
 }

core/lib/merkle_tree/src/lib.rs

@@ -241,6 +241,19 @@ impl<DB: Database, H: HashTree> MerkleTree<DB, H> {
     }
 }
 
+impl<DB: PruneDatabase> MerkleTree<DB> {
+    /// Returns the first retained version of the tree.
+    pub fn first_retained_version(&self) -> Option<u64> {
+        match self.db.min_stale_key_version() {
+            // Min stale key version is next after the first retained version since at least
+            // the root is updated on each version.
+            Some(version) => version.checked_sub(1),
+            // No stale keys means all past versions of the tree have been pruned
+            None => self.latest_version(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

core/lib/merkle_tree/src/metrics.rs

@@ -334,7 +334,7 @@ pub struct PruningStats {
 }
 
 impl PruningStats {
-    pub fn report(self) {
+    pub fn report(&self) {
         PRUNING_METRICS
             .target_retained_version
             .set(self.target_retained_version);