optimize contract storage read method

Author: Pana

crates/client/src/state_dump.rs

@@ -4,14 +4,13 @@ use cfx_rpc_eth_types::{AccountState, StateDump, EOA_STORAGE_ROOT_H256};
 use cfx_rpc_primitives::Bytes;
 use cfx_statedb::{StateDbExt, StateDbGeneric};
 use cfx_storage::state_manager::StateManagerTrait;
-use cfx_types::{Address, Space, H256};
+use cfx_types::{Address, Space, H256, U256};
 use cfxcore::NodeType;
 use chrono::Utc;
 use keccak_hash::{keccak, KECCAK_EMPTY};
 use parking_lot::{Condvar, Mutex};
 use primitives::{
-    Account, SkipInputCheck, SpaceStorageFilter, StorageKey,
-    StorageKeyWithSpace, StorageValue,
+    Account, SkipInputCheck, StorageKey, StorageKeyWithSpace, StorageValue,
 };
 use rlp::Rlp;
 use std::{
@@ -260,7 +259,7 @@ pub fn export_space_accounts_with_callback<F: Fn(AccountState)>(
     let mut core_space_key_count: u64 = 0;
     let mut total_key_count: u64 = 0;
 
-    for i in 0..=255 {
+    for i in 198..=255 {
         let prefix = [i];
         let start_key = StorageKey::AddressPrefixKey(&prefix).with_space(space);
 
@@ -295,18 +294,15 @@ pub fn export_space_accounts_with_callback<F: Fn(AccountState)>(
             }
         };
 
-        state.read_all_with_callback(
-            start_key,
-            &mut inner_callback,
-            Some(SpaceStorageFilter::from(space)),
-        )?;
+        state.read_all_with_callback(start_key, &mut inner_callback, true)?;
 
         if account_states.len() > 0 {
             println("Start to read account code and storage data...");
         }
 
         for (_address, account) in account_states {
-            let account_state = get_account_state(state, &account, config)?;
+            let account_state =
+                get_account_state(state, &account, config, space)?;
             callback(account_state);
             found_accounts += 1;
             if config.limit > 0 && found_accounts >= config.limit as usize {
@@ -321,6 +317,7 @@ pub fn export_space_accounts_with_callback<F: Fn(AccountState)>(
 #[allow(unused)]
 fn get_account_state(
     state: &mut StateDbGeneric, account: &Account, config: &StateDumpConfig,
+    space: Space,
 ) -> Result<AccountState, Box<dyn std::error::Error>> {
     let address = account.address();
 
@@ -335,7 +332,7 @@ fn get_account_state(
     };
 
     let storage = if is_contract && !config.no_storage {
-        let storage = state.get_account_storage_entries(&address, None)?;
+        let storage = get_contract_storage(state, &address.address, space)?;
         Some(storage)
     } else {
         None
@@ -358,6 +355,36 @@ fn get_account_state(
     })
 }
 
+fn get_contract_storage(
+    state: &mut StateDbGeneric, address: &Address, space: Space,
+) -> Result<BTreeMap<H256, U256>, Box<dyn std::error::Error>> {
+    let mut storage: BTreeMap<H256, U256> = Default::default();
+
+    let mut inner_callback = |(key, value): (Vec<u8>, Box<[u8]>)| {
+        let storage_key_with_space =
+            StorageKeyWithSpace::from_key_bytes::<SkipInputCheck>(&key);
+        if storage_key_with_space.space != space {
+            return;
+        }
+
+        if let StorageKey::StorageKey {
+            address_bytes: _,
+            storage_key,
+        } = storage_key_with_space.key
+        {
+            let h256_storage_key = H256::from_slice(storage_key);
+            let storage_value_with_owner: StorageValue =
+                rlp::decode(&value).expect("Failed to decode storage value");
+            storage.insert(h256_storage_key, storage_value_with_owner.value);
+        };
+    };
+
+    let start_key = StorageKey::new_storage_root_key(address).with_space(space);
+    state.read_all_with_callback(start_key, &mut inner_callback, false)?;
+
+    Ok(storage)
+}
+
 fn println(message: &str) {
     println!("[{}] {}", Utc::now().format("%Y-%m-%d %H:%M:%S"), message);
 }

crates/dbs/statedb/src/lib.rs

@@ -12,7 +12,7 @@ pub mod global_params;
 mod in_memory_storage;
 mod statedb_ext;
 use cfx_types::H256;
-use primitives::{SpaceStorageFilter, StorageValue};
+use primitives::StorageValue;
 
 use cfx_db_errors::statedb as error;
 
@@ -190,14 +190,13 @@ mod impls {
 
         pub fn read_all_with_callback(
             &mut self, access_key_prefix: StorageKeyWithSpace,
-            callback: &mut dyn FnMut(MptKeyValue),
-            space_storage_filter: Option<SpaceStorageFilter>,
+            callback: &mut dyn FnMut(MptKeyValue), only_account_key: bool,
         ) -> Result<()> {
             self.storage
                 .read_all_with_callback(
                     access_key_prefix,
                     callback,
-                    space_storage_filter,
+                    only_account_key,
                 )
                 .map_err(|err| err.into())
         }

crates/dbs/storage/src/impls/delta_mpt/cow_node_ref.rs

@@ -622,16 +622,19 @@ impl CowNodeRef {
         Ok(())
     }
 
-    // space_storage_filter can be used to filter the wanted space storage keys
+    // only_account_key can be used to filter the only account key
     pub fn iterate_internal_with_callback(
         &self, owned_node_set: &OwnedNodeSet, trie: &DeltaMpt,
         guarded_trie_node: GuardedMaybeOwnedTrieNodeAsCowCallParam,
         key_prefix: CompressedPathRaw, db: &mut DeltaDbOwnedReadTraitObj,
         callback: &mut dyn FnMut(MptKeyValue), is_delta_mpt: bool,
-        space_storage_filter: Option<SpaceStorageFilter>,
+        only_account_key: bool,
     ) -> Result<()> {
-        if let Some(filter) = space_storage_filter {
-            if filter.is_filtered(is_delta_mpt, key_prefix.path_slice()) {
+        // filter out all the key that is longer than the account key
+        if only_account_key {
+            let key_len =
+                SpaceStorageFilter::space_flag_index(is_delta_mpt) + 1;
+            if key_prefix.path_slice().len() > key_len {
                 return Ok(());
             }
         }
@@ -675,7 +678,7 @@ impl CowNodeRef {
                 db,
                 callback,
                 is_delta_mpt,
-                space_storage_filter,
+                only_account_key,
             )?;
         }
 

crates/dbs/storage/src/impls/delta_mpt/subtrie_visitor.rs

@@ -637,7 +637,7 @@ impl<'trie, 'db: 'trie> SubTrieVisitor<'trie, 'db> {
     pub fn traversal_with_callback(
         mut self, key: KeyPart, key_remaining: KeyPart,
         callback: &mut dyn FnMut(MptKeyValue), is_delta_mpt: bool,
-        space_storage_filter: Option<SpaceStorageFilter>,
+        only_account_key: bool,
     ) -> Result<()> {
         let node_memory_manager = self.node_memory_manager();
         let allocator = node_memory_manager.get_allocator();
@@ -684,7 +684,7 @@ impl<'trie, 'db: 'trie> SubTrieVisitor<'trie, 'db> {
                         key_remaining,
                         callback,
                         is_delta_mpt,
-                        space_storage_filter,
+                        only_account_key,
                     )?;
                 return Ok(());
             }
@@ -699,7 +699,7 @@ impl<'trie, 'db: 'trie> SubTrieVisitor<'trie, 'db> {
             &mut *self.db.get_mut().to_owned_read()?,
             callback,
             is_delta_mpt,
-            space_storage_filter,
+            only_account_key,
         )?;
         Ok(())
     }
@@ -928,5 +928,5 @@ use super::{
     ChildrenTableDeltaMpt, DeltaMpt, *,
 };
 use parking_lot::MutexGuard;
-use primitives::{MerkleHash, MptValue, SpaceStorageFilter, MERKLE_NULL_NODE};
+use primitives::{MerkleHash, MptValue, MERKLE_NULL_NODE};
 use std::marker::PhantomData;

crates/dbs/storage/src/impls/recording_storage.rs

@@ -54,7 +54,7 @@ impl<Storage: StateTrait + StateTraitExt> StateTrait
             fn compute_state_root(&mut self) -> Result<StateRootWithAuxInfo>;
             fn get_state_root(&self) -> Result<StateRootWithAuxInfo>;
             fn commit(&mut self, epoch_id: EpochId) -> Result<StateRootWithAuxInfo>;
-            fn read_all_with_callback(&mut self, access_key_prefix: StorageKeyWithSpace, callback: &mut dyn FnMut(MptKeyValue), space_storage_filter: Option<SpaceStorageFilter>) -> Result<()>;
+            fn read_all_with_callback(&mut self, access_key_prefix: StorageKeyWithSpace, callback: &mut dyn FnMut(MptKeyValue), only_account_key: bool) -> Result<()>;
         }
     }
 
@@ -102,6 +102,4 @@ use crate::{
 use cfx_internal_common::StateRootWithAuxInfo;
 use delegate::delegate;
 use parking_lot::Mutex;
-use primitives::{
-    CheckInput, EpochId, SpaceStorageFilter, StorageKeyWithSpace,
-};
+use primitives::{CheckInput, EpochId, StorageKeyWithSpace};

crates/dbs/storage/src/impls/replicated_state.rs

@@ -6,8 +6,7 @@ use cfx_internal_common::StateRootWithAuxInfo;
 use cfx_types::Space;
 use parking_lot::Mutex;
 use primitives::{
-    EpochId, NodeMerkleTriplet, SpaceStorageFilter, StaticBool, StorageKey,
-    StorageKeyWithSpace,
+    EpochId, NodeMerkleTriplet, StaticBool, StorageKey, StorageKeyWithSpace,
 };
 use std::{
     sync::mpsc::{channel, Sender},
@@ -316,13 +315,12 @@ impl<Main: StateTrait> StateTrait for ReplicatedState<Main> {
 
     fn read_all_with_callback(
         &mut self, access_key_prefix: StorageKeyWithSpace,
-        callback: &mut dyn FnMut(MptKeyValue),
-        space_storage_filter: Option<SpaceStorageFilter>,
+        callback: &mut dyn FnMut(MptKeyValue), only_account_key: bool,
     ) -> Result<()> {
         self.state.read_all_with_callback(
             access_key_prefix,
             callback,
-            space_storage_filter,
+            only_account_key,
         )
     }
 

crates/dbs/storage/src/impls/single_mpt_state.rs

@@ -7,8 +7,8 @@ use crate::{
 };
 use cfx_internal_common::{StateRootAuxInfo, StateRootWithAuxInfo};
 use primitives::{
-    EpochId, MerkleHash, MptValue, SpaceStorageFilter, StateRoot,
-    StorageKeyWithSpace, MERKLE_NULL_NODE,
+    EpochId, MerkleHash, MptValue, StateRoot, StorageKeyWithSpace,
+    MERKLE_NULL_NODE,
 };
 use rustc_hex::ToHex;
 use std::{cell::UnsafeCell, sync::Arc};
@@ -309,8 +309,7 @@ impl SingleMptState {
 
     fn read_all_with_callback_impl(
         &mut self, access_key_prefix: StorageKeyWithSpace,
-        callback: &mut dyn FnMut(MptKeyValue),
-        space_storage_filter: Option<SpaceStorageFilter>,
+        callback: &mut dyn FnMut(MptKeyValue), only_account_key: bool,
     ) -> Result<()> {
         self.ensure_temp_slab_for_db_load();
 
@@ -342,7 +341,7 @@ impl SingleMptState {
             &key_prefix,
             &mut inner_callback,
             false,
-            space_storage_filter,
+            only_account_key,
         )?;
 
         Ok(())
@@ -412,13 +411,12 @@ impl StateTrait for SingleMptState {
 
     fn read_all_with_callback(
         &mut self, access_key_prefix: StorageKeyWithSpace,
-        callback: &mut dyn FnMut(MptKeyValue),
-        space_storage_filter: Option<SpaceStorageFilter>,
+        callback: &mut dyn FnMut(MptKeyValue), only_account_key: bool,
     ) -> Result<()> {
         self.read_all_with_callback_impl(
             access_key_prefix,
             callback,
-            space_storage_filter,
+            only_account_key,
         )
     }
 

crates/dbs/storage/src/impls/state.rs

@@ -311,13 +311,12 @@ impl StateTrait for State {
 
     fn read_all_with_callback(
         &mut self, access_key_prefix: StorageKeyWithSpace,
-        callback: &mut dyn FnMut(MptKeyValue),
-        space_storage_filter: Option<SpaceStorageFilter>,
+        callback: &mut dyn FnMut(MptKeyValue), only_account_key: bool,
     ) -> Result<()> {
         self.read_all_with_callback_impl(
             access_key_prefix,
             callback,
-            space_storage_filter,
+            only_account_key,
         )
     }
 
@@ -994,8 +993,7 @@ impl State {
 
     pub fn read_all_with_callback_impl(
         &mut self, access_key_prefix: StorageKeyWithSpace,
-        callback: &mut dyn FnMut(MptKeyValue),
-        space_storage_filter: Option<SpaceStorageFilter>,
+        callback: &mut dyn FnMut(MptKeyValue), only_account_key: bool,
     ) -> Result<()> {
         self.ensure_temp_slab_for_db_load();
 
@@ -1040,7 +1038,7 @@ impl State {
                 &delta_mpt_key_prefix,
                 &mut inner_callback,
                 true,
-                space_storage_filter,
+                only_account_key,
             )?;
         };
 
@@ -1082,7 +1080,7 @@ impl State {
                     &intermediate_mpt_key_prefix,
                     &mut inner_callback,
                     true,
-                    space_storage_filter,
+                    only_account_key,
                 )?;
             }
         }
@@ -1131,8 +1129,8 @@ use cfx_types::AddressWithSpace;
 use fallible_iterator::FallibleIterator;
 use primitives::{
     DeltaMptKeyPadding, EpochId, MerkleHash, MptValue, NodeMerkleTriplet,
-    SkipInputCheck, SpaceStorageFilter, StateRoot, StaticBool, StorageKey,
-    StorageKeyWithSpace, StorageRoot, MERKLE_NULL_NODE, NULL_EPOCH,
+    SkipInputCheck, StateRoot, StaticBool, StorageKey, StorageKeyWithSpace,
+    StorageRoot, MERKLE_NULL_NODE, NULL_EPOCH,
 };
 use rustc_hex::ToHex;
 use std::{

crates/dbs/storage/src/state.rs

@@ -39,14 +39,9 @@ pub trait StateTrait: Sync + Send {
         &mut self, access_key_prefix: StorageKeyWithSpace,
     ) -> Result<Option<Vec<MptKeyValue>>>;
 
-    /// space_storage_filter can be used to filter only the wanted contract
-    /// storage key if you specify SpaceStorageFilter::Native, it will
-    /// filter out the ethereum space storage key/value, by this way, we can
-    /// speedup space key/value traversal
     fn read_all_with_callback(
         &mut self, _access_key_prefix: StorageKeyWithSpace,
-        _callback: &mut dyn FnMut(MptKeyValue),
-        _space_storage_filter: Option<SpaceStorageFilter>,
+        _callback: &mut dyn FnMut(MptKeyValue), _only_account_key: bool,
     ) -> Result<()> {
         Err(Error::Msg("Not implemented".into()))
     }
@@ -95,6 +90,5 @@ use super::{
 };
 use crate::StorageRootProof;
 use primitives::{
-    EpochId, NodeMerkleTriplet, SpaceStorageFilter, StaticBool,
-    StorageKeyWithSpace, StorageRoot,
+    EpochId, NodeMerkleTriplet, StaticBool, StorageKeyWithSpace, StorageRoot,
 };

crates/primitives/src/storage_key.rs

@@ -764,7 +764,7 @@ impl SpaceStorageFilter {
     // return the flag index according the trie type
     // if is_delta_mpt is true, then the space flag is at the 32th index
     // otherwise, the space flag is at the 20th index
-    fn flag_index(&self, is_delta_mpt: bool) -> usize {
+    pub fn space_flag_index(is_delta_mpt: bool) -> usize {
         if is_delta_mpt {
             delta_mpt_storage_key::KEY_PADDING_BYTES
         } else {
@@ -774,7 +774,7 @@ impl SpaceStorageFilter {
 
     // return true if the key is filtered out
     pub fn is_filtered(&self, is_delta_mpt: bool, key: &[u8]) -> bool {
-        let flag_index = self.flag_index(is_delta_mpt);
+        let flag_index = Self::space_flag_index(is_delta_mpt);
         if key.len() > flag_index {
             match self.0 {
                 Space::Native => {