Load memtries tracked shards at startup

chain/chain/src/chain.rs

@@ -390,6 +390,7 @@ impl Chain {
         chain_config: ChainConfig,
         snapshot_callbacks: Option<SnapshotCallbacks>,
         apply_chunks_spawner: Arc<dyn AsyncComputationSpawner>,
+        validator_account_id: Option<&AccountId>,
     ) -> Result<Chain, Error> {
         // Get runtime initial state and create genesis block out of it.
         let state_roots = get_genesis_state_roots(runtime_adapter.store())?
@@ -508,7 +509,19 @@ impl Chain {
         let tip = chain_store.head()?;
         let shard_uids: Vec<_> =
             epoch_manager.get_shard_layout(&tip.epoch_id)?.shard_uids().collect();
-        runtime_adapter.load_mem_tries_on_startup(&shard_uids)?;
+        let tracked_shards: Vec<_> = shard_uids
+            .iter()
+            .filter(|shard_uid| {
+                shard_tracker.care_about_shard(
+                    validator_account_id,
+                    &tip.prev_block_hash,
+                    shard_uid.shard_id(),
+                    true,
+                )
+            })
+            .cloned()
+            .collect();
+        runtime_adapter.load_mem_tries_on_startup(&shard_uids, &tracked_shards)?;
 
         info!(target: "chain", "Init: header head @ #{} {}; block head @ #{} {}",
               header_head.height, header_head.last_block_hash,

chain/chain/src/runtime/mod.rs

@@ -1230,8 +1230,12 @@ impl RuntimeAdapter for NightshadeRuntime {
         Ok(epoch_manager.will_shard_layout_change(parent_hash)?)
     }
 
-    fn load_mem_tries_on_startup(&self, shard_uids: &[ShardUId]) -> Result<(), StorageError> {
-        self.tries.load_mem_tries_for_enabled_shards(shard_uids)
+    fn load_mem_tries_on_startup(
+        &self,
+        all_shards: &[ShardUId],
+        tracked_shards: &[ShardUId],
+    ) -> Result<(), StorageError> {
+        self.tries.load_mem_tries_for_enabled_shards(all_shards, tracked_shards)
     }
 
     fn load_mem_trie_on_catchup(

chain/chain/src/runtime/tests.rs

@@ -1549,6 +1549,7 @@ fn get_test_env_with_chain_and_pool() -> (TestEnv, Chain, TransactionPool) {
         ChainConfig::test(),
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();
 

chain/chain/src/store_validator.rs

@@ -418,6 +418,7 @@ mod tests {
             ChainConfig::test(),
             None,
             Arc::new(RayonAsyncComputationSpawner),
+            None,
         )
         .unwrap();
         (

chain/chain/src/test_utils.rs

@@ -75,6 +75,7 @@ pub fn get_chain_with_epoch_length_and_num_shards(
         ChainConfig::test(),
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap()
 }
@@ -164,6 +165,7 @@ pub fn setup_with_tx_validity_period(
         ChainConfig::test(),
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();
 

chain/chain/src/test_utils/kv_runtime.rs

@@ -1453,7 +1453,11 @@ impl RuntimeAdapter for KeyValueRuntime {
         Ok(vec![])
     }
 
-    fn load_mem_tries_on_startup(&self, _shard_uids: &[ShardUId]) -> Result<(), StorageError> {
+    fn load_mem_tries_on_startup(
+        &self,
+        _all_shards: &[ShardUId],
+        _tracked_shards: &[ShardUId],
+    ) -> Result<(), StorageError> {
         Ok(())
     }
 

chain/chain/src/types.rs

@@ -514,7 +514,11 @@ pub trait RuntimeAdapter: Send + Sync {
     /// Loads in-memory tries upon startup. The given shard_uids are possible candidates to load,
     /// but which exact shards to load depends on configuration. This may only be called when flat
     /// storage is ready.
-    fn load_mem_tries_on_startup(&self, shard_uids: &[ShardUId]) -> Result<(), StorageError>;
+    fn load_mem_tries_on_startup(
+        &self,
+        all_shards: &[ShardUId],
+        tracked_shards: &[ShardUId],
+    ) -> Result<(), StorageError>;
 
     /// Loads in-memory trie upon catchup, if it is enabled.
     /// Requires state root because `ChunkExtra` is not available at the time mem-trie is being loaded.

chain/client/src/client.rs

@@ -276,6 +276,7 @@ impl Client {
             chain_config.clone(),
             snapshot_callbacks,
             async_computation_spawner.clone(),
+            validator_signer.as_ref().map(|x| x.validator_id()),
         )?;
         // Create flat storage or initiate migration to flat storage.
         let flat_storage_creator = FlatStorageCreator::new(

chain/client/src/info.rs

@@ -924,6 +924,7 @@ mod tests {
             ChainConfig::test(),
             None,
             Arc::new(RayonAsyncComputationSpawner),
+            None,
         )
         .unwrap();
 

chain/client/src/test_utils/setup.rs

@@ -121,6 +121,7 @@ pub fn setup(
         },
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();
     let genesis_block = chain.get_block(&chain.genesis().hash().clone()).unwrap();
@@ -259,6 +260,7 @@ pub fn setup_only_view(
         },
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();
 
@@ -1030,6 +1032,7 @@ pub fn setup_synchronous_shards_manager(
         }, // irrelevant
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();
     let chain_head = chain.head().unwrap();

core/store/src/test_utils.rs

@@ -147,7 +147,7 @@ impl TestTriesBuilder {
             }
         }
         if self.enable_in_memory_tries {
-            tries.load_mem_tries_for_enabled_shards(&shard_uids).unwrap();
+            tries.load_mem_tries_for_enabled_shards(&shard_uids, Default::default()).unwrap();
         }
         tries
     }

core/store/src/trie/shard_tries.rs

@@ -419,15 +419,18 @@ impl ShardTries {
     /// Should be called upon startup to load in-memory tries for enabled shards.
     pub fn load_mem_tries_for_enabled_shards(
         &self,
-        shard_uids: &[ShardUId],
+        all_shards: &[ShardUId],
+        tracked_shards: &[ShardUId],
     ) -> Result<(), StorageError> {
         let trie_config = &self.0.trie_config;
-        let shard_uids_to_load = shard_uids
+        let shard_uids_to_load = all_shards
             .iter()
             .copied()
             .filter(|shard_uid| {
                 trie_config.load_mem_tries_for_all_shards
                     || trie_config.load_mem_tries_for_shards.contains(shard_uid)
+                    || (trie_config.load_mem_tries_for_tracked_shards
+                        && tracked_shards.contains(shard_uid))
             })
             .collect::<Vec<_>>();
 

integration-tests/src/genesis_helpers.rs

@@ -38,6 +38,7 @@ pub fn genesis_header(genesis: &Genesis) -> BlockHeader {
         ChainConfig::test(),
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();
     chain.genesis().clone()
@@ -63,6 +64,7 @@ pub fn genesis_block(genesis: &Genesis) -> Block {
         ChainConfig::test(),
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();
     chain.get_block(&chain.genesis().hash().clone()).unwrap()

tools/fork-network/src/cli.rs

@@ -320,7 +320,7 @@ impl ForkNetworkCommand {
         let runtime =
             NightshadeRuntime::from_config(home_dir, store.clone(), &near_config, epoch_manager)
                 .context("could not create the transaction runtime")?;
-        runtime.load_mem_tries_on_startup(&all_shard_uids).unwrap();
+        runtime.load_mem_tries_on_startup(&all_shard_uids, Default::default()).unwrap();
 
         let make_storage_mutator: MakeSingleShardStorageMutatorFn =
             Arc::new(move |prev_state_root| {

tools/speedy_sync/src/main.rs

@@ -255,6 +255,7 @@ fn load_snapshot(load_cmd: LoadCmd) {
         },
         None,
         Arc::new(RayonAsyncComputationSpawner),
+        None,
     )
     .unwrap();