Tree states optimization using EpochCache (#4429)

* Relocate epoch cache to BeaconState

* Optimize per block processing by pulling previous epoch & current epoch calculation up.

* Revert `get_cow` change (no performance improvement)

* Initialize `EpochCache` in epoch processing and load it from state when getting base rewards.

* Initialize `EpochCache` at start of block processing if required.

* Initialize `EpochCache` in `transition_blocks` if `exclude_cache_builds` is enabled

* Fix epoch cache initialization logic

* Remove FIXME comment.

* Cache previous & current epochs in `consensus_context.rs`.

* Move `get_base_rewards` from `ConsensusContext` to `BeaconState`.

* Update Milhouse version
This commit is contained in:
Jimmy Chen
2023-06-30 11:25:51 +10:00
committed by GitHub
parent 160bbde8a2
commit 2df714e2cd
19 changed files with 237 additions and 196 deletions

View File

@@ -1,7 +1,6 @@
use crate::common::get_indexed_attestation;
use crate::per_block_processing::errors::{AttestationInvalid, BlockOperationError};
use crate::{EpochCache, EpochCacheError};
use std::borrow::Cow;
use crate::EpochCacheError;
use std::collections::{hash_map::Entry, HashMap};
use std::marker::PhantomData;
use tree_hash::TreeHash;
@@ -14,12 +13,14 @@ use types::{
pub struct ConsensusContext<T: EthSpec> {
/// Slot to act as an identifier/safeguard
slot: Slot,
/// Previous epoch of the `slot` precomputed for optimization purpose.
pub(crate) previous_epoch: Epoch,
/// Current epoch of the `slot` precomputed for optimization purpose.
pub(crate) current_epoch: Epoch,
/// Proposer index of the block at `slot`.
proposer_index: Option<u64>,
/// Block root of the block at `slot`.
current_block_root: Option<Hash256>,
/// Epoch cache of values that are useful for block processing that are static over an epoch.
epoch_cache: Option<EpochCache>,
/// Cache of indexed attestations constructed during block processing.
indexed_attestations:
HashMap<(AttestationData, BitList<T::MaxValidatorsPerCommittee>), IndexedAttestation<T>>,
@@ -48,11 +49,14 @@ impl From<EpochCacheError> for ContextError {
impl<T: EthSpec> ConsensusContext<T> {
pub fn new(slot: Slot) -> Self {
let current_epoch = slot.epoch(T::slots_per_epoch());
let previous_epoch = current_epoch.saturating_sub(1u64);
Self {
slot,
previous_epoch,
current_epoch,
proposer_index: None,
current_block_root: None,
epoch_cache: None,
indexed_attestations: HashMap::new(),
_phantom: PhantomData,
}
@@ -145,31 +149,6 @@ impl<T: EthSpec> ConsensusContext<T> {
}
}
pub fn set_epoch_cache(mut self, epoch_cache: EpochCache) -> Self {
self.epoch_cache = Some(epoch_cache);
self
}
pub fn get_base_reward(
&mut self,
state: &BeaconState<T>,
validator_index: usize,
spec: &ChainSpec,
) -> Result<u64, ContextError> {
self.check_slot(state.slot())?;
// Build epoch cache if not already built.
let epoch_cache = if let Some(ref cache) = self.epoch_cache {
Cow::Borrowed(cache)
} else {
let cache = EpochCache::new(state, spec)?;
self.epoch_cache = Some(cache.clone());
Cow::Owned(cache)
};
Ok(epoch_cache.get_base_reward(validator_index)?)
}
pub fn get_indexed_attestation(
&mut self,
state: &BeaconState<T>,

View File

@@ -1,137 +1,55 @@
use crate::common::{
altair::{self, BaseRewardPerIncrement},
base::{self, SqrtTotalActiveBalance},
};
use safe_arith::ArithError;
use std::sync::Arc;
use types::{BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec, Hash256, Slot};
use crate::common::altair::BaseRewardPerIncrement;
use crate::common::base::SqrtTotalActiveBalance;
use crate::common::{altair, base};
use types::epoch_cache::{EpochCache, EpochCacheError, EpochCacheKey};
use types::{BeaconState, ChainSpec, Epoch, EthSpec, Hash256};
/// Cache of values which are uniquely determined at the start of an epoch.
///
/// The values are fixed with respect to the last block of the _prior_ epoch, which we refer
/// to as the "decision block". This cache is very similar to the `BeaconProposerCache` in that
/// beacon proposers are determined at exactly the same time as the values in this cache, so
/// the keys for the two caches are identical.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct EpochCache {
inner: Arc<Inner>,
}
pub fn initialize_epoch_cache<E: EthSpec>(
state: &mut BeaconState<E>,
epoch: Epoch,
spec: &ChainSpec,
) -> Result<(), EpochCacheError> {
let epoch_cache: &EpochCache = state.epoch_cache();
let decision_block_root = state
.proposer_shuffling_decision_root(Hash256::zero())
.map_err(EpochCacheError::BeaconState)?;
#[derive(Debug, PartialEq, Eq, Clone)]
struct Inner {
/// Unique identifier for this cache, which can be used to check its validity before use
/// with any `BeaconState`.
key: EpochCacheKey,
/// Base reward for every validator in this epoch.
base_rewards: Vec<u64>,
}
#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
pub struct EpochCacheKey {
pub epoch: Epoch,
pub decision_block_root: Hash256,
}
#[derive(Debug, PartialEq, Clone)]
pub enum EpochCacheError {
IncorrectEpoch { cache: Epoch, state: Epoch },
IncorrectDecisionBlock { cache: Hash256, state: Hash256 },
ValidatorIndexOutOfBounds { validator_index: usize },
InvalidSlot { slot: Slot },
Arith(ArithError),
BeaconState(BeaconStateError),
}
impl From<BeaconStateError> for EpochCacheError {
fn from(e: BeaconStateError) -> Self {
Self::BeaconState(e)
}
}
impl From<ArithError> for EpochCacheError {
fn from(e: ArithError) -> Self {
Self::Arith(e)
}
}
impl EpochCache {
pub fn new<E: EthSpec>(
state: &BeaconState<E>,
spec: &ChainSpec,
) -> Result<Self, EpochCacheError> {
let epoch = state.current_epoch();
let decision_block_root = state
.proposer_shuffling_decision_root(Hash256::zero())
.map_err(EpochCacheError::BeaconState)?;
// The cache should never be constructed at slot 0 because it should only be used for
// block processing (which implies slot > 0) or epoch processing (which implies slot >= 32).
/* FIXME(sproul): EF tests like this
if decision_block_root.is_zero() {
return Err(EpochCacheError::InvalidSlot { slot: state.slot() });
}
*/
// Compute base rewards.
let total_active_balance = state.get_total_active_balance()?;
let sqrt_total_active_balance = SqrtTotalActiveBalance::new(total_active_balance);
let base_reward_per_increment = BaseRewardPerIncrement::new(total_active_balance, spec)?;
let mut base_rewards = Vec::with_capacity(state.validators().len());
for validator in state.validators().iter() {
let effective_balance = validator.effective_balance();
let base_reward = if spec
.altair_fork_epoch
.map_or(false, |altair_epoch| epoch < altair_epoch)
{
base::get_base_reward(effective_balance, sqrt_total_active_balance, spec)?
} else {
altair::get_base_reward(effective_balance, base_reward_per_increment, spec)?
};
base_rewards.push(base_reward);
}
Ok(Self {
inner: Arc::new(Inner {
key: EpochCacheKey {
epoch,
decision_block_root,
},
base_rewards,
}),
})
if epoch_cache
.check_validity::<E>(epoch, decision_block_root)
.is_ok()
{
// `EpochCache` has already been initialized and is valid, no need to initialize.
return Ok(());
}
pub fn check_validity<E: EthSpec>(
&self,
state: &BeaconState<E>,
) -> Result<(), EpochCacheError> {
if self.inner.key.epoch != state.current_epoch() {
return Err(EpochCacheError::IncorrectEpoch {
cache: self.inner.key.epoch,
state: state.current_epoch(),
});
}
let state_decision_root = state
.proposer_shuffling_decision_root(Hash256::zero())
.map_err(EpochCacheError::BeaconState)?;
if self.inner.key.decision_block_root != state_decision_root {
return Err(EpochCacheError::IncorrectDecisionBlock {
cache: self.inner.key.decision_block_root,
state: state_decision_root,
});
}
Ok(())
// Compute base rewards.
let total_active_balance = state.get_total_active_balance_at_epoch(epoch)?;
let sqrt_total_active_balance = SqrtTotalActiveBalance::new(total_active_balance);
let base_reward_per_increment = BaseRewardPerIncrement::new(total_active_balance, spec)?;
let mut base_rewards = Vec::with_capacity(state.validators().len());
for validator in state.validators().iter() {
let effective_balance = validator.effective_balance();
let base_reward = if spec
.altair_fork_epoch
.map_or(false, |altair_epoch| epoch < altair_epoch)
{
base::get_base_reward(effective_balance, sqrt_total_active_balance, spec)?
} else {
altair::get_base_reward(effective_balance, base_reward_per_increment, spec)?
};
base_rewards.push(base_reward);
}
#[inline]
pub fn get_base_reward(&self, validator_index: usize) -> Result<u64, EpochCacheError> {
self.inner
.base_rewards
.get(validator_index)
.copied()
.ok_or(EpochCacheError::ValidatorIndexOutOfBounds { validator_index })
}
*state.epoch_cache_mut() = EpochCache::new(
EpochCacheKey {
epoch,
decision_block_root,
},
base_rewards,
);
Ok(())
}

View File

@@ -30,7 +30,6 @@ pub mod verify_operation;
pub use block_replayer::{BlockReplayError, BlockReplayer, StateProcessingStrategy};
pub use consensus_context::{ConsensusContext, ContextError};
pub use epoch_cache::{EpochCache, EpochCacheError, EpochCacheKey};
pub use genesis::{
eth2_genesis_time, initialize_beacon_state_from_eth1, is_valid_genesis_state,
process_activations,
@@ -43,4 +42,5 @@ pub use per_epoch_processing::{
errors::EpochProcessingError, process_epoch as per_epoch_processing,
};
pub use per_slot_processing::{per_slot_processing, Error as SlotProcessingError};
pub use types::{EpochCache, EpochCacheError, EpochCacheKey};
pub use verify_operation::{SigVerifiedOp, VerifyOperation, VerifyOperationAt};

View File

@@ -41,6 +41,7 @@ mod verify_proposer_slashing;
use crate::common::decrease_balance;
use crate::StateProcessingStrategy;
use crate::epoch_cache::initialize_epoch_cache;
#[cfg(feature = "arbitrary-fuzz")]
use arbitrary::Arbitrary;
@@ -114,6 +115,9 @@ pub fn per_block_processing<T: EthSpec, Payload: AbstractExecPayload<T>>(
.fork_name(spec)
.map_err(BlockProcessingError::InconsistentStateFork)?;
// Build epoch cache if it hasn't already been built, or if it is no longer valid
initialize_epoch_cache(state, state.current_epoch(), spec)?;
let verify_signatures = match block_signature_strategy {
BlockSignatureStrategy::VerifyBulk => {
// Verify all signatures in the block at once.

View File

@@ -113,6 +113,7 @@ pub mod altair {
})
}
#[allow(clippy::too_many_arguments)]
pub fn process_attestation<T: EthSpec>(
state: &mut BeaconState<T>,
attestation: &Attestation<T>,
@@ -149,18 +150,22 @@ pub mod altair {
let index = *index as usize;
for (flag_index, &weight) in PARTICIPATION_FLAG_WEIGHTS.iter().enumerate() {
let epoch_participation = state.get_epoch_participation_mut(data.target.epoch)?;
let validator_participation = epoch_participation
.get_mut(index)
.ok_or(BeaconStateError::ParticipationOutOfBounds(index))?;
let epoch_participation = state.get_epoch_participation_mut(
data.target.epoch,
ctxt.previous_epoch,
ctxt.current_epoch,
)?;
if participation_flag_indices.contains(&flag_index)
&& !validator_participation.has_flag(flag_index)?
{
validator_participation.add_flag(flag_index)?;
proposer_reward_numerator.safe_add_assign(
ctxt.get_base_reward(state, index, spec)?.safe_mul(weight)?,
)?;
if participation_flag_indices.contains(&flag_index) {
let validator_participation = epoch_participation
.get_mut(index)
.ok_or(BeaconStateError::ParticipationOutOfBounds(index))?;
if !validator_participation.has_flag(flag_index)? {
validator_participation.add_flag(flag_index)?;
proposer_reward_numerator
.safe_add_assign(state.get_base_reward(index)?.safe_mul(weight)?)?;
}
}
}
}

View File

@@ -1,4 +1,5 @@
use super::{process_registry_updates, process_slashings, EpochProcessingSummary, Error};
use crate::epoch_cache::initialize_epoch_cache;
use crate::per_epoch_processing::{
effective_balance_updates::process_effective_balance_updates,
historical_roots_update::process_historical_roots_update,
@@ -75,6 +76,7 @@ pub fn process_epoch<T: EthSpec>(
// Rotate the epoch caches to suit the epoch transition.
state.advance_caches(spec)?;
initialize_epoch_cache(state, state.next_epoch()?, spec)?;
Ok(EpochProcessingSummary::Altair {
participation_cache,

View File

@@ -1,4 +1,5 @@
use super::{process_registry_updates, process_slashings, EpochProcessingSummary, Error};
use crate::epoch_cache::initialize_epoch_cache;
use crate::per_epoch_processing::{
effective_balance_updates::process_effective_balance_updates,
historical_roots_update::process_historical_roots_update,
@@ -69,6 +70,7 @@ pub fn process_epoch<T: EthSpec>(
// Rotate the epoch caches to suit the epoch transition.
state.advance_caches(spec)?;
initialize_epoch_cache(state, state.next_epoch()?, spec)?;
Ok(EpochProcessingSummary::Base {
total_balances: validator_statuses.total_balances,

View File

@@ -11,6 +11,7 @@ use crate::per_epoch_processing::{
};
use types::{BeaconState, ChainSpec, EthSpec, RelativeEpoch};
use crate::epoch_cache::initialize_epoch_cache;
pub use historical_summaries_update::process_historical_summaries_update;
mod historical_summaries_update;
@@ -71,6 +72,7 @@ pub fn process_epoch<T: EthSpec>(
// Rotate the epoch caches to suit the epoch transition.
state.advance_caches(spec)?;
initialize_epoch_cache(state, state.next_epoch()?, spec)?;
Ok(EpochProcessingSummary::Altair {
participation_cache,

View File

@@ -1,5 +1,5 @@
use crate::per_epoch_processing::altair::participation_cache::Error as ParticipationCacheError;
use types::{milhouse, BeaconStateError, InconsistentFork};
use types::{milhouse, BeaconStateError, EpochCacheError, InconsistentFork};
#[derive(Debug, PartialEq)]
pub enum EpochProcessingError {
@@ -26,6 +26,7 @@ pub enum EpochProcessingError {
InvalidFlagIndex(usize),
ParticipationCache(ParticipationCacheError),
MilhouseError(milhouse::Error),
EpochCache(EpochCacheError),
}
impl From<InclusionError> for EpochProcessingError {
@@ -64,6 +65,12 @@ impl From<milhouse::Error> for EpochProcessingError {
}
}
impl From<EpochCacheError> for EpochProcessingError {
fn from(e: EpochCacheError) -> Self {
EpochProcessingError::EpochCache(e)
}
}
#[derive(Debug, PartialEq)]
pub enum InclusionError {
/// The validator did not participate in an attestation in this period.

View File

@@ -2,8 +2,8 @@ use crate::common::{get_attestation_participation_flag_indices, get_attesting_in
use std::mem;
use std::sync::Arc;
use types::{
BeaconState, BeaconStateAltair, BeaconStateError as Error, ChainSpec, EthSpec, Fork,
ParticipationFlags, PendingAttestation, RelativeEpoch, SyncCommittee, VList,
BeaconState, BeaconStateAltair, BeaconStateError as Error, ChainSpec, EpochCache, EthSpec,
Fork, ParticipationFlags, PendingAttestation, RelativeEpoch, SyncCommittee, VList,
};
/// Translate the participation information from the epoch prior to the fork into Altair's format.
@@ -104,6 +104,7 @@ pub fn upgrade_to_altair<E: EthSpec>(
committee_caches: mem::take(&mut pre.committee_caches),
pubkey_cache: mem::take(&mut pre.pubkey_cache),
exit_cache: mem::take(&mut pre.exit_cache),
epoch_cache: EpochCache::default(),
});
// Fill in previous epoch participation from the pre state's pending attestations.

View File

@@ -1,6 +1,7 @@
use std::mem;
use types::{
BeaconState, BeaconStateCapella, BeaconStateError as Error, ChainSpec, EthSpec, Fork, VList,
BeaconState, BeaconStateCapella, BeaconStateError as Error, ChainSpec, EpochCache, EthSpec,
Fork, VList,
};
/// Transform a `Merge` state into an `Capella` state.
@@ -66,6 +67,7 @@ pub fn upgrade_to_capella<E: EthSpec>(
committee_caches: mem::take(&mut pre.committee_caches),
pubkey_cache: mem::take(&mut pre.pubkey_cache),
exit_cache: mem::take(&mut pre.exit_cache),
epoch_cache: EpochCache::default(),
});
*pre_state = post;

View File

@@ -1,6 +1,6 @@
use std::mem;
use types::{
BeaconState, BeaconStateError as Error, BeaconStateMerge, ChainSpec, EthSpec,
BeaconState, BeaconStateError as Error, BeaconStateMerge, ChainSpec, EpochCache, EthSpec,
ExecutionPayloadHeaderMerge, Fork,
};
@@ -63,6 +63,7 @@ pub fn upgrade_to_bellatrix<E: EthSpec>(
committee_caches: mem::take(&mut pre.committee_caches),
pubkey_cache: mem::take(&mut pre.pubkey_cache),
exit_cache: mem::take(&mut pre.exit_cache),
epoch_cache: EpochCache::default(),
});
*pre_state = post;