From 2fc80bcdcc2c75b60044f5a3b55d4980e7f2d6f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Chuda=C5=9B?= Date: Mon, 28 Oct 2024 16:06:40 +0100 Subject: [PATCH] Address feedback --- chain/chain/src/resharding/manager.rs | 31 +++++++-- .../src/test_loop/tests/resharding_v3.rs | 68 +++++++------------ 2 files changed, 51 insertions(+), 48 deletions(-) diff --git a/chain/chain/src/resharding/manager.rs b/chain/chain/src/resharding/manager.rs index 9272916c1fc..c2c782dbe4a 100644 --- a/chain/chain/src/resharding/manager.rs +++ b/chain/chain/src/resharding/manager.rs @@ -85,10 +85,33 @@ impl ReshardingManager { let resharding_event_type = ReshardingEventType::from_shard_layout(&next_shard_layout, *block_hash, *prev_hash)?; - let Some(ReshardingEventType::SplitShard(split_shard_event)) = resharding_event_type else { - tracing::debug!(target: "resharding", ?resharding_event_type, "resharding event type is not split shard, skipping"); - return Ok(()); + match resharding_event_type { + Some(ReshardingEventType::SplitShard(split_shard_event)) => { + self.split_shard( + chain_store_update, + block, + shard_uid, + tries, + split_shard_event, + next_shard_layout, + )?; + } + None => { + tracing::debug!(target: "resharding", ?resharding_event_type, "unsupported resharding event type, skipping"); + } }; + Ok(()) + } + + fn split_shard( + &mut self, + chain_store_update: ChainStoreUpdate, + block: &Block, + shard_uid: ShardUId, + tries: ShardTries, + split_shard_event: ReshardingSplitShardParams, + next_shard_layout: ShardLayout, + ) -> Result<(), Error> { if split_shard_event.parent_shard != shard_uid { let parent_shard = split_shard_event.parent_shard; tracing::debug!(target: "resharding", ?parent_shard, "shard uid does not match event parent shard, skipping"); @@ -124,7 +147,7 @@ impl ReshardingManager { ) -> io::Result<()> { let mut store_update = self.store.trie_store().store_update(); let parent_shard_uid = split_shard_event.parent_shard; - // TODO(reshardingV3) Leave tracked shards only? + // TODO(reshardingV3) No need to set the mapping for children shards that we won't track just after resharding? let children_shard_uids = [split_shard_event.left_child_shard, split_shard_event.right_child_shard]; diff --git a/integration-tests/src/test_loop/tests/resharding_v3.rs b/integration-tests/src/test_loop/tests/resharding_v3.rs index 899384cbb49..58a34b0af17 100644 --- a/integration-tests/src/test_loop/tests/resharding_v3.rs +++ b/integration-tests/src/test_loop/tests/resharding_v3.rs @@ -1,3 +1,4 @@ +use borsh::BorshDeserialize; use itertools::Itertools; use near_async::test_loop::data::TestLoopData; use near_async::time::Duration; @@ -9,11 +10,11 @@ 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::trie_key::TrieKey; use near_primitives::types::{AccountId, ShardId}; use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION}; use near_store::adapter::StoreAdapter; -use near_store::ShardUId; +use near_store::db::refcount::decode_value_with_rc; +use near_store::{DBCol, ShardUId}; use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; @@ -63,32 +64,8 @@ fn print_and_assert_shard_accounts(client: &Client) { } } -/// Get hash of the node associated with given `account_id` in current trie for `shard_uid`. -fn get_account_node_hash( - client: &Client, - shard_uid: ShardUId, - account_id: &AccountId, -) -> CryptoHash { - let tip = client.chain.head().unwrap(); - let chunk_extra = client.chain.get_chunk_extra(&tip.prev_block_hash, &shard_uid).unwrap(); - let trie = client - .chain - .runtime_adapter - .get_tries() - .get_view_trie_for_shard(shard_uid, *chunk_extra.state_root()); - let trie_key = TrieKey::Account { account_id: account_id.clone() }.to_vec(); - let account_node_ref = - trie.get_optimized_ref(&trie_key, near_store::KeyLookupMode::FlatStorage).unwrap().unwrap(); - account_node_ref.into_value_ref().hash -} - -/// Assert that `node_hash` is accessible from parent and children shards. -/// It should be the case, according to State mapping strategy in Resharding V3. -fn check_state_shard_uid_mapping_after_resharding( - client: &Client, - parent_shard_uid: ShardUId, - node_hash: CryptoHash, -) { +/// Asserts that all parent shard State is accessible via parent and children shards. +fn check_state_shard_uid_mapping_after_resharding(client: &Client, parent_shard_uid: ShardUId) { let tip = client.chain.head().unwrap(); let epoch_id = tip.epoch_id; let epoch_config = client.epoch_manager.get_epoch_config(&epoch_id).unwrap(); @@ -97,13 +74,24 @@ fn check_state_shard_uid_mapping_after_resharding( assert_eq!(children_shard_uids.len(), 2); let store = client.chain.chain_store.store().trie_store(); - let parent_value = store.get(parent_shard_uid, &node_hash).unwrap(); - - // Because of the mapping strategy for the State column, - // all parent shard data is available to both children shards. - for child_shard_uid in children_shard_uids { - let child_value = store.get(child_shard_uid, &node_hash); - assert_eq!(child_value.unwrap(), parent_value); + for kv in store.store().iter_raw_bytes(DBCol::State) { + let (key, value) = kv.unwrap(); + let shard_uid = ShardUId::try_from_slice(&key[0..8]).unwrap(); + // Just after resharding, no State data must be keyed using children shard UIDs. + assert!(!children_shard_uids.contains(&shard_uid)); + if shard_uid != parent_shard_uid { + continue; + } + let node_hash = CryptoHash::try_from_slice(&key[8..]).unwrap(); + let (value, _) = decode_value_with_rc(&value); + let parent_value = store.get(parent_shard_uid, &node_hash); + // Parent shard data must still be accessible using parent shard UID. + assert_eq!(&parent_value.unwrap()[..], value.unwrap()); + // All parent shard data is available via both children shards. + for child_shard_uid in &children_shard_uids { + let child_value = store.get(*child_shard_uid, &node_hash); + assert_eq!(&child_value.unwrap()[..], value.unwrap()); + } } } @@ -179,9 +167,6 @@ fn test_resharding_v3_base(chunk_ranges_to_drop: HashMap bool { let client = &test_loop_data.get(&client_handle).client; @@ -214,12 +199,7 @@ fn test_resharding_v3_base(chunk_ranges_to_drop: HashMap