feat: n historical states

packages/api/package.json

@@ -70,7 +70,7 @@
   },
   "dependencies": {
     "@chainsafe/persistent-merkle-tree": "^0.5.0",
-    "@chainsafe/ssz": "^0.13.0",
+    "@chainsafe/ssz": "../ssz/packages/ssz",
     "@lodestar/config": "^1.11.1",
     "@lodestar/params": "^1.11.1",
     "@lodestar/types": "^1.11.1",

packages/api/src/beacon/routes/lodestar.ts

@@ -68,6 +68,7 @@ export type StateCacheItem = {
   /** Unix timestamp (ms) of the last read */
   lastRead: number;
   checkpointState: boolean;
+  filePath?: string;
 };
 
 export type LodestarNodePeer = NodePeer & {

packages/beacon-node/package.json

@@ -104,7 +104,7 @@
     "@chainsafe/libp2p-noise": "^13.0.0",
     "@chainsafe/persistent-merkle-tree": "^0.5.0",
     "@chainsafe/prometheus-gc-stats": "^1.0.0",
-    "@chainsafe/ssz": "^0.13.0",
+    "@chainsafe/ssz": "../ssz/packages/ssz",
     "@chainsafe/threads": "^1.11.1",
     "@ethersproject/abi": "^5.7.0",
     "@fastify/bearer-auth": "^9.0.0",

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

@@ -5,6 +5,7 @@ import {computeEpochAtSlot, computeStartSlotAtEpoch} from "@lodestar/state-trans
 import {CheckpointWithHex} from "@lodestar/fork-choice";
 import {IBeaconDb} from "../../db/index.js";
 import {IStateRegenerator} from "../regen/interface.js";
+import {getStateSlotFromBytes} from "../../util/multifork.js";
 
 /**
  * Minimum number of epochs between single temp archived states
@@ -83,13 +84,22 @@ export class StatesArchiver {
    * Only the new finalized state is stored to disk
    */
   async archiveState(finalized: CheckpointWithHex): Promise<void> {
-    const finalizedState = this.regen.getCheckpointStateSync(finalized);
-    if (!finalizedState) {
-      throw Error("No state in cache for finalized checkpoint state epoch #" + finalized.epoch);
+    // the finalized state could be from to disk
+    const finalizedStateOrBytes = await this.regen.getCheckpointStateOrBytes(finalized);
+    const {rootHex} = finalized;
+    if (!finalizedStateOrBytes) {
+      throw Error(`No state in cache for finalized checkpoint state epoch #${finalized.epoch} root ${rootHex}`);
+    }
+    if (finalizedStateOrBytes instanceof Uint8Array) {
+      const slot = getStateSlotFromBytes(finalizedStateOrBytes);
+      await this.db.stateArchive.putBinary(slot, finalizedStateOrBytes);
+      this.logger.verbose("Archived finalized state bytes", {finalizedEpoch: finalized.epoch, slot, root: rootHex});
+    } else {
+      // state
+      await this.db.stateArchive.put(finalizedStateOrBytes.slot, finalizedStateOrBytes);
+      this.logger.verbose("Archived finalized state", {epoch: finalized.epoch, root: rootHex});
     }
-    await this.db.stateArchive.put(finalizedState.slot, finalizedState);
     // don't delete states before the finalized state, auto-prune will take care of it
-    this.logger.verbose("Archived finalized state", {finalizedEpoch: finalized.epoch});
   }
 }
 

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

@@ -54,13 +54,11 @@ export class Archiver {
 
     if (!opts.disableArchiveOnCheckpoint) {
       this.chain.emitter.on(ChainEvent.forkChoiceFinalized, this.onFinalizedCheckpoint);
-      this.chain.emitter.on(ChainEvent.checkpoint, this.onCheckpoint);
 
       signal.addEventListener(
         "abort",
         () => {
           this.chain.emitter.off(ChainEvent.forkChoiceFinalized, this.onFinalizedCheckpoint);
-          this.chain.emitter.off(ChainEvent.checkpoint, this.onCheckpoint);
         },
         {once: true}
       );
@@ -76,15 +74,6 @@ export class Archiver {
     return this.jobQueue.push(finalized);
   };
 
-  private onCheckpoint = (): void => {
-    const headStateRoot = this.chain.forkChoice.getHead().stateRoot;
-    this.chain.regen.pruneOnCheckpoint(
-      this.chain.forkChoice.getFinalizedCheckpoint().epoch,
-      this.chain.forkChoice.getJustifiedCheckpoint().epoch,
-      headStateRoot
-    );
-  };
-
   private processFinalizedCheckpoint = async (finalized: CheckpointWithHex): Promise<void> => {
     try {
       const finalizedEpoch = finalized.epoch;

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

@@ -15,7 +15,7 @@ import {ZERO_HASH_HEX} from "../../constants/index.js";
 import {toCheckpointHex} from "../stateCache/index.js";
 import {isOptimisticBlock} from "../../util/forkChoice.js";
 import {isQueueErrorAborted} from "../../util/queue/index.js";
-import {ChainEvent, ReorgEventData} from "../emitter.js";
+import {ReorgEventData} from "../emitter.js";
 import {REPROCESS_MIN_TIME_TO_NEXT_SLOT_SEC} from "../reprocess.js";
 import type {BeaconChain} from "../chain.js";
 import {FullyVerifiedBlock, ImportBlockOpts, AttestationImportOpt} from "./types.js";
@@ -62,6 +62,7 @@ export async function importBlock(
   const blockRootHex = toHexString(blockRoot);
   const currentEpoch = computeEpochAtSlot(this.forkChoice.getTime());
   const blockEpoch = computeEpochAtSlot(block.message.slot);
+  const parentEpoch = computeEpochAtSlot(parentBlockSlot);
   const prevFinalizedEpoch = this.forkChoice.getFinalizedCheckpoint().epoch;
   const blockDelaySec = (fullyVerifiedBlock.seenTimestampSec - postState.genesisTime) % this.config.SECONDS_PER_SLOT;
 
@@ -202,16 +203,15 @@ export async function importBlock(
     }
   }
 
-  // 5. Compute head. If new head, immediately stateCache.setHeadState()
+  // 5. Compute head, always add to state cache so that it'll not be pruned soon
 
   const oldHead = this.forkChoice.getHead();
   const newHead = this.recomputeForkChoiceHead();
   const currFinalizedEpoch = this.forkChoice.getFinalizedCheckpoint().epoch;
 
+  // always set head state so it'll never be pruned from state cache
+  this.regen.updateHeadState(newHead.stateRoot, postState);
   if (newHead.blockRoot !== oldHead.blockRoot) {
-    // Set head state as strong reference
-    this.regen.updateHeadState(newHead.stateRoot, postState);
-
     this.emitter.emit(routes.events.EventType.head, {
       block: newHead.blockRoot,
       epochTransition: computeStartSlotAtEpoch(computeEpochAtSlot(newHead.slot)) === newHead.slot,
@@ -331,12 +331,15 @@ export async function importBlock(
     this.logger.verbose("After importBlock caching postState without SSZ cache", {slot: postState.slot});
   }
 
-  if (block.message.slot % SLOTS_PER_EPOCH === 0) {
-    // Cache state to preserve epoch transition work
+  if (parentEpoch < blockEpoch) {
+    this.shufflingCache.processState(postState);
+    this.logger.verbose("Processed shuffling for next epoch", {parentEpoch, blockEpoch, slot: block.message.slot});
+
+    // This is the real check point state per spec because the root is in current epoch
+    // it's important to add this to cache, when chain is finalized we'll query this state later
     const checkpointState = postState;
     const cp = getCheckpointFromState(checkpointState);
     this.regen.addCheckpointState(cp, checkpointState);
-    this.emitter.emit(ChainEvent.checkpoint, cp, checkpointState);
 
     // Note: in-lined code from previos handler of ChainEvent.checkpoint
     this.logger.verbose("Checkpoint processed", toCheckpointHex(cp));

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

@@ -75,6 +75,7 @@ import {BlockAttributes, produceBlockBody} from "./produceBlock/produceBlockBody
 import {computeNewStateRoot} from "./produceBlock/computeNewStateRoot.js";
 import {BlockInput} from "./blocks/types.js";
 import {SeenAttestationDatas} from "./seenCache/seenAttestationData.js";
+import {ShufflingCache} from "./shufflingCache.js";
 
 /**
  * Arbitrary constants, blobs should be consumed immediately in the same slot they are produced.
@@ -130,6 +131,7 @@ export class BeaconChain implements IBeaconChain {
 
   readonly beaconProposerCache: BeaconProposerCache;
   readonly checkpointBalancesCache: CheckpointBalancesCache;
+  readonly shufflingCache: ShufflingCache;
   // TODO DENEB: Prune data structure every time period, for both old entries
   /** Map keyed by executionPayload.blockHash of the block for those blobs */
   readonly producedBlobSidecarsCache = new Map<BlockHash, {blobSidecars: deneb.BlobSidecars; slot: Slot}>();
@@ -211,6 +213,7 @@ export class BeaconChain implements IBeaconChain {
 
     this.beaconProposerCache = new BeaconProposerCache(opts);
     this.checkpointBalancesCache = new CheckpointBalancesCache();
+    this.shufflingCache = new ShufflingCache();
 
     // Restore state caches
     // anchorState may already by a CachedBeaconState. If so, don't create the cache again, since deserializing all
@@ -230,12 +233,14 @@ export class BeaconChain implements IBeaconChain {
     this.pubkey2index = cachedState.epochCtx.pubkey2index;
     this.index2pubkey = cachedState.epochCtx.index2pubkey;
 
-    const stateCache = new StateContextCache({metrics});
-    const checkpointStateCache = new CheckpointStateCache({metrics});
+    const stateCache = new StateContextCache(this.opts, {metrics});
+    const checkpointStateCache = new CheckpointStateCache(
+      {metrics, logger, clock, shufflingCache: this.shufflingCache, getHeadState: this.getHeadState.bind(this)},
+      this.opts
+    );
 
     const {checkpoint} = computeAnchorCheckpoint(config, anchorState);
     stateCache.add(cachedState);
-    stateCache.setHeadState(cachedState);
     checkpointStateCache.add(checkpoint, cachedState);
 
     const forkChoice = initializeForkChoice(
@@ -841,15 +846,19 @@ 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);
+    // the finalized state could be from disk
+    const finalizedStateOrBytes = await this.regen.getCheckpointStateOrBytes(cp);
+    if (!finalizedStateOrBytes) {
+      throw Error("No state in cache for finalized checkpoint state epoch #" + cp.epoch);
+    }
     if (headState) {
-      this.opPool.pruneAll(headState, finalizedState);
+      this.opPool.pruneAll(headState, finalizedStateOrBytes);
     }
   }
 

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

@@ -14,8 +14,13 @@ import {
   BLS_WITHDRAWAL_PREFIX,
 } from "@lodestar/params";
 import {Epoch, phase0, capella, ssz, ValidatorIndex} from "@lodestar/types";
+import {ChainForkConfig} from "@lodestar/config";
 import {IBeaconDb} from "../../db/index.js";
 import {SignedBLSToExecutionChangeVersioned} from "../../util/types.js";
+import {
+  getValidatorsBytesFromStateBytes,
+  getWithdrawalCredentialFirstByteFromValidatorBytes,
+} from "../../util/sszBytes.js";
 import {isValidBlsToExecutionChangeForBlockInclusion} from "./utils.js";
 
 type HexRoot = string;
@@ -270,11 +275,11 @@ export class OpPool {
   /**
    * Prune all types of transactions given the latest head state
    */
-  pruneAll(headState: CachedBeaconStateAllForks, finalizedState: CachedBeaconStateAllForks | null): void {
+  pruneAll(headState: CachedBeaconStateAllForks, finalizedState: CachedBeaconStateAllForks | Uint8Array): void {
     this.pruneAttesterSlashings(headState);
     this.pruneProposerSlashings(headState);
     this.pruneVoluntaryExits(headState);
-    this.pruneBlsToExecutionChanges(headState, finalizedState);
+    this.pruneBlsToExecutionChanges(headState.config, finalizedState);
   }
 
   /**
@@ -344,17 +349,34 @@ export class OpPool {
    * credentials
    */
   private pruneBlsToExecutionChanges(
-    headState: CachedBeaconStateAllForks,
-    finalizedState: CachedBeaconStateAllForks | null
+    config: ChainForkConfig,
+    finalizedStateOrBytes: CachedBeaconStateAllForks | Uint8Array
   ): void {
+    const validatorBytes =
+      finalizedStateOrBytes instanceof Uint8Array
+        ? getValidatorsBytesFromStateBytes(config, finalizedStateOrBytes)
+        : null;
+
     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);
-        if (validator.withdrawalCredentials[0] !== BLS_WITHDRAWAL_PREFIX) {
-          this.blsToExecutionChanges.delete(key);
+      // there are at least finalied state bytes
+      let withDrawableCredentialFirstByte: number | null;
+      const validatorIndex = blsToExecutionChange.data.message.validatorIndex;
+      if (finalizedStateOrBytes instanceof Uint8Array) {
+        if (!validatorBytes) {
+          throw Error(
+            "Not able to extract validator bytes from finalized state bytes with length " + finalizedStateOrBytes.length
+          );
         }
+        withDrawableCredentialFirstByte = getWithdrawalCredentialFirstByteFromValidatorBytes(
+          validatorBytes,
+          validatorIndex
+        );
+      } else {
+        const validator = finalizedStateOrBytes.validators.getReadonly(validatorIndex);
+        withDrawableCredentialFirstByte = validator.withdrawalCredentials[0];
+      }
+      if (withDrawableCredentialFirstByte !== BLS_WITHDRAWAL_PREFIX) {
+        this.blsToExecutionChanges.delete(key);
       }
     }
   }

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

@@ -3,12 +3,16 @@ import {defaultOptions as defaultValidatorOptions} from "@lodestar/validator";
 import {ArchiverOpts} from "./archiver/index.js";
 import {ForkChoiceOpts} from "./forkChoice/index.js";
 import {LightClientServerOpts} from "./lightClient/index.js";
+import {CheckpointStateCacheOpts} from "./stateCache/stateContextCheckpointsCache.js";
+import {StateContextCacheOpts} from "./stateCache/stateContextCache.js";
 
 export type IChainOptions = BlockProcessOpts &
   PoolOpts &
   SeenCacheOpts &
   ForkChoiceOpts &
   ArchiverOpts &
+  StateContextCacheOpts &
+  CheckpointStateCacheOpts &
   LightClientServerOpts & {
     blsVerifyAllMainThread?: boolean;
     blsVerifyAllMultiThread?: boolean;
@@ -88,4 +92,7 @@ export const defaultChainOptions: IChainOptions = {
   // batching too much may block the I/O thread so if useWorker=false, suggest this value to be 32
   // since this batch attestation work is designed to work with useWorker=true, make this the lowest value
   minSameMessageSignatureSetsToBatch: 2,
+  // since Sep 2023, only cache up to 32 states by default. If a big reorg happens it'll load checkpoint state from disk and regen from there.
+  maxStates: 32,
+  maxEpochsInMemory: 2,
 };

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

@@ -173,6 +173,12 @@ export class PrepareNextSlotScheduler {
           this.chain.emitter.emit(routes.events.EventType.payloadAttributes, {data, version: fork});
         }
       }
+
+      if (clockSlot % SLOTS_PER_EPOCH === 0) {
+        // Don't let the checkpoint state cache to prune on its own, prune at the last 1/3 slot of slot 0 of each epoch
+        const pruneCount = this.chain.regen.pruneCheckpointStateCache();
+        this.logger.verbose("Pruned checkpoint state cache", {clockSlot, nextEpoch, pruneCount});
+      }
     } catch (e) {
       if (!isErrorAborted(e) && !isQueueErrorAborted(e)) {
         this.metrics?.precomputeNextEpochTransition.count.inc({result: "error"}, 1);

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

@@ -35,12 +35,13 @@ export interface IStateRegenerator extends IStateRegeneratorInternal {
   dropCache(): void;
   dumpCacheSummary(): routes.lodestar.StateCacheItem[];
   getStateSync(stateRoot: RootHex): CachedBeaconStateAllForks | null;
+  getCheckpointStateOrBytes(cp: CheckpointHex): Promise<CachedBeaconStateAllForks | Uint8Array | null>;
   getCheckpointStateSync(cp: CheckpointHex): CachedBeaconStateAllForks | null;
   getClosestHeadState(head: ProtoBlock): CachedBeaconStateAllForks | null;
-  pruneOnCheckpoint(finalizedEpoch: Epoch, justifiedEpoch: Epoch, headStateRoot: RootHex): void;
   pruneOnFinalized(finalizedEpoch: Epoch): void;
   addPostState(postState: CachedBeaconStateAllForks): void;
   addCheckpointState(cp: phase0.Checkpoint, item: CachedBeaconStateAllForks): void;
+  pruneCheckpointStateCache(): number;
   updateHeadState(newHeadStateRoot: RootHex, maybeHeadState: CachedBeaconStateAllForks): void;
   updatePreComputedCheckpoint(rootHex: RootHex, epoch: Epoch): number | null;
 }

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

@@ -70,6 +70,10 @@ export class QueuedStateRegenerator implements IStateRegenerator {
     return this.stateCache.get(stateRoot);
   }
 
+  async getCheckpointStateOrBytes(cp: CheckpointHex): Promise<CachedBeaconStateAllForks | Uint8Array | null> {
+    return this.checkpointStateCache.getStateOrBytes(cp);
+  }
+
   getCheckpointStateSync(cp: CheckpointHex): CachedBeaconStateAllForks | null {
     return this.checkpointStateCache.get(cp);
   }
@@ -78,11 +82,6 @@ export class QueuedStateRegenerator implements IStateRegenerator {
     return this.checkpointStateCache.getLatest(head.blockRoot, Infinity) || this.stateCache.get(head.stateRoot);
   }
 
-  pruneOnCheckpoint(finalizedEpoch: Epoch, justifiedEpoch: Epoch, headStateRoot: RootHex): void {
-    this.checkpointStateCache.prune(finalizedEpoch, justifiedEpoch);
-    this.stateCache.prune(headStateRoot);
-  }
-
   pruneOnFinalized(finalizedEpoch: number): void {
     this.checkpointStateCache.pruneFinalized(finalizedEpoch);
     this.stateCache.deleteAllBeforeEpoch(finalizedEpoch);
@@ -96,22 +95,27 @@ export class QueuedStateRegenerator implements IStateRegenerator {
     this.checkpointStateCache.add(cp, item);
   }
 
+  pruneCheckpointStateCache(): number {
+    return this.checkpointStateCache.pruneFromMemory();
+  }
+
   updateHeadState(newHeadStateRoot: RootHex, maybeHeadState: CachedBeaconStateAllForks): void {
     const headState =
       newHeadStateRoot === toHexString(maybeHeadState.hashTreeRoot())
         ? maybeHeadState
         : this.stateCache.get(newHeadStateRoot);
 
     if (headState) {
-      this.stateCache.setHeadState(headState);
+      // this move the headState to the front of the queue so it'll not be pruned right away
+      this.stateCache.add(headState);
     } else {
       // Trigger regen on head change if necessary
       this.logger.warn("Head state not available, triggering regen", {stateRoot: newHeadStateRoot});
-      // head has changed, so the existing cached head state is no longer useful. Set strong reference to null to free
-      // up memory for regen step below. During regen, node won't be functional but eventually head will be available
-      this.stateCache.setHeadState(null);
-      this.regen.getState(newHeadStateRoot, RegenCaller.processBlock).then(
-        (headStateRegen) => this.stateCache.setHeadState(headStateRegen),
+      // it's important to reload state to regen head state here
+      const shouldReload = true;
+      this.regen.getState(newHeadStateRoot, RegenCaller.processBlock, shouldReload).then(
+        // this move the headState to the front of the queue so it'll not be pruned right away
+        (headStateRegen) => this.stateCache.add(headStateRegen),
         (e) => this.logger.error("Error on head state regen", {}, e)
       );
     }