[epoch] Move EpochInfoAggregator to core/primitives (#12855)

This is required for a new epoch store adapter.

chain/chain/src/runtime/tests.rs

@@ -7,7 +7,7 @@ use assert_matches::assert_matches;
 use near_chain_configs::test_utils::{TESTING_INIT_BALANCE, TESTING_INIT_STAKE};
 use near_epoch_manager::shard_assignment::shard_id_to_uid;
 use near_epoch_manager::shard_tracker::ShardTracker;
-use near_epoch_manager::{EpochManager, RngSeed};
+use near_epoch_manager::EpochManager;
 use near_pool::{
     InsertTransactionResult, PoolIteratorWrapper, TransactionGroupIteratorWrapper, TransactionPool,
 };
@@ -16,6 +16,7 @@ use near_primitives::apply::ApplyChunkReason;
 use near_primitives::bandwidth_scheduler::BlockBandwidthRequests;
 use near_primitives::congestion_info::{BlockCongestionInfo, ExtendedCongestionInfo};
 use near_primitives::epoch_block_info::BlockInfo;
+use near_primitives::epoch_info::RngSeed;
 use near_primitives::receipt::{ActionReceipt, ReceiptV1};
 use near_primitives::state::PartialState;
 use near_primitives::stateless_validation::ChunkProductionKey;

chain/chain/src/test_utils/kv_runtime.rs

@@ -11,7 +11,7 @@ use near_async::time::Duration;
 use near_chain_configs::{ProtocolConfig, DEFAULT_GC_NUM_EPOCHS_TO_KEEP};
 use near_chain_primitives::Error;
 use near_crypto::{KeyType, PublicKey, SecretKey};
-use near_epoch_manager::{EpochManagerAdapter, RngSeed};
+use near_epoch_manager::EpochManagerAdapter;
 use near_parameters::RuntimeConfig;
 use near_pool::types::TransactionGroupIterator;
 use near_primitives::account::{AccessKey, Account};
@@ -20,7 +20,7 @@ use near_primitives::bandwidth_scheduler::BandwidthRequests;
 use near_primitives::block::Tip;
 use near_primitives::congestion_info::{CongestionInfo, ExtendedCongestionInfo};
 use near_primitives::epoch_block_info::BlockInfo;
-use near_primitives::epoch_info::EpochInfo;
+use near_primitives::epoch_info::{EpochInfo, RngSeed};
 use near_primitives::epoch_manager::EpochConfig;
 use near_primitives::epoch_manager::ShardConfig;
 use near_primitives::errors::{EpochError, InvalidTxError};

chain/client/src/client_actor.rs

@@ -48,7 +48,7 @@ use near_client_primitives::types::{
     StateSyncStatus, Status, StatusError, StatusSyncInfo, SyncStatus,
 };
 use near_epoch_manager::shard_tracker::ShardTracker;
-use near_epoch_manager::{EpochManagerAdapter, RngSeed};
+use near_epoch_manager::EpochManagerAdapter;
 use near_network::client::{
     BlockApproval, BlockHeadersResponse, BlockResponse, ChunkEndorsementMessage, ProcessTxRequest,
     ProcessTxResponse, RecvChallenge, SetNetworkInfo, StateResponseReceived,
@@ -61,6 +61,7 @@ use near_performance_metrics;
 use near_performance_metrics_macros::perf;
 use near_primitives::block::Tip;
 use near_primitives::block_header::ApprovalType;
+use near_primitives::epoch_info::RngSeed;
 use near_primitives::hash::CryptoHash;
 use near_primitives::network::{AnnounceAccount, PeerId};
 use near_primitives::types::{AccountId, BlockHeight};

chain/epoch-manager/src/lib.rs

@@ -6,7 +6,7 @@ use near_cache::SyncLruCache;
 use near_chain_configs::{Genesis, GenesisConfig};
 use near_primitives::block::{BlockHeader, Tip};
 use near_primitives::epoch_block_info::{BlockInfo, SlashState};
-use near_primitives::epoch_info::EpochInfo;
+use near_primitives::epoch_info::{EpochInfo, RngSeed};
 use near_primitives::epoch_manager::{
     AllEpochConfig, EpochConfig, EpochConfigStore, EpochSummary, AGGREGATOR_KEY,
 };
@@ -28,6 +28,7 @@ use near_primitives::views::{
     CurrentEpochValidatorInfo, EpochValidatorInfo, NextEpochValidatorInfo, ValidatorKickoutView,
 };
 use near_store::adapter::StoreAdapter;
+use near_store::epoch_info_aggregator::EpochInfoAggregator;
 use near_store::{DBCol, Store, StoreUpdate, HEADER_HEAD_KEY};
 use num_rational::BigRational;
 use primitive_types::U256;
@@ -45,7 +46,6 @@ pub use crate::adapter::EpochManagerAdapter;
 pub use crate::proposals::proposals_to_epoch_info;
 pub use crate::reward_calculator::RewardCalculator;
 pub use crate::reward_calculator::NUM_SECONDS_IN_A_YEAR;
-pub use crate::types::{EpochInfoAggregator, RngSeed};
 pub use near_primitives::shard_layout::ShardInfo;
 
 mod adapter;
@@ -57,7 +57,6 @@ pub mod shard_tracker;
 pub mod test_utils;
 #[cfg(test)]
 mod tests;
-pub mod types;
 mod validator_selection;
 mod validator_stats;
 
@@ -958,7 +957,7 @@ impl EpochManager {
         height: BlockHeight,
     ) -> Result<ValidatorStake, EpochError> {
         let epoch_info = self.get_epoch_info(epoch_id)?;
-        let validator_id = Self::block_producer_from_info(&epoch_info, height);
+        let validator_id = epoch_info.sample_block_producer(height);
         Ok(epoch_info.get_validator(validator_id))
     }
 
@@ -1132,14 +1131,16 @@ impl EpochManager {
         &self,
         key: &ChunkProductionKey,
     ) -> Result<ValidatorStake, EpochError> {
-        let epoch_info = self.get_epoch_info(&key.epoch_id)?;
-        let shard_layout = self.get_shard_layout(&key.epoch_id)?;
-        let validator_id = Self::chunk_producer_from_info(
-            &epoch_info,
-            &shard_layout,
-            key.shard_id,
-            key.height_created,
-        )?;
+        let ChunkProductionKey { epoch_id, shard_id, height_created } = key;
+        let epoch_info = self.get_epoch_info(epoch_id)?;
+        let shard_layout = self.get_shard_layout(epoch_id)?;
+        let validator_id = epoch_info
+            .sample_chunk_producer(&shard_layout, *shard_id, *height_created)
+            .ok_or_else(|| {
+                EpochError::ChunkProducerSelectionError(format!(
+                    "Invalid shard {shard_id} for height {height_created}"
+                ))
+            })?;
         Ok(epoch_info.get_validator(validator_id))
     }
 
@@ -1614,28 +1615,6 @@ impl EpochManager {
 
 /// Private utilities for EpochManager.
 impl EpochManager {
-    #[inline]
-    pub(crate) fn block_producer_from_info(
-        epoch_info: &EpochInfo,
-        height: BlockHeight,
-    ) -> ValidatorId {
-        epoch_info.sample_block_producer(height)
-    }
-
-    #[inline]
-    pub(crate) fn chunk_producer_from_info(
-        epoch_info: &EpochInfo,
-        shard_layout: &ShardLayout,
-        shard_id: ShardId,
-        height: BlockHeight,
-    ) -> Result<ValidatorId, EpochError> {
-        epoch_info.sample_chunk_producer(shard_layout, shard_id, height).ok_or_else(|| {
-            EpochError::ChunkProducerSelectionError(format!(
-                "Invalid shard {shard_id} for height {height}"
-            ))
-        })
-    }
-
     /// Returns true, if given current block info, next block supposed to be in the next epoch.
     fn is_next_block_in_next_epoch(&self, block_info: &BlockInfo) -> Result<bool, EpochError> {
         if block_info.is_genesis() {

chain/epoch-manager/src/proposals.rs

@@ -85,7 +85,7 @@ mod old_validator_selection {
     use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
     use std::iter;
 
-    use near_primitives::epoch_info::EpochInfo;
+    use near_primitives::epoch_info::{EpochInfo, RngSeed};
     use near_primitives::epoch_manager::EpochConfig;
     use near_primitives::errors::EpochError;
     use near_primitives::types::validator_stake::ValidatorStake;
@@ -98,7 +98,6 @@ mod old_validator_selection {
     use rand_hc::Hc128Rng;
 
     use crate::proposals::find_threshold;
-    use crate::types::RngSeed;
 
     pub fn proposals_to_epoch_info(
         epoch_config: &EpochConfig,

chain/epoch-manager/src/tests/mod.rs

@@ -212,7 +212,7 @@ fn test_fork_finalization() {
             let height = i as u64;
             let epoch_id = epoch_manager.get_epoch_id_from_prev_block(&prev_block).unwrap();
             let epoch_info = epoch_manager.get_epoch_info(&epoch_id).unwrap().clone();
-            let block_producer_id = EpochManager::block_producer_from_info(&epoch_info, height);
+            let block_producer_id = epoch_info.sample_block_producer(height);
             let block_producer = epoch_info.get_validator(block_producer_id);
             let account_id = block_producer.account_id();
             if validator_accounts.iter().any(|v| *v == account_id) {
@@ -1040,7 +1040,7 @@ fn test_expected_chunks() {
         let height = i as u64;
         let epoch_id = epoch_manager.get_epoch_id_from_prev_block(&prev_block).unwrap();
         let epoch_info = epoch_manager.get_epoch_info(&epoch_id).unwrap().clone();
-        let block_producer = EpochManager::block_producer_from_info(&epoch_info, height);
+        let block_producer = epoch_info.sample_block_producer(height);
         // test1 does not produce blocks during first epoch
         if block_producer == 0 && epoch_id == initial_epoch_id {
             expected += 1;
@@ -1108,16 +1108,12 @@ fn test_expected_chunks_prev_block_not_produced() {
         let epoch_id = epoch_manager.get_epoch_id_from_prev_block(&prev_block).unwrap();
         let epoch_info = epoch_manager.get_epoch_info(&epoch_id).unwrap().clone();
         let shard_layout = epoch_manager.get_shard_layout(&epoch_id).unwrap();
-        let block_producer = EpochManager::block_producer_from_info(&epoch_info, height);
+        let block_producer = epoch_info.sample_block_producer(height);
         let prev_block_info = epoch_manager.get_block_info(&prev_block).unwrap();
         let prev_height = prev_block_info.height();
-        let expected_chunk_producer = EpochManager::chunk_producer_from_info(
-            &epoch_info,
-            &shard_layout,
-            ShardId::new(0),
-            prev_height + 1,
-        )
-        .unwrap();
+        let expected_chunk_producer = epoch_info
+            .sample_chunk_producer(&shard_layout, ShardId::new(0), prev_height + 1)
+            .unwrap();
         // test1 does not produce blocks during first epoch
         if block_producer == 0 && epoch_id == initial_epoch_id {
             expected += 1;
@@ -1167,11 +1163,11 @@ fn update_tracker(
     produced_heights: &[BlockHeight],
     tracker: &mut HashMap<ValidatorId, ValidatorStats>,
 ) {
-    for h in heights {
-        let block_producer = EpochManager::block_producer_from_info(epoch_info, h);
+    for height in heights {
+        let block_producer = epoch_info.sample_block_producer(height);
         let entry =
             tracker.entry(block_producer).or_insert(ValidatorStats { produced: 0, expected: 0 });
-        if produced_heights.contains(&h) {
+        if produced_heights.contains(&height) {
             entry.produced += 1;
         }
         entry.expected += 1;
@@ -1223,7 +1219,7 @@ fn test_rewards_with_kickouts() {
 
         let epoch_id = em.get_epoch_id_from_prev_block(&prev_hash).unwrap();
         let epoch_info = em.get_epoch_info(&epoch_id).unwrap().clone();
-        let validator_id = EpochManager::block_producer_from_info(&epoch_info, height);
+        let validator_id = epoch_info.sample_block_producer(height);
         let block_producer = epoch_info.validator_account_id(validator_id);
 
         // don't produce blocks for test2 so we can see it in the kickouts
@@ -1522,13 +1518,8 @@ fn test_chunk_producer_kickout() {
                     return true;
                 }
                 let shard_id = shard_layout.get_shard_id(shard_index).unwrap();
-                let chunk_producer = EpochManager::chunk_producer_from_info(
-                    &epoch_info,
-                    &shard_layout,
-                    shard_id,
-                    height,
-                )
-                .unwrap();
+                let chunk_producer =
+                    epoch_info.sample_chunk_producer(&shard_layout, shard_id, height).unwrap();
                 // test1 skips chunks
                 if chunk_producer == 0 {
                     expected += 1;
@@ -2166,7 +2157,7 @@ fn test_all_kickout_edge_case() {
         let height = height as u64;
         let epoch_id = epoch_manager.get_epoch_id_from_prev_block(&prev_block).unwrap();
         let epoch_info = epoch_manager.get_epoch_info(&epoch_id).unwrap().clone();
-        let block_producer = EpochManager::block_producer_from_info(&epoch_info, height);
+        let block_producer = epoch_info.sample_block_producer(height);
         let block_producer = epoch_info.validator_account_id(block_producer);
         if height < EPOCH_LENGTH {
             // kickout test2 during first epoch

chain/epoch-manager/src/types.rs → core/store/src/epoch_info_aggregator.rs

@@ -1,3 +1,5 @@
+use std::collections::{BTreeMap, HashMap};
+
 use borsh::{BorshDeserialize, BorshSerialize};
 use itertools::Itertools;
 use near_primitives::epoch_block_info::BlockInfo;
@@ -11,13 +13,6 @@ use near_primitives::types::{
 use near_primitives::version::ProtocolVersion;
 use near_schema_checker_lib::ProtocolSchema;
 
-use std::collections::{BTreeMap, HashMap};
-use tracing::{debug, debug_span};
-
-use crate::EpochManager;
-
-pub type RngSeed = [u8; 32];
-
 /// Aggregator of information needed for validator computation at the end of the epoch.
 #[derive(
     Clone, BorshSerialize, BorshDeserialize, Debug, Default, serde::Serialize, ProtocolSchema,
@@ -73,14 +68,14 @@ impl EpochInfoAggregator {
         shard_layout: &ShardLayout,
         prev_block_height: BlockHeight,
     ) {
-        let _span =
-            debug_span!(target: "epoch_tracker", "update_tail", prev_block_height).entered();
+        let _span = tracing::debug_span!(target: "epoch_tracker", "update_tail", prev_block_height)
+            .entered();
 
         // Step 1: update block tracer (block-production stats)
 
         let block_info_height = block_info.height();
         for height in prev_block_height + 1..=block_info_height {
-            let block_producer_id = EpochManager::block_producer_from_info(epoch_info, height);
+            let block_producer_id = epoch_info.sample_block_producer(height);
             let entry = self.block_tracker.entry(block_producer_id);
             if height == block_info_height {
                 entry
@@ -90,7 +85,7 @@ impl EpochInfoAggregator {
                     })
                     .or_insert(ValidatorStats { produced: 1, expected: 1 });
             } else {
-                debug!(
+                tracing::debug!(
                     target: "epoch_tracker",
                     block_producer = ?epoch_info.validator_account_id(block_producer_id),
                     block_height = height, "Missed block");
@@ -103,27 +98,21 @@ impl EpochInfoAggregator {
         }
 
         // Step 2: update shard tracker (chunk production/endorsement stats)
-
-        // TODO(#11900): Call EpochManager::get_chunk_validator_assignments to access the cached validator assignments.
         let chunk_validator_assignment = epoch_info.sample_chunk_validators(prev_block_height + 1);
 
         for (shard_index, mask) in block_info.chunk_mask().iter().enumerate() {
             let shard_id = shard_layout.get_shard_id(shard_index).unwrap();
-            let chunk_producer_id = EpochManager::chunk_producer_from_info(
-                epoch_info,
-                shard_layout,
-                shard_id,
-                prev_block_height + 1,
-            )
-            .unwrap();
+            let chunk_producer_id = epoch_info
+                .sample_chunk_producer(shard_layout, shard_id, prev_block_height + 1)
+                .unwrap();
             let tracker = self.shard_tracker.entry(shard_id).or_insert_with(HashMap::new);
             tracker
                 .entry(chunk_producer_id)
                 .and_modify(|stats| {
                     if *mask {
                         *stats.produced_mut() += 1;
                     } else {
-                        debug!(
+                        tracing::debug!(
                             target: "epoch_tracker",
                             chunk_validator = ?epoch_info.validator_account_id(chunk_producer_id),
                             ?shard_id,
@@ -183,8 +172,7 @@ impl EpochInfoAggregator {
         }
 
         // Step 3: update version tracker
-        let block_producer_id =
-            EpochManager::block_producer_from_info(epoch_info, block_info_height);
+        let block_producer_id = epoch_info.sample_block_producer(block_info_height);
         self.version_tracker
             .entry(block_producer_id)
             .or_insert_with(|| *block_info.latest_protocol_version());

core/store/src/lib.rs

@@ -51,6 +51,7 @@ mod columns;
 pub mod config;
 pub mod contract;
 pub mod db;
+pub mod epoch_info_aggregator;
 pub mod flat;
 pub mod genesis;
 pub mod metadata;

nearcore/src/entity_debug.rs

@@ -5,7 +5,6 @@ use borsh::BorshDeserialize;
 use near_chain::types::{LatestKnown, RuntimeAdapter};
 use near_chain::{Block, BlockHeader};
 use near_epoch_manager::shard_assignment::{account_id_to_shard_id, shard_id_to_uid};
-use near_epoch_manager::types::EpochInfoAggregator;
 use near_epoch_manager::EpochManagerAdapter;
 use near_jsonrpc_primitives::errors::RpcError;
 use near_jsonrpc_primitives::types::entity_debug::{
@@ -38,6 +37,7 @@ use near_primitives::views::{
 use near_store::adapter::flat_store::encode_flat_state_db_key;
 use near_store::adapter::StoreAdapter;
 use near_store::db::GENESIS_CONGESTION_INFO_KEY;
+use near_store::epoch_info_aggregator::EpochInfoAggregator;
 use near_store::flat::delta::KeyForFlatStateDelta;
 use near_store::flat::{FlatStateChanges, FlatStateDeltaMetadata, FlatStorageStatus};
 use near_store::{

tools/protocol-schema-check/src/main.rs

@@ -16,9 +16,9 @@ use near_primitives::*;
 use near_store::*;
 use near_vm_runner::*;
 
-use near_epoch_manager::types::EpochInfoAggregator;
 use near_schema_checker_lib::{FieldName, FieldTypeInfo, ProtocolSchema, ProtocolSchemaInfo};
 use near_stable_hasher::StableHasher;
+use near_store::epoch_info_aggregator::EpochInfoAggregator;
 use std::any::TypeId;
 use std::collections::{BTreeMap, HashSet};
 use std::fs;

tools/speedy_sync/src/main.rs

@@ -5,7 +5,6 @@ use near_chain::types::{ChainConfig, Tip};
 use near_chain::{Chain, ChainGenesis, DoomslugThresholdMode};
 use near_chain_configs::{GenesisValidationMode, MutableConfigValue, ReshardingConfig};
 use near_epoch_manager::shard_tracker::{ShardTracker, TrackedConfig};
-use near_epoch_manager::types::EpochInfoAggregator;
 use near_epoch_manager::EpochManager;
 use near_primitives::block::Block;
 use near_primitives::block_header::BlockHeader;
@@ -16,6 +15,7 @@ use near_primitives::hash::CryptoHash;
 use near_primitives::merkle::PartialMerkleTree;
 use near_primitives::types::EpochId;
 use near_primitives::utils::index_to_bytes;
+use near_store::epoch_info_aggregator::EpochInfoAggregator;
 use near_store::HEADER_HEAD_KEY;
 use near_store::{DBCol, Mode, NodeStorage, Store, StoreUpdate};
 use near_time::Clock;