diff --git a/chain/chain/src/runtime/mod.rs b/chain/chain/src/runtime/mod.rs index 051047714cb..8ae8ecb7ea4 100644 --- a/chain/chain/src/runtime/mod.rs +++ b/chain/chain/src/runtime/mod.rs @@ -1239,7 +1239,7 @@ impl RuntimeAdapter for NightshadeRuntime { shard_uid: &ShardUId, state_root: &StateRoot, ) -> Result<(), StorageError> { - if !self.get_tries().trie_config().load_mem_trie_for_tracked_shards { + if !self.get_tries().trie_config().load_mem_tries_for_tracked_shards { return Ok(()); } // It should not happen that memtrie is already loaded for a shard diff --git a/chain/chunks/src/logic.rs b/chain/chunks/src/logic.rs index 35793e9623a..41201889abf 100644 --- a/chain/chunks/src/logic.rs +++ b/chain/chunks/src/logic.rs @@ -1,4 +1,6 @@ -use near_chain::{types::EpochManagerAdapter, validate::validate_chunk_proofs, Chain, ChainStore}; +use near_chain::{ + types::EpochManagerAdapter, validate::validate_chunk_proofs, BlockHeader, Chain, ChainStore, +}; use near_chunks_primitives::Error; use near_epoch_manager::shard_tracker::ShardTracker; use near_primitives::{ @@ -46,6 +48,31 @@ pub fn cares_about_shard_this_or_next_epoch( || shard_tracker.will_care_about_shard(account_id, parent_hash, shard_id, is_me) } +pub fn get_shards_cares_about_this_or_next_epoch( + account_id: Option<&AccountId>, + is_me: bool, + block_header: &BlockHeader, + shard_tracker: &ShardTracker, + epoch_manager: &dyn EpochManagerAdapter, +) -> Vec { + let parent_hash = *block_header.prev_hash(); + let epoch_id = block_header.epoch_id().clone(); + epoch_manager + .shard_ids(&epoch_id) + .unwrap() + .into_iter() + .filter(|&shard_id| { + cares_about_shard_this_or_next_epoch( + account_id, + &parent_hash, + shard_id, + is_me, + shard_tracker, + ) + }) + .collect() +} + pub fn chunk_needs_to_be_fetched_from_archival( chunk_prev_block_hash: &CryptoHash, header_head: &CryptoHash, diff --git a/chain/client/src/client.rs b/chain/client/src/client.rs index 60420cc8592..39498e73442 100644 --- a/chain/client/src/client.rs +++ b/chain/client/src/client.rs @@ -43,7 +43,7 @@ use near_chain_configs::{ClientConfig, LogSummaryStyle, UpdateableClientConfig}; use near_chunks::adapter::ShardsManagerRequestFromClient; use near_chunks::client::ShardedTransactionPool; use near_chunks::logic::{ - cares_about_shard_this_or_next_epoch, decode_encoded_chunk, persist_chunk, + cares_about_shard_this_or_next_epoch, decode_encoded_chunk, get_shards_cares_about_this_or_next_epoch, persist_chunk }; use near_chunks::ShardsManager; use near_client_primitives::debug::ChunkProduction; @@ -2340,13 +2340,15 @@ impl Client { let _span = debug_span!(target: "sync", "run_catchup").entered(); let mut notify_state_sync = false; let me = &self.validator_signer.as_ref().map(|x| x.validator_id().clone()); + for (sync_hash, state_sync_info) in self.chain.chain_store().iterate_state_sync_infos()? { assert_eq!(sync_hash, state_sync_info.epoch_tail_hash); let network_adapter = self.network_adapter.clone(); let shards_to_split = self.get_shards_to_split(sync_hash, &state_sync_info, me)?; let state_sync_timeout = self.config.state_sync_timeout; - let epoch_id = self.chain.get_block(&sync_hash)?.header().epoch_id().clone(); + let block_header = self.chain.get_block(&sync_hash)?.header().clone(); + let epoch_id = block_header.epoch_id(); let (state_sync, shards_to_split, blocks_catch_up_state) = self.catchup_state_syncs.entry(sync_hash).or_insert_with(|| { @@ -2374,9 +2376,29 @@ impl Client { state_sync_info.shards.iter().map(|tuple| tuple.0).collect(); // Notify each shard to sync. if notify_state_sync { - for shard_id in &tracking_shards { - let shard_uid = - self.epoch_manager.shard_id_to_uid(*shard_id, &epoch_id).unwrap(); + let shard_layout = self + .epoch_manager + .get_shard_layout(&epoch_id) + .expect("Cannot get shard layout"); + + // Make sure mem-tries for shards we do not care about are unloaded before we start a new state sync. + let shards_cares_this_or_next_epoch = get_shards_cares_about_this_or_next_epoch( + me.as_ref(), + true, + &block_header, + &self.shard_tracker, + self.epoch_manager.as_ref(), + ); + let shard_uids: Vec<_> = shards_cares_this_or_next_epoch + .iter() + .map(|id| { + self.epoch_manager.shard_id_to_uid(*id, &epoch_id).unwrap() + }) + .collect(); + self.runtime_adapter.retain_mem_tries(&shard_uids); + + for &shard_id in &tracking_shards { + let shard_uid = ShardUId::from_shard_id_and_layout(shard_id, &shard_layout); match self.state_sync_adapter.clone().read() { Ok(sync_adapter) => sync_adapter.send( shard_uid, diff --git a/chain/client/src/client_actions.rs b/chain/client/src/client_actions.rs index 298771eb21d..465515f041d 100644 --- a/chain/client/src/client_actions.rs +++ b/chain/client/src/client_actions.rs @@ -33,7 +33,7 @@ use near_chain::{ use near_chain_configs::{ClientConfig, LogSummaryStyle}; use near_chain_primitives::error::EpochErrorResultToChainError; use near_chunks::client::ShardsManagerResponse; -use near_chunks::logic::cares_about_shard_this_or_next_epoch; +use near_chunks::logic::get_shards_cares_about_this_or_next_epoch; use near_client_primitives::types::{ Error, GetClientConfig, GetClientConfigError, GetNetworkInfo, NetworkInfoResponse, StateSyncStatus, Status, StatusError, StatusSyncInfo, SyncStatus, @@ -1543,30 +1543,13 @@ impl ClientActions { unwrap_and_report!(self.client.chain.get_block_header(&sync_hash)); let prev_hash = *block_header.prev_hash(); let epoch_id = block_header.epoch_id().clone(); - let shards_to_sync: Vec<_> = self - .client - .epoch_manager - .shard_ids(&epoch_id) - .unwrap() - .into_iter() - .filter(|&shard_id| { - cares_about_shard_this_or_next_epoch( - me.as_ref(), - &prev_hash, - shard_id, - true, - &self.client.shard_tracker, - ) - }) - .collect(); - let shard_uids: Vec<_> = shards_to_sync - .iter() - .map(|id| { - self.client.epoch_manager.shard_id_to_uid(*id, &epoch_id).unwrap() - }) - .collect(); - // Make sure mem-tries can be loaded only for shards we care about this or next epoch. - self.client.runtime_adapter.retain_mem_tries(&shard_uids); + let shards_to_sync = get_shards_cares_about_this_or_next_epoch( + me.as_ref(), + true, + &block_header, + &self.client.shard_tracker, + self.client.epoch_manager.as_ref(), + ); let use_colour = matches!(self.client.config.log_summary_style, LogSummaryStyle::Colored); diff --git a/core/store/src/config.rs b/core/store/src/config.rs index 64832a304ae..43981bc5a06 100644 --- a/core/store/src/config.rs +++ b/core/store/src/config.rs @@ -66,7 +66,7 @@ pub struct StoreConfig { /// If true, load mem tries for all shards; this has priority over `load_mem_tries_for_shards`. pub load_mem_tries_for_all_shards: bool, /// If true, load mem trie for each shard being tracked. - pub load_mem_trie_for_tracked_shards: bool, + pub load_mem_tries_for_tracked_shards: bool, /// Path where to create RocksDB checkpoints during database migrations or /// `false` to disable that feature. @@ -261,7 +261,7 @@ impl Default for StoreConfig { // requires more RAM and takes several minutes on startup. load_mem_tries_for_shards: Default::default(), load_mem_tries_for_all_shards: false, - load_mem_trie_for_tracked_shards: false, + load_mem_tries_for_tracked_shards: false, migration_snapshot: Default::default(), diff --git a/core/store/src/trie/config.rs b/core/store/src/trie/config.rs index bd4d80229dc..91344310587 100644 --- a/core/store/src/trie/config.rs +++ b/core/store/src/trie/config.rs @@ -35,7 +35,7 @@ pub struct TrieConfig { pub load_mem_tries_for_shards: Vec, pub load_mem_tries_for_all_shards: bool, /// Whether mem-trie should be loaded for each tracked shard. - pub load_mem_trie_for_tracked_shards: bool, + pub load_mem_tries_for_tracked_shards: bool, } impl TrieConfig { @@ -61,7 +61,7 @@ impl TrieConfig { } this.load_mem_tries_for_shards = config.load_mem_tries_for_shards.clone(); this.load_mem_tries_for_all_shards = config.load_mem_tries_for_all_shards; - this.load_mem_trie_for_tracked_shards = config.load_mem_trie_for_tracked_shards; + this.load_mem_tries_for_tracked_shards = config.load_mem_tries_for_tracked_shards; this } diff --git a/core/store/src/trie/shard_tries.rs b/core/store/src/trie/shard_tries.rs index f031e78394d..dcebca102a5 100644 --- a/core/store/src/trie/shard_tries.rs +++ b/core/store/src/trie/shard_tries.rs @@ -385,12 +385,16 @@ impl ShardTries { /// Remove trie from memory for shards not included in the given list. pub fn retain_mem_tries(&self, shard_uids: &[ShardUId]) { + info!(target: "memtrie", "Retaining memtries for shards {:?}...", shard_uids); self.0.mem_tries.write().unwrap().retain(|shard_uid, _| shard_uids.contains(shard_uid)); + info!(target: "memtrie", "Memtries retaining complete for shards {:?}", shard_uids); } /// Remove trie from memory for given shard. pub fn unload_mem_trie(&self, shard_uid: &ShardUId) { + info!(target: "memtrie", "Unloading trie from memory for shard {:?}...", shard_uid); self.0.mem_tries.write().unwrap().remove(shard_uid); + info!(target: "memtrie", "Memtrie unloading complete for shard {:?}", shard_uid); } /// Loads in-memory-trie for given shard and state root (if given). @@ -399,8 +403,10 @@ impl ShardTries { shard_uid: &ShardUId, state_root: Option, ) -> Result<(), StorageError> { + info!(target: "memtrie", "Loading trie to memory for shard {:?}...", shard_uid); let mem_tries = load_trie_from_flat_state_and_delta(&self.0.store, *shard_uid, state_root)?; self.0.mem_tries.write().unwrap().insert(*shard_uid, Arc::new(RwLock::new(mem_tries))); + info!(target: "memtrie", "Memtrie loading complete for shard {:?}", shard_uid); Ok(()) } @@ -725,7 +731,7 @@ mod test { sweat_prefetch_senders: Vec::new(), load_mem_tries_for_shards: Vec::new(), load_mem_tries_for_all_shards: false, - load_mem_trie_for_tracked_shards: false, + load_mem_tries_for_tracked_shards: false, }; let shard_uids = Vec::from([ShardUId::single_shard()]); ShardTries::new( @@ -846,7 +852,7 @@ mod test { sweat_prefetch_senders: Vec::new(), load_mem_tries_for_shards: Vec::new(), load_mem_tries_for_all_shards: false, - load_mem_trie_for_tracked_shards: false, + load_mem_tries_for_tracked_shards: false, }; let shard_uids = Vec::from([ShardUId { shard_id: 0, version: 0 }]); let shard_uid = *shard_uids.first().unwrap(); diff --git a/integration-tests/src/tests/client/state_snapshot.rs b/integration-tests/src/tests/client/state_snapshot.rs index 50b9ca14a0d..bdfe62b8ea4 100644 --- a/integration-tests/src/tests/client/state_snapshot.rs +++ b/integration-tests/src/tests/client/state_snapshot.rs @@ -45,7 +45,7 @@ impl StateSnaptshotTestEnv { sweat_prefetch_senders: Vec::new(), load_mem_tries_for_shards: Vec::new(), load_mem_tries_for_all_shards: false, - load_mem_trie_for_tracked_shards: false, + load_mem_tries_for_tracked_shards: false, }; let flat_storage_manager = FlatStorageManager::new(store.clone()); let shard_uids = [ShardUId::single_shard()];