Rewrite to avoid lock contention during computation.

Author: jimmygchen

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::cmp::Ordering;
 use std::num::NonZeroUsize;
+use std::sync::Arc;
 use types::non_zero_usize::new_non_zero_usize;
 use types::{
     BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec, Fork, Hash256, Slot, Unsigned,
@@ -41,21 +42,21 @@ pub struct Proposer {
 /// their signatures.
 pub struct EpochBlockProposers {
     /// The epoch to which the proposers pertain.
-    epoch: Epoch,
+    pub epoch: Epoch,
     /// The fork that should be used to verify proposer signatures.
-    fork: Fork,
+    pub fork: Fork,
     /// A list of length `T::EthSpec::slots_per_epoch()`, representing the proposers for each slot
     /// in that epoch.
     ///
     /// E.g., if `self.epoch == 1`, then `self.proposers[0]` contains the proposer for slot `32`.
-    proposers: SmallVec<[usize; TYPICAL_SLOTS_PER_EPOCH]>,
+    pub proposers: SmallVec<[usize; TYPICAL_SLOTS_PER_EPOCH]>,
 }
 
 /// A cache to store the proposers for some epoch.
 ///
 /// See the module-level documentation for more information.
 pub struct BeaconProposerCache {
-    cache: LruCache<(Epoch, Hash256), OnceCell<EpochBlockProposers>>,
+    cache: LruCache<(Epoch, Hash256), Arc<OnceCell<EpochBlockProposers>>>,
 }
 
 impl Default for BeaconProposerCache {
@@ -111,6 +112,17 @@ impl BeaconProposerCache {
             .and_then(|cache_once_cell| cache_once_cell.get().map(|proposers| &proposers.proposers))
     }
 
+    pub fn get_or_insert_key(
+        &mut self,
+        epoch: Epoch,
+        shuffling_decision_block: Hash256,
+    ) -> Arc<OnceCell<EpochBlockProposers>> {
+        let key = (epoch, shuffling_decision_block);
+        self.cache
+            .get_or_insert(key, || Arc::new(OnceCell::new()))
+            .clone()
+    }
+
     /// Insert the proposers into the cache.
     ///
     /// See `Self::get` for a description of `shuffling_decision_block`.
@@ -125,61 +137,16 @@ impl BeaconProposerCache {
     ) -> Result<(), BeaconStateError> {
         let key = (epoch, shuffling_decision_block);
         if !self.cache.contains(&key) {
-            self.cache.put(
-                key,
-                OnceCell::with_value(EpochBlockProposers {
-                    epoch,
-                    fork,
-                    proposers: proposers.into(),
-                }),
-            );
-        }
-
-        Ok(())
-    }
-
-    pub fn get_slot_or_insert_with<E: EthSpec, F, Err: BlockBlobError>(
-        &mut self,
-        shuffling_decision_block: Hash256,
-        slot: Slot,
-        compute_proposers_and_fork_fn: F,
-    ) -> Result<Option<Proposer>, Err>
-    where
-        F: FnOnce() -> Result<(Vec<usize>, Fork), Err>,
-    {
-        let epoch = slot.epoch(E::slots_per_epoch());
-        let (proposers, fork) = self.get_epoch_or_insert_with(
-            shuffling_decision_block,
-            epoch,
-            compute_proposers_and_fork_fn,
-        )?;
-
-        Ok(proposers
-            .get(slot.as_usize() % E::SlotsPerEpoch::to_usize())
-            .map(|&index| Proposer { index, fork }))
-    }
-
-    pub fn get_epoch_or_insert_with<F, Err: BlockBlobError>(
-        &mut self,
-        shuffling_decision_block: Hash256,
-        epoch: Epoch,
-        compute_proposers_and_fork_fn: F,
-    ) -> Result<(&SmallVec<[usize; TYPICAL_SLOTS_PER_EPOCH]>, Fork), Err>
-    where
-        F: FnOnce() -> Result<(Vec<usize>, Fork), Err>,
-    {
-        let key = (epoch, shuffling_decision_block);
-        let once_cell = self.cache.get_or_insert_mut(key, OnceCell::new);
-        let epoch_block_proposers = once_cell.get_or_try_init(|| {
-            let (proposers, fork) = compute_proposers_and_fork_fn()?;
-            Ok::<EpochBlockProposers, Err>(EpochBlockProposers {
+            let epoch_proposers = EpochBlockProposers {
                 epoch,
                 fork,
                 proposers: proposers.into(),
-            })
-        })?;
+            };
+            self.cache
+                .put(key, Arc::new(OnceCell::with_value(epoch_proposers)));
+        }
 
-        Ok((&epoch_block_proposers.proposers, epoch_block_proposers.fork))
+        Ok(())
     }
 }
 

beacon_node/beacon_chain/src/data_column_verification.rs

@@ -1,4 +1,4 @@
-use crate::beacon_proposer_cache::Proposer;
+use crate::beacon_proposer_cache::{EpochBlockProposers, Proposer};
 use crate::block_verification::{
     cheap_state_advance_to_obtain_committees, get_validator_pubkey_cache, process_block_slash_info,
     BlockSlashInfo,
@@ -20,7 +20,7 @@ use tracing::debug;
 use types::data_column_sidecar::{ColumnIndex, DataColumnIdentifier};
 use types::{
     BeaconStateError, ChainSpec, DataColumnSidecar, DataColumnSubnetId, EthSpec, Hash256,
-    RuntimeVariableList, SignedBeaconBlockHeader, Slot,
+    RuntimeVariableList, SignedBeaconBlockHeader, Slot, Unsigned,
 };
 
 /// An error occurred while validating a gossip data column.
@@ -572,45 +572,57 @@ fn verify_proposer_and_signature<T: BeaconChainTypes>(
             parent_block.root
         };
 
-    let Proposer {
-        index: proposer_index,
-        fork,
-    } = chain
+    // We lock the cache briefly to get or insert a OnceCell, then drop the lock
+    // before doing proposer shuffling calculation via `OnceCell::get_or_try_init`. This avoids
+    // holding the lock during the computation, while still ensuring the result is cached and
+    // initialised only once.
+    //
+    // This approach unfortunately exposes the cache internals (`OnceCell` & `EpochBlockProposers`)
+    //  as a trade-off for avoiding lock contention.
+    let epoch_proposers_cell = chain
         .beacon_proposer_cache
         .lock()
-        .get_slot_or_insert_with::<E, _, GossipDataColumnError>(
-            proposer_shuffling_root,
+        .get_or_insert_key(column_epoch, proposer_shuffling_root);
+
+    let epoch_proposers = epoch_proposers_cell.get_or_try_init(move || {
+        debug!(
+            %block_root,
+            index = %column_index,
+            "Proposer shuffling cache miss for column verification"
+        );
+        let (parent_state_root, mut parent_state) = chain
+            .store
+            .get_advanced_hot_state(block_parent_root, column_slot, parent_block.state_root)
+            .map_err(|e| GossipDataColumnError::BeaconChainError(e.into()))?
+            .ok_or_else(|| {
+                BeaconChainError::DBInconsistent(format!(
+                    "Missing state for parent block {block_parent_root:?}",
+                ))
+            })?;
+
+        let state = cheap_state_advance_to_obtain_committees::<_, GossipDataColumnError>(
+            &mut parent_state,
+            Some(parent_state_root),
             column_slot,
-            move || {
-                debug!(
-                    %block_root,
-                    index = %column_index,
-                    "Proposer shuffling cache miss for column verification"
-                );
-                let (parent_state_root, mut parent_state) = chain
-                    .store
-                    .get_advanced_hot_state(block_parent_root, column_slot, parent_block.state_root)
-                    .map_err(|e| GossipDataColumnError::BeaconChainError(e.into()))?
-                    .ok_or_else(|| {
-                        BeaconChainError::DBInconsistent(format!(
-                            "Missing state for parent block {block_parent_root:?}",
-                        ))
-                    })?;
-
-                let state = cheap_state_advance_to_obtain_committees::<_, GossipDataColumnError>(
-                    &mut parent_state,
-                    Some(parent_state_root),
-                    column_slot,
-                    &chain.spec,
-                )?;
-
-                let proposers = state.get_beacon_proposer_indices(&chain.spec)?;
-                // Prime the proposer shuffling cache with the newly-learned value.
-                Ok((proposers, state.fork()))
-            },
-        )?
+            &chain.spec,
+        )?;
+
+        let proposers = state.get_beacon_proposer_indices(&chain.spec)?;
+        // Prime the proposer shuffling cache with the newly-learned value.
+        Ok::<EpochBlockProposers, Err>(EpochBlockProposers {
+            epoch: column_epoch,
+            fork: state.fork(),
+            proposers: proposers.into(),
+        })
+    })?;
+
+    let proposer_index = epoch_proposers
+        .proposers
+        .get(column_slot.as_usize() % E::SlotsPerEpoch::to_usize())
         .ok_or_else(|| BeaconChainError::NoProposerForSlot(column_slot))?;
 
+    let fork = epoch_proposers.fork;
+
     // Signature verify the signed block header.
     let signature_is_valid = {
         let pubkey_cache = get_validator_pubkey_cache(chain)