support iterate dump mode

Author: Pana

Cargo.lock

@@ -1,12 +1,12 @@
-use cfx_types::Address;
+use cfx_types::parse_hex_string;
 use clap::{ArgMatches, Args};
 use client::{
     configuration::Configuration,
-    state_dump::{dump_whole_state, StateDumpConfig},
+    state_dump::{dump_whole_state, iterate_dump_whole_state, StateDumpConfig},
 };
 use parking_lot::{Condvar, Mutex};
 use serde_json;
-use std::{collections::HashMap, fs, path::Path, str::FromStr, sync::Arc};
+use std::{collections::HashMap, fs, path::Path, sync::Arc};
 
 #[derive(Args, Debug)]
 pub struct DumpCommand {
@@ -68,8 +68,7 @@ impl DumpCommand {
     }
 
     fn get_state_dump_config(&self) -> Result<StateDumpConfig, String> {
-        let address_str = self.start.strip_prefix("0x").unwrap_or(&self.start);
-        let start_address = Address::from_str(address_str)
+        let start_address = parse_hex_string(&self.start)
             .map_err(|e| format!("Invalid address: {}", e))?;
         Ok(StateDumpConfig {
             start_address,
@@ -97,58 +96,69 @@ impl DumpCommand {
         }
 
         let exit = Arc::new((Mutex::new(false), Condvar::new()));
-
         let config = self.get_state_dump_config()?;
-        let state = dump_whole_state(conf, exit, &config)?;
-        let total_accounts = state.accounts.len();
 
-        if self.multi_file {
+        let _total_accounts = if self.multi_file {
             // Write to multiple files
-            for (address, account_state) in state.accounts {
-                // Create filename using address (without 0x prefix)
-                let filename = format!("{}.json", address);
-                let file_path = Path::new(output_path).join(&filename);
+            let state_root = iterate_dump_whole_state(
+                conf,
+                exit,
+                &config,
+                |account_state| {
+                    let address =
+                        account_state.address.expect("address is not set");
+                    let filename = format!("{:?}.json", address);
+                    let file_path = Path::new(output_path).join(&filename);
 
-                // Serialize account_state to JSON
-                let json_content = serde_json::to_string_pretty(&account_state)
-                    .map_err(|e| {
-                        format!(
+                    // Serialize account_state to JSON
+                    let json_content =
+                        serde_json::to_string_pretty(&account_state)
+                            .map_err(|e| {
+                                format!(
                             "Failed to serialize account state for {}: {}",
                             address, e
                         )
-                    })?;
+                            })
+                            .expect("Failed to serialize account state");
 
-                // Write to file
-                fs::write(&file_path, json_content).map_err(|e| {
-                    format!(
-                        "Failed to write file {}: {}",
-                        file_path.display(),
-                        e
-                    )
-                })?;
-            }
+                    // Write to file
+                    fs::write(&file_path, json_content)
+                        .map_err(|e| {
+                            format!(
+                                "Failed to write file {}: {}",
+                                file_path.display(),
+                                e
+                            )
+                        })
+                        .expect("Failed to write file");
+                },
+            )?;
 
             // Write meta info
-            let meta_file_path = Path::new(output_path).join("meta.json");
             let mut meta_info = HashMap::new();
-            meta_info.insert("root".to_string(), state.root);
+            meta_info.insert("root".to_string(), state_root);
+            let meta_file_path = Path::new(output_path).join("meta.json");
             let meta_content = serde_json::to_string_pretty(&meta_info)
                 .map_err(|e| format!("Failed to serialize state: {}", e))?;
             fs::write(&meta_file_path, meta_content)
                 .map_err(|e| format!("Failed to write meta file: {}", e))?;
+            0
         } else {
+            let state = dump_whole_state(conf, exit, &config)?;
+            let total_accounts = state.accounts.len();
             // Write to a single file
             let file_path = Path::new(output_path).join("state.json");
             let json_content = serde_json::to_string_pretty(&state)
                 .map_err(|e| format!("Failed to serialize state: {}", e))?;
             fs::write(&file_path, json_content).map_err(|e| {
                 format!("Failed to write file {}: {}", file_path.display(), e)
             })?;
-        }
+            total_accounts
+        };
 
         Ok(format!(
-            "Dumped {} account state to output directory: {}",
-            total_accounts, output_path
+            "Dumped account state to output directory: {}",
+            output_path
         ))
     }
 }

bins/conflux/src/command/dump.rs

@@ -79,6 +79,7 @@ cfx-parity-trace-types = { workspace = true }
 cfx-tasks = { workspace = true }
 cfx-config = { workspace = true }
 cfxcore-types = { workspace = true }
+fallible-iterator = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }

crates/client/Cargo.toml

@@ -3,17 +3,21 @@ use cfx_config::Configuration;
 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_storage::{
+    state_manager::StateManagerTrait, utils::to_key_prefix_iter_upper_bound,
+    KeyValueDbIterableTrait,
+};
 use cfx_types::{Address, Space, H256};
 use cfxcore::NodeType;
+use fallible_iterator::FallibleIterator;
 use keccak_hash::{keccak, KECCAK_EMPTY};
 use parking_lot::{Condvar, Mutex};
 use primitives::{
     Account, SkipInputCheck, StorageKey, StorageKeyWithSpace, StorageValue,
 };
 use rlp::Rlp;
 use std::{
-    collections::{BTreeMap, HashMap},
+    collections::{BTreeMap, HashMap, HashSet},
     ops::Deref,
     sync::Arc,
     thread,
@@ -32,6 +36,44 @@ pub fn dump_whole_state(
     conf: &mut Configuration, exit_cond_var: Arc<(Mutex<bool>, Condvar)>,
     config: &StateDumpConfig,
 ) -> Result<StateDump, String> {
+    let (mut state_db, state_root) =
+        prepare_state_db(conf, exit_cond_var, config)?;
+
+    let accounts =
+        export_space_accounts(&mut state_db, Space::Ethereum, config)
+            .map_err(|e| e.to_string())?;
+
+    let state_dump = StateDump {
+        root: state_root,
+        accounts,
+        next: None,
+    };
+
+    Ok(state_dump)
+}
+
+pub fn iterate_dump_whole_state<F: Fn(AccountState)>(
+    conf: &mut Configuration, exit_cond_var: Arc<(Mutex<bool>, Condvar)>,
+    config: &StateDumpConfig, callback: F,
+) -> Result<H256, String> {
+    let (mut state_db, state_root) =
+        prepare_state_db(conf, exit_cond_var, config)?;
+
+    export_space_accounts_with_iterator(
+        &mut state_db,
+        Space::Ethereum,
+        config,
+        callback,
+    )
+    .map_err(|e| e.to_string())?;
+
+    Ok(state_root)
+}
+
+fn prepare_state_db(
+    conf: &mut Configuration, exit_cond_var: Arc<(Mutex<bool>, Condvar)>,
+    config: &StateDumpConfig,
+) -> Result<(StateDbGeneric, H256), String> {
     println!("Preparing state...");
     let (
         data_man,
@@ -71,8 +113,6 @@ pub fn dump_whole_state(
         None => consensus.latest_confirmed_epoch_number(),
     };
 
-    println!("Start to dump state at epoch: {:?}", target_height);
-
     let epoch_hash = consensus
         .get_hash_from_epoch_number(target_height.into())
         .map_err(|e| e.to_string())?;
@@ -92,19 +132,9 @@ pub fn dump_whole_state(
         .map_err(|e| e.to_string())?
         .ok_or("Failed to get state")?;
 
-    let mut state_db = StateDbGeneric::new(state);
-
-    let accounts =
-        export_space_accounts(&mut state_db, Space::Ethereum, config)
-            .map_err(|e| e.to_string())?;
-
-    let state_dump = StateDump {
-        root: state_root.clone(),
-        accounts,
-        next: None,
-    };
+    let state_db = StateDbGeneric::new(state);
 
-    Ok(state_dump)
+    Ok((state_db, state_root.clone()))
 }
 
 fn export_space_accounts(
@@ -177,12 +207,18 @@ fn export_space_accounts(
         let code = if is_contract {
             codes_map.get(&address).cloned()
         } else {
+            if let Some(code) = codes_map.get(&address) {
+                println!("no-contract account have code: {:?}", code);
+            }
             None
         };
 
         let storage = if is_contract {
             storage_map.get(&address).cloned()
         } else {
+            if let Some(storage) = storage_map.get(&address) {
+                println!("no-contract account have storage: {:?}", storage);
+            }
             None
         };
 
@@ -207,6 +243,93 @@ fn export_space_accounts(
     Ok(accounts)
 }
 
+fn export_space_accounts_with_iterator<F: Fn(AccountState)>(
+    state: &mut StateDbGeneric, space: Space, config: &StateDumpConfig,
+    callback: F,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let empty_key = StorageKey::EmptyKey.with_space(space);
+    let (kvs, maybe_kv_iterator) = state.read_all_iterator(empty_key)?;
+
+    let mut deleted_keys = HashSet::new();
+    let mut found_accounts = 0;
+
+    // Iterate key value pairs from delta trie and intermediate trie
+    for (k, v) in kvs {
+        let storage_key = StorageKeyWithSpace::from_delta_mpt_key(&k);
+        let key = storage_key.to_key_bytes();
+        deleted_keys.insert(key.clone());
+
+        let storage_key_with_space =
+            StorageKeyWithSpace::from_key_bytes::<SkipInputCheck>(&key);
+        if storage_key_with_space.space != space {
+            continue;
+        }
+
+        if let StorageKey::AccountKey(address_bytes) =
+            storage_key_with_space.key
+        {
+            let address = Address::from_slice(address_bytes);
+            println!("Find account: {:?}", address);
+            let account = Account::new_from_rlp(address, &Rlp::new(&v))?;
+
+            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;
+            }
+        } else {
+            continue;
+        }
+    }
+
+    let lower_bound_incl = empty_key.to_key_bytes();
+    let upper_bound_excl = to_key_prefix_iter_upper_bound(&lower_bound_incl);
+
+    if let Some(mut kv_iterator) = maybe_kv_iterator {
+        let mut kvs = kv_iterator
+            .iter_range(
+                lower_bound_incl.as_slice(),
+                upper_bound_excl.as_ref().map(|v| &**v),
+            )?
+            .take();
+
+        while let Some((key, value)) = kvs.next()? {
+            if deleted_keys.contains(&key) {
+                continue;
+            }
+
+            let storage_key_with_space =
+                StorageKeyWithSpace::from_key_bytes::<SkipInputCheck>(&key);
+            if storage_key_with_space.space != space {
+                continue;
+            }
+
+            if let StorageKey::AccountKey(address_bytes) =
+                storage_key_with_space.key
+            {
+                let address = Address::from_slice(address_bytes);
+                println!("Find account: {:?}", address);
+                let account =
+                    Account::new_from_rlp(address, &Rlp::new(&value))?;
+
+                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;
+                }
+            } else {
+                continue;
+            }
+        }
+    }
+
+    Ok(())
+}
+
 #[allow(unused)]
 fn get_account_state(
     state: &mut StateDbGeneric, account: &Account, config: &StateDumpConfig,

crates/client/src/state_dump.rs

@@ -188,6 +188,15 @@ mod impls {
             self.delete_all::<access_mode::Read>(key_prefix, debug_record)
         }
 
+        pub fn read_all_iterator(
+            &mut self, access_key_prefix: StorageKeyWithSpace,
+        ) -> Result<(Vec<MptKeyValue>, Option<KvdbSqliteSharded<Box<[u8]>>>)>
+        {
+            self.storage
+                .read_all_iterator(access_key_prefix)
+                .map_err(|err| err.into())
+        }
+
         pub fn delete_all<AM: access_mode::AccessMode>(
             &mut self, key_prefix: StorageKeyWithSpace,
             debug_record: Option<&mut ComputeEpochDebugRecord>,
@@ -535,7 +544,7 @@ mod impls {
     };
     use cfx_storage::{
         utils::{access_mode, to_key_prefix_iter_upper_bound},
-        MptKeyValue, StorageStateTrait,
+        KvdbSqliteSharded, MptKeyValue, StorageStateTrait,
     };
     use cfx_types::{
         address_util::AddressUtil, Address, AddressWithSpace, Space,

crates/dbs/statedb/src/lib.rs

@@ -54,6 +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_iterator(&mut self, access_key_prefix: StorageKeyWithSpace) -> Result<(Vec<MptKeyValue>, Option<KvdbSqliteSharded<Box<[u8]>>>)>;
         }
     }
 
@@ -94,6 +95,7 @@ impl<Storage: StateTrait + StateTraitExt> StateTrait
 use crate::{
     impls::{
         errors::*, merkle_patricia_trie::MptKeyValue, state_proof::StateProof,
+        storage_db::kvdb_sqlite_sharded::KvdbSqliteSharded,
     },
     state::*,
     StateProofMerger,