Skip to content

Commit

Permalink
feat(resharding) - Make shard ids non-contiguous - part 2 (#12209)
Browse files Browse the repository at this point in the history
More of the same. I added a rust feature to quickly switch between the
old and new shard ids to speed up work.

No interesting changes here, just syntax everywhere. I hope.
  • Loading branch information
wacban authored Oct 11, 2024
1 parent 16fd9e8 commit 0c135f2
Show file tree
Hide file tree
Showing 19 changed files with 313 additions and 215 deletions.
8 changes: 4 additions & 4 deletions chain/chain/src/runtime/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,8 @@ use near_primitives::state_part::PartId;
use near_primitives::transaction::SignedTransaction;
use near_primitives::trie_key::TrieKey;
use near_primitives::types::{
new_shard_id_tmp, AccountId, Balance, BlockHeight, EpochHeight, EpochId, EpochInfoProvider,
Gas, MerkleHash, ShardId, StateChangeCause, StateRoot, StateRootNode,
shard_id_as_u32, AccountId, Balance, BlockHeight, EpochHeight, EpochId, EpochInfoProvider, Gas,
MerkleHash, ShardId, StateChangeCause, StateRoot, StateRootNode,
};
use near_primitives::version::{ProtocolFeature, ProtocolVersion};
use near_primitives::views::{
Expand Down Expand Up @@ -223,7 +223,7 @@ impl NightshadeRuntime {
epoch_manager.get_epoch_id_from_prev_block(prev_hash).map_err(Error::from)?;
let shard_version =
epoch_manager.get_shard_layout(&epoch_id).map_err(Error::from)?.version();
Ok(ShardUId { version: shard_version, shard_id: new_shard_id_tmp(shard_id) as u32 })
Ok(ShardUId { version: shard_version, shard_id: shard_id_as_u32(shard_id) })
}

fn get_shard_uid_from_epoch_id(
Expand All @@ -234,7 +234,7 @@ impl NightshadeRuntime {
let epoch_manager = self.epoch_manager.read();
let shard_version =
epoch_manager.get_shard_layout(epoch_id).map_err(Error::from)?.version();
Ok(ShardUId { version: shard_version, shard_id: new_shard_id_tmp(shard_id) as u32 })
Ok(ShardUId { version: shard_version, shard_id: shard_id_as_u32(shard_id) })
}

fn account_id_to_shard_uid(
Expand Down
3 changes: 2 additions & 1 deletion chain/chain/src/runtime/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@ use near_primitives::block::Tip;
use near_primitives::challenge::{ChallengesResult, PartialState, SlashedValidator};
use near_primitives::transaction::{Action, DeleteAccountAction, StakeAction, TransferAction};
use near_primitives::types::{
BlockHeightDelta, Nonce, ValidatorId, ValidatorInfoIdentifier, ValidatorKickoutReason,
new_shard_id_tmp, BlockHeightDelta, Nonce, ValidatorId, ValidatorInfoIdentifier,
ValidatorKickoutReason,
};
use near_primitives::validator_signer::ValidatorSigner;
use near_primitives::views::{
Expand Down
2 changes: 2 additions & 0 deletions chain/chunks/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use itertools::Itertools;
use near_pool::types::TransactionGroupIterator;
use near_pool::{InsertTransactionResult, PoolIteratorWrapper, TransactionPool};
use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout, ShardUId};
use near_primitives::types::shard_id_as_u16;
use near_primitives::{
epoch_info::RngSeed,
sharding::{EncodedShardChunk, PartialEncodedChunk, ShardChunk, ShardChunkHeader},
Expand Down Expand Up @@ -74,6 +75,7 @@ impl ShardedTransactionPool {
/// For better security we want the seed to different in each shard.
/// For testing purposes we want it to be the reproducible and derived from the `self.rng_seed` and `shard_id`
fn random_seed(base_seed: &RngSeed, shard_id: ShardId) -> RngSeed {
let shard_id = shard_id_as_u16(shard_id);
let mut res = *base_seed;
res[0] = shard_id as u8;
res[1] = (shard_id / 256) as u8;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ mod tests {
use near_primitives::hash::{hash, CryptoHash};
use near_primitives::sharding::{ShardChunkHeader, ShardChunkHeaderInner};
use near_primitives::stateless_validation::state_witness::ChunkStateWitness;
use near_primitives::types::{new_shard_id_tmp, BlockHeight, ShardId};
use near_primitives::types::{new_shard_id_tmp, shard_id_max, BlockHeight, ShardId};

use super::OrphanStateWitnessPool;

Expand Down Expand Up @@ -336,7 +336,7 @@ mod tests {
fn large_shard_id() {
let mut pool = OrphanStateWitnessPool::new(10);

let large_shard_id = ShardId::MAX;
let large_shard_id = shard_id_max();
let witness = make_witness(101, large_shard_id.into(), block(99), 0);
pool.add_orphan_state_witness(witness.clone(), 0);

Expand Down
3 changes: 2 additions & 1 deletion chain/client/src/tests/bug_repros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use futures::FutureExt;
use near_async::messaging::CanSend;
use near_async::time::Clock;
use near_network::shards_manager::ShardsManagerRequestFromNetwork;
use near_primitives::types::shard_id_as_usize;
use rand::{thread_rng, Rng};

use crate::test_utils::{setup_mock_all_validators, ActorHandlesForTesting};
Expand Down Expand Up @@ -112,7 +113,7 @@ fn repro_1183() {
// This test uses the V0 shard layout so it's ok to
// cast ShardId to ShardIndex.
let shard_id = account_id_to_shard_id(&from, 4);
let shard_index = shard_id as usize;
let shard_index = shard_id_as_usize(shard_id);
connectors1.write().unwrap()[shard_index].client_actor.do_send(
ProcessTxRequest {
transaction: SignedTransaction::send_money(
Expand Down
12 changes: 6 additions & 6 deletions chain/client/src/tests/cross_shard_tx.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use near_o11y::testonly::init_integration_logger;
use near_o11y::WithSpanContextExt;
use near_primitives::hash::CryptoHash;
use near_primitives::transaction::SignedTransaction;
use near_primitives::types::{AccountId, BlockId, BlockReference};
use near_primitives::types::{shard_id_as_usize, AccountId, BlockId, BlockReference};
use near_primitives::views::QueryResponseKind::ViewAccount;
use near_primitives::views::{QueryRequest, QueryResponse};
use std::collections::HashSet;
Expand Down Expand Up @@ -192,7 +192,7 @@ fn test_cross_shard_tx_callback(
// This test uses the V0 shard layout so it's ok to cast ShardId to
// ShardIndex.
let shard_id = account_id_to_shard_id(&account_id, 8);
let shard_index = shard_id as usize;
let shard_index = shard_id_as_usize(shard_id);
let actor = &connectors_[shard_index + (*presumable_epoch.read().unwrap() * 8) % 24]
.view_client_actor;
let actor = actor.send(
Expand Down Expand Up @@ -260,7 +260,7 @@ fn test_cross_shard_tx_callback(
// This test uses the V0 shard layout so it's ok to cast ShardId to
// ShardIndex.
let shard_id = account_id_to_shard_id(&validators[from], 8);
let shard_index = shard_id as usize;
let shard_index = shard_id_as_usize(shard_id);

send_tx(
validators.len(),
Expand Down Expand Up @@ -299,7 +299,7 @@ fn test_cross_shard_tx_callback(
// This test uses the V0 shard layout so it's ok to cast ShardId to
// ShardIndex.
let shard_id = account_id_to_shard_id(&validators[i], 8);
let shard_index = shard_id as usize;
let shard_index = shard_id_as_usize(shard_id);

let actor = &connectors_
[shard_index + (*presumable_epoch.read().unwrap() * 8) % 24]
Expand Down Expand Up @@ -359,7 +359,7 @@ fn test_cross_shard_tx_callback(
// This test uses the V0 shard layout so it's ok to cast ShardId to
// ShardIndex.
let shard_id = account_id_to_shard_id(&account_id, 8);
let shard_index = shard_id as usize;
let shard_index = shard_id_as_usize(shard_id);

let actor = &connectors_[shard_index + (*presumable_epoch.read().unwrap() * 8) % 24]
.view_client_actor;
Expand Down Expand Up @@ -521,7 +521,7 @@ fn test_cross_shard_tx_common(
// This test uses the V0 shard layout so it's ok to cast ShardId to
// ShardIndex.
let shard_id = account_id_to_shard_id(&validators[i], 8);
let shard_index = shard_id as usize;
let shard_index = shard_id_as_usize(shard_id);

let actor =
&connectors_[shard_index + *presumable_epoch.read().unwrap() * 8].view_client_actor;
Expand Down
12 changes: 9 additions & 3 deletions chain/network/src/peer_manager/tests/snapshot_hosts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ use near_o11y::WithSpanContextExt;
use near_primitives::hash::CryptoHash;
use near_primitives::network::PeerId;
use near_primitives::types::new_shard_id_tmp;
use near_primitives::types::shard_id_as_u64;
use near_primitives::types::shard_id_max;
use near_primitives::types::EpochHeight;
use near_primitives::types::ShardId;
use peer_manager::testonly::FDS_PER_PEER;
Expand Down Expand Up @@ -371,12 +373,16 @@ async fn large_shard_id_in_cache() {
let peer1 = pm.start_inbound(chain.clone(), peer1_config.clone()).await.handshake(clock).await;

tracing::info!(target:"test", "Send a SnapshotHostInfo message with very large shard ids.");
let max_shard_id: ShardId = ShardId::MAX;
let max_shard_id = shard_id_max();
let max_shard_id_minus_one = shard_id_as_u64(max_shard_id) - 1;
let max_shard_id_minus_one = new_shard_id_tmp(max_shard_id_minus_one);
let big_shard_info = Arc::new(SnapshotHostInfo::new(
peer1_config.node_id(),
CryptoHash::hash_borsh(1234_u64),
1234,
vec![0, 1232232, max_shard_id - 1, max_shard_id].into_iter().map(Into::into).collect(),
vec![new_shard_id_tmp(0), new_shard_id_tmp(1232232), max_shard_id_minus_one, max_shard_id]
.into_iter()
.collect(),
&peer1_config.node_key,
));

Expand Down Expand Up @@ -447,7 +453,7 @@ async fn too_many_shards_truncate() {
assert_eq!(info.shards.len(), MAX_SHARDS_PER_SNAPSHOT_HOST_INFO);
for &shard_id in &info.shards {
// Shard ids are taken from the original vector
assert!(shard_id < 2 * MAX_SHARDS_PER_SNAPSHOT_HOST_INFO as u64);
assert!(shard_id_as_u64(shard_id) < 2 * MAX_SHARDS_PER_SNAPSHOT_HOST_INFO as u64);
}
// The shard_ids are sorted and unique (no two elements are equal, hence the < condition instead of <=)
assert!(info.shards.windows(2).all(|twoelems| twoelems[0] < twoelems[1]));
Expand Down
4 changes: 4 additions & 0 deletions core/primitives-core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@ protocol_feature_fix_contract_loading_cost = []
protocol_feature_reject_blocks_with_outdated_protocol_version = []
protocol_feature_nonrefundable_transfer_nep491 = []

# TODO(wacban) remove after the transition is done
# default = ["new_shard_id"] # DO NOT COMMIT THIS
new_shard_id = []

nightly = [
"nightly_protocol",
"protocol_feature_fix_contract_loading_cost",
Expand Down
Loading

0 comments on commit 0c135f2

Please sign in to comment.