tests(resharding): add a resharding test with single-shard tracking a…

…nd shard shuffling (#12454)

When testing resharding and state sync integration, the existing
testloop tests with some modifications are very helpful. So this makes a
few edits to the basic resharding tests and adds a test that has nodes
tracking only one shard and enables shard shuffling. For now we
`#[ignore]` it because it fails, but we can enable it once everything is
fixed.

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

@@ -6,11 +6,12 @@ use near_chain::ChainStoreAccess;
 use near_chain_configs::test_genesis::TestGenesisBuilder;
 use near_client::Client;
 use near_o11y::testonly::init_test_logger;
+use near_primitives::block::Tip;
 use near_primitives::epoch_manager::EpochConfigStore;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout};
 use near_primitives::state_record::StateRecord;
-use near_primitives::types::{AccountId, BlockHeightDelta};
+use near_primitives::types::{AccountId, BlockHeightDelta, ShardId};
 use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
 use near_store::adapter::StoreAdapter;
 use near_store::db::refcount::decode_value_with_rc;
@@ -20,14 +21,33 @@ use std::sync::Arc;
 
 use crate::test_loop::builder::TestLoopBuilder;
 use crate::test_loop::env::TestLoopEnv;
+use crate::test_loop::utils::transactions::get_smallest_height_head;
 use crate::test_loop::utils::ONE_NEAR;
 use near_client::client_actor::ClientActorInner;
 
-fn print_and_assert_shard_accounts(client: &Client) {
-    let tip = client.chain.head().unwrap();
-    let epoch_id = tip.epoch_id;
-    let epoch_config = client.epoch_manager.get_epoch_config(&epoch_id).unwrap();
+fn client_tracking_shard<'a>(clients: &'a [&Client], tip: &Tip, shard_id: ShardId) -> &'a Client {
+    for client in clients {
+        let signer = client.validator_signer.get();
+        let cares_about_shard = client.shard_tracker.care_about_shard(
+            signer.as_ref().map(|s| s.validator_id()),
+            &tip.prev_block_hash,
+            shard_id,
+            true,
+        );
+        if cares_about_shard {
+            return client;
+        }
+    }
+    panic!(
+        "client_tracking_shard() could not find client tracking shard {} at {} #{}",
+        shard_id, &tip.last_block_hash, tip.height
+    );
+}
+
+fn print_and_assert_shard_accounts(clients: &[&Client], tip: &Tip) {
+    let epoch_config = clients[0].epoch_manager.get_epoch_config(&tip.epoch_id).unwrap();
     for shard_uid in epoch_config.shard_layout.shard_uids() {
+        let client = client_tracking_shard(clients, tip, shard_uid.shard_id());
         let chunk_extra = client.chain.get_chunk_extra(&tip.prev_block_hash, &shard_uid).unwrap();
         let trie = client
             .runtime_adapter
@@ -90,8 +110,14 @@ struct TestReshardingParameters {
     block_and_chunk_producers: Vec<AccountId>,
     initial_balance: u128,
     epoch_length: BlockHeightDelta,
+    shuffle_shard_assignment_for_chunk_producers: bool,
+    track_all_shards: bool,
     /// Custom behavior executed at every iteration of test loop.
     loop_action: Option<Box<dyn Fn(&mut TestLoopData, TestLoopDataHandle<ClientActorInner>)>>,
+    // When enabling shard shuffling with a short epoch length, sometimes a node might not finish
+    // catching up by the end of the epoch, and then misses a chunk. This can be fixed by using a longer
+    // epoch length, but it's good to also check what happens with shorter ones.
+    all_chunks_expected: bool,
 }
 
 impl TestReshardingParameters {
@@ -103,6 +129,8 @@ impl TestReshardingParameters {
         let num_accounts = 8;
         let initial_balance = 1_000_000 * ONE_NEAR;
         let epoch_length = 6;
+        let track_all_shards = true;
+        let all_chunks_expected = true;
 
         // #12195 prevents number of BPs bigger than `epoch_length`.
         assert!(num_clients > 0 && num_clients <= epoch_length);
@@ -137,6 +165,8 @@ impl TestReshardingParameters {
             block_and_chunk_producers,
             initial_balance,
             epoch_length,
+            track_all_shards,
+            all_chunks_expected,
             ..Default::default()
         }
     }
@@ -169,6 +199,21 @@ impl TestReshardingParameters {
         self.loop_action = loop_action;
         self
     }
+
+    fn shuffle_shard_assignment(mut self) -> Self {
+        self.shuffle_shard_assignment_for_chunk_producers = true;
+        self
+    }
+
+    fn single_shard_tracking(mut self) -> Self {
+        self.track_all_shards = false;
+        self
+    }
+
+    fn chunk_miss_possible(mut self) -> Self {
+        self.all_chunks_expected = false;
+        self
+    }
 }
 
 // Returns a callable function that, when invoked inside a test loop iteration, can force the creation of a chain fork.
@@ -238,15 +283,15 @@ fn test_resharding_v3_base(params: TestReshardingParameters) {
     }
 
     init_test_logger();
-    let builder = TestLoopBuilder::new();
+    let mut builder = TestLoopBuilder::new();
 
     // Prepare shard split configuration.
     let base_epoch_config_store = EpochConfigStore::for_chain_id("mainnet", None).unwrap();
     let base_protocol_version = ProtocolFeature::SimpleNightshadeV4.protocol_version() - 1;
     let mut base_epoch_config =
         base_epoch_config_store.get_config(base_protocol_version).as_ref().clone();
     base_epoch_config.validator_selection_config.shuffle_shard_assignment_for_chunk_producers =
-        false;
+        params.shuffle_shard_assignment_for_chunk_producers;
     if !params.chunk_ranges_to_drop.is_empty() {
         base_epoch_config.block_producer_kickout_threshold = 0;
         base_epoch_config.chunk_producer_kickout_threshold = 0;
@@ -290,6 +335,9 @@ fn test_resharding_v3_base(params: TestReshardingParameters) {
     }
     let (genesis, _) = genesis_builder.build();
 
+    if params.track_all_shards {
+        builder = builder.track_all_shards();
+    }
     let TestLoopEnv { mut test_loop, datas: node_datas, tempdir } = builder
         .genesis(genesis)
         .epoch_config_store(epoch_config_store)
@@ -298,27 +346,31 @@ fn test_resharding_v3_base(params: TestReshardingParameters) {
             base_protocol_version + 1,
             params.chunk_ranges_to_drop.clone(),
         )
-        .track_all_shards()
         .build();
 
-    let client_handle = node_datas[0].client_sender.actor_handle();
+    let client_handles =
+        node_datas.iter().map(|data| data.client_sender.actor_handle()).collect_vec();
+
     let latest_block_height = std::cell::Cell::new(0u64);
     let success_condition = |test_loop_data: &mut TestLoopData| -> bool {
-        params.loop_action.as_ref().map(|action| action(test_loop_data, client_handle.clone()));
+        params.loop_action.as_ref().map(|action| action(test_loop_data, client_handles[0].clone()));
 
-        let client = &test_loop_data.get(&client_handle).client;
-        let tip = client.chain.head().unwrap();
+        let clients =
+            client_handles.iter().map(|handle| &test_loop_data.get(handle).client).collect_vec();
+        let client = &clients[0];
+
+        let tip = get_smallest_height_head(&clients);
 
         // Check that all chunks are included.
         let block_header = client.chain.get_block_header(&tip.last_block_hash).unwrap();
         if latest_block_height.get() < tip.height {
             if latest_block_height.get() == 0 {
                 println!("State before resharding:");
-                print_and_assert_shard_accounts(client);
+                print_and_assert_shard_accounts(&clients, &tip);
             }
             latest_block_height.set(tip.height);
             println!("block: {} chunks: {:?}", tip.height, block_header.chunk_mask());
-            if params.chunk_ranges_to_drop.is_empty() {
+            if params.all_chunks_expected && params.chunk_ranges_to_drop.is_empty() {
                 assert!(block_header.chunk_mask().iter().all(|chunk_bit| *chunk_bit));
             }
         }
@@ -335,7 +387,7 @@ fn test_resharding_v3_base(params: TestReshardingParameters) {
         }
 
         println!("State after resharding:");
-        print_and_assert_shard_accounts(client);
+        print_and_assert_shard_accounts(&clients, &tip);
         check_state_shard_uid_mapping_after_resharding(&client, parent_shard_uid);
         return true;
     };
@@ -415,3 +467,14 @@ fn test_resharding_v3_double_sign_resharding_block() {
             .loop_action(Some(fork_before_resharding_block(true))),
     );
 }
+
+// TODO(resharding): fix nearcore and un-ignore this test
+#[test]
+#[ignore]
+fn test_resharding_v3_shard_shuffling() {
+    let params = TestReshardingParameters::new()
+        .shuffle_shard_assignment()
+        .single_shard_tracking()
+        .chunk_miss_possible();
+    test_resharding_v3_base(params);
+}

integration-tests/src/test_loop/utils/transactions.rs

@@ -8,6 +8,7 @@ use near_client::test_utils::test_loop::ClientQueries;
 use near_client::{Client, ProcessTxResponse};
 use near_crypto::Signer;
 use near_network::client::ProcessTxRequest;
+use near_primitives::block::Tip;
 use near_primitives::errors::InvalidTxError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::test_utils::create_user_test_signer;
@@ -34,19 +35,20 @@ pub(crate) struct BalanceMismatchError {
     pub actual: u128,
 }
 
+// Returns the head with the smallest height
+pub(crate) fn get_smallest_height_head(clients: &[&Client]) -> Tip {
+    clients
+        .iter()
+        .map(|client| client.chain.head().unwrap())
+        .min_by_key(|head| head.height)
+        .unwrap()
+}
+
 // Transactions have to be built on top of some block in chain. To make
 // sure all clients accept them, we select the head of the client with
 // the smallest height.
 pub(crate) fn get_anchor_hash(clients: &[&Client]) -> CryptoHash {
-    let (_, anchor_hash) = clients
-        .iter()
-        .map(|client| {
-            let head = client.chain.head().unwrap();
-            (head.height, head.last_block_hash)
-        })
-        .min_by_key(|&(height, _)| height)
-        .unwrap();
-    anchor_hash
+    get_smallest_height_head(clients).last_block_hash
 }
 
 /// Get next available nonce for the account's public key.