From 52dfb61de08eb89c493fbf6b19444cececa1d7f8 Mon Sep 17 00:00:00 2001 From: Paul Hauner Date: Fri, 17 Jan 2020 14:53:06 +1100 Subject: [PATCH] Fix justified balances bug --- beacon_node/beacon_chain/src/fork_choice.rs | 2 + .../src/fork_choice/checkpoint_manager.rs | 97 +++++++++++++++++-- 2 files changed, 91 insertions(+), 8 deletions(-) diff --git a/beacon_node/beacon_chain/src/fork_choice.rs b/beacon_node/beacon_chain/src/fork_choice.rs index 22cbff6621..d7410a8f86 100644 --- a/beacon_node/beacon_chain/src/fork_choice.rs +++ b/beacon_node/beacon_chain/src/fork_choice.rs @@ -27,6 +27,8 @@ pub enum Error { StoreError(StoreError), BeaconChainError(Box), UnknownBlockSlot(Hash256), + UnknownJustifiedBlock(Hash256), + UnknownJustifiedState(Hash256), } pub struct ForkChoice { diff --git a/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs b/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs index 9e59da6ea4..00b4513e55 100644 --- a/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs +++ b/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs @@ -4,10 +4,63 @@ use proto_array_fork_choice::ProtoArrayForkChoice; use ssz_derive::{Decode, Encode}; use types::{BeaconState, Checkpoint, Epoch, EthSpec, Hash256, Slot}; +const MAX_BALANCE_CACHE_SIZE: usize = 4; + +#[derive(PartialEq, Clone, Encode, Decode)] +struct CacheItem { + block_root: Hash256, + balances: Vec, +} + +#[derive(PartialEq, Clone, Default, Encode, Decode)] +struct BalancesCache { + items: Vec, +} + +impl BalancesCache { + pub fn process_state(&mut self, state: &BeaconState) -> Result<(), Error> { + let epoch_boundary_slot = state.current_epoch().start_slot(E::slots_per_epoch()); + let epoch_boundary_root = *state.get_block_root(epoch_boundary_slot)?; + + if self.position(epoch_boundary_root).is_none() { + let item = CacheItem { + block_root: epoch_boundary_root, + balances: state.balances.clone().into(), + }; + + if self.items.len() == MAX_BALANCE_CACHE_SIZE { + self.items.remove(0); + } + + self.items.push(item); + } + + Ok(()) + } + + fn position(&self, block_root: Hash256) -> Option { + self.items + .iter() + .position(|item| item.block_root == block_root) + } + + pub fn get(&mut self, block_root: Hash256) -> Option> { + let i = self.position(block_root)?; + Some(self.items.remove(i).balances) + } +} + +/// A `types::Checkpoint` that also stores the validator balances from a `BeaconState`. +/// +/// Useful because we need to track the justified checkpoint balances. #[derive(PartialEq, Clone, Encode, Decode)] pub struct CheckpointWithBalances { pub epoch: Epoch, pub root: Hash256, + /// These are the balances of the state with `self.root`. + /// + /// Importantly, these are _not_ the balances of the first state that we saw with a matching + /// `state.current_justified_checkpoint`. pub balances: Vec, } @@ -31,6 +84,7 @@ pub struct CheckpointManager { pub current: FFGCheckpoints, best: FFGCheckpoints, update_at: Option, + balances_cache: BalancesCache, } impl CheckpointManager { @@ -43,6 +97,7 @@ impl CheckpointManager { current: ffg_checkpoint.clone(), best: ffg_checkpoint, update_at: None, + balances_cache: BalancesCache::default(), } } @@ -53,12 +108,14 @@ impl CheckpointManager { match self.update_at { None => { - if Self::compute_slots_since_epoch_start::(current_slot) - < chain.spec.safe_slots_to_update_justified - { - self.current = self.best.clone(); - } else { - self.update_at = Some(current_epoch + 1) + if self.best.justified.epoch > self.current.justified.epoch { + if Self::compute_slots_since_epoch_start::(current_slot) + < chain.spec.safe_slots_to_update_justified + { + self.current = self.best.clone(); + } else { + self.update_at = Some(current_epoch + 1) + } } } Some(epoch) if epoch <= current_epoch => { @@ -82,7 +139,7 @@ impl CheckpointManager { chain: &BeaconChain, proto_array: &ProtoArrayForkChoice, ) -> Result<(), Error> { - // Only proceeed if the new checkpoint is better than our current checkpoint. + // Only proceed if the new checkpoint is better than our current checkpoint. if state.current_justified_checkpoint.epoch > self.current.justified.epoch && state.finalized_checkpoint.epoch >= self.current.finalized.epoch { @@ -90,7 +147,8 @@ impl CheckpointManager { justified: CheckpointWithBalances { epoch: state.current_justified_checkpoint.epoch, root: state.current_justified_checkpoint.root, - balances: state.balances.clone().into(), + balances: self + .get_balances_for_block(state.current_justified_checkpoint.root, chain)?, }, finalized: state.finalized_checkpoint.clone(), }; @@ -129,11 +187,34 @@ impl CheckpointManager { // Always update the best checkpoint, if it's better. self.best = candidate; } + + // Add the state's balances to the balances cache to avoid a state read later. + self.balances_cache.process_state(state)?; } Ok(()) } + fn get_balances_for_block( + &mut self, + block_root: Hash256, + chain: &BeaconChain, + ) -> Result, Error> { + if let Some(balances) = self.balances_cache.get(block_root) { + Ok(balances) + } else { + let block = chain + .get_block_caching(&block_root)? + .ok_or_else(|| Error::UnknownJustifiedBlock(block_root))?; + + let state = chain + .get_state_caching_only_with_committee_caches(&block.state_root, Some(block.slot))? + .ok_or_else(|| Error::UnknownJustifiedState(block.state_root))?; + + Ok(state.balances.into()) + } + } + /// Attempts to get the block root for the given `slot`. /// /// First, the `state` is used to see if the slot is within the distance of its historical