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;
+  persistentKey?: 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/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,20 @@ 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) {
+    // current epoch and previous epoch are likely cached in previous states
+    this.shufflingCache.processState(postState, postState.epochCtx.nextShuffling.epoch);
+    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);
+    // add Current Root Checkpoint State to the checkpoint state cache
+    // this could be the justified/finalized checkpoint state later according to https://github.com/ethereum/consensus-specs/blob/v1.4.0-beta.2/specs/phase0/beacon-chain.md
+    if (block.message.slot % SLOTS_PER_EPOCH === 0) {
+      this.regen.addCheckpointState(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

@@ -39,7 +39,7 @@ import {IExecutionEngine, IExecutionBuilder} from "../execution/index.js";
 import {Clock, ClockEvent, IClock} from "../util/clock.js";
 import {ensureDir, writeIfNotExist} from "../util/file.js";
 import {isOptimisticBlock} from "../util/forkChoice.js";
-import {CheckpointStateCache, StateContextCache} from "./stateCache/index.js";
+import {CHECKPOINT_STATES_FOLDER, PersistentCheckpointStateCache, LRUBlockStateCache} from "./stateCache/index.js";
 import {BlockProcessor, ImportBlockOpts} from "./blocks/index.js";
 import {ChainEventEmitter, ChainEvent} from "./emitter.js";
 import {IBeaconChain, ProposerPreparationData, BlockHash, StateGetOpts} from "./interface.js";
@@ -75,6 +75,11 @@ 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";
+import {MemoryCheckpointStateCache} from "./stateCache/memoryCheckpointsCache.js";
+import {FilePersistentApis} from "./stateCache/persistent/file.js";
+import {DbPersistentApis} from "./stateCache/persistent/db.js";
+import {StateContextCache} from "./stateCache/stateContextCache.js";
 
 /**
  * Arbitrary constants, blobs should be consumed immediately in the same slot they are produced.
@@ -130,6 +135,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 +217,7 @@ export class BeaconChain implements IBeaconChain {
 
     this.beaconProposerCache = new BeaconProposerCache(opts);
     this.checkpointBalancesCache = new CheckpointBalancesCache();
+    this.shufflingCache = new ShufflingCache(metrics);
 
     // Restore state caches
     // anchorState may already by a CachedBeaconState. If so, don't create the cache again, since deserializing all
@@ -225,16 +232,37 @@ export class BeaconChain implements IBeaconChain {
             pubkey2index: new PubkeyIndexMap(),
             index2pubkey: [],
           });
+    this.shufflingCache.processState(cachedState, cachedState.epochCtx.previousShuffling.epoch);
+    this.shufflingCache.processState(cachedState, cachedState.epochCtx.currentShuffling.epoch);
+    this.shufflingCache.processState(cachedState, cachedState.epochCtx.nextShuffling.epoch);
 
     // Persist single global instance of state caches
     this.pubkey2index = cachedState.epochCtx.pubkey2index;
     this.index2pubkey = cachedState.epochCtx.index2pubkey;
 
-    const stateCache = new StateContextCache({metrics});
-    const checkpointStateCache = new CheckpointStateCache({metrics});
+    const stateCache = this.opts.nHistoricalStates
+      ? new LRUBlockStateCache(this.opts, {metrics})
+      : new StateContextCache({metrics});
+    const persistentApis = this.opts.persistCheckpointStatesToFile
+      ? new FilePersistentApis(CHECKPOINT_STATES_FOLDER)
+      : new DbPersistentApis(this.db);
+    const checkpointStateCache = this.opts.nHistoricalStates
+      ? new PersistentCheckpointStateCache(
+          {
+            metrics,
+            logger,
+            clock,
+            shufflingCache: this.shufflingCache,
+            getHeadState: this.getHeadState.bind(this),
+            persistentApis,
+          },
+          this.opts
+        )
+      : new MemoryCheckpointStateCache({metrics});
 
     const {checkpoint} = computeAnchorCheckpoint(config, anchorState);
     stateCache.add(cachedState);
+    // TODO: remove once we go with n-historical states
     stateCache.setHeadState(cachedState);
     checkpointStateCache.add(checkpoint, cachedState);
 
@@ -841,15 +869,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/interface.ts

@@ -36,6 +36,7 @@ import {CheckpointBalancesCache} from "./balancesCache.js";
 import {IChainOptions} from "./options.js";
 import {AssembledBlockType, BlockAttributes, BlockType} from "./produceBlock/produceBlockBody.js";
 import {SeenAttestationDatas} from "./seenCache/seenAttestationData.js";
+import {ShufflingCache} from "./shufflingCache.js";
 
 export {BlockType, AssembledBlockType};
 export {ProposerPreparationData};
@@ -93,6 +94,7 @@ export interface IBeaconChain {
 
   readonly beaconProposerCache: BeaconProposerCache;
   readonly checkpointBalancesCache: CheckpointBalancesCache;
+  readonly shufflingCache: ShufflingCache;
   readonly producedBlobSidecarsCache: Map<BlockHash, {blobSidecars: deneb.BlobSidecars; slot: Slot}>;
   readonly producedBlindedBlobSidecarsCache: Map<BlockHash, {blobSidecars: deneb.BlindedBlobSidecars; slot: Slot}>;
   readonly producedBlockRoot: Set<RootHex>;

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 {PersistentCheckpointStateCacheOpts} from "./stateCache/types.js";
+import {LRUBlockStateCacheOpts} from "./stateCache/lruBlockStateCache.js";
 
 export type IChainOptions = BlockProcessOpts &
   PoolOpts &
   SeenCacheOpts &
   ForkChoiceOpts &
   ArchiverOpts &
+  LRUBlockStateCacheOpts &
+  PersistentCheckpointStateCacheOpts &
   LightClientServerOpts & {
     blsVerifyAllMainThread?: boolean;
     blsVerifyAllMultiThread?: boolean;
@@ -27,6 +31,9 @@ export type IChainOptions = BlockProcessOpts &
     trustedSetup?: string;
     broadcastValidationStrictness?: string;
     minSameMessageSignatureSetsToBatch: number;
+    nHistoricalStates?: boolean;
+    /** by default persist checkpoint state to db */
+    persistCheckpointStatesToFile?: boolean;
   };
 
 export type BlockProcessOpts = {
@@ -88,4 +95,14 @@ 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,
+  // TODO: change to false, leaving here to ease testing
+  nHistoricalStates: true,
+  // by default, persist checkpoint states to db
+  persistCheckpointStatesToFile: false,
+
+  // 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.
+  // TODO: change to 128 which is the old StateCache config, only change back to 32 when we enable n-historical state, leaving here to ease testing
+  maxStates: 32,
+  // only used when persistentCheckpointStateCache = true
+  maxEpochsInMemory: 2,
 };

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

@@ -83,7 +83,8 @@ export class PrepareNextSlotScheduler {
           headSlot,
           clockSlot,
         });
-
+        // It's important to still do this to get through Holesky unfinality time of low resouce nodes
+        await this.prunePerSlot(clockSlot);
         return;
       }
 
@@ -173,11 +174,28 @@ export class PrepareNextSlotScheduler {
           this.chain.emitter.emit(routes.events.EventType.payloadAttributes, {data, version: fork});
         }
       }
+
+      // do this after all as it's the lowest priority task
+      await this.prunePerSlot(clockSlot);
     } catch (e) {
       if (!isErrorAborted(e) && !isQueueErrorAborted(e)) {
         this.metrics?.precomputeNextEpochTransition.count.inc({result: "error"}, 1);
         this.logger.error("Failed to run prepareForNextSlot", {nextEpoch, isEpochTransition, prepareSlot}, e as Error);
       }
     }
   };
+
+  /**
+   * Pruning at the last 1/3 slot of first slot of epoch is the safest time because all epoch transitions already use the checkpoint states cached
+   * one down side of this is when `inMemoryEpochs = 0` and gossip block hasn't come yet then we have to reload state we added 2/3 slot ago
+   * However, it's not likely `inMemoryEpochs` is configured as 0, and this scenario rarely happen
+   * since we only use `inMemoryEpochs = 0` for testing, if it happens it's a good thing because it helps us test the reload flow
+   */
+  private prunePerSlot = async (clockSlot: Slot): Promise<void> => {
+    // a contabo vpss can have 10-12 holesky epoch transitions per epoch when syncing, stronger node may have more
+    // it's better to prune at the last 1/3 of every slot in order not to cache a lot of checkpoint states
+    // at synced time, it's likely we only prune at the 1st slot of epoch, all other prunes are no-op
+    const pruneCount = await this.chain.regen.pruneCheckpointStateCache();
+    this.logger.verbose("Pruned checkpoint state cache", {clockSlot, pruneCount});
+  };
 }