diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs index 20a93e31e8d..ea1c5089a36 100644 --- a/beacon_node/beacon_chain/src/beacon_chain.rs +++ b/beacon_node/beacon_chain/src/beacon_chain.rs @@ -4166,21 +4166,14 @@ impl BeaconChain { (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 BeaconChain { 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)> { + // 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> { + 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 BeaconChain { // 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 BeaconChain { "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!(