feat(resharding) - Make shard ids non-contiguous - part 3

chain/chain/src/chain.rs

@@ -3847,7 +3847,7 @@ impl Chain {
             let prev_hash = head.prev_block_hash;
             let epoch_height = self.epoch_manager.get_epoch_height_from_prev_block(&prev_hash)?;
             let shard_layout = &self.epoch_manager.get_shard_layout_from_prev_block(&prev_hash)?;
-            let shard_uids = shard_layout.shard_uids().collect();
+            let shard_uids = shard_layout.shard_uids().enumerate().collect();
             let last_block = self.get_block(&head.last_block_hash)?;
             let make_snapshot_callback = &snapshot_callbacks.make_snapshot_callback;
             make_snapshot_callback(prev_hash, epoch_height, shard_uids, last_block);

chain/chain/src/runtime/mod.rs

@@ -117,7 +117,7 @@ impl NightshadeRuntime {
             let epoch_manager = epoch_manager.read();
             let epoch_id = epoch_manager.get_epoch_id(&prev_block_hash)?;
             let shard_layout = epoch_manager.get_shard_layout(&epoch_id)?;
-            Ok(shard_layout.shard_uids().collect())
+            Ok(shard_layout.shard_uids().enumerate().collect())
         }) {
             tracing::debug!(target: "runtime", ?err, "The state snapshot is not available.");
         }

chain/chain/src/state_snapshot_actor.rs

@@ -5,7 +5,7 @@ use near_performance_metrics_macros::perf;
 use near_primitives::block::Block;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::ShardUId;
-use near_primitives::types::EpochHeight;
+use near_primitives::types::{EpochHeight, ShardIndex};
 use near_store::flat::FlatStorageManager;
 use near_store::ShardTries;
 use std::sync::Arc;
@@ -49,7 +49,7 @@ pub struct CreateSnapshotRequest {
     /// epoch height associated with prev_block_hash
     epoch_height: EpochHeight,
     /// Shards that need to be present in the snapshot.
-    shard_uids: Vec<ShardUId>,
+    shard_indexes_and_uids: Vec<(ShardIndex, ShardUId)>,
     /// Last block of the prev epoch.
     block: Block,
 }
@@ -74,13 +74,15 @@ impl StateSnapshotActor {
     pub fn handle_create_snapshot_request(&mut self, msg: CreateSnapshotRequest) {
         tracing::debug!(target: "state_snapshot", ?msg);
 
-        let CreateSnapshotRequest { prev_block_hash, epoch_height, shard_uids, block } = msg;
-        let res = self.tries.create_state_snapshot(prev_block_hash, &shard_uids, &block);
+        let CreateSnapshotRequest { prev_block_hash, epoch_height, shard_indexes_and_uids, block } =
+            msg;
+        let res =
+            self.tries.create_state_snapshot(prev_block_hash, &shard_indexes_and_uids, &block);
 
         // Unlocking flat state head can be done asynchronously in state_snapshot_actor.
         // The next flat storage update will bring flat storage to latest head.
         if !self.flat_storage_manager.set_flat_state_updates_mode(true) {
-            tracing::error!(target: "state_snapshot", ?prev_block_hash, ?shard_uids, "Failed to unlock flat state updates");
+            tracing::error!(target: "state_snapshot", ?prev_block_hash, ?shard_indexes_and_uids, "Failed to unlock flat state updates");
         }
         match res {
             Ok(res_shard_uids) => {
@@ -125,8 +127,12 @@ pub struct StateSnapshotSenderForStateSnapshot {
 #[derive(Clone, MultiSend, MultiSenderFrom)]
 pub struct StateSnapshotSenderForClient(Sender<DeleteAndMaybeCreateSnapshotRequest>);
 
-type MakeSnapshotCallback =
-    Arc<dyn Fn(CryptoHash, EpochHeight, Vec<ShardUId>, Block) -> () + Send + Sync + 'static>;
+type MakeSnapshotCallback = Arc<
+    dyn Fn(CryptoHash, EpochHeight, Vec<(ShardIndex, ShardUId)>, Block) -> ()
+        + Send
+        + Sync
+        + 'static,
+>;
 
 type DeleteSnapshotCallback = Arc<dyn Fn() -> () + Send + Sync + 'static>;
 
@@ -140,20 +146,20 @@ pub fn get_make_snapshot_callback(
     sender: StateSnapshotSenderForClient,
     flat_storage_manager: FlatStorageManager,
 ) -> MakeSnapshotCallback {
-    Arc::new(move |prev_block_hash, epoch_height, shard_uids, block| {
+    Arc::new(move |prev_block_hash, epoch_height, shard_indexes_and_uids, block| {
         tracing::info!(
             target: "state_snapshot",
             ?prev_block_hash,
-            ?shard_uids,
+            ?shard_indexes_and_uids,
             "make_snapshot_callback sends `DeleteAndMaybeCreateSnapshotRequest` to state_snapshot_addr");
         // We need to stop flat head updates synchronously in the client thread.
         // Async update in state_snapshot_actor and potentially lead to flat head progressing beyond prev_block_hash
         if !flat_storage_manager.set_flat_state_updates_mode(false) {
-            tracing::error!(target: "state_snapshot", ?prev_block_hash, ?shard_uids, "Failed to lock flat state updates");
+            tracing::error!(target: "state_snapshot", ?prev_block_hash, ?shard_indexes_and_uids, "Failed to lock flat state updates");
             return;
         }
         let create_snapshot_request =
-            CreateSnapshotRequest { prev_block_hash, epoch_height, shard_uids, block };
+            CreateSnapshotRequest { prev_block_hash, epoch_height, shard_indexes_and_uids, block };
         sender.send(DeleteAndMaybeCreateSnapshotRequest {
             create_snapshot_request: Some(create_snapshot_request),
         });

chain/client/src/test_utils/test_env_builder.rs

@@ -19,7 +19,7 @@ use near_parameters::RuntimeConfigStore;
 use near_primitives::epoch_info::RngSeed;
 use near_primitives::epoch_manager::{AllEpochConfigTestOverrides, EpochConfig, EpochConfigStore};
 use near_primitives::test_utils::create_test_signer;
-use near_primitives::types::{AccountId, NumShards};
+use near_primitives::types::{AccountId, NumShards, ShardIndex};
 use near_store::config::StateSnapshotType;
 use near_store::test_utils::create_test_store;
 use near_store::{NodeStorage, ShardUId, Store, StoreConfig, TrieConfig};
@@ -588,7 +588,7 @@ impl TestEnvBuilder {
                         None => TEST_SEED,
                     };
                     let tries = runtime.get_tries();
-                    let make_snapshot_callback = Arc::new(move |prev_block_hash, _epoch_height, shard_uids: Vec<ShardUId>, block| {
+                    let make_snapshot_callback = Arc::new(move |prev_block_hash, _epoch_height, shard_uids: Vec<(ShardIndex, ShardUId)>, block| {
                         tracing::info!(target: "state_snapshot", ?prev_block_hash, "make_snapshot_callback");
                         tries.delete_state_snapshot();
                         tries.create_state_snapshot(prev_block_hash, &shard_uids, &block).unwrap();

core/primitives-core/src/types.rs

@@ -1,5 +1,7 @@
 // #[cfg(feature = "new_shard_id")]
 // use std::fmt::Display;
+// #[cfg(feature = "new_shard_id")]
+// use std::ops::Add;
 
 use crate::hash::CryptoHash;
 
@@ -266,4 +268,16 @@ impl Into<u16> for ShardId {
     }
 }
 
+#[cfg(feature = "new_shard_id")]
+impl<T> Add<T> for ShardId
+where
+    T: Add<u64, Output = u64>,
+{
+    type Output = Self;
+
+    fn add(self, rhs: T) -> Self::Output {
+        Self(T::add(rhs, self.0))
+    }
+}
+
 */

core/store/src/trie/state_snapshot.rs

@@ -11,6 +11,7 @@ use near_primitives::errors::EpochError;
 use near_primitives::errors::StorageError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::ShardUId;
+use near_primitives::types::ShardIndex;
 use std::error::Error;
 use std::io;
 use std::path::{Path, PathBuf};
@@ -83,32 +84,32 @@ impl StateSnapshot {
         store: TrieStoreAdapter,
         prev_block_hash: CryptoHash,
         flat_storage_manager: FlatStorageManager,
-        requested_shard_uids: &[ShardUId],
+        shard_indexes_and_uids: &[(ShardIndex, ShardUId)],
         block: Option<&Block>,
     ) -> Self {
-        tracing::debug!(target: "state_snapshot", ?requested_shard_uids, ?prev_block_hash, "new StateSnapshot");
+        tracing::debug!(target: "state_snapshot", ?shard_indexes_and_uids, ?prev_block_hash, "new StateSnapshot");
         let mut included_shard_uids = vec![];
-        for shard_uid in requested_shard_uids {
-            if let Err(err) = flat_storage_manager.create_flat_storage_for_shard(*shard_uid) {
+        for &(shard_index, shard_uid) in shard_indexes_and_uids {
+            if let Err(err) = flat_storage_manager.create_flat_storage_for_shard(shard_uid) {
                 tracing::warn!(target: "state_snapshot", ?err, ?shard_uid, "Failed to create a flat storage for snapshot shard");
                 continue;
             }
             if let Some(block) = block {
                 let flat_storage =
-                    flat_storage_manager.get_flat_storage_for_shard(*shard_uid).unwrap();
+                    flat_storage_manager.get_flat_storage_for_shard(shard_uid).unwrap();
                 let current_flat_head = flat_storage.get_head_hash();
                 tracing::debug!(target: "state_snapshot", ?shard_uid, ?current_flat_head, block_hash = ?block.header().hash(), block_height = block.header().height(), "Moving FlatStorage head of the snapshot");
                 let _timer = metrics::MOVE_STATE_SNAPSHOT_FLAT_HEAD_ELAPSED
                     .with_label_values(&[&shard_uid.shard_id.to_string()])
                     .start_timer();
-                if let Some(chunk) = block.chunks().get(shard_uid.shard_id as usize) {
+                if let Some(chunk) = block.chunks().get(shard_index) {
                     // Flat state snapshot needs to be at a height that lets it
                     // replay the last chunk of the shard.
                     let desired_flat_head = chunk.prev_block_hash();
                     match flat_storage.update_flat_head(desired_flat_head) {
                         Ok(_) => {
                             tracing::debug!(target: "state_snapshot", ?shard_uid, ?current_flat_head, ?desired_flat_head, "Successfully moved FlatStorage head of the snapshot");
-                            included_shard_uids.push(*shard_uid);
+                            included_shard_uids.push(shard_uid);
                         }
                         Err(err) => {
                             tracing::error!(target: "state_snapshot", ?shard_uid, ?err, ?current_flat_head, ?desired_flat_head, "Failed to move FlatStorage head of the snapshot");
@@ -155,6 +156,8 @@ pub const STATE_SNAPSHOT_COLUMNS: &[DBCol] = &[
     DBCol::FlatStorageStatus,
 ];
 
+type ShardIndexesAndUIds = Vec<(ShardIndex, ShardUId)>;
+
 impl ShardTries {
     pub fn get_state_snapshot(
         &self,
@@ -177,7 +180,7 @@ impl ShardTries {
     pub fn create_state_snapshot(
         &self,
         prev_block_hash: CryptoHash,
-        shard_uids: &[ShardUId],
+        shard_indexes_and_uids: &[(ShardIndex, ShardUId)],
         block: &Block,
     ) -> Result<Option<Vec<ShardUId>>, anyhow::Error> {
         metrics::HAS_STATE_SNAPSHOT.set(0);
@@ -224,7 +227,7 @@ impl ShardTries {
             store,
             prev_block_hash,
             flat_storage_manager,
-            shard_uids,
+            shard_indexes_and_uids,
             Some(block),
         ));
 
@@ -317,7 +320,9 @@ impl ShardTries {
     /// we don't deal with multiple snapshots here because we will deal with it whenever a new snapshot is created and saved to file system
     pub fn maybe_open_state_snapshot(
         &self,
-        get_shard_uids_fn: impl FnOnce(CryptoHash) -> Result<Vec<ShardUId>, EpochError>,
+        get_shard_indexes_and_uids_fn: impl FnOnce(
+            CryptoHash,
+        ) -> Result<ShardIndexesAndUIds, EpochError>,
     ) -> Result<(), anyhow::Error> {
         let _span =
             tracing::info_span!(target: "state_snapshot", "maybe_open_state_snapshot").entered();
@@ -346,10 +351,15 @@ impl ShardTries {
         let store = storage.get_hot_store().trie_store();
         let flat_storage_manager = FlatStorageManager::new(store.flat_store());
 
-        let shard_uids = get_shard_uids_fn(snapshot_hash)?;
+        let shard_indexes_and_uids = get_shard_indexes_and_uids_fn(snapshot_hash)?;
         let mut guard = self.state_snapshot().write().unwrap();
-        *guard =
-            Some(StateSnapshot::new(store, snapshot_hash, flat_storage_manager, &shard_uids, None));
+        *guard = Some(StateSnapshot::new(
+            store,
+            snapshot_hash,
+            flat_storage_manager,
+            &shard_indexes_and_uids,
+            None,
+        ));
         metrics::HAS_STATE_SNAPSHOT.set(1);
         tracing::info!(target: "runtime", ?snapshot_hash, ?snapshot_path, "Detected and opened a state snapshot.");
         Ok(())

integration-tests/src/tests/client/block_corruption.rs

@@ -8,6 +8,7 @@ use near_crypto::{InMemorySigner, KeyType};
 use near_o11y::testonly::init_test_logger;
 use near_primitives::sharding::{ShardChunkHeader, ShardChunkHeaderInner};
 use near_primitives::transaction::SignedTransaction;
+use near_primitives::types::new_shard_id_tmp;
 use near_primitives::validator_signer::InMemoryValidatorSigner;
 use near_primitives_core::types::BlockHeight;
 use nearcore::test_utils::TestEnvNightshadeSetupExt;
@@ -61,7 +62,7 @@ fn change_shard_id_to_invalid() {
     let mut block = env.clients[0].produce_block(2).unwrap().unwrap();
 
     // 1. Corrupt chunks
-    let bad_shard_id = 100;
+    let bad_shard_id = new_shard_id_tmp(100);
     let mut new_chunks = vec![];
     for chunk in block.chunks().iter() {
         let mut new_chunk = chunk.clone();

integration-tests/src/tests/client/challenges.rs

@@ -22,7 +22,7 @@ use near_primitives::stateless_validation::chunk_endorsement::ChunkEndorsementV1
 use near_primitives::test_utils::create_test_signer;
 use near_primitives::transaction::SignedTransaction;
 use near_primitives::types::chunk_extra::ChunkExtra;
-use near_primitives::types::{AccountId, ShardId};
+use near_primitives::types::{new_shard_id_tmp, AccountId, ShardId};
 use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
 use near_store::Trie;
 use nearcore::test_utils::TestEnvNightshadeSetupExt;
@@ -312,12 +312,16 @@ fn challenge(
     chunk: Box<MaybeEncodedShardChunk>,
     block: &Block,
 ) -> Result<(CryptoHash, Vec<AccountId>), Error> {
+    let epoch_id = block.header().epoch_id();
+    let shard_layout = env.clients[0].chain.epoch_manager.get_shard_layout(epoch_id).unwrap();
+    let shard_index = shard_layout.get_shard_index(shard_id);
+
     let merkle_paths = Block::compute_chunk_headers_root(block.chunks().iter()).1;
     let valid_challenge = Challenge::produce(
         ChallengeBody::ChunkProofs(ChunkProofs {
             block_header: borsh::to_vec(&block.header()).unwrap(),
             chunk,
-            merkle_proof: merkle_paths[shard_id as usize].clone(),
+            merkle_proof: merkle_paths[shard_index].clone(),
         }),
         &*env.clients[0].validator_signer.get().unwrap(),
     );
@@ -372,7 +376,7 @@ fn test_verify_chunk_invalid_state_challenge() {
         Trie::EMPTY_ROOT,
         CryptoHash::default(),
         last_block.header().height() + 1,
-        0,
+        new_shard_id_tmp(0),
         0,
         1_000,
         0,

integration-tests/src/tests/client/features/adversarial_behaviors.rs

@@ -236,6 +236,7 @@ fn test_banning_chunk_producer_when_seeing_invalid_chunk_base(
             )
             .unwrap();
         let block_producer = epoch_manager.get_block_producer(&epoch_id, height).unwrap();
+        let shard_layout = epoch_manager.get_shard_layout(&epoch_id).unwrap();
 
         let block = test.env.client(&block_producer).produce_block(height).unwrap().unwrap();
         assert_eq!(block.header().height(), height);
@@ -248,7 +249,7 @@ fn test_banning_chunk_producer_when_seeing_invalid_chunk_base(
             } else {
                 assert_eq!(block.header().prev_height().unwrap(), height - 1);
             }
-            for shard_id in 0..4 {
+            for shard_id in shard_layout.shard_ids() {
                 let chunk_producer = epoch_manager
                     .get_chunk_producer(
                         &epoch_id,
@@ -286,20 +287,25 @@ fn test_banning_chunk_producer_when_seeing_invalid_chunk_base(
         if height > 1 {
             let prev_block =
                 test.env.clients[0].chain.get_block(&block.header().prev_hash()).unwrap();
-            for i in 0..4 {
-                if invalid_chunks_in_this_block.contains(&(i as ShardId))
-                    && !this_block_should_be_skipped
+            for shard_id in shard_layout.shard_ids() {
+                let shard_index = shard_layout.get_shard_index(shard_id);
+                if invalid_chunks_in_this_block.contains(&shard_id) && !this_block_should_be_skipped
                 {
-                    assert_eq!(block.chunks()[i].chunk_hash(), prev_block.chunks()[i].chunk_hash());
+                    assert_eq!(
+                        block.chunks()[shard_index].chunk_hash(),
+                        prev_block.chunks()[shard_index].chunk_hash()
+                    );
                 } else {
                     // TODO: mysteriously we might miss a chunk around epoch boundaries.
                     // Figure out why...
-                    assert!(
-                        block.chunks()[i].height_created() == prev_block.header().height() + 1
-                            || (height % EPOCH_LENGTH == 1
-                                && block.chunks()[i].chunk_hash()
-                                    == prev_block.chunks()[i].chunk_hash())
-                    );
+                    let is_epoch_boundary = height % EPOCH_LENGTH == 1;
+                    let chunk_header = &block.chunks()[shard_index];
+                    let prev_chunk_header = &prev_block.chunks()[shard_index];
+                    let is_new_chunk =
+                        chunk_header.height_created() == prev_block.header().height() + 1;
+                    let is_old_chunk_same_as_prev =
+                        chunk_header.chunk_hash() == prev_chunk_header.chunk_hash();
+                    assert!(is_new_chunk || (is_epoch_boundary && is_old_chunk_same_as_prev));
                 }
             }
         }