From f93844e63bcaeb4328fbf7d9beb9898145a6d9da Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Fri, 8 Mar 2024 16:15:28 +1100
Subject: [PATCH] Optimise concurrent block production (#5368)

* Optimise concurrent block production
---
 beacon_node/beacon_chain/src/beacon_chain.rs | 63 +++++++++++++-------
 1 file changed, 41 insertions(+), 22 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index 20a93e31e..ea1c5089a 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -4166,21 +4166,14 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 (re_org_state.pre_state, re_org_state.state_root)
             }
             // Normal case: proposing a block atop the current head using the cache.
-            else if let Some((_, cached_state)) = self
-                .block_production_state
-                .lock()
-                .take()
-                .filter(|(cached_block_root, _)| *cached_block_root == head_block_root)
+            else if let Some((_, cached_state)) =
+                self.get_state_from_block_production_cache(head_block_root)
             {
                 (cached_state.pre_state, cached_state.state_root)
             }
             // Fall back to a direct read of the snapshot cache.
-            else if let Some(pre_state) = self
-                .snapshot_cache
-                .try_read_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
-                .and_then(|snapshot_cache| {
-                    snapshot_cache.get_state_for_block_production(head_block_root)
-                })
+            else if let Some(pre_state) =
+                self.get_state_from_snapshot_cache_for_block_production(head_block_root)
             {
                 warn!(
                     self.log,
@@ -4221,6 +4214,40 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         Ok((state, state_root_opt))
     }
 
+    /// Get the state cached for block production *if* it matches `head_block_root`.
+    ///
+    /// This will clear the cache regardless of whether the block root matches, so only call this if
+    /// you think the `head_block_root` is likely to match!
+    fn get_state_from_block_production_cache(
+        &self,
+        head_block_root: Hash256,
+    ) -> Option<(Hash256, BlockProductionPreState<T::EthSpec>)> {
+        // Take care to drop the lock as quickly as possible.
+        let mut lock = self.block_production_state.lock();
+        let result = lock
+            .take()
+            .filter(|(cached_block_root, _)| *cached_block_root == head_block_root);
+        drop(lock);
+        result
+    }
+
+    /// Get a state for block production from the snapshot cache.
+    fn get_state_from_snapshot_cache_for_block_production(
+        &self,
+        head_block_root: Hash256,
+    ) -> Option<BlockProductionPreState<T::EthSpec>> {
+        if let Some(lock) = self
+            .snapshot_cache
+            .try_read_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
+        {
+            let result = lock.get_state_for_block_production(head_block_root);
+            drop(lock);
+            result
+        } else {
+            None
+        }
+    }
+
     /// Fetch the beacon state to use for producing a block if a 1-slot proposer re-org is viable.
     ///
     /// This function will return `None` if proposer re-orgs are disabled.
@@ -4313,12 +4340,8 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         // Only attempt a re-org if we hit the block production cache or snapshot cache.
         let pre_state = self
-            .block_production_state
-            .lock()
-            .take()
-            .and_then(|(cached_block_root, state)| {
-                (cached_block_root == re_org_parent_block).then_some(state)
-            })
+            .get_state_from_block_production_cache(re_org_parent_block)
+            .map(|(_, state)| state)
             .or_else(|| {
                 warn!(
                     self.log,
@@ -4327,11 +4350,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     "slot" => slot,
                     "block_root" => ?re_org_parent_block
                 );
-                self.snapshot_cache
-                    .try_read_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
-                    .and_then(|snapshot_cache| {
-                        snapshot_cache.get_state_for_block_production(re_org_parent_block)
-                    })
+                self.get_state_from_snapshot_cache_for_block_production(re_org_parent_block)
             })
             .or_else(|| {
                 debug!(