Cache next-epoch shuffling on block processing

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();

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());

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,

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.
         //