feat(resharding) - Make shard ids non-contiguous - part 2 (#12209)

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.

chain/chain/src/runtime/mod.rs

@@ -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::{
@@ -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(
@@ -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(

chain/chain/src/runtime/tests.rs

@@ -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::{

chain/chunks/src/client.rs

@@ -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},
@@ -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;

chain/client/src/stateless_validation/chunk_validator/orphan_witness_pool.rs

@@ -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;
 
@@ -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);
 

chain/client/src/tests/bug_repros.rs

@@ -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};
@@ -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(

chain/client/src/tests/cross_shard_tx.rs

@@ -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;
@@ -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(
@@ -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(),
@@ -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]
@@ -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;
@@ -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;

chain/network/src/peer_manager/tests/snapshot_hosts.rs

@@ -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;
@@ -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,
     ));
 
@@ -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]));

core/primitives-core/Cargo.toml

@@ -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",