[resharding] Enable snapshot iterator for resharding (#9607)

chain/chain/src/resharding.rs

@@ -13,8 +13,6 @@ use near_primitives::errors::StorageError::StorageInconsistentState;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout};
 use near_primitives::state::FlatStateValue;
-use near_primitives::state_part::PartId;
-use near_primitives::state_sync::get_num_state_parts;
 use near_primitives::types::chunk_extra::ChunkExtra;
 use near_primitives::types::{AccountId, ShardId, StateRoot};
 use near_store::flat::{
@@ -216,50 +214,8 @@ impl Chain {
         StateSplitResponse { shard_id, sync_hash, new_state_roots }
     }
 
-    // TODO(#9446) remove function when shifting to flat storage iteration for resharding
     fn build_state_for_split_shards_impl(
         state_split_request: StateSplitRequest,
-    ) -> Result<HashMap<ShardUId, StateRoot>, Error> {
-        let StateSplitRequest { tries, shard_uid, state_root, next_epoch_shard_layout, .. } =
-            state_split_request;
-        let trie = tries.get_view_trie_for_shard(shard_uid, state_root);
-        let shard_id = shard_uid.shard_id();
-        let new_shards = next_epoch_shard_layout
-            .get_split_shard_uids(shard_id)
-            .ok_or(Error::InvalidShardId(shard_id))?;
-        let mut state_roots: HashMap<_, _> =
-            new_shards.iter().map(|shard_uid| (*shard_uid, Trie::EMPTY_ROOT)).collect();
-        let checked_account_id_to_shard_uid =
-            get_checked_account_id_to_shard_uid_fn(shard_uid, new_shards, next_epoch_shard_layout);
-
-        let state_root_node = trie.retrieve_root_node()?;
-        let num_parts = get_num_state_parts(state_root_node.memory_usage);
-        debug!(target: "resharding", "splitting state for shard {} to {} parts to build new states", shard_id, num_parts);
-        for part_id in 0..num_parts {
-            let trie_items = trie.get_trie_items_for_part(PartId::new(part_id, num_parts))?;
-            let (store_update, new_state_roots) = tries.add_values_to_split_states(
-                &state_roots,
-                trie_items.into_iter().map(|(key, value)| (key, Some(value))).collect(),
-                &checked_account_id_to_shard_uid,
-            )?;
-            state_roots = new_state_roots;
-            store_update.commit()?;
-        }
-        state_roots = apply_delayed_receipts(
-            &tries,
-            shard_uid,
-            state_root,
-            state_roots,
-            &checked_account_id_to_shard_uid,
-        )?;
-
-        Ok(state_roots)
-    }
-
-    // TODO(#9446) After implementing iterator at specific head, shift to build_state_for_split_shards_impl_v2
-    #[allow(dead_code)]
-    fn build_state_for_split_shards_impl_v2(
-        state_split_request: StateSplitRequest,
     ) -> Result<HashMap<ShardUId, StateRoot>, Error> {
         let StateSplitRequest {
             tries,

chain/client/src/test_utils/test_env_builder.rs

@@ -16,7 +16,7 @@ use near_network::test_utils::MockPeerManagerAdapter;
 use near_primitives::epoch_manager::RngSeed;
 use near_primitives::types::{AccountId, NumShards};
 use near_store::test_utils::create_test_store;
-use near_store::{NodeStorage, Store};
+use near_store::{NodeStorage, Store, StoreConfig};
 
 use super::setup::{setup_client_with_runtime, setup_synchronous_shards_manager};
 use super::test_env::TestEnv;
@@ -157,7 +157,13 @@ impl TestEnvBuilder {
             .unwrap()
             .iter()
             .map(|home_dir| {
-                NodeStorage::opener(home_dir.as_path(), false, &Default::default(), None)
+                // The max number of open files across all RocksDB instances is INT_MAX i.e. 65,535
+                // The default value of max_open_files is 10,000 which only allows upto 6 RocksDB
+                // instance to open at a time. This is problematic in testing resharding. To overcome
+                // this limit, we set the max_open_files config to 1000.
+                let mut store_config = StoreConfig::default();
+                store_config.max_open_files = 1000;
+                NodeStorage::opener(home_dir.as_path(), false, &store_config, None)
                     .open()
                     .unwrap()
                     .get_hot_store()

integration-tests/src/tests/client/resharding.rs

@@ -194,9 +194,11 @@ impl TestReshardingEnv {
         let env = TestEnv::builder(chain_genesis)
             .clients_count(num_clients)
             .validator_seats(num_validators)
+            .real_stores()
             .real_epoch_managers(&genesis.config)
             .nightshade_runtimes(&genesis)
             .track_all_shards()
+            .use_state_snapshots()
             .build();
         assert_eq!(env.validators.len(), num_validators);
         Self {

tools/state-viewer/src/state_dump.rs

@@ -295,7 +295,7 @@ mod test {
     use std::path::Path;
     use std::sync::Arc;
 
-    use near_chain::{ChainGenesis, Provenance};
+    use near_chain::{ChainGenesis, ChainStoreAccess, Provenance};
     use near_chain_configs::genesis_validate::validate_genesis;
     use near_chain_configs::{Genesis, GenesisChangeConfig};
     use near_client::test_utils::TestEnv;
@@ -315,6 +315,7 @@ mod test {
     use nearcore::config::GenesisExt;
     use nearcore::config::TESTING_INIT_STAKE;
     use nearcore::config::{Config, NearConfig};
+    use nearcore::test_utils::TestEnvNightshadeSetupExt;
     use nearcore::NightshadeRuntime;
 
     use crate::state_dump::state_dump;
@@ -324,33 +325,31 @@ mod test {
     fn setup(
         epoch_length: NumBlocks,
         protocol_version: ProtocolVersion,
-        use_production_config: bool,
+        test_resharding: bool,
     ) -> (Store, Genesis, TestEnv, NearConfig) {
         let mut genesis =
             Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
         genesis.config.num_block_producer_seats = 2;
         genesis.config.num_block_producer_seats_per_shard = vec![2];
         genesis.config.epoch_length = epoch_length;
         genesis.config.protocol_version = protocol_version;
-        genesis.config.use_production_config = use_production_config;
-        let store = create_test_store();
-        initialize_genesis_state(store.clone(), &genesis, None);
-        let epoch_manager = EpochManager::new_arc_handle(store.clone(), &genesis.config);
-        let nightshade_runtime = NightshadeRuntime::test(
-            Path::new("."),
-            store.clone(),
-            &genesis.config,
-            epoch_manager.clone(),
-        );
-        let mut chain_genesis = ChainGenesis::test();
-        chain_genesis.epoch_length = epoch_length;
-        chain_genesis.gas_limit = genesis.config.gas_limit;
-        let env = TestEnv::builder(chain_genesis)
-            .validator_seats(2)
-            .stores(vec![store.clone()])
-            .epoch_managers(vec![epoch_manager])
-            .runtimes(vec![nightshade_runtime])
-            .build();
+        genesis.config.use_production_config = test_resharding;
+
+        let env = if test_resharding {
+            TestEnv::builder(ChainGenesis::new(&genesis))
+                .validator_seats(2)
+                .real_stores()
+                .real_epoch_managers(&genesis.config)
+                .nightshade_runtimes(&genesis)
+                .use_state_snapshots()
+                .build()
+        } else {
+            TestEnv::builder(ChainGenesis::new(&genesis))
+                .validator_seats(2)
+                .real_epoch_managers(&genesis.config)
+                .nightshade_runtimes(&genesis)
+                .build()
+        };
 
         let near_config = NearConfig::new(
             Config::default(),
@@ -367,6 +366,7 @@ mod test {
         )
         .unwrap();
 
+        let store = env.clients[0].chain.store().store().clone();
         (store, genesis, env, near_config)
     }