From f3accf0db2b5f5ded7e44a9229429aef92a3150a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Chuda=C5=9B?= Date: Thu, 23 Jan 2025 18:21:47 +0100 Subject: [PATCH] address PR comments --- core/primitives/src/shard_layout.rs | 4 ++++ .../src/test_loop/tests/resharding_v3.rs | 15 +++++++-------- integration-tests/src/test_loop/utils/setups.rs | 2 ++ integration-tests/src/test_loop/utils/sharding.rs | 14 ++++---------- 4 files changed, 17 insertions(+), 18 deletions(-) diff --git a/core/primitives/src/shard_layout.rs b/core/primitives/src/shard_layout.rs index 1211cf069f1..82b5271055d 100644 --- a/core/primitives/src/shard_layout.rs +++ b/core/primitives/src/shard_layout.rs @@ -94,6 +94,10 @@ type ShardsSplitMapV2 = BTreeMap>; /// A mapping from the child shard to the parent shard. type ShardsParentMapV2 = BTreeMap; +pub fn shard_uids_to_ids(shard_uids: &[ShardUId]) -> Vec { + shard_uids.iter().map(|shard_uid| shard_uid.shard_id()).collect_vec() +} + fn new_shard_ids_vec(shard_ids: Vec) -> Vec { shard_ids.into_iter().map(Into::into).collect() } diff --git a/integration-tests/src/test_loop/tests/resharding_v3.rs b/integration-tests/src/test_loop/tests/resharding_v3.rs index 825f677a1ee..ecaae15af93 100644 --- a/integration-tests/src/test_loop/tests/resharding_v3.rs +++ b/integration-tests/src/test_loop/tests/resharding_v3.rs @@ -5,7 +5,7 @@ use near_async::time::Duration; use near_chain_configs::test_genesis::{TestGenesisBuilder, ValidatorsSpec}; use near_o11y::testonly::init_test_logger; use near_primitives::epoch_manager::EpochConfigStore; -use near_primitives::shard_layout::ShardLayout; +use near_primitives::shard_layout::{shard_uids_to_ids, ShardLayout}; use near_primitives::types::{AccountId, BlockHeightDelta, ShardId, ShardIndex}; use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION}; use std::cell::Cell; @@ -29,8 +29,7 @@ use crate::test_loop::utils::setups::{ derive_new_epoch_config_from_boundary, two_upgrades_voting_schedule, }; use crate::test_loop::utils::sharding::{ - get_shards_needs_for_next_epoch, get_tracked_shards, print_and_assert_shard_accounts, - shard_uids_to_ids, + get_shards_will_care_about, get_tracked_shards, print_and_assert_shard_accounts, }; use crate::test_loop::utils::transactions::{ check_txs, create_account, deploy_contract, get_smallest_height_head, @@ -539,14 +538,14 @@ fn test_resharding_v3_base(params: TestReshardingParameters) { tip.epoch_id.0, ); for (client_index, client) in clients.iter().enumerate() { - let tracked_shards = - shard_uids_to_ids(&get_tracked_shards(client, &tip.last_block_hash)); + let tracked_shards = get_tracked_shards(client, &tip.last_block_hash); + let tracked_shards = shard_uids_to_ids(&tracked_shards); // That's not accurate in case of tracked shard schedule: it won't return parent shard before resharding boundary, if we track child after resharding. - let shards_needs_for_next_epoch = - shard_uids_to_ids(&get_shards_needs_for_next_epoch(client, &tip.last_block_hash)); + let shards_will_care_about = &get_shards_will_care_about(client, &tip.last_block_hash); + let shards_will_care_about = shard_uids_to_ids(shards_will_care_about); let signer = client.validator_signer.get().unwrap(); let account_id = signer.validator_id().as_str(); - println!("client_{client_index}: id={account_id:?} tracks={tracked_shards:?}\tneeds_for_next_epoch={shards_needs_for_next_epoch:?}"); + println!("client_{client_index}: id={account_id:?} tracks={tracked_shards:?}\twill_care_about={shards_will_care_about:?}"); } // Check that all chunks are included. diff --git a/integration-tests/src/test_loop/utils/setups.rs b/integration-tests/src/test_loop/utils/setups.rs index 14195100159..fa5117db24a 100644 --- a/integration-tests/src/test_loop/utils/setups.rs +++ b/integration-tests/src/test_loop/utils/setups.rs @@ -73,6 +73,8 @@ pub fn derive_new_epoch_config_from_boundary( epoch_config } +/// Two protocol upgrades would happen as soon as possible, +/// usually in two consecutive epochs, unless upgrade voting decides differently. pub fn two_upgrades_voting_schedule( target_protocol_version: ProtocolVersion, ) -> ProtocolUpgradeVotingSchedule { diff --git a/integration-tests/src/test_loop/utils/sharding.rs b/integration-tests/src/test_loop/utils/sharding.rs index 0e69749df32..7e4ef35be71 100644 --- a/integration-tests/src/test_loop/utils/sharding.rs +++ b/integration-tests/src/test_loop/utils/sharding.rs @@ -1,4 +1,3 @@ -use itertools::Itertools; use near_chain::types::Tip; use near_client::Client; use near_epoch_manager::EpochManagerAdapter; @@ -148,17 +147,16 @@ pub fn get_tracked_shards(client: &Client, block_hash: &CryptoHash) -> Vec Vec { +pub fn get_shards_will_care_about(client: &Client, block_hash: &CryptoHash) -> Vec { let signer = client.validator_signer.get(); let account_id = signer.as_ref().map(|s| s.validator_id()); - let prev_block_hash = *client.chain.get_block_header(block_hash).unwrap().prev_hash(); - let shard_layout = - client.epoch_manager.get_shard_layout_from_prev_block(&prev_block_hash).unwrap(); + let block_header = client.chain.get_block_header(block_hash).unwrap(); + let shard_layout = client.epoch_manager.get_shard_layout(&block_header.epoch_id()).unwrap(); let mut shards_needs_for_next_epoch = vec![]; for shard_uid in shard_layout.shard_uids() { if client.shard_tracker.will_care_about_shard( account_id, - &prev_block_hash, + &block_header.prev_hash(), shard_uid.shard_id(), true, ) { @@ -167,7 +165,3 @@ pub fn get_shards_needs_for_next_epoch(client: &Client, block_hash: &CryptoHash) } shards_needs_for_next_epoch } - -pub fn shard_uids_to_ids(shard_uids: &[ShardUId]) -> Vec { - shard_uids.iter().map(|shard_uid| shard_uid.shard_id()).collect_vec() -}