split state dump into multiple task

Author: Pana

crates/client/src/state_dump.rs

@@ -276,45 +276,52 @@ pub fn export_space_accounts_with_callback<F: Fn(AccountState)>(
         "[{}] Start to iterate state...",
         Utc::now().format("%Y-%m-%d %H:%M:%S")
     );
-    let mut account_states = BTreeMap::new();
 
-    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;
-        }
+    for i in 0..=255 {
+        let prefix = [i];
+        let start_key = StorageKey::AddressPrefixKey(&prefix).with_space(space);
 
-        if let StorageKey::AccountKey(address_bytes) =
-            storage_key_with_space.key
-        {
-            let address = Address::from_slice(address_bytes);
-            println!(
-                "[{}] Find account: {:?}",
-                Utc::now().format("%Y-%m-%d %H:%M:%S"),
-                address
-            );
-            let account = Account::new_from_rlp(address, &Rlp::new(&value))
-                .expect("Failed to decode account");
+        let mut account_states = BTreeMap::new();
 
-            account_states.insert(address, account);
-        }
-    };
+        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;
+            }
 
-    let empty_key = StorageKey::EmptyKey.with_space(space);
-    state.read_all_with_callback(empty_key, &mut inner_callback)?;
+            if let StorageKey::AccountKey(address_bytes) =
+                storage_key_with_space.key
+            {
+                let address = Address::from_slice(address_bytes);
+                println!(
+                    "[{}] Find account: {:?}",
+                    Utc::now().format("%Y-%m-%d %H:%M:%S"),
+                    address
+                );
+                let account = Account::new_from_rlp(address, &Rlp::new(&value))
+                    .expect("Failed to decode account");
 
-    println!(
-        "[{}] Start to read account code and storage data...",
-        Utc::now().format("%Y-%m-%d %H:%M:%S")
-    );
-    let mut found_accounts = 0;
-    for (_address, account) in account_states {
-        let account_state = get_account_state(state, &account, config)?;
-        callback(account_state);
-        found_accounts += 1;
-        if config.limit > 0 && found_accounts >= config.limit as usize {
-            break;
+                account_states.insert(address, account);
+            }
+        };
+
+        state.read_all_with_callback(start_key, &mut inner_callback)?;
+
+        if account_states.len() > 0 {
+            println!(
+                "[{}] Start to read account code and storage data...",
+                Utc::now().format("%Y-%m-%d %H:%M:%S")
+            );
+        }
+        let mut found_accounts = 0;
+        for (_address, account) in account_states {
+            let account_state = get_account_state(state, &account, config)?;
+            callback(account_state);
+            found_accounts += 1;
+            if config.limit > 0 && found_accounts >= config.limit as usize {
+                break;
+            }
         }
     }
 

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

@@ -161,6 +161,7 @@ enum OwnedStorageKey {
     DepositListKey(Vec<u8>),
     VoteListKey(Vec<u8>),
     EmptyKey,
+    AddressPrefixKey(Vec<u8>),
 }
 
 impl OwnedStorageKey {
@@ -196,6 +197,9 @@ impl OwnedStorageKey {
                 StorageKey::VoteListKey(k.as_slice())
             }
             OwnedStorageKey::EmptyKey => StorageKey::EmptyKey,
+            OwnedStorageKey::AddressPrefixKey(k) => {
+                StorageKey::AddressPrefixKey(k.as_slice())
+            }
         }
     }
 }
@@ -248,6 +252,9 @@ impl<'a> From<StorageKey<'a>> for OwnedStorageKey {
                 OwnedStorageKey::VoteListKey(k.to_vec())
             }
             StorageKey::EmptyKey => OwnedStorageKey::EmptyKey,
+            StorageKey::AddressPrefixKey(k) => {
+                OwnedStorageKey::AddressPrefixKey(k.to_vec())
+            }
         }
     }
 }

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

@@ -909,13 +909,30 @@ impl State {
             snapshot_kvs.push((key, value));
         }
 
+        let is_address_search_prefix =
+            if let StorageKey::AddressPrefixKey(prefix) = access_key_prefix.key
+            {
+                Some(prefix)
+            } else {
+                None
+            };
+
         let mut result = Vec::new();
         // This is used to keep track of the deleted keys.
         let mut deleted_keys = HashSet::new();
         if let Some(kvs) = delta_trie_kvs {
             for (k, v) in kvs {
                 let storage_key = StorageKeyWithSpace::from_delta_mpt_key(&k);
                 let k = storage_key.to_key_bytes();
+
+                // If it's an address search prefix, and k is not start with
+                // prefix, skip the key.
+                if let Some(prefix) = is_address_search_prefix {
+                    if !k.starts_with(prefix) {
+                        continue;
+                    }
+                }
+
                 deleted_keys.insert(k.clone());
                 if v.len() > 0 {
                     result.push((k, v));
@@ -926,11 +943,21 @@ impl State {
         if let Some(kvs) = intermediate_trie_kvs {
             for (k, v) in kvs {
                 let storage_key = StorageKeyWithSpace::from_delta_mpt_key(&k);
+                let k = storage_key.to_key_bytes();
+
+                // If it's an address search prefix, and k is not start with
+                // prefix, skip the key.
+                if let Some(prefix) = is_address_search_prefix {
+                    if !k.starts_with(prefix) {
+                        continue;
+                    }
+                }
+
                 // Only delete non-empty keys.
                 if v.len() > 0 && !AM::READ_ONLY {
                     self.delete(storage_key)?;
                 }
-                let k = storage_key.to_key_bytes();
+
                 if !deleted_keys.contains(&k) {
                     deleted_keys.insert(k.clone());
                     if v.len() > 0 {
@@ -966,6 +993,14 @@ impl State {
     ) -> Result<()> {
         self.ensure_temp_slab_for_db_load();
 
+        let is_address_search_prefix =
+            if let StorageKey::AddressPrefixKey(prefix) = access_key_prefix.key
+            {
+                Some(prefix)
+            } else {
+                None
+            };
+
         // This is used to keep track of the deleted keys.
         let mut deleted_keys = HashSet::new();
 
@@ -974,6 +1009,14 @@ impl State {
             let mut inner_callback = |(k, v): MptKeyValue| {
                 let storage_key = StorageKeyWithSpace::from_delta_mpt_key(&k);
                 let k = storage_key.to_key_bytes();
+
+                // If it's an address search prefix, and k is not start with
+                // prefix, skip the key.
+                if let Some(prefix) = is_address_search_prefix {
+                    if !k.starts_with(prefix) {
+                        return;
+                    }
+                }
                 deleted_keys.insert(k.clone());
                 if v.len() > 0 {
                     callback((k, v));
@@ -998,6 +1041,15 @@ impl State {
             let mut inner_callback = |(k, v): MptKeyValue| {
                 let storage_key = StorageKeyWithSpace::from_delta_mpt_key(&k);
                 let k = storage_key.to_key_bytes();
+
+                // If it's an address search prefix, and k is not start with
+                // prefix, skip the key.
+                if let Some(prefix) = is_address_search_prefix {
+                    if !k.starts_with(prefix) {
+                        return;
+                    }
+                }
+
                 if !deleted_keys.contains(&k) {
                     deleted_keys.insert(k.clone());
                     if v.len() > 0 {

crates/primitives/src/storage_key.rs

@@ -60,6 +60,9 @@ pub enum StorageKey<'a> {
     VoteListKey(&'a [u8]),
     // Empty key is used to traverse all key and value pairs.
     EmptyKey,
+    // Address prefix key is used to search all keys with the same address
+    // prefix, eg [1, 2](0x0102) will search all keys with prefix 0x0102
+    AddressPrefixKey(&'a [u8]),
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -200,6 +203,13 @@ impl<'a> StorageKeyWithSpace<'a> {
             StorageKey::EmptyKey => {
                 return vec![];
             }
+            StorageKey::AddressPrefixKey(_address_bytes) => {
+                // delta mpt trie does not support address prefix key search
+                // so we search all keys and filter them by address prefix
+                // due to delta mpt trie won't be very big, so the performance
+                // impact is not very big
+                return vec![];
+            }
         };
 
         return if self.space == Space::Native {
@@ -292,6 +302,11 @@ impl<'a> StorageKeyWithSpace<'a> {
             StorageKey::EmptyKey => {
                 return vec![];
             }
+            StorageKey::AddressPrefixKey(address_bytes) => {
+                let mut key = Vec::with_capacity(address_bytes.len());
+                key.extend_from_slice(address_bytes);
+                return key;
+            }
         };
 
         return if self.space == Space::Native {