From f2f384898defc3bab54604313ec66ded10043348 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Mon, 10 Aug 2020 11:53:37 +1000
Subject: [PATCH] Cache next-epoch shuffling on block processing

---
 beacon_node/beacon_chain/src/beacon_chain.rs  | 33 ++++++++++---------
 .../beacon_chain/src/shuffling_cache.rs       |  4 +++
 consensus/fork_choice/src/fork_choice.rs      |  5 +--
 consensus/types/src/shuffling_id.rs           |  3 +-
 4 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index 639490ed7f7..447e1ac614a 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -1419,7 +1419,6 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         let block = &signed_block.message;
         let block_root = fully_verified_block.block_root;
         let state = fully_verified_block.state;
-        let parent_block = fully_verified_block.parent_block;
         let current_slot = self.slot()?;
         let mut ops = fully_verified_block.intermediate_states;
 
@@ -1451,22 +1450,24 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .ok_or_else(|| Error::ValidatorPubkeyCacheLockTimeout)?
             .import_new_pubkeys(&state)?;
 
-        // If the imported block is in the previous or current epochs (according to the
-        // wall-clock), check to see if this is the first block of the epoch. If so, add the
-        // committee to the shuffling cache.
-        if state.current_epoch() + 1 >= self.epoch()?
-            && parent_block.slot().epoch(T::EthSpec::slots_per_epoch()) != state.current_epoch()
-        {
-            let mut shuffling_cache = self
-                .shuffling_cache
-                .try_write_for(ATTESTATION_CACHE_LOCK_TIMEOUT)
-                .ok_or_else(|| Error::AttestationCacheLockTimeout)?;
-
-            let committee_cache = state.committee_cache(RelativeEpoch::Current)?;
+        // For the current and next epoch of this state, ensure we have the shuffling from this
+        // block in our cache.
+        for relative_epoch in &[RelativeEpoch::Current, RelativeEpoch::Next] {
+            let shuffling_id = ShufflingId::new(block_root, &state, *relative_epoch)?;
 
-            let shuffling_id = ShufflingId::new(block_root, &state)?;
-
-            shuffling_cache.insert(shuffling_id, committee_cache);
+            let shuffling_is_cached = self
+                .shuffling_cache
+                .try_read_for(ATTESTATION_CACHE_LOCK_TIMEOUT)
+                .ok_or_else(|| Error::AttestationCacheLockTimeout)?
+                .contains(&shuffling_id);
+
+            if !shuffling_is_cached {
+                let committee_cache = state.committee_cache(RelativeEpoch::Current)?;
+                self.shuffling_cache
+                    .try_write_for(ATTESTATION_CACHE_LOCK_TIMEOUT)
+                    .ok_or_else(|| Error::AttestationCacheLockTimeout)?
+                    .insert(shuffling_id, committee_cache);
+            }
         }
 
         let mut fork_choice = self.fork_choice.write();
diff --git a/beacon_node/beacon_chain/src/shuffling_cache.rs b/beacon_node/beacon_chain/src/shuffling_cache.rs
index 2da7c38cece..e29106291ac 100644
--- a/beacon_node/beacon_chain/src/shuffling_cache.rs
+++ b/beacon_node/beacon_chain/src/shuffling_cache.rs
@@ -36,6 +36,10 @@ impl ShufflingCache {
         opt
     }
 
+    pub fn contains(&self, key: &ShufflingId) -> bool {
+        self.cache.contains(key)
+    }
+
     pub fn insert(&mut self, key: ShufflingId, committee_cache: &CommitteeCache) {
         if !self.cache.contains(&key) {
             self.cache.put(key, committee_cache.clone());
diff --git a/consensus/fork_choice/src/fork_choice.rs b/consensus/fork_choice/src/fork_choice.rs
index 1604d467b83..c63660fbeda 100644
--- a/consensus/fork_choice/src/fork_choice.rs
+++ b/consensus/fork_choice/src/fork_choice.rs
@@ -4,7 +4,7 @@ use proto_array::{Block as ProtoBlock, ProtoArrayForkChoice};
 use ssz_derive::{Decode, Encode};
 use types::{
     BeaconBlock, BeaconState, BeaconStateError, Epoch, EthSpec, Hash256, IndexedAttestation,
-    ShufflingId, Slot,
+    RelativeEpoch, ShufflingId, Slot,
 };
 
 use crate::ForkChoiceStore;
@@ -534,7 +534,8 @@ where
             root: block_root,
             parent_root: Some(block.parent_root),
             target_root,
-            shuffling_id: ShufflingId::new(block_root, state).map_err(Error::BeaconStateError)?,
+            shuffling_id: ShufflingId::new(block_root, state, RelativeEpoch::Current)
+                .map_err(Error::BeaconStateError)?,
             state_root: block.state_root,
             justified_epoch: state.current_justified_checkpoint.epoch,
             finalized_epoch: state.finalized_checkpoint.epoch,
diff --git a/consensus/types/src/shuffling_id.rs b/consensus/types/src/shuffling_id.rs
index 68a07950e04..ace574bd0da 100644
--- a/consensus/types/src/shuffling_id.rs
+++ b/consensus/types/src/shuffling_id.rs
@@ -13,8 +13,9 @@ impl ShufflingId {
     pub fn new<E: EthSpec>(
         block_root: Hash256,
         state: &BeaconState<E>,
+        relative_epoch: RelativeEpoch,
     ) -> Result<Self, BeaconStateError> {
-        let shuffling_epoch = state.current_epoch();
+        let shuffling_epoch = relative_epoch.into_epoch(state.current_epoch());
 
         // Taking advantage of saturating subtraction on slot and epoch.
         //