New Storage Write API REDO

.changelog/unreleased/improvements/2438-new-storage-write.md

@@ -0,0 +1,2 @@
+- Skip writing some MASP and IBC storage keys to merkle tree and DB diffs.
+  ([\#2438](https://github.com/anoma/namada/pull/2438))

crates/apps/src/lib/node/ledger/shell/finalize_block.rs

@@ -2850,7 +2850,7 @@ mod test_finalize_block {
 
     /// Test that replay protection keys are not added to the merkle tree
     #[test]
-    fn test_replay_keys_not_merkelized() {
+    fn test_replay_keys_not_merklized() {
         let (mut shell, _, _, _) = setup();
 
         let (wrapper_tx, processed_tx) =

crates/apps/src/lib/node/ledger/shell/init_chain.rs

@@ -215,9 +215,9 @@ where
             self.update_eth_oracle(&Default::default());
         } else {
             self.wl_storage
-                .write_bytes(
+                .write(
                     &namada::eth_bridge::storage::active_key(),
-                    EthBridgeStatus::Disabled.serialize_to_vec(),
+                    EthBridgeStatus::Disabled,
                 )
                 .unwrap();
         }

crates/apps/src/lib/node/ledger/shell/mod.rs

@@ -363,6 +363,13 @@ where
     event_log: EventLog,
 }
 
+/// Merkle tree storage key filter. Return `false` for keys that shouldn't be
+/// merklized.
+pub fn is_merklized_storage_key(key: &namada_sdk::types::storage::Key) -> bool {
+    !token::storage_key::is_masp_key(key)
+        && !namada::ibc::storage::is_ibc_counter_key(key)
+}
+
 /// Channels for communicating with an Ethereum oracle.
 #[derive(Debug)]
 pub struct EthereumOracleChannels {
@@ -431,6 +438,7 @@ where
             native_token,
             db_cache,
             config.shell.storage_read_past_height_limit,
+            is_merklized_storage_key,
         );
         storage
             .load_last_state()
@@ -2041,10 +2049,7 @@ mod test_utils {
         use namada::eth_bridge::storage::eth_bridge_queries::EthBridgeStatus;
         shell
             .wl_storage
-            .write_bytes(
-                &active_key(),
-                EthBridgeStatus::Disabled.serialize_to_vec(),
-            )
+            .write(&active_key(), EthBridgeStatus::Disabled)
             .expect("Test failed");
     }
 

crates/apps/src/lib/node/ledger/storage/mod.rs

@@ -52,14 +52,17 @@ fn new_blake2b() -> Blake2b {
 mod tests {
     use std::collections::HashMap;
 
+    use borsh::BorshDeserialize;
     use itertools::Itertools;
     use namada::eth_bridge::storage::proof::BridgePoolRootProof;
     use namada::ledger::eth_bridge::storage::bridge_pool;
     use namada::ledger::gas::STORAGE_ACCESS_GAS_PER_BYTE;
     use namada::ledger::ibc::storage::ibc_key;
     use namada::ledger::parameters::{EpochDuration, Parameters};
     use namada::state::write_log::WriteLog;
-    use namada::state::{self, StorageWrite, StoreType, WlStorage};
+    use namada::state::{
+        self, StorageRead, StorageWrite, StoreType, WlStorage, DB,
+    };
     use namada::token::conversion::update_allowed_conversions;
     use namada::types::chain::ChainId;
     use namada::types::ethereum_events::Uint;
@@ -75,6 +78,7 @@ mod tests {
     use tempfile::TempDir;
 
     use super::*;
+    use crate::node::ledger::shell::is_merklized_storage_key;
 
     #[test]
     fn test_crud_value() {
@@ -86,6 +90,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
         let key = Key::parse("key").expect("cannot parse the key string");
         let value: u64 = 1;
@@ -138,6 +143,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
         storage
             .begin_block(BlockHash::default(), BlockHeight(100))
@@ -239,6 +245,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
         storage
             .load_last_state()
@@ -263,6 +270,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
         storage
             .begin_block(BlockHash::default(), BlockHeight(100))
@@ -309,6 +317,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
         storage
             .begin_block(BlockHash::default(), BlockHeight(100))
@@ -376,6 +385,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
 
         // 1. For each `blocks_write_value`, write the current block height if
@@ -469,6 +479,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
 
         let num_keys = 5;
@@ -587,6 +598,7 @@ mod tests {
             address::nam(),
             None,
             Some(5),
+            is_merklized_storage_key,
         );
         let new_epoch_start = BlockHeight(1);
         let signed_root_key = bridge_pool::get_signed_root_key();
@@ -692,6 +704,7 @@ mod tests {
             address::nam(),
             None,
             None,
+            is_merklized_storage_key,
         );
         let mut storage = WlStorage {
             storage,
@@ -783,4 +796,204 @@ mod tests {
             .map(Result::unwrap);
         itertools::assert_equal(iter, expected);
     }
+
+    fn test_key_1() -> Key {
+        Key::parse("testing1").unwrap()
+    }
+
+    fn test_key_2() -> Key {
+        Key::parse("testing2").unwrap()
+    }
+
+    fn merkle_tree_key_filter(key: &Key) -> bool {
+        key == &test_key_1()
+    }
+
+    #[test]
+    fn test_persistent_storage_writing_without_merklizing_or_diffs() {
+        let db_path =
+            TempDir::new().expect("Unable to create a temporary DB directory");
+        let storage = PersistentStorage::open(
+            db_path.path(),
+            ChainId::default(),
+            address::nam(),
+            None,
+            None,
+            merkle_tree_key_filter,
+        );
+        let mut wls = WlStorage {
+            storage,
+            write_log: Default::default(),
+        };
+        // Start the first block
+        let first_height = BlockHeight::first();
+        wls.storage.block.height = first_height;
+
+        let key1 = test_key_1();
+        let val1 = 1u64;
+        let key2 = test_key_2();
+        let val2 = 2u64;
+
+        // Standard write of key-val-1
+        wls.write(&key1, val1).unwrap();
+
+        // Read from WlStorage should return val1
+        let res = wls.read::<u64>(&key1).unwrap().unwrap();
+        assert_eq!(res, val1);
+
+        // Read from Storage shouldn't return val1 because the block hasn't been
+        // committed
+        let (res, _) = wls.storage.read(&key1).unwrap();
+        assert!(res.is_none());
+
+        // Write key-val-2 without merklizing or diffs
+        wls.write(&key2, val2).unwrap();
+
+        // Read from WlStorage should return val2
+        let res = wls.read::<u64>(&key2).unwrap().unwrap();
+        assert_eq!(res, val2);
+
+        // Commit block and storage changes
+        wls.commit_block().unwrap();
+        wls.storage.block.height = wls.storage.block.height.next_height();
+        let second_height = wls.storage.block.height;
+
+        // Read key1 from Storage should return val1
+        let (res1, _) = wls.storage.read(&key1).unwrap();
+        let res1 = u64::try_from_slice(&res1.unwrap()).unwrap();
+        assert_eq!(res1, val1);
+
+        // Check merkle tree inclusion of key-val-1 explicitly
+        let is_merklized1 = wls.storage.block.tree.has_key(&key1).unwrap();
+        assert!(is_merklized1);
+
+        // Key2 should be in storage. Confirm by reading from
+        // WlStorage and also by reading Storage subspace directly
+        let res2 = wls.read::<u64>(&key2).unwrap().unwrap();
+        assert_eq!(res2, val2);
+        let res2 = wls.storage.db.read_subspace_val(&key2).unwrap().unwrap();
+        let res2 = u64::try_from_slice(&res2).unwrap();
+        assert_eq!(res2, val2);
+
+        // Check explicitly that key-val-2 is not in merkle tree
+        let is_merklized2 = wls.storage.block.tree.has_key(&key2).unwrap();
+        assert!(!is_merklized2);
+
+        // Check that the proper diffs exist for key-val-1
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, true)
+            .unwrap();
+        assert!(res1.is_none());
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, false)
+            .unwrap()
+            .unwrap();
+        let res1 = u64::try_from_slice(&res1).unwrap();
+        assert_eq!(res1, val1);
+
+        // Check that there are diffs for key-val-2 in block 0, since all keys
+        // need to have diffs for at least 1 block for rollback purposes
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, true)
+            .unwrap();
+        assert!(res2.is_none());
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, false)
+            .unwrap()
+            .unwrap();
+        let res2 = u64::try_from_slice(&res2).unwrap();
+        assert_eq!(res2, val2);
+
+        // Delete the data then commit the block
+        wls.delete(&key1).unwrap();
+        wls.delete(&key2).unwrap();
+        wls.commit_block().unwrap();
+        wls.storage.block.height = wls.storage.block.height.next_height();
+
+        // Check the key-vals are removed from the storage subspace
+        let res1 = wls.read::<u64>(&key1).unwrap();
+        let res2 = wls.read::<u64>(&key2).unwrap();
+        assert!(res1.is_none() && res2.is_none());
+        let res1 = wls.storage.db.read_subspace_val(&key1).unwrap();
+        let res2 = wls.storage.db.read_subspace_val(&key2).unwrap();
+        assert!(res1.is_none() && res2.is_none());
+
+        // Check that the key-vals don't exist in the merkle tree anymore
+        let is_merklized1 = wls.storage.block.tree.has_key(&key1).unwrap();
+        let is_merklized2 = wls.storage.block.tree.has_key(&key2).unwrap();
+        assert!(!is_merklized1 && !is_merklized2);
+
+        // Check that key-val-1 diffs are properly updated for blocks 0 and 1
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, true)
+            .unwrap();
+        assert!(res1.is_none());
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, first_height, false)
+            .unwrap()
+            .unwrap();
+        let res1 = u64::try_from_slice(&res1).unwrap();
+        assert_eq!(res1, val1);
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, second_height, true)
+            .unwrap()
+            .unwrap();
+        let res1 = u64::try_from_slice(&res1).unwrap();
+        assert_eq!(res1, val1);
+
+        let res1 = wls
+            .storage
+            .db
+            .read_diffs_val(&key1, second_height, false)
+            .unwrap();
+        assert!(res1.is_none());
+
+        // Check that key-val-2 diffs don't exist for block 0 anymore
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, true)
+            .unwrap();
+        assert!(res2.is_none());
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, first_height, false)
+            .unwrap();
+        assert!(res2.is_none());
+
+        // Check that the block 1 diffs for key-val-2 include an "old" value of
+        // val2 and no "new" value
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, second_height, true)
+            .unwrap()
+            .unwrap();
+        let res2 = u64::try_from_slice(&res2).unwrap();
+        assert_eq!(res2, val2);
+        let res2 = wls
+            .storage
+            .db
+            .read_diffs_val(&key2, second_height, false)
+            .unwrap();
+        assert!(res2.is_none());
+    }
 }