fix: allow loading of memtrie during catchup, for shards pending resh…

…arding (#12679)

This PR contains a small code change to fix a failure in the tests
`test_resharding_v3_load_mem_trie_`.

In the test a client that doesn't track the parent shard is assigned to
track one of the children in the next epoch after resharding. However,
in order to track the children shard, the client must state sync the
state of the parent and perform the shard split.
This fix allows the catchup routine to build the memtrie even with
`load_mem_tries_for_tracked_shards=false`, only if the shard is "pending
resharding in the future".

The first two commits are updates of code comments.

chain/chain/src/chain.rs

@@ -2024,7 +2024,6 @@ impl Chain {
             tracing::debug!(target: "chain", ?shard_id, need_storage_update, "Updating storage");
 
             if need_storage_update {
-                // TODO(resharding): consider adding to catchup flow.
                 self.resharding_manager.start_resharding(
                     self.chain_store.store_update(),
                     &block,
@@ -3127,8 +3126,8 @@ impl Chain {
                         }
                         blocks_catch_up_state.done_blocks.push(queued_block);
                     }
-                    Err(_) => {
-                        error!("Error processing block during catch up, retrying");
+                    Err(err) => {
+                        error!(target: "chain", ?err, "Error processing block during catch up, retrying");
                         blocks_catch_up_state.pending_blocks.push(queued_block);
                     }
                 }

chain/client/src/sync/state/shard.rs

@@ -2,7 +2,6 @@ use super::downloader::StateSyncDownloader;
 use super::task_tracker::TaskTracker;
 use crate::metrics;
 use crate::sync::state::chain_requests::ChainFinalizationRequest;
-use crate::sync::state::util::query_epoch_id_and_height_for_block;
 use futures::{StreamExt, TryStreamExt};
 use near_async::futures::{FutureSpawner, FutureSpawnerExt};
 use near_async::messaging::AsyncSender;
@@ -15,6 +14,7 @@ use near_primitives::sharding::ShardChunk;
 use near_primitives::state_part::PartId;
 use near_primitives::state_sync::StatePartKey;
 use near_primitives::types::{EpochId, ShardId};
+use near_primitives::version::PROTOCOL_VERSION;
 use near_store::adapter::{StoreAdapter, StoreUpdateAdapter};
 use near_store::flat::{FlatStorageReadyStatus, FlatStorageStatus};
 use near_store::{DBCol, ShardUId, Store};
@@ -162,8 +162,6 @@ pub(super) async fn run_state_sync_for_shard(
     return_if_cancelled!(cancel);
     // Create flat storage.
     {
-        let (epoch_id, _) = query_epoch_id_and_height_for_block(&store, sync_hash)?;
-        let shard_uid = epoch_manager.shard_id_to_uid(shard_id, &epoch_id)?;
         let chunk = header.cloned_chunk();
         let block_hash = chunk.prev_block();
 
@@ -180,8 +178,15 @@ pub(super) async fn run_state_sync_for_shard(
     // Load memtrie.
     {
         let handle = computation_task_tracker.get_handle(&format!("shard {}", shard_id)).await;
+        let head_protocol_version = epoch_manager.get_epoch_protocol_version(&epoch_id)?;
+        let shard_uids_pending_resharding = epoch_manager
+            .get_shard_uids_pending_resharding(head_protocol_version, PROTOCOL_VERSION)?;
         handle.set_status("Loading memtrie");
-        runtime.get_tries().load_mem_trie_on_catchup(&shard_uid, &state_root)?;
+        runtime.get_tries().load_mem_trie_on_catchup(
+            &shard_uid,
+            &state_root,
+            &shard_uids_pending_resharding,
+        )?;
     }
 
     return_if_cancelled!(cancel);

core/store/src/trie/shard_tries.rs

@@ -429,12 +429,16 @@ impl ShardTries {
 
     /// 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.
+    /// Mem-tries of shards that are pending resharding must be loaded in any case.
     pub fn load_mem_trie_on_catchup(
         &self,
         shard_uid: &ShardUId,
         state_root: &StateRoot,
+        shard_uids_pending_resharding: &HashSet<ShardUId>,
     ) -> Result<(), StorageError> {
-        if !self.0.trie_config.load_mem_tries_for_tracked_shards {
+        if !self.0.trie_config.load_mem_tries_for_tracked_shards
+            && !shard_uids_pending_resharding.contains(shard_uid)
+        {
             return Ok(());
         }
         // It should not happen that memtrie is already loaded for a shard

integration-tests/src/test_loop/tests/resharding_v3.rs

@@ -660,7 +660,9 @@ fn test_resharding_v3_resharding_block_in_fork() {
 #[test]
 // TODO(resharding): fix nearcore and un-ignore this test
 // TODO(resharding): duplicate this test so that in one case resharding is performed on block
-//                   B(height=13) and in another case resharding is performed on block B'(height=13)
+//                   B(height=13) and in another case resharding is performed on block B'(height=13).
+//                   In the current scenario the real resharding happens on block B'. Low priority TODO
+//                   since it's a very rare corner case.
 #[ignore]
 #[cfg(feature = "test_features")]
 fn test_resharding_v3_double_sign_resharding_block() {
@@ -866,8 +868,6 @@ fn test_resharding_v3_load_mem_trie_v1() {
     let params = TestReshardingParametersBuilder::default()
         .base_shard_layout_version(1)
         .load_mem_tries_for_tracked_shards(false)
-        // TODO(resharding): should it work without tracking all shards?
-        .track_all_shards(true)
         .build();
     test_resharding_v3_base(params);
 }
@@ -877,8 +877,6 @@ fn test_resharding_v3_load_mem_trie_v2() {
     let params = TestReshardingParametersBuilder::default()
         .base_shard_layout_version(2)
         .load_mem_tries_for_tracked_shards(false)
-        // TODO(resharding): should it work without tracking all shards?
-        .track_all_shards(true)
         .build();
     test_resharding_v3_base(params);
 }