chore: enable nHistoricalStates e2e test

packages/beacon-node/test/e2e/chain/stateCache/nHistoricalStates.test.ts

@@ -12,7 +12,7 @@ import {waitForEvent} from "../../../utils/events/resolver.js";
 import {ChainEvent, ReorgEventData} from "../../../../src/chain/emitter.js";
 import {connect, onPeerConnect} from "../../../utils/network.js";
 import {CacheItemType} from "../../../../src/chain/stateCache/types.js";
-import {ReorgedForkChoice} from "../../../mocks/forkchoice.js";
+import {ReorgedForkChoice} from "../../../mocks/fork-choice/reorg.js";
 
 /**
  * Test different reorg scenarios to make sure the StateCache implementations are correct.
@@ -249,8 +249,6 @@ describe(
         // chain is not finalized, epoch 4 is in-memory so CP state at epoch 0 1 2 3 are persisted
         numEpochsPersisted: 4,
         // chain is NOT finalized end of test
-        // TODO: remove this after proposer boost reorg is fully implemented
-        skip: true,
       },
     ];
 
@@ -295,6 +293,8 @@ describe(
             chain: {
               blsVerifyAllMainThread: true,
               forkchoiceConstructor: ReorgedForkChoice,
+              // this node does not need to reload state
+              nHistoricalStates: false,
               proposerBoost: true,
             },
           },
@@ -315,6 +315,7 @@ describe(
             chain: {
               blsVerifyAllMainThread: true,
               forkchoiceConstructor: ReorgedForkChoice,
+              // this node can follow with nHistoricalStates flag and it has to reload state
               nHistoricalStates: true,
               maxBlockStates,
               maxCPStateEpochsInMemory,
@@ -347,9 +348,16 @@ describe(
         afterEachCallbacks.push(() => Promise.all(validators.map((v) => v.close())));
 
         // wait for checkpoint 3 at slot 24, both nodes should reach same checkpoint
+        const cpEpoch = 3;
+        const cpSlot = 3 * SLOTS_PER_EPOCH;
         const checkpoints = await Promise.all(
           [reorgedBn, followupBn].map((bn) =>
-            waitForEvent<phase0.Checkpoint>(bn.chain.emitter, ChainEvent.checkpoint, 240000, (cp) => cp.epoch === 3)
+            waitForEvent<phase0.Checkpoint>(
+              bn.chain.emitter,
+              ChainEvent.checkpoint,
+              (cpSlot + genesisSlotsDelay + 1) * testParams.SECONDS_PER_SLOT * 1000,
+              (cp) => cp.epoch === cpEpoch
+            )
           )
         );
         expect(checkpoints[0]).toEqual(checkpoints[1]);
@@ -369,7 +377,8 @@ describe(
             waitForEvent<ReorgEventData>(
               bn.chain.emitter,
               routes.events.EventType.chainReorg,
-              240000,
+              // reorged event happens at reorgedSlot + 1
+              (reorgedSlot + 1 - cpSlot + 1) * testParams.SECONDS_PER_SLOT * 1000,
               (reorgData) => reorgData.slot === reorgedSlot + 1
             )
           )

packages/beacon-node/test/mocks/forkchoice.ts → packages/beacon-node/test/mocks/fork-choice/reorg.ts

@@ -1,5 +1,6 @@
 import {ChainForkConfig} from "@lodestar/config";
 import {ForkChoice, ForkChoiceOpts, IForkChoiceStore, ProtoArray, ProtoBlock} from "@lodestar/fork-choice";
+import {NotReorgedReason} from "@lodestar/fork-choice/lib/forkChoice/interface.js";
 import {Slot} from "@lodestar/types";
 
 /**
@@ -22,10 +23,6 @@ export class ReorgedForkChoice extends ForkChoice {
   reorgedSlot: Slot | undefined;
   reorgDistance: number | undefined;
   private readonly _fcStore: IForkChoiceStore;
-  // these flags to mark if the current call of getHead() is to produce a block
-  // the other way to check this is to check the n-th call of getHead() in the same slot, but this is easier
-  private calledUpdateHead = false;
-  private calledUpdateTime = false;
 
   constructor(
     config: ChainForkConfig,
@@ -38,33 +35,40 @@ export class ReorgedForkChoice extends ForkChoice {
     this._fcStore = fcStore;
   }
 
+  /**
+   * Override to trigger reorged event at `reorgedSlot + 1`
+   */
+  getProposerHead(
+    headBlock: ProtoBlock,
+    secFromSlot: number,
+    slot: Slot
+  ): {proposerHead: ProtoBlock; isHeadTimely: boolean; notReorgedReason?: NotReorgedReason} {
+    const currentSlot = this._fcStore.currentSlot;
+    if (this.reorgedSlot !== undefined && this.reorgDistance !== undefined && currentSlot === this.reorgedSlot + 1) {
+      const nodes = super.getAllNodes();
+      const headSlot = currentSlot - this.reorgDistance;
+      const headNode = nodes.find((node) => node.slot === headSlot);
+      if (headNode !== undefined) {
+        return {proposerHead: headNode, isHeadTimely: true};
+      }
+    }
+
+    return super.getProposerHead(headBlock, secFromSlot, slot);
+  }
+
   /**
    * Override the getHead() method
-   * - produceBlock: to reorg at a given slot and distance.
    * - produceAttestation: to build on the latest node after the reorged slot
    * - importBlock: to return the old branch at the reorged slot to produce the reorg event
    */
   getHead = (): ProtoBlock => {
     const currentSlot = this._fcStore.currentSlot;
-    const producingBlock = this.calledUpdateHead && this.calledUpdateTime;
     if (this.reorgedSlot === undefined || this.reorgDistance === undefined) {
       return super.getHead();
     }
 
-    this.calledUpdateTime = false;
-    this.calledUpdateHead = false;
-
-    // produceBlock: at reorgedSlot + 1, build new branch
-    if (currentSlot === this.reorgedSlot + 1 && producingBlock) {
-      const nodes = super.getAllNodes();
-      const headSlot = currentSlot - this.reorgDistance;
-      const headNode = nodes.find((node) => node.slot === headSlot);
-      if (headNode !== undefined) {
-        return headNode;
-      }
-    }
-
     // this is mainly for producing attestations + produceBlock for latter slots
+    // at `reorgedSlot + 1` should return the old head to trigger reorg event
     if (currentSlot > this.reorgedSlot + 1) {
       // from now on build on latest node which reorged at the given slot
       const nodes = super.getAllNodes();
@@ -75,12 +79,6 @@ export class ReorgedForkChoice extends ForkChoice {
     return super.getHead();
   };
 
-  updateTime(currentSlot: Slot): void {
-    // set flag to signal produceBlock flow
-    this.calledUpdateTime = true;
-    super.updateTime(currentSlot);
-  }
-
   /**
    * Override this function to:
    * - produceBlock flow: mark flags to indicate that the current call of getHead() is to produce a block
@@ -90,10 +88,6 @@ export class ReorgedForkChoice extends ForkChoice {
     if (this.reorgedSlot === undefined || this.reorgDistance === undefined) {
       return super.updateHead();
     }
-    // in all produce blocks flow, it always call updateTime() first then recomputeForkChoiceHead()
-    if (this.calledUpdateTime) {
-      this.calledUpdateHead = true;
-    }
     const currentSlot = this._fcStore.currentSlot;
     if (currentSlot <= this.reorgedSlot) {
       return super.updateHead();