[stateless_validation] Add cache for generated chunk_validators (#10442)

This PR adds a layer or caching for calls to `get_chunk_validators` in epoch_manager
near · Jan 17, 2024 · 82e9e4e · 82e9e4e
1 parent db20a66
commit 82e9e4e
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 24 deletions.
diff --git a/chain/chain/src/test_utils/kv_runtime.rs b/chain/chain/src/test_utils/kv_runtime.rs
@@ -13,7 +13,7 @@ use near_epoch_manager::{EpochManagerAdapter, RngSeed};
 use near_pool::types::PoolIterator;
 use near_primitives::account::{AccessKey, Account};
 use near_primitives::block_header::{Approval, ApprovalInner};
-use near_primitives::chunk_validation::ChunkEndorsement;
+use near_primitives::chunk_validation::{ChunkEndorsement, ChunkValidators};
 use near_primitives::epoch_manager::block_info::BlockInfo;
 use near_primitives::epoch_manager::epoch_info::EpochInfo;
 use near_primitives::epoch_manager::EpochConfig;
@@ -35,7 +35,6 @@ use near_primitives::types::{
     AccountId, ApprovalStake, Balance, BlockHeight, EpochHeight, EpochId, Gas, Nonce, NumShards,
     ShardId, StateChangesForResharding, StateRoot, StateRootNode, ValidatorInfoIdentifier,
 };
-use near_primitives::validator_mandates::AssignmentWeight;
 use near_primitives::version::{ProtocolVersion, PROTOCOL_VERSION};
 use near_primitives::views::{
     AccessKeyInfoView, AccessKeyList, CallResult, ContractCodeView, EpochValidatorInfo,
@@ -712,8 +711,8 @@ impl EpochManagerAdapter for MockEpochManager {
         _epoch_id: &EpochId,
         _shard_id: ShardId,
         _height: BlockHeight,
-    ) -> Result<HashMap<AccountId, AssignmentWeight>, EpochError> {
-        Ok(HashMap::new())
+    ) -> Result<Arc<ChunkValidators>, EpochError> {
+        Ok(Arc::new(HashMap::new()))
     }
 
     fn get_validator_by_account_id(

diff --git a/chain/client/src/chunk_validation.rs b/chain/client/src/chunk_validation.rs
@@ -528,7 +528,7 @@ impl Client {
             chunk_validators.keys(),
         );
         self.network_adapter.send(PeerManagerMessageRequest::NetworkRequests(
-            NetworkRequests::ChunkStateWitness(chunk_validators.into_keys().collect(), witness),
+            NetworkRequests::ChunkStateWitness(chunk_validators.keys().cloned().collect(), witness),
         ));
         Ok(())
     }

diff --git a/chain/epoch-manager/src/adapter.rs b/chain/epoch-manager/src/adapter.rs
@@ -5,7 +5,7 @@ use crate::EpochManagerHandle;
 use near_chain_primitives::Error;
 use near_crypto::Signature;
 use near_primitives::block_header::{Approval, ApprovalInner, BlockHeader};
-use near_primitives::chunk_validation::ChunkEndorsement;
+use near_primitives::chunk_validation::{ChunkEndorsement, ChunkValidators};
 use near_primitives::epoch_manager::block_info::BlockInfo;
 use near_primitives::epoch_manager::epoch_info::EpochInfo;
 use near_primitives::epoch_manager::EpochConfig;
@@ -19,11 +19,11 @@ use near_primitives::types::{
     AccountId, ApprovalStake, Balance, BlockHeight, EpochHeight, EpochId, ShardId,
     ValidatorInfoIdentifier,
 };
-use near_primitives::validator_mandates::AssignmentWeight;
 use near_primitives::version::ProtocolVersion;
 use near_primitives::views::EpochValidatorInfo;
 use near_store::{ShardUId, StoreUpdate};
 use std::cmp::Ordering;
+#[cfg(feature = "new_epoch_sync")]
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -194,7 +194,7 @@ pub trait EpochManagerAdapter: Send + Sync {
         epoch_id: &EpochId,
         shard_id: ShardId,
         height: BlockHeight,
-    ) -> Result<HashMap<AccountId, AssignmentWeight>, EpochError>;
+    ) -> Result<Arc<ChunkValidators>, EpochError>;
 
     fn get_validator_by_account_id(
         &self,
@@ -663,7 +663,7 @@ impl EpochManagerAdapter for EpochManagerHandle {
         epoch_id: &EpochId,
         shard_id: ShardId,
         height: BlockHeight,
-    ) -> Result<HashMap<AccountId, AssignmentWeight>, EpochError> {
+    ) -> Result<Arc<ChunkValidators>, EpochError> {
         let epoch_manager = self.read();
         epoch_manager.get_chunk_validators(epoch_id, shard_id, height)
     }

diff --git a/chain/epoch-manager/src/lib.rs b/chain/epoch-manager/src/lib.rs
@@ -3,6 +3,7 @@ use crate::types::EpochInfoAggregator;
 use near_cache::SyncLruCache;
 use near_chain_configs::GenesisConfig;
 use near_primitives::checked_feature;
+use near_primitives::chunk_validation::ChunkValidators;
 use near_primitives::epoch_manager::block_info::BlockInfo;
 use near_primitives::epoch_manager::epoch_info::{EpochInfo, EpochSummary};
 use near_primitives::epoch_manager::{
@@ -18,7 +19,6 @@ use near_primitives::types::{
     EpochInfoProvider, NumBlocks, NumSeats, ShardId, ValidatorId, ValidatorInfoIdentifier,
     ValidatorKickoutReason, ValidatorStats,
 };
-use near_primitives::validator_mandates::AssignmentWeight;
 use near_primitives::version::{ProtocolVersion, UPGRADABILITY_FIX_PROTOCOL_VERSION};
 use near_primitives::views::{
     CurrentEpochValidatorInfo, EpochValidatorInfo, NextEpochValidatorInfo, ValidatorKickoutView,
@@ -142,6 +142,8 @@ pub struct EpochManager {
     epoch_info_aggregator: EpochInfoAggregator,
     /// Largest final height. Monotonically increasing.
     largest_final_height: BlockHeight,
+    /// Cache for chunk_validators
+    chunk_validators_cache: SyncLruCache<(EpochId, ShardId, BlockHeight), Arc<ChunkValidators>>,
 
     /// Counts loop iterations inside of aggregate_epoch_info_upto method.
     /// Used for tests as a bit of white-box testing.
@@ -235,6 +237,7 @@ impl EpochManager {
             epoch_validators_ordered: SyncLruCache::new(EPOCH_CACHE_SIZE),
             epoch_validators_ordered_unique: SyncLruCache::new(EPOCH_CACHE_SIZE),
             epoch_chunk_producers_unique: SyncLruCache::new(EPOCH_CACHE_SIZE),
+            chunk_validators_cache: SyncLruCache::new(BLOCK_CACHE_SIZE),
             epoch_info_aggregator,
             #[cfg(test)]
             epoch_info_aggregator_loop_counter: Default::default(),
@@ -923,22 +926,31 @@ impl EpochManager {
         epoch_id: &EpochId,
         shard_id: ShardId,
         height: BlockHeight,
-    ) -> Result<HashMap<AccountId, AssignmentWeight>, EpochError> {
+    ) -> Result<Arc<ChunkValidators>, EpochError> {
+        let cache_key = (epoch_id.clone(), shard_id, height);
+        if let Some(chunk_validators) = self.chunk_validators_cache.get(&cache_key) {
+            return Ok(chunk_validators);
+        }
+
         let epoch_info = self.get_epoch_info(epoch_id)?;
         let chunk_validators_per_shard = epoch_info.sample_chunk_validators(height);
-        let chunk_validators =
-            chunk_validators_per_shard.get(shard_id as usize).ok_or_else(|| {
-                EpochError::ChunkValidatorSelectionError(format!(
-                    "Invalid shard ID {} for height {}, epoch {:?} for chunk validation",
-                    shard_id, height, epoch_id,
-                ))
-            })?;
-        Ok(chunk_validators
-            .iter()
-            .map(|(validator_id, seats)| {
-                (epoch_info.get_validator(*validator_id).take_account_id(), seats.clone())
-            })
-            .collect())
+        for (shard_id, chunk_validators) in chunk_validators_per_shard.iter().enumerate() {
+            let chunk_validators = chunk_validators
+                .iter()
+                .map(|(validator_id, seats)| {
+                    (epoch_info.get_validator(*validator_id).take_account_id(), seats.clone())
+                })
+                .collect();
+            let cache_key = (epoch_id.clone(), shard_id as ShardId, height);
+            self.chunk_validators_cache.put(cache_key, Arc::new(chunk_validators));
+        }
+
+        self.chunk_validators_cache.get(&cache_key).ok_or_else(|| {
+            EpochError::ChunkValidatorSelectionError(format!(
+                "Invalid shard ID {} for height {}, epoch {:?} for chunk validation",
+                shard_id, height, epoch_id,
+            ))
+        })
     }
 
     /// get_heuristic_block_approvers_ordered: block producers for epoch

diff --git a/core/primitives/src/chunk_validation.rs b/core/primitives/src/chunk_validation.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 use crate::challenge::PartialState;
 use crate::sharding::{ChunkHash, ReceiptProof, ShardChunkHeader};
 use crate::transaction::SignedTransaction;
+use crate::validator_mandates::AssignmentWeight;
 use crate::validator_signer::ValidatorSigner;
 use borsh::{BorshDeserialize, BorshSerialize};
 use near_crypto::{PublicKey, Signature};
@@ -156,3 +157,5 @@ pub struct StoredChunkStateTransitionData {
     /// to ease debugging of why a state witness may be incorrect.
     pub receipts_hash: CryptoHash,
 }
+
+pub type ChunkValidators = HashMap<AccountId, AssignmentWeight>;