diff --git a/tools/fork-network/src/cli.rs b/tools/fork-network/src/cli.rs index 204ed70a446..1dc2db75575 100644 --- a/tools/fork-network/src/cli.rs +++ b/tools/fork-network/src/cli.rs @@ -25,8 +25,7 @@ use near_primitives::types::{ }; use near_primitives::version::PROTOCOL_VERSION; use near_store::db::RocksDB; -use near_store::flat::{store_helper, BlockInfo}; -use near_store::flat::{FlatStorageManager, FlatStorageStatus}; +use near_store::flat::{store_helper, BlockInfo, FlatStorageManager, FlatStorageStatus}; use near_store::{ checkpoint_hot_storage_and_cleanup_columns, DBCol, Store, TrieDBStorage, TrieStorage, FINAL_HEAD_KEY, @@ -283,6 +282,7 @@ impl ForkNetworkCommand { let mut store_update = store.store_update(); store_update.set_ser(DBCol::Misc, b"FORK_TOOL_EPOCH_ID", epoch_id)?; + store_update.set_ser(DBCol::Misc, b"FORK_TOOL_BLOCK_HASH", &desired_block_hash)?; store_update.set(DBCol::Misc, b"FORK_TOOL_BLOCK_HEIGHT", &block_height.to_le_bytes()); for (shard_id, state_root) in state_roots.iter().enumerate() { store_update.set_ser( @@ -307,8 +307,9 @@ impl ForkNetworkCommand { let storage = open_storage(&home_dir, near_config).unwrap(); let store = storage.get_hot_store(); - let (prev_state_roots, epoch_id, _block_height) = + let (prev_state_roots, prev_hash, epoch_id, _block_height) = self.get_state_roots_and_hash(store.clone())?; + tracing::info!(?prev_state_roots, ?epoch_id, ?prev_hash); let epoch_manager = EpochManager::new_arc_handle(store.clone(), &near_config.genesis.config); @@ -321,7 +322,12 @@ impl ForkNetworkCommand { let make_storage_mutator: MakeSingleShardStorageMutatorFn = Arc::new(move |shard_id, prev_state_root| { - SingleShardStorageMutator::new(shard_id, &runtime.clone(), prev_state_root) + SingleShardStorageMutator::new( + shard_id, + &runtime.clone(), + prev_hash, + prev_state_root, + ) }); let new_state_roots = self.prepare_state( @@ -350,7 +356,7 @@ impl ForkNetworkCommand { let storage = open_storage(&home_dir, near_config).unwrap(); let store = storage.get_hot_store(); - let (prev_state_roots, epoch_id, block_height) = + let (prev_state_roots, prev_hash, epoch_id, block_height) = self.get_state_roots_and_hash(store.clone())?; let epoch_manager = @@ -366,6 +372,7 @@ impl ForkNetworkCommand { epoch_manager.clone(), &runtime, epoch_id.clone(), + prev_hash, prev_state_roots, )?; let (new_state_roots, new_validator_accounts) = @@ -410,8 +417,9 @@ impl ForkNetworkCommand { fn get_state_roots_and_hash( &self, store: Store, - ) -> anyhow::Result<(Vec, EpochId, BlockHeight)> { + ) -> anyhow::Result<(Vec, CryptoHash, EpochId, BlockHeight)> { let epoch_id = EpochId(store.get_ser(DBCol::Misc, b"FORK_TOOL_EPOCH_ID")?.unwrap()); + let block_hash = store.get_ser(DBCol::Misc, b"FORK_TOOL_BLOCK_HASH")?.unwrap(); let block_height = store.get(DBCol::Misc, b"FORK_TOOL_BLOCK_HEIGHT")?.unwrap(); let block_height = u64::from_le_bytes(block_height.as_slice().try_into().unwrap()); let mut state_roots = vec![]; @@ -424,8 +432,8 @@ impl ForkNetworkCommand { assert_eq!(key, format!("FORK_TOOL_SHARD_ID:{shard_id}")); state_roots.push(state_root); } - tracing::info!(?state_roots, ?epoch_id, block_height); - Ok((state_roots, epoch_id, block_height)) + tracing::info!(?state_roots, ?block_hash, ?epoch_id, block_height); + Ok((state_roots, block_hash, epoch_id, block_height)) } /// Checks that `~/.near/data/fork-snapshot/data` exists. diff --git a/tools/fork-network/src/single_shard_storage_mutator.rs b/tools/fork-network/src/single_shard_storage_mutator.rs index 57ea95ea285..01b1752a084 100644 --- a/tools/fork-network/src/single_shard_storage_mutator.rs +++ b/tools/fork-network/src/single_shard_storage_mutator.rs @@ -23,15 +23,10 @@ impl SingleShardStorageMutator { pub(crate) fn new( shard_id: ShardId, runtime: &NightshadeRuntime, + prev_block_hash: CryptoHash, state_root: StateRoot, ) -> anyhow::Result { - let block_hash = CryptoHash::new(); - let trie = runtime.get_trie_for_shard( - shard_id, - &block_hash, // unused because of use_flat_storage=false - state_root, - false, - )?; + let trie = runtime.get_trie_for_shard(shard_id, &prev_block_hash, state_root, false)?; let trie_update = TrieUpdate::new(trie); Ok(Self { trie_update, shard_tries: runtime.get_tries(), num_changes: 0 }) } diff --git a/tools/fork-network/src/storage_mutator.rs b/tools/fork-network/src/storage_mutator.rs index be4527615f2..743aad4aa63 100644 --- a/tools/fork-network/src/storage_mutator.rs +++ b/tools/fork-network/src/storage_mutator.rs @@ -2,7 +2,8 @@ use crate::single_shard_storage_mutator::SingleShardStorageMutator; use near_crypto::PublicKey; use near_epoch_manager::EpochManagerAdapter; use near_primitives::account::{AccessKey, Account}; -use near_primitives::types::{AccountId, EpochId, StateRoot}; +use near_primitives::hash::CryptoHash; +use near_primitives::types::{AccountId, EpochId, ShardId, StateRoot}; use nearcore::NightshadeRuntime; use std::sync::Arc; @@ -19,17 +20,19 @@ impl StorageMutator { epoch_manager: Arc, runtime: &NightshadeRuntime, epoch_id: EpochId, + prev_block_hash: CryptoHash, state_roots: Vec, ) -> anyhow::Result { let shard_layout = epoch_manager.get_shard_layout(&epoch_id)?; - let num_shards = shard_layout.num_shards(); + assert_eq!(shard_layout.num_shards(), state_roots.len() as u64); let mut mutators = vec![]; - for shard_id in 0..num_shards { + for (shard_id, state_root) in state_roots.iter().enumerate() { mutators.push(SingleShardStorageMutator::new( - shard_id, + shard_id as ShardId, runtime, - state_roots[shard_id as usize], + prev_block_hash, + *state_root, )?); } Ok(Self { epoch_manager, epoch_id, mutators })