From f3c97aeddcbcdc0ac3c710eaff3d3105e8936b73 Mon Sep 17 00:00:00 2001 From: wacban Date: Fri, 11 Oct 2024 15:49:26 +0100 Subject: [PATCH] congestion info, shard layout, test utils, trie_key --- core/primitives-core/Cargo.toml | 4 +- core/primitives/benches/serialization.rs | 4 +- core/primitives/src/congestion_info.rs | 2 +- core/primitives/src/shard_layout.rs | 79 +++++++++++++----------- core/primitives/src/test_utils.rs | 6 +- core/primitives/src/trie_key.rs | 7 ++- 6 files changed, 57 insertions(+), 45 deletions(-) diff --git a/core/primitives-core/Cargo.toml b/core/primitives-core/Cargo.toml index 62be2a02172..41be9c49194 100644 --- a/core/primitives-core/Cargo.toml +++ b/core/primitives-core/Cargo.toml @@ -33,8 +33,8 @@ insta.workspace = true expect-test.workspace = true [features] -# default = ["new_shard_id"] # DO NOT COMMIT THIS -default = [] +default = ["new_shard_id"] # DO NOT COMMIT THIS +# default = [] protocol_feature_fix_staking_threshold = [] protocol_feature_fix_contract_loading_cost = [] protocol_feature_reject_blocks_with_outdated_protocol_version = [] diff --git a/core/primitives/benches/serialization.rs b/core/primitives/benches/serialization.rs index 5236e2b0124..a77eda9e7cc 100644 --- a/core/primitives/benches/serialization.rs +++ b/core/primitives/benches/serialization.rs @@ -13,7 +13,7 @@ use near_primitives::test_utils::account_new; use near_primitives::transaction::{ Action, SignedTransaction, Transaction, TransactionV0, TransferAction, }; -use near_primitives::types::{EpochId, StateRoot}; +use near_primitives::types::{new_shard_id_vec_tmp, EpochId, StateRoot}; use near_primitives::validator_signer::InMemoryValidatorSigner; use near_primitives::version::PROTOCOL_VERSION; use near_primitives_core::types::MerkleHash; @@ -39,7 +39,7 @@ fn create_transaction() -> SignedTransaction { } fn create_block() -> Block { - let shard_ids = vec![0]; + let shard_ids = new_shard_id_vec_tmp(&[0]); let genesis_chunks = genesis_chunks( vec![StateRoot::new()], vec![Default::default(); shard_ids.len()], diff --git a/core/primitives/src/congestion_info.rs b/core/primitives/src/congestion_info.rs index 991e9091b1d..bee31ab84a2 100644 --- a/core/primitives/src/congestion_info.rs +++ b/core/primitives/src/congestion_info.rs @@ -844,7 +844,7 @@ mod tests { // Full congestion - only the allowed shard should be able to send something. for shard in all_shards { - if shard == control.info.allowed_shard() as u64 { + if shard_id_as_u16(shard) == control.info.allowed_shard() { assert_eq!(control.outgoing_gas_limit(shard), config.allowed_shard_outgoing_gas); } else { assert_eq!(control.outgoing_gas_limit(shard), 0); diff --git a/core/primitives/src/shard_layout.rs b/core/primitives/src/shard_layout.rs index 9d9a6df6418..a11914d9d55 100644 --- a/core/primitives/src/shard_layout.rs +++ b/core/primitives/src/shard_layout.rs @@ -752,7 +752,7 @@ mod tests { ShardLayoutV1, ShardUId, }; use itertools::Itertools; - use near_primitives_core::types::{new_shard_id_tmp, ProtocolVersion}; + use near_primitives_core::types::{new_shard_id_tmp, shard_id_as_u64, ProtocolVersion}; use near_primitives_core::types::{AccountId, ShardId}; use near_primitives_core::version::{ProtocolFeature, PROTOCOL_VERSION}; use rand::distributions::Alphanumeric; @@ -829,16 +829,25 @@ mod tests { let s = String::from_utf8(s).unwrap(); let account_id = s.to_lowercase().parse().unwrap(); let shard_id = account_id_to_shard_id(&account_id, &shard_layout); - assert!(shard_id < num_shards); - *shard_id_distribution.get_mut(&shard_id.into()).unwrap() += 1; + assert!(shard_id_as_u64(shard_id) < num_shards); + *shard_id_distribution.get_mut(&shard_id).unwrap() += 1; } - let expected_distribution: HashMap<_, _> = - [(0, 247), (1, 268), (2, 233), (3, 252)].into_iter().collect(); + let expected_distribution: HashMap = [ + (new_shard_id_tmp(0), 247), + (new_shard_id_tmp(1), 268), + (new_shard_id_tmp(2), 233), + (new_shard_id_tmp(3), 252), + ] + .into_iter() + .collect(); assert_eq!(shard_id_distribution, expected_distribution); } #[test] fn test_shard_layout_v1() { + let aid = |s: &str| s.parse().unwrap(); + let sid = |s: u64| new_shard_id_tmp(s); + let shard_layout = ShardLayout::v1( parse_account_ids(&["aurora", "bar", "foo", "foo.baz", "paz"]), Some(new_shards_split_map(vec![vec![0, 1, 2], vec![3, 4, 5]])), @@ -853,27 +862,26 @@ mod tests { (3..6).map(|x| ShardUId { version: 1, shard_id: x }).collect::>() ); for x in 0..3 { - assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(x)).unwrap(), 0); - assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(x + 3)).unwrap(), 1); + assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(x)).unwrap(), sid(0)); + assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(x + 3)).unwrap(), sid(1)); } - let aid = |s: &str| s.parse().unwrap(); - assert_eq!(account_id_to_shard_id(&aid("aurora"), &shard_layout), 1); - assert_eq!(account_id_to_shard_id(&aid("foo.aurora"), &shard_layout), 3); - assert_eq!(account_id_to_shard_id(&aid("bar.foo.aurora"), &shard_layout), 2); - assert_eq!(account_id_to_shard_id(&aid("bar"), &shard_layout), 2); - assert_eq!(account_id_to_shard_id(&aid("bar.bar"), &shard_layout), 2); - assert_eq!(account_id_to_shard_id(&aid("foo"), &shard_layout), 3); - assert_eq!(account_id_to_shard_id(&aid("baz.foo"), &shard_layout), 2); - assert_eq!(account_id_to_shard_id(&aid("foo.baz"), &shard_layout), 4); - assert_eq!(account_id_to_shard_id(&aid("a.foo.baz"), &shard_layout), 0); - - assert_eq!(account_id_to_shard_id(&aid("aaa"), &shard_layout), 0); - assert_eq!(account_id_to_shard_id(&aid("abc"), &shard_layout), 0); - assert_eq!(account_id_to_shard_id(&aid("bbb"), &shard_layout), 2); - assert_eq!(account_id_to_shard_id(&aid("foo.goo"), &shard_layout), 4); - assert_eq!(account_id_to_shard_id(&aid("goo"), &shard_layout), 4); - assert_eq!(account_id_to_shard_id(&aid("zoo"), &shard_layout), 5); + assert_eq!(account_id_to_shard_id(&aid("aurora"), &shard_layout), sid(1)); + assert_eq!(account_id_to_shard_id(&aid("foo.aurora"), &shard_layout), sid(3)); + assert_eq!(account_id_to_shard_id(&aid("bar.foo.aurora"), &shard_layout), sid(2)); + assert_eq!(account_id_to_shard_id(&aid("bar"), &shard_layout), sid(2)); + assert_eq!(account_id_to_shard_id(&aid("bar.bar"), &shard_layout), sid(2)); + assert_eq!(account_id_to_shard_id(&aid("foo"), &shard_layout), sid(3)); + assert_eq!(account_id_to_shard_id(&aid("baz.foo"), &shard_layout), sid(2)); + assert_eq!(account_id_to_shard_id(&aid("foo.baz"), &shard_layout), sid(4)); + assert_eq!(account_id_to_shard_id(&aid("a.foo.baz"), &shard_layout), sid(0)); + + assert_eq!(account_id_to_shard_id(&aid("aaa"), &shard_layout), sid(0)); + assert_eq!(account_id_to_shard_id(&aid("abc"), &shard_layout), sid(0)); + assert_eq!(account_id_to_shard_id(&aid("bbb"), &shard_layout), sid(2)); + assert_eq!(account_id_to_shard_id(&aid("foo.goo"), &shard_layout), sid(4)); + assert_eq!(account_id_to_shard_id(&aid("goo"), &shard_layout), sid(4)); + assert_eq!(account_id_to_shard_id(&aid("zoo"), &shard_layout), sid(5)); } // check that after removing the fixed shards from the shard layout v1 @@ -905,18 +913,19 @@ mod tests { #[test] fn test_shard_layout_v2() { + let sid = |s: u64| new_shard_id_tmp(s); let shard_layout = get_test_shard_layout_v2(); // check accounts mapping in the middle of each range - assert_eq!(account_id_to_shard_id(&"aaa".parse().unwrap(), &shard_layout), 3); - assert_eq!(account_id_to_shard_id(&"ddd".parse().unwrap(), &shard_layout), 8); - assert_eq!(account_id_to_shard_id(&"mmm".parse().unwrap(), &shard_layout), 4); - assert_eq!(account_id_to_shard_id(&"rrr".parse().unwrap(), &shard_layout), 7); + assert_eq!(account_id_to_shard_id(&"aaa".parse().unwrap(), &shard_layout), sid(3)); + assert_eq!(account_id_to_shard_id(&"ddd".parse().unwrap(), &shard_layout), sid(8)); + assert_eq!(account_id_to_shard_id(&"mmm".parse().unwrap(), &shard_layout), sid(4)); + assert_eq!(account_id_to_shard_id(&"rrr".parse().unwrap(), &shard_layout), sid(7)); // check accounts mapping for the boundary accounts - assert_eq!(account_id_to_shard_id(&"ccc".parse().unwrap(), &shard_layout), 8); - assert_eq!(account_id_to_shard_id(&"kkk".parse().unwrap(), &shard_layout), 4); - assert_eq!(account_id_to_shard_id(&"ppp".parse().unwrap(), &shard_layout), 7); + assert_eq!(account_id_to_shard_id(&"ccc".parse().unwrap(), &shard_layout), sid(8)); + assert_eq!(account_id_to_shard_id(&"kkk".parse().unwrap(), &shard_layout), sid(4)); + assert_eq!(account_id_to_shard_id(&"ppp".parse().unwrap(), &shard_layout), sid(7)); // check shard ids assert_eq!(shard_layout.shard_ids().collect_vec(), new_shard_ids_vec(vec![3, 8, 4, 7])); @@ -927,10 +936,10 @@ mod tests { assert_eq!(shard_layout.shard_uids().collect_vec(), vec![u(3), u(8), u(4), u(7)]); // check parent - assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(3)).unwrap(), 3); - assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(8)).unwrap(), 1); - assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(4)).unwrap(), 4); - assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(7)).unwrap(), 1); + assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(3)).unwrap(), sid(3)); + assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(8)).unwrap(), sid(1)); + assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(4)).unwrap(), sid(4)); + assert_eq!(shard_layout.get_parent_shard_id(new_shard_id_tmp(7)).unwrap(), sid(1)); // check child assert_eq!( diff --git a/core/primitives/src/test_utils.rs b/core/primitives/src/test_utils.rs index a15b63fc08c..e9d6c043f64 100644 --- a/core/primitives/src/test_utils.rs +++ b/core/primitives/src/test_utils.rs @@ -20,7 +20,9 @@ use crate::version::PROTOCOL_VERSION; use crate::views::{ExecutionStatusView, FinalExecutionOutcomeView, FinalExecutionStatus}; use near_crypto::vrf::Value; use near_crypto::{EmptySigner, PublicKey, SecretKey, Signature, Signer}; -use near_primitives_core::types::{BlockHeight, MerkleHash, ProtocolVersion, ShardId}; +use near_primitives_core::types::{ + new_shard_id_tmp, BlockHeight, MerkleHash, ProtocolVersion, ShardId, +}; use std::collections::HashMap; use std::sync::Arc; @@ -1011,7 +1013,7 @@ impl EpochInfoProvider for MockEpochInfoProvider { _account_id: &AccountId, _epoch_id: &EpochId, ) -> Result { - Ok(0) + Ok(new_shard_id_tmp(0)) } } diff --git a/core/primitives/src/trie_key.rs b/core/primitives/src/trie_key.rs index 6e17aec60a4..415053e69b7 100644 --- a/core/primitives/src/trie_key.rs +++ b/core/primitives/src/trie_key.rs @@ -2,7 +2,7 @@ use crate::hash::CryptoHash; use crate::types::AccountId; use borsh::{BorshDeserialize, BorshSerialize}; use near_crypto::PublicKey; -use near_primitives_core::types::ShardId; +use near_primitives_core::types::{shard_id_as_u16, shard_id_as_u64, ShardId}; use near_schema_checker_lib::ProtocolSchema; use std::mem::size_of; @@ -292,10 +292,11 @@ impl TrieKey { } TrieKey::BufferedReceiptIndices => buf.push(col::BUFFERED_RECEIPT_INDICES), TrieKey::BufferedReceipt { index, receiving_shard } => { + let receiving_shard = *receiving_shard; buf.push(col::BUFFERED_RECEIPT); // Use u16 for shard id to reduce depth in trie. - assert!(*receiving_shard <= u16::MAX as u64, "Shard ID too big."); - buf.extend(&(*receiving_shard as u16).to_le_bytes()); + assert!(shard_id_as_u64(receiving_shard) <= u16::MAX as u64, "Shard ID too big."); + buf.extend(&shard_id_as_u16(receiving_shard).to_le_bytes()); buf.extend(&index.to_le_bytes()); } };