Optimise payload attributes calculation and add SSE (#4027)

## Issue Addressed

Closes #3896
Closes #3998
Closes #3700

## Proposed Changes

- Optimise the calculation of withdrawals for payload attributes by avoiding state clones, avoiding unnecessary state advances and reading from the snapshot cache if possible.
- Use the execution layer's payload attributes cache to avoid re-calculating payload attributes. I actually implemented a new LRU cache just for withdrawals but it had the exact same key and most of the same data as the existing payload attributes cache, so I deleted it.
- Add a new SSE event that fires when payloadAttributes are calculated. This is useful for block builders, a la ethereum/beacon-APIs#244.
- Add a new CLI flag `--always-prepare-payload` which forces payload attributes to be sent with every fcU regardless of connected proposers. This is intended for use by builders/relays.

For maximum effect, the flags I've been using to run Lighthouse in "payload builder mode" are:

```
--always-prepare-payload \
--prepare-payload-lookahead 12000 \
--suggested-fee-recipient 0x0000000000000000000000000000000000000000
```

The fee recipient is required so Lighthouse has something to pack in the payload attributes (it can be ignored by the builder). The lookahead causes fcU to be sent at the start of every slot rather than at 8s. As usual, fcU will also be sent after each change of head block. I think this combination is sufficient for builders to build on all viable heads. Often there will be two fcU (and two payload attributes) sent for the same slot: one sent at the start of the slot with the head from `n - 1` as the parent, and one sent after the block arrives with `n` as the parent.

Example usage of the new event stream:

```bash
curl -N "http://localhost:5052/eth/v1/events?topics=payload_attributes"
```

## Additional Info

- [x] Tests added by updating the proposer re-org tests. This has the benefit of testing the proposer re-org code paths with withdrawals too, confirming that the new changes don't interact poorly.
- [ ] Benchmarking with `blockdreamer` on devnet-7 showed promising results but I'm yet to do a comparison to `unstable`.


Co-authored-by: Michael Sproul <[email protected]>

Author: michaelsproul

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -57,7 +57,7 @@ use crate::validator_monitor::{
 };
 use crate::validator_pubkey_cache::ValidatorPubkeyCache;
 use crate::{metrics, BeaconChainError, BeaconForkChoiceStore, BeaconSnapshot, CachedHead};
-use eth2::types::{EventKind, SseBlock, SyncDuty};
+use eth2::types::{EventKind, SseBlock, SseExtendedPayloadAttributes, SyncDuty};
 use execution_layer::{
     BlockProposalContents, BuilderParams, ChainHealth, ExecutionLayer, FailedCondition,
     PayloadAttributes, PayloadStatus,
@@ -89,6 +89,7 @@ use state_processing::{
     state_advance::{complete_state_advance, partial_state_advance},
     BlockSignatureStrategy, ConsensusContext, SigVerifiedOp, VerifyBlockRoot, VerifyOperation,
 };
+use std::borrow::Cow;
 use std::cmp::Ordering;
 use std::collections::HashMap;
 use std::collections::HashSet;
@@ -3878,6 +3879,75 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         }))
     }
 
+    pub fn get_expected_withdrawals(
+        &self,
+        forkchoice_update_params: &ForkchoiceUpdateParameters,
+        proposal_slot: Slot,
+    ) -> Result<Withdrawals<T::EthSpec>, Error> {
+        let cached_head = self.canonical_head.cached_head();
+        let head_state = &cached_head.snapshot.beacon_state;
+
+        let parent_block_root = forkchoice_update_params.head_root;
+
+        let (unadvanced_state, unadvanced_state_root) =
+            if cached_head.head_block_root() == parent_block_root {
+                (Cow::Borrowed(head_state), cached_head.head_state_root())
+            } else if let Some(snapshot) = self
+                .snapshot_cache
+                .try_read_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
+                .ok_or(Error::SnapshotCacheLockTimeout)?
+                .get_cloned(parent_block_root, CloneConfig::none())
+            {
+                debug!(
+                    self.log,
+                    "Hit snapshot cache during withdrawals calculation";
+                    "slot" => proposal_slot,
+                    "parent_block_root" => ?parent_block_root,
+                );
+                let state_root = snapshot.beacon_state_root();
+                (Cow::Owned(snapshot.beacon_state), state_root)
+            } else {
+                info!(
+                    self.log,
+                    "Missed snapshot cache during withdrawals calculation";
+                    "slot" => proposal_slot,
+                    "parent_block_root" => ?parent_block_root
+                );
+                let block = self
+                    .get_blinded_block(&parent_block_root)?
+                    .ok_or(Error::MissingBeaconBlock(parent_block_root))?;
+                let state = self
+                    .get_state(&block.state_root(), Some(block.slot()))?
+                    .ok_or(Error::MissingBeaconState(block.state_root()))?;
+                (Cow::Owned(state), block.state_root())
+            };
+
+        // Parent state epoch is the same as the proposal, we don't need to advance because the
+        // list of expected withdrawals can only change after an epoch advance or a
+        // block application.
+        let proposal_epoch = proposal_slot.epoch(T::EthSpec::slots_per_epoch());
+        if head_state.current_epoch() == proposal_epoch {
+            return get_expected_withdrawals(&unadvanced_state, &self.spec)
+                .map_err(Error::PrepareProposerFailed);
+        }
+
+        // Advance the state using the partial method.
+        debug!(
+            self.log,
+            "Advancing state for withdrawals calculation";
+            "proposal_slot" => proposal_slot,
+            "parent_block_root" => ?parent_block_root,
+        );
+        let mut advanced_state = unadvanced_state.into_owned();
+        partial_state_advance(
+            &mut advanced_state,
+            Some(unadvanced_state_root),
+            proposal_epoch.start_slot(T::EthSpec::slots_per_epoch()),
+            &self.spec,
+        )?;
+        get_expected_withdrawals(&advanced_state, &self.spec).map_err(Error::PrepareProposerFailed)
+    }
+
     /// Determine whether a fork choice update to the execution layer should be overridden.
     ///
     /// This is *only* necessary when proposer re-orgs are enabled, because we have to prevent the
@@ -4664,7 +4734,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         // Nothing to do if there are no proposers registered with the EL, exit early to avoid
         // wasting cycles.
-        if !execution_layer.has_any_proposer_preparation_data().await {
+        if !self.config.always_prepare_payload
+            && !execution_layer.has_any_proposer_preparation_data().await
+        {
             return Ok(());
         }
 
@@ -4721,71 +4793,84 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         // If the execution layer doesn't have any proposer data for this validator then we assume
         // it's not connected to this BN and no action is required.
         let proposer = pre_payload_attributes.proposer_index;
-        if !execution_layer
-            .has_proposer_preparation_data(proposer)
-            .await
+        if !self.config.always_prepare_payload
+            && !execution_layer
+                .has_proposer_preparation_data(proposer)
+                .await
         {
             return Ok(());
         }
 
-        let withdrawals = match self.spec.fork_name_at_slot::<T::EthSpec>(prepare_slot) {
-            ForkName::Base | ForkName::Altair | ForkName::Merge => None,
-            ForkName::Capella => {
-                // We must use the advanced state because balances can change at epoch boundaries
-                // and balances affect withdrawals.
-                // FIXME(mark)
-                //   Might implement caching here in the future..
-                let prepare_state = self
-                    .state_at_slot(prepare_slot, StateSkipConfig::WithoutStateRoots)
-                    .map_err(|e| {
-                        error!(self.log, "State advance for withdrawals failed"; "error" => ?e);
-                        e
-                    })?;
-                Some(get_expected_withdrawals(&prepare_state, &self.spec))
-            }
-        }
-        .transpose()
-        .map_err(|e| {
-            error!(self.log, "Error preparing beacon proposer"; "error" => ?e);
-            e
-        })
-        .map(|withdrawals_opt| withdrawals_opt.map(|w| w.into()))
-        .map_err(Error::PrepareProposerFailed)?;
-
+        // Fetch payoad attributes from the execution layer's cache, or compute them from scratch
+        // if no matching entry is found. This saves recomputing the withdrawals which can take
+        // considerable time to compute if a state load is required.
         let head_root = forkchoice_update_params.head_root;
-        let payload_attributes = PayloadAttributes::new(
-            self.slot_clock
-                .start_of(prepare_slot)
-                .ok_or(Error::InvalidSlot(prepare_slot))?
-                .as_secs(),
-            pre_payload_attributes.prev_randao,
-            execution_layer.get_suggested_fee_recipient(proposer).await,
-            withdrawals,
-        );
+        let payload_attributes = if let Some(payload_attributes) = execution_layer
+            .payload_attributes(prepare_slot, head_root)
+            .await
+        {
+            payload_attributes
+        } else {
+            let withdrawals = match self.spec.fork_name_at_slot::<T::EthSpec>(prepare_slot) {
+                ForkName::Base | ForkName::Altair | ForkName::Merge => None,
+                ForkName::Capella => {
+                    let chain = self.clone();
+                    self.spawn_blocking_handle(
+                        move || {
+                            chain.get_expected_withdrawals(&forkchoice_update_params, prepare_slot)
+                        },
+                        "prepare_beacon_proposer_withdrawals",
+                    )
+                    .await?
+                    .map(Some)?
+                }
+            };
 
-        debug!(
-            self.log,
-            "Preparing beacon proposer";
-            "payload_attributes" => ?payload_attributes,
-            "prepare_slot" => prepare_slot,
-            "validator" => proposer,
-            "parent_root" => ?head_root,
-        );
+            let payload_attributes = PayloadAttributes::new(
+                self.slot_clock
+                    .start_of(prepare_slot)
+                    .ok_or(Error::InvalidSlot(prepare_slot))?
+                    .as_secs(),
+                pre_payload_attributes.prev_randao,
+                execution_layer.get_suggested_fee_recipient(proposer).await,
+                withdrawals.map(Into::into),
+            );
 
-        let already_known = execution_layer
-            .insert_proposer(prepare_slot, head_root, proposer, payload_attributes)
-            .await;
+            execution_layer
+                .insert_proposer(
+                    prepare_slot,
+                    head_root,
+                    proposer,
+                    payload_attributes.clone(),
+                )
+                .await;
 
-        // Only push a log to the user if this is the first time we've seen this proposer for this
-        // slot.
-        if !already_known {
+            // Only push a log to the user if this is the first time we've seen this proposer for
+            // this slot.
             info!(
                 self.log,
                 "Prepared beacon proposer";
                 "prepare_slot" => prepare_slot,
                 "validator" => proposer,
                 "parent_root" => ?head_root,
             );
+            payload_attributes
+        };
+
+        // Push a server-sent event (probably to a block builder or relay).
+        if let Some(event_handler) = &self.event_handler {
+            if event_handler.has_payload_attributes_subscribers() {
+                event_handler.register(EventKind::PayloadAttributes(ForkVersionedResponse {
+                    data: SseExtendedPayloadAttributes {
+                        proposal_slot: prepare_slot,
+                        proposer_index: proposer,
+                        parent_block_root: head_root,
+                        parent_block_hash: forkchoice_update_params.head_hash.unwrap_or_default(),
+                        payload_attributes: payload_attributes.into(),
+                    },
+                    version: Some(self.spec.fork_name_at_slot::<T::EthSpec>(prepare_slot)),
+                }));
+            }
         }
 
         let till_prepare_slot =
@@ -4808,7 +4893,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         // If we are close enough to the proposal slot, send an fcU, which will have payload
         // attributes filled in by the execution layer cache we just primed.
-        if till_prepare_slot <= self.config.prepare_payload_lookahead {
+        if self.config.always_prepare_payload
+            || till_prepare_slot <= self.config.prepare_payload_lookahead
+        {
             debug!(
                 self.log,
                 "Sending forkchoiceUpdate for proposer prep";

beacon_node/beacon_chain/src/chain_config.rs

@@ -67,6 +67,10 @@ pub struct ChainConfig {
     pub prepare_payload_lookahead: Duration,
     /// Use EL-free optimistic sync for the finalized part of the chain.
     pub optimistic_finalized_sync: bool,
+    /// Whether to send payload attributes every slot, regardless of connected proposers.
+    ///
+    /// This is useful for block builders and testing.
+    pub always_prepare_payload: bool,
 }
 
 impl Default for ChainConfig {
@@ -93,6 +97,7 @@ impl Default for ChainConfig {
             prepare_payload_lookahead: Duration::from_secs(4),
             // This value isn't actually read except in tests.
             optimistic_finalized_sync: true,
+            always_prepare_payload: false,
         }
     }
 }

beacon_node/beacon_chain/src/events.rs

@@ -14,6 +14,7 @@ pub struct ServerSentEventHandler<T: EthSpec> {
     exit_tx: Sender<EventKind<T>>,
     chain_reorg_tx: Sender<EventKind<T>>,
     contribution_tx: Sender<EventKind<T>>,
+    payload_attributes_tx: Sender<EventKind<T>>,
     late_head: Sender<EventKind<T>>,
     block_reward_tx: Sender<EventKind<T>>,
     log: Logger,
@@ -32,6 +33,7 @@ impl<T: EthSpec> ServerSentEventHandler<T> {
         let (exit_tx, _) = broadcast::channel(capacity);
         let (chain_reorg_tx, _) = broadcast::channel(capacity);
         let (contribution_tx, _) = broadcast::channel(capacity);
+        let (payload_attributes_tx, _) = broadcast::channel(capacity);
         let (late_head, _) = broadcast::channel(capacity);
         let (block_reward_tx, _) = broadcast::channel(capacity);
 
@@ -43,35 +45,63 @@ impl<T: EthSpec> ServerSentEventHandler<T> {
             exit_tx,
             chain_reorg_tx,
             contribution_tx,
+            payload_attributes_tx,
             late_head,
             block_reward_tx,
             log,
         }
     }
 
     pub fn register(&self, kind: EventKind<T>) {
-        let result = match kind {
-            EventKind::Attestation(attestation) => self
+        let log_count = |name, count| {
+            trace!(
+                self.log,
+                "Registering server-sent event";
+                "kind" => name,
+                "receiver_count" => count
+            );
+        };
+        let result = match &kind {
+            EventKind::Attestation(_) => self
                 .attestation_tx
-                .send(EventKind::Attestation(attestation))
-                .map(|count| trace!(self.log, "Registering server-sent attestation event"; "receiver_count" => count)),
-            EventKind::Block(block) => self.block_tx.send(EventKind::Block(block))
-                .map(|count| trace!(self.log, "Registering server-sent block event"; "receiver_count" => count)),
-            EventKind::FinalizedCheckpoint(checkpoint) => self.finalized_tx
-                .send(EventKind::FinalizedCheckpoint(checkpoint))
-                .map(|count| trace!(self.log, "Registering server-sent finalized checkpoint event"; "receiver_count" => count)),
-            EventKind::Head(head) => self.head_tx.send(EventKind::Head(head))
-                .map(|count| trace!(self.log, "Registering server-sent head event"; "receiver_count" => count)),
-            EventKind::VoluntaryExit(exit) => self.exit_tx.send(EventKind::VoluntaryExit(exit))
-                .map(|count| trace!(self.log, "Registering server-sent voluntary exit event"; "receiver_count" => count)),
-            EventKind::ChainReorg(reorg) => self.chain_reorg_tx.send(EventKind::ChainReorg(reorg))
-                .map(|count| trace!(self.log, "Registering server-sent chain reorg event"; "receiver_count" => count)),
-            EventKind::ContributionAndProof(contribution_and_proof) => self.contribution_tx.send(EventKind::ContributionAndProof(contribution_and_proof))
-                .map(|count| trace!(self.log, "Registering server-sent contribution and proof event"; "receiver_count" => count)),
-            EventKind::LateHead(late_head) => self.late_head.send(EventKind::LateHead(late_head))
-                .map(|count| trace!(self.log, "Registering server-sent late head event"; "receiver_count" => count)),
-            EventKind::BlockReward(block_reward) => self.block_reward_tx.send(EventKind::BlockReward(block_reward))
-                .map(|count| trace!(self.log, "Registering server-sent contribution and proof event"; "receiver_count" => count)),
+                .send(kind)
+                .map(|count| log_count(count, "attestation")),
+            EventKind::Block(_) => self
+                .block_tx
+                .send(kind)
+                .map(|count| log_count(count, "block")),
+            EventKind::FinalizedCheckpoint(_) => self
+                .finalized_tx
+                .send(kind)
+                .map(|count| log_count(count, "finalized checkpoint")),
+            EventKind::Head(_) => self
+                .head_tx
+                .send(kind)
+                .map(|count| log_count(count, "head")),
+            EventKind::VoluntaryExit(_) => self
+                .exit_tx
+                .send(kind)
+                .map(|count| log_count(count, "exit")),
+            EventKind::ChainReorg(_) => self
+                .chain_reorg_tx
+                .send(kind)
+                .map(|count| log_count(count, "chain reorg")),
+            EventKind::ContributionAndProof(_) => self
+                .contribution_tx
+                .send(kind)
+                .map(|count| log_count(count, "contribution and proof")),
+            EventKind::PayloadAttributes(_) => self
+                .payload_attributes_tx
+                .send(kind)
+                .map(|count| log_count(count, "payload attributes")),
+            EventKind::LateHead(_) => self
+                .late_head
+                .send(kind)
+                .map(|count| log_count(count, "late head")),
+            EventKind::BlockReward(_) => self
+                .block_reward_tx
+                .send(kind)
+                .map(|count| log_count(count, "block reward")),
         };
         if let Err(SendError(event)) = result {
             trace!(self.log, "No receivers registered to listen for event"; "event" => ?event);
@@ -106,6 +136,10 @@ impl<T: EthSpec> ServerSentEventHandler<T> {
         self.contribution_tx.subscribe()
     }
 
+    pub fn subscribe_payload_attributes(&self) -> Receiver<EventKind<T>> {
+        self.payload_attributes_tx.subscribe()
+    }
+
     pub fn subscribe_late_head(&self) -> Receiver<EventKind<T>> {
         self.late_head.subscribe()
     }
@@ -142,6 +176,10 @@ impl<T: EthSpec> ServerSentEventHandler<T> {
         self.contribution_tx.receiver_count() > 0
     }
 
+    pub fn has_payload_attributes_subscribers(&self) -> bool {
+        self.payload_attributes_tx.receiver_count() > 0
+    }
+
     pub fn has_late_head_subscribers(&self) -> bool {
         self.late_head.receiver_count() > 0
     }