fix: no state id finalized

packages/beacon-node/src/api/impl/beacon/state/utils.ts

@@ -38,15 +38,13 @@ async function resolveStateIdOrNull(
   }
 
   if (stateId === "finalized") {
-    const block = chain.forkChoice.getFinalizedBlock();
-    const state = await chain.getStateByStateRoot(block.stateRoot, opts);
-    return state && {state: state.state, executionOptimistic: isOptimisticBlock(block)};
+    const checkpoint = chain.forkChoice.getFinalizedCheckpoint();
+    return chain.getStateByCheckpoint(checkpoint);
   }
 
   if (stateId === "justified") {
-    const block = chain.forkChoice.getJustifiedBlock();
-    const state = await chain.getStateByStateRoot(block.stateRoot, opts);
-    return state && {state: state.state, executionOptimistic: isOptimisticBlock(block)};
+    const checkpoint = chain.forkChoice.getJustifiedCheckpoint();
+    return chain.getStateByCheckpoint(checkpoint);
   }
 
   if (typeof stateId === "string" && stateId.startsWith("0x")) {

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

@@ -82,7 +82,7 @@ import {BlockRewards, computeBlockRewards} from "./rewards/blockRewards.js";
 import {ShufflingCache} from "./shufflingCache.js";
 import {StateContextCache} from "./stateCache/stateContextCache.js";
 import {SeenGossipBlockInput} from "./seenCache/index.js";
-import {InMemoryCheckpointStateCache} from "./stateCache/stateContextCheckpointsCache.js";
+import {InMemoryCheckpointStateCache} from "./stateCache/inMemoryCheckpointsCache.js";
 import {FIFOBlockStateCache} from "./stateCache/fifoBlockStateCache.js";
 import {PersistentCheckpointStateCache} from "./stateCache/persistentCheckpointsCache.js";
 import {DbCPStateDatastore} from "./stateCache/datastore/db.js";
@@ -463,6 +463,19 @@ export class BeaconChain implements IBeaconChain {
     return data && {state: data, executionOptimistic: false};
   }
 
+  getStateByCheckpoint(
+    checkpoint: CheckpointWithHex
+  ): {state: BeaconStateAllForks; executionOptimistic: boolean} | null {
+    // TODO: this is not guaranteed to work with new state caches, should work on this before we turn n-historical state on
+    const cachedStateCtx = this.regen.getCheckpointStateSync(checkpoint);
+    if (cachedStateCtx) {
+      const block = this.forkChoice.getBlock(cachedStateCtx.latestBlockHeader.hashTreeRoot());
+      return {state: cachedStateCtx, executionOptimistic: block != null && isOptimisticBlock(block)};
+    }
+
+    return null;
+  }
+
   async getCanonicalBlockAtSlot(
     slot: Slot
   ): Promise<{block: allForks.SignedBeaconBlock; executionOptimistic: boolean} | null> {

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

@@ -23,7 +23,7 @@ import {
 import {BeaconConfig} from "@lodestar/config";
 import {Logger} from "@lodestar/utils";
 
-import {IForkChoice, ProtoBlock} from "@lodestar/fork-choice";
+import {CheckpointWithHex, IForkChoice, ProtoBlock} from "@lodestar/fork-choice";
 import {IEth1ForBlockProduction} from "../eth1/index.js";
 import {IExecutionEngine, IExecutionBuilder} from "../execution/index.js";
 import {Metrics} from "../metrics/metrics.js";
@@ -144,6 +144,10 @@ export interface IBeaconChain {
     stateRoot: RootHex,
     opts?: StateGetOpts
   ): Promise<{state: BeaconStateAllForks; executionOptimistic: boolean} | null>;
+  /** Returns a cached state by checkpoint */
+  getStateByCheckpoint(
+    checkpoint: CheckpointWithHex
+  ): {state: BeaconStateAllForks; executionOptimistic: boolean} | null;
 
   /**
    * Since we can have multiple parallel chains,

packages/beacon-node/src/chain/stateCache/stateContextCheckpointsCache.ts → packages/beacon-node/src/chain/stateCache/inMemoryCheckpointsCache.ts

@@ -21,11 +21,15 @@ export class InMemoryCheckpointStateCache implements CheckpointStateCache {
   private readonly cache: MapTracker<string, CachedBeaconStateAllForks>;
   /** Epoch -> Set<blockRoot> */
   private readonly epochIndex = new MapDef<Epoch, Set<string>>(() => new Set<string>());
+  /**
+   * Max number of epochs allowed in the cache
+   */
+  private readonly maxEpochs: number;
   private readonly metrics: Metrics["cpStateCache"] | null | undefined;
   private preComputedCheckpoint: string | null = null;
   private preComputedCheckpointHits: number | null = null;
 
-  constructor({metrics}: {metrics?: Metrics | null}) {
+  constructor({maxEpochs = MAX_EPOCHS, metrics}: {maxEpochs?: number; metrics?: Metrics | null}) {
     this.cache = new MapTracker(metrics?.cpStateCache);
     if (metrics) {
       this.metrics = metrics.cpStateCache;
@@ -36,6 +40,7 @@ export class InMemoryCheckpointStateCache implements CheckpointStateCache {
         metrics.cpStateCache.epochSize.set({type: CacheItemType.inMemory}, this.epochIndex.size)
       );
     }
+    this.maxEpochs = maxEpochs;
   }
 
   async getOrReload(cp: CheckpointHex, opts?: StateCloneOpts): Promise<CachedBeaconStateAllForks | null> {
@@ -130,8 +135,8 @@ export class InMemoryCheckpointStateCache implements CheckpointStateCache {
     const epochs = Array.from(this.epochIndex.keys()).filter(
       (epoch) => epoch !== finalizedEpoch && epoch !== justifiedEpoch
     );
-    if (epochs.length > MAX_EPOCHS) {
-      for (const epoch of epochs.slice(0, epochs.length - MAX_EPOCHS)) {
+    if (epochs.length > this.maxEpochs) {
+      for (const epoch of epochs.slice(0, epochs.length - this.maxEpochs)) {
         this.deleteAllEpochItems(epoch);
       }
     }

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

@@ -1,3 +1,3 @@
 export * from "./stateContextCache.js";
-export * from "./stateContextCheckpointsCache.js";
+export * from "./inMemoryCheckpointsCache.js";
 export * from "./fifoBlockStateCache.js";

packages/beacon-node/test/perf/chain/stateCache/stateContextCheckpointsCache.test.ts → packages/beacon-node/test/perf/chain/stateCache/inMemoryCheckpointsCache.test.ts

@@ -4,7 +4,7 @@ import {ssz, phase0} from "@lodestar/types";
 import {generateCachedState} from "../../../utils/state.js";
 import {InMemoryCheckpointStateCache, toCheckpointHex} from "../../../../src/chain/stateCache/index.js";
 
-describe("CheckpointStateCache perf tests", function () {
+describe("InMemoryCheckpointStateCache perf tests", function () {
   setBenchOpts({noThreshold: true});
 
   let state: CachedBeaconStateAllForks;
@@ -17,7 +17,7 @@ describe("CheckpointStateCache perf tests", function () {
     checkpoint = ssz.phase0.Checkpoint.defaultValue();
   });
 
-  itBench("CheckpointStateCache - add get delete", () => {
+  itBench("InMemoryCheckpointStateCache - add get delete", () => {
     checkpointStateCache.add(checkpoint, state);
     checkpointStateCache.get(toCheckpointHex(checkpoint));
     checkpointStateCache.delete(checkpoint);

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

@@ -0,0 +1,143 @@
+import {describe, beforeAll, it, expect, beforeEach} from "vitest";
+import {CachedBeaconStateAllForks, computeEpochAtSlot, computeStartSlotAtEpoch} from "@lodestar/state-transition";
+import {phase0} from "@lodestar/types";
+import {SLOTS_PER_EPOCH, SLOTS_PER_HISTORICAL_ROOT} from "@lodestar/params";
+import {
+  CheckpointHex,
+  InMemoryCheckpointStateCache,
+  toCheckpointHex,
+} from "../../../../src/chain/stateCache/inMemoryCheckpointsCache.js";
+import {generateCachedState} from "../../../utils/state.js";
+
+describe("InMemoryCheckpointStateCache", function () {
+  let root0a: Buffer, root0b: Buffer, root1: Buffer, root2: Buffer;
+  let cp0a: phase0.Checkpoint, cp0b: phase0.Checkpoint, cp1: phase0.Checkpoint, cp2: phase0.Checkpoint;
+  let cp0aHex: CheckpointHex, cp0bHex: CheckpointHex, cp1Hex: CheckpointHex, cp2Hex: CheckpointHex;
+  let states: Record<"cp0a" | "cp0b" | "cp1" | "cp2", CachedBeaconStateAllForks>;
+
+  let cache: InMemoryCheckpointStateCache;
+
+  const startSlotEpoch20 = computeStartSlotAtEpoch(20);
+  const startSlotEpoch21 = computeStartSlotAtEpoch(21);
+  const startSlotEpoch22 = computeStartSlotAtEpoch(22);
+
+  beforeAll(() => {
+    root0a = Buffer.alloc(32);
+    root0b = Buffer.alloc(32, 1);
+    root1 = Buffer.alloc(32, 2);
+    root2 = Buffer.alloc(32, 3);
+    root0b[31] = 1;
+    //     epoch: 19         20           21         22          23
+    //            |-----------|-----------|-----------|-----------|
+    //                       ^^           ^           ^
+    //                       ||           |           |
+    //                       |0b--------root1--------root2
+    //                       |
+    //                       0a
+    // root0a is of the last slot of epoch 19
+    cp0a = {epoch: 20, root: root0a};
+    // root0b is of the first slot of epoch 20
+    cp0b = {epoch: 20, root: root0b};
+    cp1 = {epoch: 21, root: root1};
+    cp2 = {epoch: 22, root: root2};
+    [cp0aHex, cp0bHex, cp1Hex, cp2Hex] = [cp0a, cp0b, cp1, cp2].map((cp) => toCheckpointHex(cp));
+    const allStates = [cp0a, cp0b, cp1, cp2]
+      .map((cp) => generateCachedState({slot: cp.epoch * SLOTS_PER_EPOCH}))
+      .map((state, i) => {
+        const stateEpoch = computeEpochAtSlot(state.slot);
+        if (stateEpoch === 20 && i === 0) {
+          // cp0a
+          state.blockRoots.set((startSlotEpoch20 - 1) % SLOTS_PER_HISTORICAL_ROOT, root0a);
+          state.blockRoots.set(startSlotEpoch20 % SLOTS_PER_HISTORICAL_ROOT, root0a);
+          return state;
+        }
+
+        // other states based on cp0b
+        state.blockRoots.set((startSlotEpoch20 - 1) % SLOTS_PER_HISTORICAL_ROOT, root0a);
+        state.blockRoots.set(startSlotEpoch20 % SLOTS_PER_HISTORICAL_ROOT, root0b);
+
+        if (stateEpoch >= 21) {
+          state.blockRoots.set(startSlotEpoch21 % SLOTS_PER_HISTORICAL_ROOT, root1);
+        }
+        if (stateEpoch >= 22) {
+          state.blockRoots.set(startSlotEpoch22 % SLOTS_PER_HISTORICAL_ROOT, root2);
+        }
+        return state;
+      });
+
+    states = {
+      // Previous Root Checkpoint State of epoch 20
+      cp0a: allStates[0],
+      // Current Root Checkpoint State of epoch 20
+      cp0b: allStates[1],
+      // Current Root Checkpoint State of epoch 21
+      cp1: allStates[2],
+      // Current Root Checkpoint State of epoch 22
+      cp2: allStates[3],
+    };
+
+    for (const state of allStates) {
+      state.hashTreeRoot();
+    }
+  });
+
+  beforeEach(() => {
+    cache = new InMemoryCheckpointStateCache({maxEpochs: 0});
+    cache.add(cp0a, states["cp0a"]);
+    cache.add(cp0b, states["cp0b"]);
+    cache.add(cp1, states["cp1"]);
+  });
+
+  it("getLatest", () => {
+    // cp0
+    expect(cache.getLatest(cp0aHex.rootHex, cp0a.epoch)?.hashTreeRoot()).toEqual(states["cp0a"].hashTreeRoot());
+    expect(cache.getLatest(cp0aHex.rootHex, cp0a.epoch + 1)?.hashTreeRoot()).toEqual(states["cp0a"].hashTreeRoot());
+    expect(cache.getLatest(cp0aHex.rootHex, cp0a.epoch - 1)?.hashTreeRoot()).toBeUndefined();
+
+    // cp1
+    expect(cache.getLatest(cp1Hex.rootHex, cp1.epoch)?.hashTreeRoot()).toEqual(states["cp1"].hashTreeRoot());
+    expect(cache.getLatest(cp1Hex.rootHex, cp1.epoch + 1)?.hashTreeRoot()).toEqual(states["cp1"].hashTreeRoot());
+    expect(cache.getLatest(cp1Hex.rootHex, cp1.epoch - 1)?.hashTreeRoot()).toBeUndefined();
+
+    // cp2
+    expect(cache.getLatest(cp2Hex.rootHex, cp2.epoch)?.hashTreeRoot()).toBeUndefined();
+  });
+
+  it("getStateOrBytes", async () => {
+    expect(((await cache.getStateOrBytes(cp0aHex)) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(
+      states["cp0a"].hashTreeRoot()
+    );
+    expect(((await cache.getStateOrBytes(cp0bHex)) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(
+      states["cp0b"].hashTreeRoot()
+    );
+    expect(((await cache.getStateOrBytes(cp1Hex)) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(
+      states["cp1"].hashTreeRoot()
+    );
+    expect(await cache.getStateOrBytes(cp2Hex)).toBeNull();
+  });
+
+  it("get", () => {
+    expect((cache.get(cp0aHex) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(states["cp0a"].hashTreeRoot());
+    expect((cache.get(cp0bHex) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(states["cp0b"].hashTreeRoot());
+    expect((cache.get(cp1Hex) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(states["cp1"].hashTreeRoot());
+    expect(cache.get(cp2Hex) as CachedBeaconStateAllForks).toBeNull();
+  });
+
+  it("pruneFinalized", () => {
+    cache.pruneFinalized(21);
+    expect(cache.get(cp0aHex) as CachedBeaconStateAllForks).toBeNull();
+    expect(cache.get(cp0bHex) as CachedBeaconStateAllForks).toBeNull();
+    expect((cache.get(cp1Hex) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(states["cp1"].hashTreeRoot());
+  });
+
+  it("prune", () => {
+    cache.add(cp2, states["cp2"]);
+    const finalizedEpoch = 21;
+    const justifiedEpoch = 22;
+    cache.prune(finalizedEpoch, justifiedEpoch);
+    expect(cache.get(cp0aHex) as CachedBeaconStateAllForks).toBeNull();
+    expect(cache.get(cp0bHex) as CachedBeaconStateAllForks).toBeNull();
+    expect((cache.get(cp1Hex) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(states["cp1"].hashTreeRoot());
+    expect((cache.get(cp2Hex) as CachedBeaconStateAllForks).hashTreeRoot()).toEqual(states["cp2"].hashTreeRoot());
+  });
+});