Implement activation queue cache

Author: michaelsproul

consensus/state_processing/src/epoch_cache.rs

@@ -2,7 +2,7 @@ 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};
+use types::{ActivationQueue, BeaconState, ChainSpec, Epoch, EthSpec, Hash256};
 
 pub fn initialize_epoch_cache<E: EthSpec>(
     state: &mut BeaconState<E>,
@@ -23,13 +23,17 @@ pub fn initialize_epoch_cache<E: EthSpec>(
     }
 
     // Compute base rewards.
+    state.build_total_active_balance_cache_at(epoch, spec)?;
     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() {
+    // Compute activation queue.
+    let mut activation_queue = ActivationQueue::default();
+
+    for (index, validator) in state.validators().iter().enumerate() {
         let effective_balance = validator.effective_balance();
 
         let base_reward = if spec
@@ -41,6 +45,9 @@ pub fn initialize_epoch_cache<E: EthSpec>(
             altair::get_base_reward(effective_balance, base_reward_per_increment, spec)?
         };
         base_rewards.push(base_reward);
+
+        // Add to speculative activation queue.
+        activation_queue.add_if_could_be_eligible_for_activation(index, validator, epoch, spec);
     }
 
     *state.epoch_cache_mut() = EpochCache::new(
@@ -49,6 +56,7 @@ pub fn initialize_epoch_cache<E: EthSpec>(
             decision_block_root,
         },
         base_rewards,
+        activation_queue,
     );
 
     Ok(())

consensus/state_processing/src/per_epoch_processing/registry_updates.rs

@@ -1,5 +1,4 @@
 use crate::{common::initiate_validator_exit, per_epoch_processing::Error};
-use itertools::Itertools;
 use safe_arith::SafeArith;
 use types::{BeaconState, ChainSpec, EthSpec, Validator};
 
@@ -40,19 +39,16 @@ pub fn process_registry_updates<T: EthSpec>(
     }
 
     // Queue validators eligible for activation and not dequeued for activation prior to finalized epoch
-    let activation_queue = state
-        .validators()
-        .iter()
-        .enumerate()
-        .filter(|(_, validator)| validator.is_eligible_for_activation(state, spec))
-        .sorted_by_key(|(index, validator)| (validator.activation_eligibility_epoch(), *index))
-        .map(|(index, _)| index)
-        .collect_vec();
-
     // Dequeue validators for activation up to churn limit
     let churn_limit = state.get_churn_limit(spec)? as usize;
+
+    let epoch_cache = state.epoch_cache().clone();
+    let activation_queue = epoch_cache
+        .activation_queue()?
+        .get_validators_eligible_for_activation(state.finalized_checkpoint().epoch, churn_limit);
+
     let delayed_activation_epoch = state.compute_activation_exit_epoch(current_epoch, spec)?;
-    for index in activation_queue.into_iter().take(churn_limit) {
+    for index in activation_queue {
         state.get_validator_mut(index)?.mutable.activation_epoch = delayed_activation_epoch;
     }
 

consensus/types/src/activation_queue.rs

@@ -0,0 +1,38 @@
+use crate::{ChainSpec, Epoch, Validator};
+use std::collections::BTreeSet;
+
+/// Activation queue computed during epoch processing for use in the *next* epoch.
+#[derive(Debug, PartialEq, Eq, Default, Clone, arbitrary::Arbitrary)]
+pub struct ActivationQueue {
+    /// Validators represented by `(activation_eligibility_epoch, index)` in sorted order.
+    queue: BTreeSet<(Epoch, usize)>,
+}
+
+impl ActivationQueue {
+    pub fn add_if_could_be_eligible_for_activation(
+        &mut self,
+        index: usize,
+        validator: &Validator,
+        next_epoch: Epoch,
+        spec: &ChainSpec,
+    ) {
+        if validator.could_be_eligible_for_activation_at(next_epoch, spec) {
+            self.queue
+                .insert((validator.activation_eligibility_epoch(), index));
+        }
+    }
+
+    pub fn get_validators_eligible_for_activation(
+        &self,
+        finalized_epoch: Epoch,
+        churn_limit: usize,
+    ) -> BTreeSet<usize> {
+        self.queue
+            .iter()
+            .filter_map(|&(eligibility_epoch, index)| {
+                (eligibility_epoch <= finalized_epoch).then_some(index)
+            })
+            .take(churn_limit)
+            .collect()
+    }
+}

consensus/types/src/epoch_cache.rs

@@ -1,4 +1,4 @@
-use crate::{BeaconStateError, Epoch, EthSpec, Hash256, Slot};
+use crate::{ActivationQueue, BeaconStateError, Epoch, EthSpec, Hash256, Slot};
 use safe_arith::ArithError;
 use std::sync::Arc;
 
@@ -20,6 +20,8 @@ struct Inner {
     key: EpochCacheKey,
     /// Base reward for every validator in this epoch.
     base_rewards: Vec<u64>,
+    /// Validator activation queue.
+    activation_queue: ActivationQueue,
 }
 
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, arbitrary::Arbitrary)]
@@ -52,9 +54,17 @@ impl From<ArithError> for EpochCacheError {
 }
 
 impl EpochCache {
-    pub fn new(key: EpochCacheKey, base_rewards: Vec<u64>) -> EpochCache {
+    pub fn new(
+        key: EpochCacheKey,
+        base_rewards: Vec<u64>,
+        activation_queue: ActivationQueue,
+    ) -> EpochCache {
         Self {
-            inner: Some(Arc::new(Inner { key, base_rewards })),
+            inner: Some(Arc::new(Inner {
+                key,
+                base_rewards,
+                activation_queue,
+            })),
         }
     }
 
@@ -92,4 +102,12 @@ impl EpochCache {
             .copied()
             .ok_or(EpochCacheError::ValidatorIndexOutOfBounds { validator_index })
     }
+
+    pub fn activation_queue(&self) -> Result<&ActivationQueue, EpochCacheError> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or(EpochCacheError::CacheNotInitialized)?;
+        Ok(&inner.activation_queue)
+    }
 }

consensus/types/src/lib.rs

@@ -73,6 +73,7 @@ pub mod validator_subscription;
 pub mod voluntary_exit;
 #[macro_use]
 pub mod slot_epoch_macros;
+pub mod activation_queue;
 pub mod config_and_preset;
 pub mod execution_block_header;
 pub mod fork_context;
@@ -99,6 +100,7 @@ pub mod sqlite;
 
 use ethereum_types::{H160, H256};
 
+pub use crate::activation_queue::ActivationQueue;
 pub use crate::aggregate_and_proof::AggregateAndProof;
 pub use crate::attestation::{Attestation, Error as AttestationError};
 pub use crate::attestation_data::AttestationData;

consensus/types/src/validator.rs

@@ -159,10 +159,26 @@ impl Validator {
         state: &BeaconState<E>,
         spec: &ChainSpec,
     ) -> bool {
+        // Has not yet been activated
+        self.activation_epoch() == spec.far_future_epoch &&
         // Placement in queue is finalized
         self.activation_eligibility_epoch() <= state.finalized_checkpoint().epoch
+    }
+
+    /// Returns `true` if the validator *could* be eligible for activation at `epoch`.
+    ///
+    /// Eligibility depends on finalization, so we assume best-possible finalization. This function
+    /// returning true is a necessary but *not sufficient* condition for a validator to activate in
+    /// the epoch transition at the end of `epoch`.
+    pub fn could_be_eligible_for_activation_at(&self, epoch: Epoch, spec: &ChainSpec) -> bool {
         // Has not yet been activated
-        && self.activation_epoch() == spec.far_future_epoch
+        self.activation_epoch() == spec.far_future_epoch
+        // Placement in queue could be finalized.
+        //
+        // NOTE: it's +1 rather than +2 because we consider the activations that occur at the *end*
+        // of `epoch`, after `process_justification_and_finalization` has already updated the
+        // state's checkpoint.
+        && self.activation_eligibility_epoch() + 1 <= epoch
     }
 
     fn tree_hash_root_internal(&self) -> Result<Hash256, tree_hash::Error> {