feat: implement new state caches (#6176)

* feat: implement LRUBlockStateCache

* feat: implement PersistentCheckpointStateCache

* feat: implement findSeedStateToReload

* fix: add missing catch()

* fix: import path in state-transition

* fix: model CacheItem and type in PersistentCheckpointStateCache

* refactor: use for loop in PersistentCheckpointStateCache.processState

* chore: move test code to beforeAll() in persistentCheckpointsCache.test.ts

* feat: do not prune persisted state when reload

* fix: fifo instead of lru BlockStateCache

* fix: do not prune the last added item in FIFOBlockStateCache

* fix: sync epochIndex and cache in PersistentCheckpointStateCache

* chore: cleanup persistent checkpoint cache types

* chore: tweak comments

* chore: tweak more comments

* chore: reword persistent apis

* chore: add type to cpStateCache size metrics

* fix: metrics labels after rebasing from unstable

---------

Co-authored-by: Cayman <caymannava@gmail.com>

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

@@ -167,6 +167,23 @@ export class ShufflingCache {
     }
   }
 
+  /**
+   * Same to get() function but synchronous.
+   */
+  getSync(shufflingEpoch: Epoch, decisionRootHex: RootHex): EpochShuffling | null {
+    const cacheItem = this.itemsByDecisionRootByEpoch.getOrDefault(shufflingEpoch).get(decisionRootHex);
+    if (cacheItem === undefined) {
+      return null;
+    }
+
+    if (isShufflingCacheItem(cacheItem)) {
+      return cacheItem.shuffling;
+    }
+
+    // ignore promise
+    return null;
+  }
+
   private add(shufflingEpoch: Epoch, decisionBlock: RootHex, cacheItem: CacheItem): void {
     this.itemsByDecisionRootByEpoch.getOrDefault(shufflingEpoch).set(decisionBlock, cacheItem);
     pruneSetToMax(this.itemsByDecisionRootByEpoch, this.maxEpochs);

packages/beacon-node/src/chain/stateCache/datastore/db.ts

@@ -0,0 +1,38 @@
+import {CachedBeaconStateAllForks} from "@lodestar/state-transition";
+import {phase0, ssz} from "@lodestar/types";
+import {IBeaconDb} from "../../../db/interface.js";
+import {CPStateDatastore, DatastoreKey} from "./types.js";
+
+/**
+ * Implementation of CPStateDatastore using db.
+ */
+export class DbCPStateDatastore implements CPStateDatastore {
+  constructor(private readonly db: IBeaconDb) {}
+
+  async write(cpKey: phase0.Checkpoint, state: CachedBeaconStateAllForks): Promise<DatastoreKey> {
+    const serializedCheckpoint = checkpointToDatastoreKey(cpKey);
+    const stateBytes = state.serialize();
+    await this.db.checkpointState.putBinary(serializedCheckpoint, stateBytes);
+    return serializedCheckpoint;
+  }
+
+  async remove(serializedCheckpoint: DatastoreKey): Promise<void> {
+    await this.db.checkpointState.delete(serializedCheckpoint);
+  }
+
+  async read(serializedCheckpoint: DatastoreKey): Promise<Uint8Array | null> {
+    return this.db.checkpointState.getBinary(serializedCheckpoint);
+  }
+
+  async readKeys(): Promise<DatastoreKey[]> {
+    return this.db.checkpointState.keys();
+  }
+}
+
+export function datastoreKeyToCheckpoint(key: DatastoreKey): phase0.Checkpoint {
+  return ssz.phase0.Checkpoint.deserialize(key);
+}
+
+export function checkpointToDatastoreKey(cp: phase0.Checkpoint): DatastoreKey {
+  return ssz.phase0.Checkpoint.serialize(cp);
+}

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

@@ -0,0 +1,2 @@
+export * from "./types.js";
+export * from "./db.js";

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

@@ -0,0 +1,13 @@
+import {CachedBeaconStateAllForks} from "@lodestar/state-transition";
+import {phase0} from "@lodestar/types";
+
+// With db implementation, persistedKey is serialized data of a checkpoint
+export type DatastoreKey = Uint8Array;
+
+// Make this generic to support testing
+export interface CPStateDatastore {
+  write: (cpKey: phase0.Checkpoint, state: CachedBeaconStateAllForks) => Promise<DatastoreKey>;
+  remove: (key: DatastoreKey) => Promise<void>;
+  read: (key: DatastoreKey) => Promise<Uint8Array | null>;
+  readKeys: () => Promise<DatastoreKey[]>;
+}

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

@@ -0,0 +1,181 @@
+import {toHexString} from "@chainsafe/ssz";
+import {RootHex} from "@lodestar/types";
+import {CachedBeaconStateAllForks} from "@lodestar/state-transition";
+import {routes} from "@lodestar/api";
+import {Metrics} from "../../metrics/index.js";
+import {LinkedList} from "../../util/array.js";
+import {MapTracker} from "./mapMetrics.js";
+import {BlockStateCache} from "./types.js";
+
+export type FIFOBlockStateCacheOpts = {
+  maxBlockStates?: number;
+};
+
+/**
+ * Regen state if there's a reorg distance > 32 slots.
+ */
+export const DEFAULT_MAX_BLOCK_STATES = 32;
+
+/**
+ * New implementation of BlockStateCache that keeps the most recent n states consistently
+ *  - Maintain a linked list (FIFO) with special handling for head state, which is always the first item in the list
+ *  - Prune per add() instead of per checkpoint so it only keeps n historical states consistently, prune from tail
+ *  - No need to prune per finalized checkpoint
+ *
+ * Given this block tree with Block 11 as head:
+ * ```
+         Block 10
+           |
+     +-----+-----+
+     |           |
+  Block 11   Block 12
+     ^           |
+     |           |
+    head       Block 13
+ * ```
+ * The maintained key order would be: 11 -> 13 -> 12 -> 10, and state 10 will be pruned first.
+ */
+export class FIFOBlockStateCache implements BlockStateCache {
+  /**
+   * Max number of states allowed in the cache
+   */
+  readonly maxStates: number;
+
+  private readonly cache: MapTracker<string, CachedBeaconStateAllForks>;
+  /**
+   * Key order to implement FIFO cache
+   */
+  private readonly keyOrder: LinkedList<string>;
+  private readonly metrics: Metrics["stateCache"] | null | undefined;
+
+  constructor(opts: FIFOBlockStateCacheOpts, {metrics}: {metrics?: Metrics | null}) {
+    this.maxStates = opts.maxBlockStates ?? DEFAULT_MAX_BLOCK_STATES;
+    this.cache = new MapTracker(metrics?.stateCache);
+    if (metrics) {
+      this.metrics = metrics.stateCache;
+      metrics.stateCache.size.addCollect(() => metrics.stateCache.size.set(this.cache.size));
+    }
+    this.keyOrder = new LinkedList();
+  }
+
+  /**
+   * Set a state as head, happens when importing a block and head block is changed.
+   */
+  setHeadState(item: CachedBeaconStateAllForks | null): void {
+    if (item !== null) {
+      this.add(item, true);
+    }
+  }
+
+  /**
+   * Get a state from this cache given a state root hex.
+   */
+  get(rootHex: RootHex): CachedBeaconStateAllForks | null {
+    this.metrics?.lookups.inc();
+    const item = this.cache.get(rootHex);
+    if (!item) {
+      return null;
+    }
+
+    this.metrics?.hits.inc();
+    this.metrics?.stateClonedCount.observe(item.clonedCount);
+
+    return item;
+  }
+
+  /**
+   * Add a state to this cache.
+   * @param isHead if true, move it to the head of the list. Otherwise add to the 2nd position.
+   * In importBlock() steps, normally it'll call add() with isHead = false first. Then call setHeadState() to set the head.
+   */
+  add(item: CachedBeaconStateAllForks, isHead = false): void {
+    const key = toHexString(item.hashTreeRoot());
+    if (this.cache.get(key) != null) {
+      if (!this.keyOrder.has(key)) {
+        throw Error(`State exists but key not found in keyOrder: ${key}`);
+      }
+      if (isHead) {
+        this.keyOrder.moveToHead(key);
+      } else {
+        this.keyOrder.moveToSecond(key);
+      }
+      // same size, no prune
+      return;
+    }
+
+    // new state
+    this.metrics?.adds.inc();
+    this.cache.set(key, item);
+    if (isHead) {
+      this.keyOrder.unshift(key);
+    } else {
+      // insert after head
+      const head = this.keyOrder.first();
+      if (head == null) {
+        // should not happen, however handle just in case
+        this.keyOrder.unshift(key);
+      } else {
+        this.keyOrder.insertAfter(head, key);
+      }
+    }
+    this.prune(key);
+  }
+
+  get size(): number {
+    return this.cache.size;
+  }
+
+  /**
+   * Prune the cache from tail to keep the most recent n states consistently.
+   * The tail of the list is the oldest state, in case regen adds back the same state,
+   * it should stay next to head so that it won't be pruned right away.
+   * The FIFO cache helps with this.
+   */
+  prune(lastAddedKey: string): void {
+    while (this.keyOrder.length > this.maxStates) {
+      const key = this.keyOrder.last();
+      // it does not make sense to prune the last added state
+      // this only happens when max state is 1 in a short period of time
+      if (key === lastAddedKey) {
+        break;
+      }
+      if (!key) {
+        // should not happen
+        throw new Error("No key");
+      }
+      this.keyOrder.pop();
+      this.cache.delete(key);
+    }
+  }
+
+  /**
+   * No need for this implementation
+   * This is only to conform to the old api
+   */
+  deleteAllBeforeEpoch(): void {}
+
+  /**
+   * ONLY FOR DEBUGGING PURPOSES. For lodestar debug API.
+   */
+  clear(): void {
+    this.cache.clear();
+  }
+
+  /** ONLY FOR DEBUGGING PURPOSES. For lodestar debug API */
+  dumpSummary(): routes.lodestar.StateCacheItem[] {
+    return Array.from(this.cache.entries()).map(([key, state]) => ({
+      slot: state.slot,
+      root: toHexString(state.hashTreeRoot()),
+      reads: this.cache.readCount.get(key) ?? 0,
+      lastRead: this.cache.lastRead.get(key) ?? 0,
+      checkpointState: false,
+    }));
+  }
+
+  /**
+   * For unit test only.
+   */
+  dumpKeyOrder(): string[] {
+    return this.keyOrder.toArray();
+  }
+}

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

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