Merge remote-tracking branch 'origin/master' into fork-net-shards

Author: marcelo-gonzalez

chain/chain/src/chain.rs

@@ -52,14 +52,15 @@ use near_async::messaging::{noop, IntoMultiSender};
 use near_async::time::{Clock, Duration, Instant};
 use near_chain_configs::{MutableConfigValue, MutableValidatorSigner};
 use near_chain_primitives::error::{BlockKnownError, Error, LogTransientStorageError};
+use near_epoch_manager::shard_assignment::shard_id_to_uid;
 use near_epoch_manager::shard_tracker::ShardTracker;
 use near_epoch_manager::EpochManagerAdapter;
 use near_primitives::bandwidth_scheduler::BandwidthRequests;
 use near_primitives::block::{genesis_chunks, Block, BlockValidityError, Chunks, MaybeNew, Tip};
 use near_primitives::block_header::BlockHeader;
 use near_primitives::challenge::{
     BlockDoubleSign, Challenge, ChallengeBody, ChallengesResult, ChunkProofs, ChunkState,
-    MaybeEncodedShardChunk, PartialState, SlashedValidator,
+    MaybeEncodedShardChunk, SlashedValidator,
 };
 use near_primitives::checked_feature;
 use near_primitives::congestion_info::CongestionInfo;
@@ -74,6 +75,7 @@ use near_primitives::sharding::{
     ChunkHash, ChunkHashHeight, EncodedShardChunk, ReceiptList, ReceiptProof, ShardChunk,
     ShardChunkHeader, ShardProof, StateSyncInfo,
 };
+use near_primitives::state::PartialState;
 use near_primitives::state_part::PartId;
 use near_primitives::state_sync::{
     get_num_state_parts, ReceiptProofResponse, RootProof, ShardStateSyncResponseHeader,
@@ -365,12 +367,8 @@ impl Chain {
     ) -> Result<Chain, Error> {
         let store = runtime_adapter.store();
         let transaction_validity_period = chain_genesis.transaction_validity_period;
-        let chain_store = ChainStore::new(
-            store.clone(),
-            chain_genesis.height,
-            save_trie_changes,
-            transaction_validity_period,
-        );
+        let chain_store =
+            ChainStore::new(store.clone(), save_trie_changes, transaction_validity_period);
         let state_roots = get_genesis_state_roots(runtime_adapter.store())?
             .expect("genesis should be initialized.");
         let (genesis, _genesis_chunks) = Self::make_genesis_block(
@@ -439,7 +437,6 @@ impl Chain {
         // Check if we have a head in the store, otherwise pick genesis block.
         let mut chain_store = ChainStore::new(
             runtime_adapter.store().clone(),
-            chain_genesis.height,
             chain_config.save_trie_changes,
             transaction_validity_period,
         );
@@ -710,7 +707,7 @@ impl Chain {
 
             store_update.save_chunk_extra(
                 genesis.hash(),
-                &epoch_manager.shard_id_to_uid(chunk_header.shard_id(), &EpochId::default())?,
+                &shard_id_to_uid(epoch_manager, chunk_header.shard_id(), &EpochId::default())?,
                 Self::create_genesis_chunk_extra(
                     state_root,
                     chain_genesis.gas_limit,
@@ -901,14 +898,8 @@ impl Chain {
     fn partial_verify_orphan_header_signature(&self, header: &BlockHeader) -> Result<bool, Error> {
         let block_producer =
             self.epoch_manager.get_block_producer(header.epoch_id(), header.height())?;
-        // DEVNOTE: we pass head which is not necessarily on block's chain, but it's only used for
-        // slashing info which we will ignore
-        let head = self.head()?;
-        let (block_producer, _slashed) = self.epoch_manager.get_validator_by_account_id(
-            header.epoch_id(),
-            &head.last_block_hash,
-            &block_producer,
-        )?;
+        let block_producer =
+            self.epoch_manager.get_validator_by_account_id(header.epoch_id(), &block_producer)?;
         Ok(header.signature().verify(header.hash().as_ref(), block_producer.public_key()))
     }
 
@@ -1206,7 +1197,6 @@ impl Chain {
         &self,
         challenges: &[Challenge],
         epoch_id: &EpochId,
-        prev_block_hash: &CryptoHash,
     ) -> Result<(ChallengesResult, Vec<CryptoHash>), Error> {
         let _span = tracing::debug_span!(
             target: "chain",
@@ -1220,7 +1210,6 @@ impl Chain {
                 self.epoch_manager.as_ref(),
                 self.runtime_adapter.as_ref(),
                 epoch_id,
-                prev_block_hash,
                 challenge,
             ) {
                 Ok((hash, account_ids)) => {
@@ -1557,7 +1546,7 @@ impl Chain {
     /// soon as possible and allow next block producer to skip invalid blocks.
     pub fn process_challenge(&mut self, challenge: &Challenge) {
         let head = unwrap_or_return!(self.head());
-        match self.verify_challenges(&[challenge.clone()], &head.epoch_id, &head.last_block_hash) {
+        match self.verify_challenges(&[challenge.clone()], &head.epoch_id) {
             Ok((_, challenged_blocks)) => {
                 let mut chain_update = self.chain_update();
                 for block_hash in challenged_blocks {
@@ -2010,7 +1999,7 @@ impl Chain {
                 true,
             );
             let care_about_shard_this_or_next_epoch = care_about_shard || will_care_about_shard;
-            let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, &epoch_id).unwrap();
+            let shard_uid = shard_id_to_uid(self.epoch_manager.as_ref(), shard_id, epoch_id)?;
             if care_about_shard_this_or_next_epoch {
                 shards_cares_this_or_next_epoch.push(shard_uid);
             }
@@ -2177,7 +2166,7 @@ impl Chain {
         shard_id: ShardId,
     ) -> Result<(), Error> {
         let epoch_id = block.header().epoch_id();
-        let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, epoch_id)?;
+        let shard_uid = shard_id_to_uid(self.epoch_manager.as_ref(), shard_id, epoch_id)?;
 
         // Update flat storage.
         let flat_storage_manager = self.runtime_adapter.get_flat_storage_manager();
@@ -2334,7 +2323,7 @@ impl Chain {
         }
 
         let (challenges_result, challenged_blocks) =
-            self.verify_challenges(block.challenges(), header.epoch_id(), header.prev_hash())?;
+            self.verify_challenges(block.challenges(), header.epoch_id())?;
 
         let prev_block = self.get_block(&prev_hash)?;
 
@@ -3297,7 +3286,7 @@ impl Chain {
                 shard_id,
                 true,
             ) {
-                let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, &epoch_id)?;
+                let shard_uid = shard_id_to_uid(self.epoch_manager.as_ref(), shard_id, epoch_id)?;
                 self.resharding_manager.start_resharding(
                     self.chain_store.store_update(),
                     &block,
@@ -3742,7 +3731,7 @@ impl Chain {
             cares_about_shard_next_epoch,
             cared_about_shard_prev_epoch,
         );
-        let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, epoch_id)?;
+        let shard_uid = shard_id_to_uid(self.epoch_manager.as_ref(), shard_id, epoch_id)?;
         Ok(ShardContext { shard_uid, should_apply_chunk })
     }
 
@@ -4483,9 +4472,16 @@ impl Chain {
     ) -> HashMap<ShardId, Vec<Receipt>> {
         let mut result = HashMap::new();
         for receipt in receipts {
-            let shard_id = shard_layout.account_id_to_shard_id(receipt.receiver_id());
-            let entry = result.entry(shard_id).or_insert_with(Vec::new);
-            entry.push(receipt)
+            if receipt.send_to_all_shards() {
+                for shard_id in shard_layout.shard_ids() {
+                    let entry = result.entry(shard_id).or_insert_with(Vec::new);
+                    entry.push(receipt.clone());
+                }
+            } else {
+                let shard_id = shard_layout.account_id_to_shard_id(receipt.receiver_id());
+                let entry = result.entry(shard_id).or_insert_with(Vec::new);
+                entry.push(receipt);
+            }
         }
         result
     }
@@ -4506,13 +4502,22 @@ impl Chain {
         }
         let mut cache = HashMap::new();
         for receipt in receipts {
-            let &mut shard_id = cache
-                .entry(receipt.receiver_id())
-                .or_insert_with(|| shard_layout.account_id_to_shard_id(receipt.receiver_id()));
-            // This unwrap should be safe as we pre-populated the map with all
-            // valid shard ids.
-            let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
-            result_map.get_mut(&shard_index).unwrap().1.push(receipt);
+            if receipt.send_to_all_shards() {
+                for shard_id in shard_layout.shard_ids() {
+                    // This unwrap should be safe as we pre-populated the map with all
+                    // valid shard ids.
+                    let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
+                    result_map.get_mut(&shard_index).unwrap().1.push(receipt);
+                }
+            } else {
+                let &mut shard_id = cache
+                    .entry(receipt.receiver_id())
+                    .or_insert_with(|| shard_layout.account_id_to_shard_id(receipt.receiver_id()));
+                // This unwrap should be safe as we pre-populated the map with all
+                // valid shard ids.
+                let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
+                result_map.get_mut(&shard_index).unwrap().1.push(receipt);
+            }
         }
 
         let mut result_vec = vec![];

chain/chain/src/chain_update.rs

@@ -11,6 +11,7 @@ use crate::update_shard::{NewChunkResult, OldChunkResult, ShardUpdateResult};
 use crate::{metrics, DoomslugThresholdMode};
 use crate::{Chain, Doomslug};
 use near_chain_primitives::error::Error;
+use near_epoch_manager::shard_assignment::shard_id_to_uid;
 use near_epoch_manager::EpochManagerAdapter;
 use near_primitives::apply::ApplyChunkReason;
 use near_primitives::block::{Block, Tip};
@@ -562,7 +563,8 @@ impl<'a> ChainUpdate<'a> {
 
         self.chain_store_update.save_chunk(chunk);
 
-        let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, block_header.epoch_id())?;
+        let shard_uid =
+            shard_id_to_uid(self.epoch_manager.as_ref(), shard_id, block_header.epoch_id())?;
         let flat_storage_manager = self.runtime_adapter.get_flat_storage_manager();
         let store_update = flat_storage_manager.save_flat_state_changes(
             *block_header.hash(),
@@ -639,7 +641,8 @@ impl<'a> ChainUpdate<'a> {
         let prev_hash = block_header.prev_hash();
         let prev_block_header = self.chain_store_update.get_block_header(prev_hash)?;
 
-        let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, block_header.epoch_id())?;
+        let shard_uid =
+            shard_id_to_uid(self.epoch_manager.as_ref(), shard_id, block_header.epoch_id())?;
         let chunk_extra = self.chain_store_update.get_chunk_extra(prev_hash, &shard_uid)?;
 
         let apply_result = self.runtime_adapter.apply_chunk(

chain/chain/src/flat_storage_init.rs

@@ -1,4 +1,5 @@
 use near_chain_primitives::Error;
+use near_epoch_manager::shard_assignment::shard_id_to_uid;
 use near_epoch_manager::EpochManagerAdapter;
 use near_primitives::block::Tip;
 use near_store::flat::{FlatStorageManager, FlatStorageStatus};
@@ -44,7 +45,7 @@ fn init_flat_storage_for_current_epoch(
 
     let shard_ids = epoch_manager.shard_ids(epoch_id)?;
     for shard_id in shard_ids {
-        let shard_uid = epoch_manager.shard_id_to_uid(shard_id, &chain_head.epoch_id)?;
+        let shard_uid = shard_id_to_uid(epoch_manager, shard_id, &chain_head.epoch_id)?;
         let status = flat_storage_manager.get_flat_storage_status(shard_uid);
         match status {
             FlatStorageStatus::Ready(_) => {
@@ -76,7 +77,7 @@ fn init_flat_storage_for_next_epoch(
 
     let shard_ids = epoch_manager.shard_ids(next_epoch_id)?;
     for shard_id in shard_ids {
-        let shard_uid = epoch_manager.shard_id_to_uid(shard_id, next_epoch_id)?;
+        let shard_uid = shard_id_to_uid(epoch_manager, shard_id, next_epoch_id)?;
         let status = flat_storage_manager.get_flat_storage_status(shard_uid);
         match status {
             FlatStorageStatus::Ready(_) => {

chain/chain/src/flat_storage_resharder.rs

@@ -1512,7 +1512,6 @@ mod tests {
         let chain_genesis = ChainGenesis::new(&genesis.config);
         let sender = Arc::new(T::new(ChainStore::new(
             store,
-            chain_genesis.height,
             false,
             chain_genesis.transaction_validity_period,
         )));

chain/chain/src/garbage_collection.rs

@@ -4,6 +4,7 @@ use std::{fmt, io};
 
 use near_chain_configs::GCConfig;
 use near_chain_primitives::Error;
+use near_epoch_manager::shard_assignment::shard_id_to_uid;
 use near_epoch_manager::shard_tracker::ShardTracker;
 use near_epoch_manager::EpochManagerAdapter;
 use near_primitives::block::Block;
@@ -739,7 +740,7 @@ impl<'a> ChainStoreUpdate<'a> {
 
         // 1. Delete shard_id-indexed data (TrieChanges, Receipts, ChunkExtra, State Headers and Parts, FlatStorage data)
         for shard_id in shard_layout.shard_ids() {
-            let shard_uid = epoch_manager.shard_id_to_uid(shard_id, epoch_id).unwrap();
+            let shard_uid = shard_id_to_uid(epoch_manager, shard_id, epoch_id).unwrap();
             let block_shard_id = get_block_shard_uid(&block_hash, &shard_uid);
 
             // delete TrieChanges

chain/chain/src/resharding/manager.rs

@@ -12,10 +12,10 @@ use near_chain_configs::{MutableConfigValue, ReshardingConfig, ReshardingHandle}
 use near_chain_primitives::Error;
 use near_epoch_manager::EpochManagerAdapter;
 use near_primitives::block::Block;
-use near_primitives::challenge::PartialState;
 use near_primitives::congestion_info::CongestionInfo;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::{get_block_shard_uid, ShardLayout};
+use near_primitives::state::PartialState;
 use near_primitives::types::chunk_extra::ChunkExtra;
 use near_store::adapter::trie_store::get_shard_uid_mapping;
 use near_store::adapter::{StoreAdapter, StoreUpdateAdapter};

chain/chain/src/resharding/resharding_actor.rs

@@ -46,14 +46,7 @@ impl Handler<MemtrieReloadRequest> for ReshardingActor {
 
 impl ReshardingActor {
     pub fn new(store: Store, genesis: &ChainGenesis) -> Self {
-        Self {
-            chain_store: ChainStore::new(
-                store,
-                genesis.height,
-                false,
-                genesis.transaction_validity_period,
-            ),
-        }
+        Self { chain_store: ChainStore::new(store, false, genesis.transaction_validity_period) }
     }
 
     fn handle_flat_storage_split_shard(

chain/chain/src/runtime/mod.rs

@@ -9,6 +9,7 @@ use errors::FromStateViewerErrors;
 use near_async::time::{Duration, Instant};
 use near_chain_configs::{GenesisConfig, ProtocolConfig, MIN_GC_NUM_EPOCHS_TO_KEEP};
 use near_crypto::PublicKey;
+use near_epoch_manager::shard_assignment::account_id_to_shard_id;
 use near_epoch_manager::{EpochManagerAdapter, EpochManagerHandle};
 use near_parameters::{ActionCosts, ExtCosts, RuntimeConfig, RuntimeConfigStore};
 use near_pool::types::TransactionGroupIterator;
@@ -1378,8 +1379,7 @@ fn congestion_control_accepts_transaction(
         return Ok(true);
     }
     let receiver_id = tx.transaction.receiver_id();
-    let receiving_shard = epoch_manager.account_id_to_shard_id(receiver_id, &epoch_id)?;
-
+    let receiving_shard = account_id_to_shard_id(epoch_manager, receiver_id, &epoch_id)?;
     let congestion_info = prev_block.congestion_info.get(&receiving_shard);
     let Some(congestion_info) = congestion_info else {
         return Ok(true);