feat: resharding - tests for missing chunks (#10052)

* generalized the missing chunk test
* added missing chunk test cases to cover the v2 resharding
* fixed the check for `get_next_block_hash_with_new_chunk` that didn't
work for the V2 resharding
* fixed validators getting kicked out of the test when missing chunks
* introduced AllEpochConfigTestOverrides - overrides that allow changing
the epoch config in tests

chain/client/src/client.rs

@@ -13,6 +13,7 @@ use crate::SyncAdapter;
 use crate::SyncMessage;
 use crate::{metrics, SyncStatus};
 use actix_rt::ArbiterHandle;
+use itertools::Itertools;
 use lru::LruCache;
 use near_async::messaging::{CanSend, Sender};
 use near_chain::chain::VerifyBlockHashAndSignatureResult;
@@ -626,8 +627,16 @@ impl Client {
         }
 
         let new_chunks = self.get_chunk_headers_ready_for_inclusion(&epoch_id, &prev_hash);
-        debug!(target: "client", "{:?} Producing block at height {}, parent {} @ {}, {} new chunks", validator_signer.validator_id(),
-               next_height, prev.height(), format_hash(head.last_block_hash), new_chunks.len());
+        debug!(
+            target: "client",
+            validator=?validator_signer.validator_id(),
+            next_height=next_height,
+            prev_height=prev.height(),
+            prev_hash=format_hash(head.last_block_hash),
+            new_chunks_count=new_chunks.len(),
+            new_chunks=?new_chunks.keys().sorted().collect_vec(),
+            "Producing block",
+        );
 
         // If we are producing empty blocks and there are no transactions.
         if !self.config.produce_empty_blocks && new_chunks.is_empty() {

chain/client/src/test_utils/test_env.rs

@@ -123,7 +123,12 @@ impl TestEnv {
 
         let mut keep_going = true;
         while keep_going {
-            for network_adapter in network_adapters.iter() {
+            // for network_adapter in network_adapters.iter() {
+            for i in 0..network_adapters.len() {
+                let network_adapter = network_adapters.get(i).unwrap();
+                let _span =
+                    tracing::debug_span!(target: "test", "process_partial_encoded_chunks", client=i).entered();
+
                 keep_going = false;
                 // process partial encoded chunks
                 while let Some(request) = network_adapter.pop() {

chain/client/src/test_utils/test_env_builder.rs

@@ -16,7 +16,7 @@ use near_chunks::test_utils::MockClientAdapterForShardsManager;
 use near_epoch_manager::shard_tracker::{ShardTracker, TrackedConfig};
 use near_epoch_manager::{EpochManager, EpochManagerAdapter, EpochManagerHandle};
 use near_network::test_utils::MockPeerManagerAdapter;
-use near_primitives::epoch_manager::RngSeed;
+use near_primitives::epoch_manager::{AllEpochConfigTestOverrides, RngSeed};
 use near_primitives::types::{AccountId, NumShards};
 use near_store::test_utils::create_test_store;
 use near_store::{NodeStorage, ShardUId, Store, StoreConfig};
@@ -238,18 +238,27 @@ impl TestEnvBuilder {
         self
     }
 
-    /// Constructs real EpochManager implementations for each instance.
     pub fn real_epoch_managers(self, genesis_config: &GenesisConfig) -> Self {
+        self.real_epoch_managers_with_test_overrides(genesis_config, None)
+    }
+
+    /// Constructs real EpochManager implementations for each instance.
+    pub fn real_epoch_managers_with_test_overrides(
+        self,
+        genesis_config: &GenesisConfig,
+        test_overrides: Option<AllEpochConfigTestOverrides>,
+    ) -> Self {
         assert!(
             self.num_shards.is_none(),
             "Cannot set both num_shards and epoch_managers at the same time"
         );
         let ret = self.ensure_stores();
         let epoch_managers = (0..ret.clients.len())
             .map(|i| {
-                EpochManager::new_arc_handle(
+                EpochManager::new_arc_handle_with_test_overrides(
                     ret.stores.as_ref().unwrap()[i].clone(),
                     genesis_config,
+                    test_overrides.clone(),
                 )
             })
             .collect();

chain/epoch-manager/src/lib.rs

@@ -6,7 +6,8 @@ use near_primitives::checked_feature;
 use near_primitives::epoch_manager::block_info::BlockInfo;
 use near_primitives::epoch_manager::epoch_info::{EpochInfo, EpochSummary};
 use near_primitives::epoch_manager::{
-    AllEpochConfig, EpochConfig, ShardConfig, SlashState, AGGREGATOR_KEY,
+    AllEpochConfig, AllEpochConfigTestOverrides, EpochConfig, ShardConfig, SlashState,
+    AGGREGATOR_KEY,
 };
 use near_primitives::errors::EpochError;
 use near_primitives::hash::CryptoHash;
@@ -151,9 +152,18 @@ impl EpochManager {
     pub fn new_from_genesis_config(
         store: Store,
         genesis_config: &GenesisConfig,
+    ) -> Result<Self, EpochError> {
+        Self::new_from_genesis_config_with_test_overrides(store, genesis_config, None)
+    }
+
+    pub fn new_from_genesis_config_with_test_overrides(
+        store: Store,
+        genesis_config: &GenesisConfig,
+        test_overrides: Option<AllEpochConfigTestOverrides>,
     ) -> Result<Self, EpochError> {
         let reward_calculator = RewardCalculator::new(genesis_config);
-        let all_epoch_config = AllEpochConfig::from(genesis_config);
+        let all_epoch_config =
+            Self::new_all_epoch_config_with_test_overrides(genesis_config, test_overrides);
         Self::new(
             store,
             all_epoch_config,
@@ -164,7 +174,37 @@ impl EpochManager {
     }
 
     pub fn new_arc_handle(store: Store, genesis_config: &GenesisConfig) -> Arc<EpochManagerHandle> {
-        Arc::new(Self::new_from_genesis_config(store, genesis_config).unwrap().into_handle())
+        Self::new_arc_handle_with_test_overrides(store, genesis_config, None)
+    }
+
+    pub fn new_arc_handle_with_test_overrides(
+        store: Store,
+        genesis_config: &GenesisConfig,
+        test_overrides: Option<AllEpochConfigTestOverrides>,
+    ) -> Arc<EpochManagerHandle> {
+        Arc::new(
+            Self::new_from_genesis_config_with_test_overrides(
+                store,
+                genesis_config,
+                test_overrides,
+            )
+            .unwrap()
+            .into_handle(),
+        )
+    }
+
+    fn new_all_epoch_config_with_test_overrides(
+        genesis_config: &GenesisConfig,
+        test_overrides: Option<AllEpochConfigTestOverrides>,
+    ) -> AllEpochConfig {
+        let initial_epoch_config = EpochConfig::from(genesis_config);
+        let epoch_config = AllEpochConfig::new_with_test_overrides(
+            genesis_config.use_production_config(),
+            initial_epoch_config,
+            &genesis_config.chain_id,
+            test_overrides,
+        );
+        epoch_config
     }
 
     pub fn new(

core/chain-configs/src/genesis_config.rs

@@ -7,7 +7,7 @@ use crate::genesis_validate::validate_genesis;
 use anyhow::Context;
 use chrono::{DateTime, Utc};
 use near_config_utils::ValidationError;
-use near_primitives::epoch_manager::{AllEpochConfig, EpochConfig};
+use near_primitives::epoch_manager::EpochConfig;
 use near_primitives::shard_layout::ShardLayout;
 use near_primitives::types::validator_stake::ValidatorStake;
 use near_primitives::types::StateRoot;
@@ -216,18 +216,6 @@ impl From<&GenesisConfig> for EpochConfig {
     }
 }
 
-impl From<&GenesisConfig> for AllEpochConfig {
-    fn from(genesis_config: &GenesisConfig) -> Self {
-        let initial_epoch_config = EpochConfig::from(genesis_config);
-        let epoch_config = Self::new(
-            genesis_config.use_production_config(),
-            initial_epoch_config,
-            &genesis_config.chain_id,
-        );
-        epoch_config
-    }
-}
-
 /// Records in storage at genesis (get split into shards at genesis creation).
 #[derive(
     Debug,

core/primitives/src/epoch_manager.rs

@@ -73,6 +73,14 @@ impl ShardConfig {
     }
 }
 
+/// Testing overrides to apply to the EpochConfig returned by the `for_protocol_version`.
+/// All fields should be optional and the default should be a no-op.
+#[derive(Clone, Default)]
+pub struct AllEpochConfigTestOverrides {
+    pub block_producer_kickout_threshold: Option<u8>,
+    pub chunk_producer_kickout_threshold: Option<u8>,
+}
+
 /// AllEpochConfig manages protocol configs that might be changing throughout epochs (hence EpochConfig).
 /// The main function in AllEpochConfig is ::for_protocol_version which takes a protocol version
 /// and returns the EpochConfig that should be used for this protocol version.
@@ -85,6 +93,9 @@ pub struct AllEpochConfig {
     genesis_epoch_config: EpochConfig,
     /// Chain Id. Some parameters are specific to certain chains.
     chain_id: String,
+
+    /// Testing overrides to apply to the EpochConfig returned by the `for_protocol_version`.
+    test_overrides: AllEpochConfigTestOverrides,
 }
 
 impl AllEpochConfig {
@@ -93,7 +104,26 @@ impl AllEpochConfig {
         genesis_epoch_config: EpochConfig,
         chain_id: &str,
     ) -> Self {
-        Self { use_production_config, genesis_epoch_config, chain_id: chain_id.to_string() }
+        Self {
+            use_production_config,
+            genesis_epoch_config,
+            chain_id: chain_id.to_string(),
+            test_overrides: AllEpochConfigTestOverrides::default(),
+        }
+    }
+
+    pub fn new_with_test_overrides(
+        use_production_config: bool,
+        genesis_epoch_config: EpochConfig,
+        chain_id: &str,
+        test_overrides: Option<AllEpochConfigTestOverrides>,
+    ) -> Self {
+        Self {
+            use_production_config,
+            genesis_epoch_config,
+            chain_id: chain_id.to_string(),
+            test_overrides: test_overrides.unwrap_or_default(),
+        }
     }
 
     pub fn for_protocol_version(&self, protocol_version: ProtocolVersion) -> EpochConfig {
@@ -108,6 +138,8 @@ impl AllEpochConfig {
 
         Self::config_max_kickout_stake(&mut config, protocol_version);
 
+        Self::config_test_overrides(&mut config, &self.test_overrides);
+
         config
     }
 
@@ -170,6 +202,23 @@ impl AllEpochConfig {
             config.validator_max_kickout_stake_perc = 30;
         }
     }
+
+    fn config_test_overrides(
+        config: &mut EpochConfig,
+        test_overrides: &AllEpochConfigTestOverrides,
+    ) {
+        if let Some(block_producer_kickout_threshold) =
+            test_overrides.block_producer_kickout_threshold
+        {
+            config.block_producer_kickout_threshold = block_producer_kickout_threshold;
+        }
+
+        if let Some(chunk_producer_kickout_threshold) =
+            test_overrides.chunk_producer_kickout_threshold
+        {
+            config.chunk_producer_kickout_threshold = chunk_producer_kickout_threshold;
+        }
+    }
 }
 
 /// Additional configuration parameters for the new validator selection

core/store/src/trie/state_snapshot.rs

@@ -297,7 +297,10 @@ impl ShardTries {
         let _span =
             tracing::info_span!(target: "state_snapshot", "delete_all_state_snapshots").entered();
         let path = home_dir.join(hot_store_path).join(state_snapshot_subdir);
-        std::fs::remove_dir_all(&path)
+        if path.exists() {
+            std::fs::remove_dir_all(&path)?
+        }
+        Ok(())
     }
 
     pub fn get_state_snapshot_base_dir(