feat: prune BlsToExecutionChange opPool with head state (ChainSafe#6252)

* feat: prune BlsToExecutionChange opPool with head state

* fix: address PR comment

packages/beacon-node/src/chain/chain.ts

@@ -906,15 +906,20 @@ export class BeaconChain implements IBeaconChain {
     this.logger.verbose("Fork choice justified", {epoch: cp.epoch, root: cp.rootHex});
   }
 
-  private onForkChoiceFinalized(this: BeaconChain, cp: CheckpointWithHex): void {
+  private async onForkChoiceFinalized(this: BeaconChain, cp: CheckpointWithHex): Promise<void> {
     this.logger.verbose("Fork choice finalized", {epoch: cp.epoch, root: cp.rootHex});
     this.seenBlockProposers.prune(computeStartSlotAtEpoch(cp.epoch));
 
     // TODO: Improve using regen here
-    const headState = this.regen.getStateSync(this.forkChoice.getHead().stateRoot);
-    const finalizedState = this.regen.getCheckpointStateSync(cp);
+    const {blockRoot, stateRoot, slot} = this.forkChoice.getHead();
+    const headState = this.regen.getStateSync(stateRoot);
+    const headBlock = await this.db.block.get(fromHexString(blockRoot));
+    if (headBlock == null) {
+      throw Error(`Head block ${slot} ${headBlock} is not available in database`);
+    }
+
     if (headState) {
-      this.opPool.pruneAll(headState, finalizedState);
+      this.opPool.pruneAll(headBlock, headState);
     }
   }
 

packages/beacon-node/src/chain/opPools/opPool.ts

@@ -13,8 +13,9 @@ import {
   MAX_BLS_TO_EXECUTION_CHANGES,
   BLS_WITHDRAWAL_PREFIX,
   MAX_ATTESTER_SLASHINGS,
+  ForkSeq,
 } from "@lodestar/params";
-import {Epoch, phase0, capella, ssz, ValidatorIndex} from "@lodestar/types";
+import {Epoch, phase0, capella, ssz, ValidatorIndex, allForks} from "@lodestar/types";
 import {IBeaconDb} from "../../db/index.js";
 import {SignedBLSToExecutionChangeVersioned} from "../../util/types.js";
 import {BlockType} from "../interface.js";
@@ -300,11 +301,11 @@ export class OpPool {
   /**
    * Prune all types of transactions given the latest head state
    */
-  pruneAll(headState: CachedBeaconStateAllForks, finalizedState: CachedBeaconStateAllForks | null): void {
+  pruneAll(headBlock: allForks.SignedBeaconBlock, headState: CachedBeaconStateAllForks): void {
     this.pruneAttesterSlashings(headState);
     this.pruneProposerSlashings(headState);
     this.pruneVoluntaryExits(headState);
-    this.pruneBlsToExecutionChanges(headState, finalizedState);
+    this.pruneBlsToExecutionChanges(headBlock, headState);
   }
 
   /**
@@ -369,19 +370,28 @@ export class OpPool {
   }
 
   /**
-   * Call after finalizing
-   * Prune blsToExecutionChanges for validators which have been set with withdrawal
-   * credentials
+   * Prune BLS to execution changes that have been applied to the state more than 1 block ago.
+   * In the worse case where head block is reorged, the same BlsToExecutionChange message can be re-added
+   * to opPool once gossipsub seen cache TTL passes.
    */
   private pruneBlsToExecutionChanges(
-    headState: CachedBeaconStateAllForks,
-    finalizedState: CachedBeaconStateAllForks | null
+    headBlock: allForks.SignedBeaconBlock,
+    headState: CachedBeaconStateAllForks
   ): void {
+    const {config} = headState;
+    const recentBlsToExecutionChanges =
+      config.getForkSeq(headBlock.message.slot) >= ForkSeq.capella
+        ? (headBlock as capella.SignedBeaconBlock).message.body.blsToExecutionChanges
+        : [];
+
+    const recentBlsToExecutionChangeIndexes = new Set(
+      recentBlsToExecutionChanges.map((blsToExecutionChange) => blsToExecutionChange.message.validatorIndex)
+    );
+
     for (const [key, blsToExecutionChange] of this.blsToExecutionChanges.entries()) {
-      // TODO CAPELLA: We need the finalizedState to safely prune BlsToExecutionChanges. Finalized state may not be
-      // available in the cache, so it can be null. Once there's a head only prunning strategy, change
-      if (finalizedState !== null) {
-        const validator = finalizedState.validators.getReadonly(blsToExecutionChange.data.message.validatorIndex);
+      const {validatorIndex} = blsToExecutionChange.data.message;
+      if (!recentBlsToExecutionChangeIndexes.has(validatorIndex)) {
+        const validator = headState.validators.getReadonly(validatorIndex);
         if (validator.withdrawalCredentials[0] !== BLS_WITHDRAWAL_PREFIX) {
           this.blsToExecutionChanges.delete(key);
         }