Fix justified balances bug

beacon_node/beacon_chain/src/fork_choice.rs

@@ -27,6 +27,8 @@ pub enum Error {
     StoreError(StoreError),
     BeaconChainError(Box<BeaconChainError>),
     UnknownBlockSlot(Hash256),
+    UnknownJustifiedBlock(Hash256),
+    UnknownJustifiedState(Hash256),
 }
 
 pub struct ForkChoice<T: BeaconChainTypes> {

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<u64>,
+}
+
+#[derive(PartialEq, Clone, Default, Encode, Decode)]
+struct BalancesCache {
+    items: Vec<CacheItem>,
+}
+
+impl BalancesCache {
+    pub fn process_state<E: EthSpec>(&mut self, state: &BeaconState<E>) -> 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<usize> {
+        self.items
+            .iter()
+            .position(|item| item.block_root == block_root)
+    }
+
+    pub fn get(&mut self, block_root: Hash256) -> Option<Vec<u64>> {
+        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<u64>,
 }
 
@@ -31,6 +84,7 @@ pub struct CheckpointManager {
     pub current: FFGCheckpoints,
     best: FFGCheckpoints,
     update_at: Option<Epoch>,
+    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::<T>(current_slot)
+                    if self.best.justified.epoch > self.current.justified.epoch {
-                        < chain.spec.safe_slots_to_update_justified
+                        if Self::compute_slots_since_epoch_start::<T>(current_slot)
-                    {
+                            < chain.spec.safe_slots_to_update_justified
-                        self.current = self.best.clone();
+                        {
-                    } else {
+                            self.current = self.best.clone();
-                        self.update_at = Some(current_epoch + 1)
+                        } else {
+                            self.update_at = Some(current_epoch + 1)
+                        }
                     }
                 }
                 Some(epoch) if epoch <= current_epoch => {
@@ -82,7 +139,7 @@ impl CheckpointManager {
         chain: &BeaconChain<T>,
         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<T: BeaconChainTypes>(
+        &mut self,
+        block_root: Hash256,
+        chain: &BeaconChain<T>,
+    ) -> Result<Vec<u64>, 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