Merge pull request #4043 from anoma/yuji/prune-more-subtree-snapshots

Prune more Merkle tree snapshots

.changelog/unreleased/improvements/4043-prune-more-subtree-snapshots.md

@@ -0,0 +1,2 @@
+- Prune old Merkle tree snapshots which are saved every block
+  ([\#4043](https://github.com/anoma/namada/issues/4043))

crates/node/src/storage/mod.rs

@@ -76,6 +76,7 @@ mod tests {
     use namada_sdk::{
         address, decode, encode, parameters, storage, token, validation,
     };
+    use namada_vp::state::MerkleTree;
     use proptest::collection::vec;
     use proptest::prelude::*;
     use proptest::test_runner::Config;
@@ -533,7 +534,8 @@ mod tests {
             }
         }
         // save the last root
-        roots.insert(state.in_mem().block.height, state.in_mem().merkle_root());
+        let last_height = state.in_mem().block.height;
+        roots.insert(last_height, state.in_mem().merkle_root());
         state.commit_block_from_batch(batch)?;
 
         let mut current_state = HashMap::new();
@@ -577,34 +579,53 @@ mod tests {
             let key = make_key(i);
             current_state.insert(key, true);
         }
+        // NoDiff subtree can be restored for the last 2 blocks
+        let prev_last_height = last_height.prev_height().unwrap();
         for (height, key, write_type) in blocks_write_type {
             if key != client_counter_key() {
                 continue;
             }
             let merkle_key =
                 Key::from(NO_DIFF_KEY_PREFIX.to_string().to_db_key())
                     .join(&key);
-            let tree =
-                state.get_merkle_tree(height, Some(StoreType::NoDiff))?;
-            // Check if the rebuilt tree's root is the same as the saved one
-            assert_eq!(tree.root().0, roots.get(&height).unwrap().0);
+            let tree = match state
+                .get_merkle_tree(height, Some(StoreType::NoDiff))
+            {
+                Ok(tree) if height >= prev_last_height => {
+                    // Check if the rebuilt tree's root is the same as the saved
+                    // one
+                    assert_eq!(tree.root().0, roots.get(&height).unwrap().0);
+                    tree
+                }
+                Ok(_) => panic!("The tree at the height should be pruned"),
+                // expected error because of no merkle subtree store
+                // set an empty merkle tree as a dummy
+                Err(_) => MerkleTree::<PersistentStorageHasher>::default(),
+            };
+            // Check the NoDiff subtree at the last height
             match write_type {
                 0 => {
                     // data was not updated
-                    if *current_state.get(&key).unwrap() {
-                        assert!(tree.has_key(&merkle_key)?);
-                    } else {
-                        assert!(!tree.has_key(&merkle_key)?);
+                    if height >= prev_last_height {
+                        if *current_state.get(&key).unwrap() {
+                            assert!(tree.has_key(&merkle_key)?);
+                        } else {
+                            assert!(!tree.has_key(&merkle_key)?);
+                        }
                     }
                 }
                 1 | 3 => {
                     // data was deleted
-                    assert!(!tree.has_key(&merkle_key)?);
+                    if height >= prev_last_height {
+                        assert!(!tree.has_key(&merkle_key)?);
+                    }
                     current_state.insert(key, false);
                 }
                 _ => {
                     // data was updated
-                    assert!(tree.has_key(&merkle_key)?);
+                    if height >= prev_last_height {
+                        assert!(tree.has_key(&merkle_key)?);
+                    }
                     current_state.insert(key, true);
                 }
             }
@@ -733,6 +754,11 @@ mod tests {
         let result =
             state.get_merkle_tree(6.into(), Some(StoreType::BridgePool));
         assert!(result.is_err(), "The bridge pool tree should be pruned");
+
+        let result = state.get_merkle_tree(10.into(), Some(StoreType::NoDiff));
+        assert!(result.is_err(), "The tree at Height 10 should be pruned");
+        let result = state.get_merkle_tree(11.into(), Some(StoreType::NoDiff));
+        assert!(result.is_ok(), "The tree at Height 11 should be restored");
     }
 
     /// Test the prefix iterator with RocksDB.

crates/node/src/storage/rocksdb.rs

@@ -1410,7 +1410,8 @@ impl DB for RocksDB {
                 format!("{key_prefix}/{MERKLE_TREE_ROOT_KEY_SEGMENT}");
             match self.read_value(block_cf, root_key)? {
                 Some(root) => merkle_tree_stores.set_root(st, root),
-                None => return Ok(None),
+                None if store_type.is_some() => return Ok(None),
+                _ => continue,
             }
 
             let store_key =
@@ -1419,7 +1420,8 @@ impl DB for RocksDB {
                 Some(bytes) => {
                     merkle_tree_stores.set_store(st.decode_store(bytes)?)
                 }
-                None => return Ok(None),
+                None if store_type.is_some() => return Ok(None),
+                _ => continue,
             }
         }
         Ok(Some(merkle_tree_stores))
@@ -1659,10 +1661,17 @@ impl DB for RocksDB {
         &mut self,
         batch: &mut Self::WriteBatch,
         store_type: &StoreType,
-        epoch: Epoch,
+        pruned_target: Either<BlockHeight, Epoch>,
     ) -> Result<()> {
         let block_cf = self.get_column_family(BLOCK_CF)?;
-        let key_prefix = tree_key_prefix_with_epoch(store_type, epoch);
+        let key_prefix = match pruned_target {
+            Either::Left(height) => {
+                tree_key_prefix_with_height(store_type, height)
+            }
+            Either::Right(epoch) => {
+                tree_key_prefix_with_epoch(store_type, epoch)
+            }
+        };
         let root_key = format!("{key_prefix}/{MERKLE_TREE_ROOT_KEY_SEGMENT}");
         batch.0.delete_cf(block_cf, root_key);
         let store_key = format!("{key_prefix}/{MERKLE_TREE_STORE_KEY_SEGMENT}");

crates/state/src/wl_state.rs

@@ -1,6 +1,7 @@
 use std::cmp::Ordering;
 use std::ops::{Deref, DerefMut};
 
+use itertools::Either;
 use namada_core::address::Address;
 use namada_core::arith::checked;
 use namada_core::borsh::BorshSerializeExt;
@@ -393,12 +394,41 @@ where
     // commit.
     fn prune_merkle_tree_stores(
         &mut self,
+        is_full_commit: bool,
         batch: &mut D::WriteBatch,
     ) -> Result<()> {
+        if let Some(prev_height) = self
+            .in_mem
+            .block
+            .height
+            .prev_height()
+            .and_then(|h| h.prev_height())
+        {
+            for st in StoreType::iter().filter(|st| st.is_stored_every_block())
+            {
+                match st {
+                    StoreType::Base => continue,
+                    _ => self.0.db.prune_merkle_tree_store(
+                        batch,
+                        st,
+                        Either::Left(prev_height),
+                    )?,
+                }
+            }
+        }
+
+        if !is_full_commit {
+            return Ok(());
+        }
+
         // Prune non-provable stores at the previous epoch
         if let Some(prev_epoch) = self.in_mem.block.epoch.prev() {
             for st in StoreType::iter_non_provable() {
-                self.0.db.prune_merkle_tree_store(batch, st, prev_epoch)?;
+                self.0.db.prune_merkle_tree_store(
+                    batch,
+                    st,
+                    Either::Right(prev_epoch),
+                )?;
             }
         }
         // Prune provable stores
@@ -411,7 +441,11 @@ where
                 self.db.prune_merkle_tree_store(
                     batch,
                     st,
-                    oldest_epoch.prev().unwrap(),
+                    Either::Right(
+                        oldest_epoch
+                            .prev()
+                            .expect("the previous epoch should exist"),
+                    ),
                 )?;
             }
 
@@ -425,7 +459,7 @@ where
                 self.db.prune_merkle_tree_store(
                     batch,
                     &StoreType::BridgePool,
-                    epoch,
+                    Either::Right(epoch),
                 )?;
             }
         }
@@ -637,10 +671,8 @@ where
             time: header.time,
         });
         self.in_mem.last_epoch = self.in_mem.block.epoch;
-        if is_full_commit {
-            // prune old merkle tree stores
-            self.prune_merkle_tree_stores(&mut batch)?;
-        }
+        // prune old merkle tree stores
+        self.prune_merkle_tree_stores(is_full_commit, &mut batch)?;
         // If there's a previous block, prune non-persisted diffs from it
         if let Some(height) = self.in_mem.block.height.prev_height() {
             self.db.prune_non_persisted_diffs(&mut batch, height)?;
@@ -996,15 +1028,14 @@ where
                 None => BlockHeight(1),
             },
         };
+        // Try to read all subtrees, but some subtrees would be stale or empty.
+        // They will be rebuild later. That's why the tree isn't validated here.
         let stores = self
             .db
             .read_merkle_tree_stores(epoch, start_height, store_type)?
             .ok_or(StateError::NoMerkleTree { height })?;
+        let mut tree = MerkleTree::<H>::new_partial(stores);
         let prefix = store_type.and_then(|st| st.provable_prefix());
-        let mut tree = match store_type {
-            Some(_) => MerkleTree::<H>::new_partial(stores),
-            None => MerkleTree::<H>::new(stores).expect("invalid stores"),
-        };
         // Restore the tree state with diffs
         let mut target_height = start_height;
         while target_height < height {

crates/storage/src/db.rs

@@ -1,6 +1,7 @@
 use std::fmt::Debug;
 use std::num::TryFromIntError;
 
+use itertools::Either;
 use namada_core::address::EstablishedAddressGen;
 use namada_core::chain::{BlockHeader, BlockHeight, Epoch, Epochs};
 use namada_core::hash::{Error as HashError, Hash};
@@ -169,7 +170,7 @@ pub trait DB: Debug {
 
     /// Read the merkle tree stores with the given epoch. If a store_type is
     /// given, it reads only the specified tree. Otherwise, it reads all
-    /// trees.
+    /// trees, but some subtrees could be empty if the stores aren't saved.
     fn read_merkle_tree_stores(
         &self,
         epoch: Epoch,
@@ -258,7 +259,7 @@ pub trait DB: Debug {
         &mut self,
         batch: &mut Self::WriteBatch,
         store_type: &StoreType,
-        pruned_epoch: Epoch,
+        pruned_target: Either<BlockHeight, Epoch>,
     ) -> Result<()>;
 
     /// Read the signed nonce of Bridge Pool

crates/storage/src/mockdb.rs

@@ -322,16 +322,16 @@ impl DB for MockDB {
                 format!("{key_prefix}/{MERKLE_TREE_ROOT_KEY_SEGMENT}");
             match self.read_value(root_key)? {
                 Some(root) => merkle_tree_stores.set_root(st, root),
-                None => return Ok(None),
+                None if store_type.is_some() => return Ok(None),
+                _ => continue,
             }
             let store_key =
                 format!("{key_prefix}/{MERKLE_TREE_STORE_KEY_SEGMENT}");
             let bytes = self.0.borrow().get(&store_key.to_string()).cloned();
             match bytes {
-                Some(b) => {
-                    merkle_tree_stores.set_store(st.decode_store(b)?);
-                }
-                None => return Ok(None),
+                Some(b) => merkle_tree_stores.set_store(st.decode_store(b)?),
+                None if store_type.is_some() => return Ok(None),
+                _ => continue,
             }
         }
         Ok(Some(merkle_tree_stores))
@@ -547,9 +547,16 @@ impl DB for MockDB {
         &mut self,
         _batch: &mut Self::WriteBatch,
         store_type: &StoreType,
-        epoch: Epoch,
+        pruned_target: Either<BlockHeight, Epoch>,
     ) -> Result<()> {
-        let key_prefix = tree_key_prefix_with_epoch(store_type, epoch);
+        let key_prefix = match pruned_target {
+            Either::Left(height) => {
+                tree_key_prefix_with_height(store_type, height)
+            }
+            Either::Right(epoch) => {
+                tree_key_prefix_with_epoch(store_type, epoch)
+            }
+        };
         let root_key = format!("{key_prefix}/{MERKLE_TREE_ROOT_KEY_SEGMENT}");
         self.0.borrow_mut().remove(&root_key);
         let store_key = format!("{key_prefix}/{MERKLE_TREE_STORE_KEY_SEGMENT}");