feat: async shuffling refactor

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

@@ -202,7 +202,14 @@ export function getBeaconStateApi({
 
       const epoch = filters.epoch ?? computeEpochAtSlot(state.slot);
       const startSlot = computeStartSlotAtEpoch(epoch);
-      const shuffling = stateCached.epochCtx.getShufflingAtEpoch(epoch);
+      const decisionRoot = stateCached.epochCtx.getShufflingDecisionRoot(epoch);
+      const shuffling = await chain.shufflingCache.get(epoch, decisionRoot);
+      if (!shuffling) {
+        throw new ApiError(
+          500,
+          `No shuffling found to calculate committees for epoch: ${epoch} and decisionRoot: ${decisionRoot}`
+        );
+      }
       const committees = shuffling.committees;
       const committeesFlat = committees.flatMap((slotCommittees, slotInEpoch) => {
         const slot = startSlot + slotInEpoch;

packages/beacon-node/src/api/impl/validator/index.ts

@@ -3,6 +3,7 @@ import {ApplicationMethods} from "@lodestar/api/server";
 import {
   CachedBeaconStateAllForks,
   computeStartSlotAtEpoch,
+  calculateCommitteeAssignments,
   proposerShufflingDecisionRoot,
   attesterShufflingDecisionRoot,
   getBlockRootAtSlot,
@@ -995,7 +996,15 @@ export function getValidatorApi(
 
       // Check that all validatorIndex belong to the state before calling getCommitteeAssignments()
       const pubkeys = getPubkeysForIndices(state.validators, indices);
-      const committeeAssignments = state.epochCtx.getCommitteeAssignments(epoch, indices);
+      const decisionRoot = state.epochCtx.getShufflingDecisionRoot(epoch);
+      const shuffling = await chain.shufflingCache.get(epoch, decisionRoot);
+      if (!shuffling) {
+        throw new ApiError(
+          500,
+          `No shuffling found to calculate committee assignments for epoch: ${epoch} and decisionRoot: ${decisionRoot}`
+        );
+      }
+      const committeeAssignments = calculateCommitteeAssignments(shuffling, indices);
       const duties: routes.validator.AttesterDuty[] = [];
       for (let i = 0, len = indices.length; i < len; i++) {
         const validatorIndex = indices[i];

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

@@ -65,7 +65,6 @@ export async function importBlock(
   const blockRootHex = toRootHex(blockRoot);
   const currentEpoch = computeEpochAtSlot(this.forkChoice.getTime());
   const blockEpoch = computeEpochAtSlot(blockSlot);
-  const parentEpoch = computeEpochAtSlot(parentBlockSlot);
   const prevFinalizedEpoch = this.forkChoice.getFinalizedCheckpoint().epoch;
   const blockDelaySec = (fullyVerifiedBlock.seenTimestampSec - postState.genesisTime) % this.config.SECONDS_PER_SLOT;
   const recvToValLatency = Date.now() / 1000 - (opts.seenTimestampSec ?? Date.now() / 1000);
@@ -336,12 +335,6 @@ export async function importBlock(
     this.logger.verbose("After importBlock caching postState without SSZ cache", {slot: postState.slot});
   }
 
-  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: blockSlot});
-  }
-
   if (blockSlot % SLOTS_PER_EPOCH === 0) {
     // Cache state to preserve epoch transition work
     const checkpointState = postState;

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

@@ -13,6 +13,7 @@ import {
   PubkeyIndexMap,
   EpochShuffling,
   computeEndSlotAtEpoch,
+  computeAnchorCheckpoint,
 } from "@lodestar/state-transition";
 import {BeaconConfig} from "@lodestar/config";
 import {
@@ -60,7 +61,6 @@ import {
 import {IChainOptions} from "./options.js";
 import {QueuedStateRegenerator, RegenCaller} from "./regen/index.js";
 import {initializeForkChoice} from "./forkChoice/index.js";
-import {computeAnchorCheckpoint} from "./initState.js";
 import {IBlsVerifier, BlsSingleThreadVerifier, BlsMultiThreadWorkerPool} from "./bls/index.js";
 import {
   SeenAttesters,
@@ -245,7 +245,6 @@ export class BeaconChain implements IBeaconChain {
 
     this.beaconProposerCache = new BeaconProposerCache(opts);
     this.checkpointBalancesCache = new CheckpointBalancesCache();
-    this.shufflingCache = new ShufflingCache(metrics, this.opts);
 
     // Restore state caches
     // anchorState may already by a CachedBeaconState. If so, don't create the cache again, since deserializing all
@@ -260,9 +259,21 @@ 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);
+
+    this.shufflingCache = cachedState.epochCtx.shufflingCache = new ShufflingCache(metrics, logger, this.opts, [
+      {
+        shuffling: cachedState.epochCtx.previousShuffling,
+        decisionRoot: cachedState.epochCtx.previousDecisionRoot,
+      },
+      {
+        shuffling: cachedState.epochCtx.currentShuffling,
+        decisionRoot: cachedState.epochCtx.currentDecisionRoot,
+      },
+      {
+        shuffling: cachedState.epochCtx.nextShuffling,
+        decisionRoot: cachedState.epochCtx.nextDecisionRoot,
+      },
+    ]);
 
     // Persist single global instance of state caches
     this.pubkey2index = cachedState.epochCtx.pubkey2index;
@@ -898,8 +909,8 @@ export class BeaconChain implements IBeaconChain {
       state = await this.regen.getState(attHeadBlock.stateRoot, regenCaller);
     }
 
-    // resolve the promise to unblock other calls of the same epoch and dependent root
-    return this.shufflingCache.processState(state, attEpoch);
+    // should always be the current epoch of the active context so no need to await a result from the ShufflingCache
+    return state.epochCtx.getShufflingAtEpoch(attEpoch);
   }
 
   /**

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

@@ -14,10 +14,10 @@ import {
   getEffectiveBalanceIncrementsZeroInactive,
   isExecutionStateType,
   isMergeTransitionComplete,
+  computeAnchorCheckpoint,
 } from "@lodestar/state-transition";
 
 import {Logger, toRootHex} from "@lodestar/utils";
-import {computeAnchorCheckpoint} from "../initState.js";
 import {ChainEventEmitter} from "../emitter.js";
 import {ChainEvent} from "../emitter.js";
 import {GENESIS_SLOT} from "../../constants/index.js";

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

@@ -1,15 +1,13 @@
 import {
-  blockToHeader,
   computeEpochAtSlot,
   BeaconStateAllForks,
   CachedBeaconStateAllForks,
-  computeCheckpointEpochAtStateSlot,
   computeStartSlotAtEpoch,
 } from "@lodestar/state-transition";
-import {SignedBeaconBlock, phase0, ssz} from "@lodestar/types";
+import {SignedBeaconBlock} from "@lodestar/types";
 import {ChainForkConfig} from "@lodestar/config";
 import {Logger, toHex, toRootHex} from "@lodestar/utils";
-import {GENESIS_SLOT, ZERO_HASH} from "../constants/index.js";
+import {GENESIS_SLOT} from "../constants/index.js";
 import {IBeaconDb} from "../db/index.js";
 import {Eth1Provider} from "../eth1/index.js";
 import {Metrics} from "../metrics/index.js";
@@ -202,35 +200,3 @@ export function initBeaconMetrics(metrics: Metrics, state: BeaconStateAllForks):
   metrics.currentJustifiedEpoch.set(state.currentJustifiedCheckpoint.epoch);
   metrics.finalizedEpoch.set(state.finalizedCheckpoint.epoch);
 }
-
-export function computeAnchorCheckpoint(
-  config: ChainForkConfig,
-  anchorState: BeaconStateAllForks
-): {checkpoint: phase0.Checkpoint; blockHeader: phase0.BeaconBlockHeader} {
-  let blockHeader;
-  let root;
-  const blockTypes = config.getForkTypes(anchorState.latestBlockHeader.slot);
-
-  if (anchorState.latestBlockHeader.slot === GENESIS_SLOT) {
-    const block = blockTypes.BeaconBlock.defaultValue();
-    block.stateRoot = anchorState.hashTreeRoot();
-    blockHeader = blockToHeader(config, block);
-    root = ssz.phase0.BeaconBlockHeader.hashTreeRoot(blockHeader);
-  } else {
-    blockHeader = ssz.phase0.BeaconBlockHeader.clone(anchorState.latestBlockHeader);
-    if (ssz.Root.equals(blockHeader.stateRoot, ZERO_HASH)) {
-      blockHeader.stateRoot = anchorState.hashTreeRoot();
-    }
-    root = ssz.phase0.BeaconBlockHeader.hashTreeRoot(blockHeader);
-  }
-
-  return {
-    checkpoint: {
-      root,
-      // the checkpoint epoch = computeEpochAtSlot(anchorState.slot) + 1 if slot is not at epoch boundary
-      // this is similar to a process_slots() call
-      epoch: computeCheckpointEpochAtStateSlot(anchorState.slot),
-    },
-    blockHeader,
-  };
-}