More proposer shuffling cleanup

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -6553,6 +6553,24 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         }
     }
 
+    /// This function provides safe and efficient multi-threaded access to the beacon proposer cache.
+    ///
+    /// The arguments are:
+    ///
+    /// - `shuffling_decision_block`: The block root of the decision block for the desired proposer
+    ///   shuffling. This should be computed using one of the methods for computing proposer
+    ///   shuffling decision roots, e.g. `BeaconState::proposer_shuffling_decision_root_at_epoch`.
+    /// - `proposal_epoch`: The epoch at which the proposer shuffling is required.
+    /// - `accessor`: A closure to run against the proposers for the selected epoch. Usually this
+    ///   closure just grabs a single proposer, or takes the vec of proposers for the epoch.
+    /// - `state_provider`: A closure to compute a state suitable for determining the shuffling.
+    ///   This closure is evaluated lazily ONLY in the case
+    ///
+    /// This function makes use of closures in order to efficiently handle concurrent accesses to
+    /// the cache.
+    ///
+    /// The error type is polymorphic, if in doubt you can use `BeaconChainError`. You might need
+    /// to use a turbofish if type inference can't work it out.
     pub fn with_proposer_cache<V, E: From<BeaconChainError> + From<BeaconStateError>>(
         &self,
         shuffling_decision_block: Hash256,

beacon_node/beacon_chain/src/beacon_proposer_cache.rs

@@ -17,6 +17,7 @@ use smallvec::SmallVec;
 use state_processing::state_advance::partial_state_advance;
 use std::num::NonZeroUsize;
 use std::sync::Arc;
+use tracing::instrument;
 use types::non_zero_usize::new_non_zero_usize;
 use types::{
     BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec, Fork, Hash256, Slot, Unsigned,
@@ -199,8 +200,7 @@ pub fn compute_proposer_duties_from_head<T: BeaconChainTypes>(
         .map_err(BeaconChainError::from)?;
 
     let dependent_root = state
-        // The only block which decides its own shuffling is the genesis block.
-        .proposer_shuffling_decision_root(chain.genesis_block_root, &chain.spec)
+        .proposer_shuffling_decision_root_at_epoch(request_epoch, head_block_root, &chain.spec)
         .map_err(BeaconChainError::from)?;
 
     Ok((indices, dependent_root, execution_status, state.fork()))
@@ -214,6 +214,7 @@ pub fn compute_proposer_duties_from_head<T: BeaconChainTypes>(
 /// - No-op if `state.current_epoch() == target_epoch`.
 /// - It must be the case that `state.canonical_root() == state_root`, but this function will not
 ///   check that.
+#[instrument(skip_all, fields(?state_root, %target_epoch, state_slot = %state.slot()), level = "debug")]
 pub fn ensure_state_can_determine_proposers_for_epoch<E: EthSpec>(
     state: &mut BeaconState<E>,
     state_root: Hash256,

beacon_node/beacon_chain/src/block_verification.rs

@@ -950,8 +950,6 @@ impl<T: BeaconChainTypes> GossipVerifiedBlock<T> {
         let proposer_shuffling_decision_block =
             parent_block.proposer_shuffling_root_for_child_block(block_epoch, &chain.spec);
 
-        // We assign to a variable instead of using `if let Some` directly to ensure we drop the
-        // write lock before trying to acquire it again in the `else` clause.
         let block_slot = block.slot();
         let mut opt_parent = None;
         let proposer = chain.with_proposer_cache::<_, BlockError>(

beacon_node/beacon_chain/src/state_advance_timer.rs

@@ -365,25 +365,40 @@ fn advance_head<T: BeaconChainTypes>(beacon_chain: &Arc<BeaconChain<T>>) -> Resu
         .build_committee_cache(RelativeEpoch::Next, &beacon_chain.spec)
         .map_err(BeaconChainError::from)?;
 
-    // If the `pre_state` is in a later epoch than `state`, pre-emptively add the proposer shuffling
-    // for the state's current epoch and the committee cache for the state's next epoch.
+    // The state root is required to prime the proposer cache AND for writing it to disk.
+    let advanced_state_root = state.update_tree_hash_cache()?;
+
+    // If the `pre_state` is in a later epoch than `state`, pre-emptively update the proposer
+    // shuffling and attester shuffling caches.
     if initial_epoch < state.current_epoch() {
-        // Update the proposer cache.
-        //
-        // We supply the `head_block_root` as the decision block since the prior `if` statement guarantees
-        // the head root is the latest block from the prior epoch.
-        beacon_chain
-            .beacon_proposer_cache
-            .lock()
-            .insert(
-                state.current_epoch(),
-                head_block_root,
-                state
-                    .get_beacon_proposer_indices(state.current_epoch(), &beacon_chain.spec)
-                    .map_err(BeaconChainError::from)?,
-                state.fork(),
-            )
-            .map_err(BeaconChainError::from)?;
+        // Include the proposer shuffling from the current epoch, which is likely to be useful
+        // pre-Fulu, and probably redundant post-Fulu (it should already have been in the cache).
+        let current_epoch_decision_root = state.proposer_shuffling_decision_root_at_epoch(
+            state.current_epoch(),
+            head_block_root,
+            &beacon_chain.spec,
+        )?;
+        beacon_chain.with_proposer_cache(
+            current_epoch_decision_root,
+            state.current_epoch(),
+            |_| Ok(()),
+            || Ok::<_, Error>((advanced_state_root, state.clone())),
+        )?;
+
+        // For epochs *greater than* the Fulu fork epoch, we have also determined the proposer
+        // shuffling for the next epoch.
+        let next_epoch = state.next_epoch()?;
+        let next_epoch_decision_root = state.proposer_shuffling_decision_root_at_epoch(
+            next_epoch,
+            head_block_root,
+            &beacon_chain.spec,
+        )?;
+        beacon_chain.with_proposer_cache(
+            next_epoch_decision_root,
+            next_epoch,
+            |_| Ok(()),
+            || Ok::<_, Error>((advanced_state_root, state.clone())),
+        )?;
 
         // Update the attester cache.
         let shuffling_id =
@@ -438,7 +453,6 @@ fn advance_head<T: BeaconChainTypes>(beacon_chain: &Arc<BeaconChain<T>>) -> Resu
     // even if we race with the deletion of this state by the finalization pruning code, the worst
     // case is we end up with a finalized state stored, that will get pruned the next time pruning
     // runs.
-    let advanced_state_root = state.update_tree_hash_cache()?;
     beacon_chain.store.put_state(&advanced_state_root, &state)?;
 
     debug!(

beacon_node/beacon_chain/tests/store_tests.rs

@@ -1273,19 +1273,34 @@ async fn proposer_shuffling_root_consistency_test(
 #[tokio::test]
 async fn proposer_shuffling_root_consistency_same_epoch() {
     let spec = test_spec::<E>();
-    proposer_shuffling_root_consistency_test(spec, 32, 39).await;
+    proposer_shuffling_root_consistency_test(
+        spec,
+        4 * E::slots_per_epoch(),
+        5 * E::slots_per_epoch() - 1,
+    )
+    .await;
 }
 
 #[tokio::test]
 async fn proposer_shuffling_root_consistency_next_epoch() {
     let spec = test_spec::<E>();
-    proposer_shuffling_root_consistency_test(spec, 32, 47).await;
+    proposer_shuffling_root_consistency_test(
+        spec,
+        4 * E::slots_per_epoch(),
+        6 * E::slots_per_epoch() - 1,
+    )
+    .await;
 }
 
 #[tokio::test]
 async fn proposer_shuffling_root_consistency_two_epochs() {
     let spec = test_spec::<E>();
-    proposer_shuffling_root_consistency_test(spec, 32, 55).await;
+    proposer_shuffling_root_consistency_test(
+        spec,
+        4 * E::slots_per_epoch(),
+        7 * E::slots_per_epoch() - 1,
+    )
+    .await;
 }
 
 #[tokio::test]

beacon_node/http_api/src/proposer_duties.rs

@@ -103,14 +103,6 @@ fn try_proposer_duties_from_cache<T: BeaconChainTypes>(
     let head_block = &head.snapshot.beacon_block;
     let head_block_root = head.head_block_root();
     let head_epoch = head_block.slot().epoch(T::EthSpec::slots_per_epoch());
-    let head_decision_root = head
-        .snapshot
-        .beacon_state
-        .proposer_shuffling_decision_root(head_block_root, &chain.spec)
-        .map_err(warp_utils::reject::beacon_state_error)?;
-    let execution_optimistic = chain
-        .is_optimistic_or_invalid_head_block(head_block)
-        .map_err(warp_utils::reject::unhandled_error)?;
 
     // This code path can't handle requests for past epochs.
     if head_epoch > request_epoch {
@@ -119,6 +111,15 @@ fn try_proposer_duties_from_cache<T: BeaconChainTypes>(
         )));
     }
 
+    let head_decision_root = head
+        .snapshot
+        .beacon_state
+        .proposer_shuffling_decision_root_at_epoch(request_epoch, head_block_root, &chain.spec)
+        .map_err(warp_utils::reject::beacon_state_error)?;
+    let execution_optimistic = chain
+        .is_optimistic_or_invalid_head_block(head_block)
+        .map_err(warp_utils::reject::unhandled_error)?;
+
     chain
         .beacon_proposer_cache
         .lock()