fix: head state regen in n-historical state (#6674)

* fix: head state regen in n-historical state

* chore: improve getSeedState()

* chore: add more comments

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

@@ -257,7 +257,7 @@ export class BeaconChain implements IBeaconChain {
             logger,
             clock,
             shufflingCache: this.shufflingCache,
-            getHeadState: this.getHeadState.bind(this),
+            blockStateCache: stateCache,
             bufferPool: new BufferPool(anchorState.type.tree_serializedSize(anchorState.node), metrics),
             datastore: fileDataStore
               ? // debug option if we want to investigate any issues with the DB

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

@@ -68,6 +68,23 @@ export class FIFOBlockStateCache implements BlockStateCache {
     }
   }
 
+  /**
+   * Get a seed state for state reload, this could be any states. The goal is to have the same
+   * base merkle tree for all BeaconState objects across application.
+   * See packages/state-transition/src/util/loadState/loadState.ts for more detail
+   */
+  getSeedState(): CachedBeaconStateAllForks {
+    const firstValue = this.cache.values().next();
+    if (firstValue.done) {
+      // should not happen
+      throw Error("No state in FIFOBlockStateCache");
+    }
+
+    const firstState = firstValue.value;
+    // don't transfer cache because consumer only use this cache to reload another state from disc
+    return firstState.clone(true);
+  }
+
   /**
    * Get a state from this cache given a state root hex.
    */

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

@@ -12,7 +12,7 @@ import {BufferPool, BufferWithKey} from "../../util/bufferPool.js";
 import {StateCloneOpts} from "../regen/interface.js";
 import {MapTracker} from "./mapMetrics.js";
 import {CPStateDatastore, DatastoreKey, datastoreKeyToCheckpoint} from "./datastore/index.js";
-import {CheckpointHex, CacheItemType, CheckpointStateCache} from "./types.js";
+import {CheckpointHex, CacheItemType, CheckpointStateCache, BlockStateCache} from "./types.js";
 
 export type PersistentCheckpointStateCacheOpts = {
   /** Keep max n states in memory, persist the rest to disk */
@@ -21,16 +21,14 @@ export type PersistentCheckpointStateCacheOpts = {
   processLateBlock?: boolean;
 };
 
-type GetHeadStateFn = () => CachedBeaconStateAllForks;
-
 type PersistentCheckpointStateCacheModules = {
   metrics?: Metrics | null;
   logger: Logger;
   clock?: IClock | null;
   signal?: AbortSignal;
   shufflingCache: ShufflingCache;
   datastore: CPStateDatastore;
-  getHeadState?: GetHeadStateFn;
+  blockStateCache: BlockStateCache;
   bufferPool?: BufferPool;
 };
 
@@ -107,7 +105,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
   private readonly processLateBlock: boolean;
   private readonly datastore: CPStateDatastore;
   private readonly shufflingCache: ShufflingCache;
-  private readonly getHeadState?: GetHeadStateFn;
+  private readonly blockStateCache: BlockStateCache;
   private readonly bufferPool?: BufferPool;
 
   constructor(
@@ -118,7 +116,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
       signal,
       shufflingCache,
       datastore,
-      getHeadState,
+      blockStateCache,
       bufferPool,
     }: PersistentCheckpointStateCacheModules,
     opts: PersistentCheckpointStateCacheOpts
@@ -158,7 +156,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
     // Specify different datastore for testing
     this.datastore = datastore;
     this.shufflingCache = shufflingCache;
-    this.getHeadState = getHeadState;
+    this.blockStateCache = blockStateCache;
     this.bufferPool = bufferPool;
   }
 
@@ -197,10 +195,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
     const logMeta = {persistedKey: toHexString(persistedKey)};
     this.logger.debug("Reload: read state successful", logMeta);
     this.metrics?.stateReloadSecFromSlot.observe(this.clock?.secFromSlot(this.clock?.currentSlot ?? 0) ?? 0);
-    const seedState = this.findSeedStateToReload(cp) ?? this.getHeadState?.();
-    if (seedState == null) {
-      throw new Error("No seed state found for cp " + toCacheKey(cp));
-    }
+    const seedState = this.findSeedStateToReload(cp);
     this.metrics?.stateReloadEpochDiff.observe(Math.abs(seedState.epochCtx.epoch - cp.epoch));
     this.logger.debug("Reload: found seed state", {...logMeta, seedSlot: seedState.slot});
 
@@ -537,9 +532,9 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
    *
    * we always reload an epoch in the past. We'll start with epoch n then (n+1) prioritizing ones with the same view of `reloadedCp`.
    *
-   * This could return null and we should get head state in that case.
+   * Use seed state from the block cache if cannot find any seed states within this cache.
    */
-  findSeedStateToReload(reloadedCp: CheckpointHex): CachedBeaconStateAllForks | null {
+  findSeedStateToReload(reloadedCp: CheckpointHex): CachedBeaconStateAllForks {
     const maxEpoch = Math.max(...Array.from(this.epochIndex.keys()));
     const reloadedCpSlot = computeStartSlotAtEpoch(reloadedCp.epoch);
     let firstState: CachedBeaconStateAllForks | null = null;
@@ -574,7 +569,9 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
       }
     }
 
-    return firstState;
+    const seedBlockState = this.blockStateCache.getSeedState();
+    this.logger.verbose("Reload: use block state as seed state", {slot: seedBlockState.slot});
+    return seedBlockState;
   }
 
   clear(): void {

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

@@ -77,6 +77,15 @@ export class StateContextCache implements BlockStateCache {
     }
   }
 
+  /**
+   * Get a seed state for state reload.
+   * This is to conform to the api only as this cache is not used in n-historical state.
+   * See ./fifoBlockStateCache.ts for implementation
+   */
+  getSeedState(): CachedBeaconStateAllForks {
+    throw Error("Not implemented for StateContextCache");
+  }
+
   clear(): void {
     this.cache.clear();
     this.epochIndex.clear();

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

@@ -24,6 +24,10 @@ export interface BlockStateCache {
   get(rootHex: RootHex, opts?: StateCloneOpts): CachedBeaconStateAllForks | null;
   add(item: CachedBeaconStateAllForks): void;
   setHeadState(item: CachedBeaconStateAllForks | null): void;
+  /**
+   * Get a seed state for state reload.
+   */
+  getSeedState(): CachedBeaconStateAllForks;
   clear(): void;
   size: number;
   prune(headStateRootHex: RootHex): void;

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

@@ -249,6 +249,7 @@ 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,
       },
     ];

packages/beacon-node/test/unit/chain/stateCache/fifoBlockStateCache.test.ts

@@ -89,6 +89,7 @@ describe("FIFOBlockStateCache", function () {
 
   for (const {name, headState, addAsHeadArr, keptStates, prunedState} of testCases) {
     it(name, () => {
+      expect(cache.getSeedState().hashTreeRoot()).toEqual(state1.hashTreeRoot());
       // move to head this state
       cache.setHeadState(headState);
       expect(cache.size).toEqualWithMessage(2, "Size must be same as initial 2");

packages/beacon-node/test/unit/chain/stateCache/persistentCheckpointsCache.test.ts

@@ -10,7 +10,7 @@ import {ShufflingCache} from "../../../../src/chain/shufflingCache.js";
 import {testLogger} from "../../../utils/logger.js";
 import {getTestDatastore} from "../../../utils/chain/stateCache/datastore.js";
 import {CheckpointHex} from "../../../../src/chain/stateCache/types.js";
-import {toCheckpointHex} from "../../../../src/chain/index.js";
+import {FIFOBlockStateCache, toCheckpointHex} from "../../../../src/chain/index.js";
 
 describe("PersistentCheckpointStateCache", function () {
   let root0a: Buffer, root0b: Buffer, root1: Buffer, root2: Buffer;
@@ -87,7 +87,12 @@ describe("PersistentCheckpointStateCache", function () {
     fileApisBuffer = new Map();
     const datastore = getTestDatastore(fileApisBuffer);
     cache = new PersistentCheckpointStateCache(
-      {datastore, logger: testLogger(), shufflingCache: new ShufflingCache()},
+      {
+        datastore,
+        logger: testLogger(),
+        shufflingCache: new ShufflingCache(),
+        blockStateCache: new FIFOBlockStateCache({}, {}),
+      },
       {maxCPStateEpochsInMemory: 2, processLateBlock: true}
     );
     cache.add(cp0a, states["cp0a"]);
@@ -157,7 +162,12 @@ describe("PersistentCheckpointStateCache", function () {
       fileApisBuffer = new Map();
       const datastore = getTestDatastore(fileApisBuffer);
       cache = new PersistentCheckpointStateCache(
-        {datastore, logger: testLogger(), shufflingCache: new ShufflingCache()},
+        {
+          datastore,
+          logger: testLogger(),
+          shufflingCache: new ShufflingCache(),
+          blockStateCache: new FIFOBlockStateCache({}, {}),
+        },
         {maxCPStateEpochsInMemory: 2, processLateBlock: true}
       );
       cache.add(cp0a, states["cp0a"]);
@@ -229,7 +239,12 @@ describe("PersistentCheckpointStateCache", function () {
       fileApisBuffer = new Map();
       const datastore = getTestDatastore(fileApisBuffer);
       cache = new PersistentCheckpointStateCache(
-        {datastore, logger: testLogger(), shufflingCache: new ShufflingCache()},
+        {
+          datastore,
+          logger: testLogger(),
+          shufflingCache: new ShufflingCache(),
+          blockStateCache: new FIFOBlockStateCache({}, {}),
+        },
         {maxCPStateEpochsInMemory: 2, processLateBlock: true}
       );
       cache.add(cp0a, states["cp0a"]);
@@ -530,7 +545,12 @@ describe("PersistentCheckpointStateCache", function () {
       fileApisBuffer = new Map();
       const datastore = getTestDatastore(fileApisBuffer);
       cache = new PersistentCheckpointStateCache(
-        {datastore, logger: testLogger(), shufflingCache: new ShufflingCache()},
+        {
+          datastore,
+          logger: testLogger(),
+          shufflingCache: new ShufflingCache(),
+          blockStateCache: new FIFOBlockStateCache({}, {}),
+        },
         {maxCPStateEpochsInMemory: 1, processLateBlock: true}
       );
       cache.add(cp0a, states["cp0a"]);
@@ -797,7 +817,12 @@ describe("PersistentCheckpointStateCache", function () {
         fileApisBuffer = new Map();
         const datastore = getTestDatastore(fileApisBuffer);
         cache = new PersistentCheckpointStateCache(
-          {datastore, logger: testLogger(), shufflingCache: new ShufflingCache()},
+          {
+            datastore,
+            logger: testLogger(),
+            shufflingCache: new ShufflingCache(),
+            blockStateCache: new FIFOBlockStateCache({}, {}),
+          },
           {maxCPStateEpochsInMemory: 0, processLateBlock: true}
         );
         cache.add(cp0a, states["cp0a"]);
@@ -883,7 +908,12 @@ describe("PersistentCheckpointStateCache", function () {
         fileApisBuffer = new Map();
         const datastore = getTestDatastore(fileApisBuffer);
         cache = new PersistentCheckpointStateCache(
-          {datastore, logger: testLogger(), shufflingCache: new ShufflingCache()},
+          {
+            datastore,
+            logger: testLogger(),
+            shufflingCache: new ShufflingCache(),
+            blockStateCache: new FIFOBlockStateCache({}, {}),
+          },
           {maxCPStateEpochsInMemory: 0, processLateBlock: true}
         );