refactor: remove build() and async shuffling calculation

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

@@ -418,6 +418,13 @@ export async function importBlock(
     this.logger.verbose("After importBlock caching postState without SSZ cache", {slot: postState.slot});
   }
 
+  // Cache shufflings when crossing an epoch boundary
+  const parentEpoch = computeEpochAtSlot(parentBlockSlot);
+  if (parentEpoch < blockEpoch) {
+    this.shufflingCache.processState(postState);
+    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

@@ -291,7 +291,8 @@ export class BeaconChain implements IBeaconChain {
     });
 
     this._earliestAvailableSlot = anchorState.slot;
-    this.shufflingCache = anchorState.epochCtx.shufflingCache = new ShufflingCache(metrics, logger, this.opts, [
+
+    this.shufflingCache = new ShufflingCache(metrics, logger, this.opts, [
       {
         shuffling: anchorState.epochCtx.previousShuffling,
         decisionRoot: anchorState.epochCtx.previousDecisionRoot,
@@ -417,6 +418,7 @@ export class BeaconChain implements IBeaconChain {
     clock.addListener(ClockEvent.epoch, this.onClockEpoch.bind(this));
     emitter.addListener(ChainEvent.forkChoiceFinalized, this.onForkChoiceFinalized.bind(this));
     emitter.addListener(ChainEvent.forkChoiceJustified, this.onForkChoiceJustified.bind(this));
+    emitter.addListener(ChainEvent.checkpoint, this.onCheckpoint.bind(this));
   }
 
   async init(): Promise<void> {
@@ -980,8 +982,8 @@ export class BeaconChain implements IBeaconChain {
       this.metrics?.gossipAttestation.useHeadBlockState.inc({caller: regenCaller});
       state = await this.regen.getState(attHeadBlock.stateRoot, regenCaller);
     }
-
-    // should always be the current epoch of the active context so no need to await a result from the ShufflingCache
+    // resolve the promise to unblock other calls of the same epoch and dependent root
+    this.shufflingCache.processState(state);
     return state.epochCtx.getShufflingAtEpoch(attEpoch);
   }
 
@@ -1165,6 +1167,10 @@ export class BeaconChain implements IBeaconChain {
     this.logger.verbose("Fork choice justified", {epoch: cp.epoch, root: cp.rootHex});
   }
 
+  private onCheckpoint(this: BeaconChain, _checkpoint: phase0.Checkpoint, state: CachedBeaconStateAllForks): void {
+    this.shufflingCache.processState(state);
+  }
+
   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));

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

@@ -117,12 +117,7 @@ export class PrepareNextSlotScheduler {
         // the slot 0 of next epoch will likely use this Previous Root Checkpoint state for state transition so we transfer cache here
         // the resulting state with cache will be cached in Checkpoint State Cache which is used for the upcoming block processing
         // for other slots dontTransferCached=true because we don't run state transition on this state
-        //
-        // Shuffling calculation will be done asynchronously when passing asyncShufflingCalculation=true.  Shuffling will be queued in
-        // beforeProcessEpoch and should theoretically be ready immediately after the synchronous epoch transition finished and the
-        // event loop is free.  In long periods of non-finality too many forks will cause the shufflingCache to throw an error for
-        // too many queued shufflings so only run async during normal epoch transition. See issue ChainSafe/lodestar#7244
-        {dontTransferCache: !isEpochTransition, asyncShufflingCalculation: true},
+        {dontTransferCache: !isEpochTransition},
         RegenCaller.precomputeEpoch
       );
 

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

@@ -31,10 +31,6 @@ export enum RegenFnName {
 
 export type StateRegenerationOpts = {
   dontTransferCache: boolean;
-  /**
-   * Do not queue shuffling calculation async. Forces sync JIT calculation in afterProcessEpoch if not passed as `true`
-   */
-  asyncShufflingCalculation?: boolean;
 };
 
 export interface IStateRegenerator extends IStateRegeneratorInternal {

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

@@ -1,11 +1,4 @@
-import {
-  BeaconStateAllForks,
-  EpochShuffling,
-  IShufflingCache,
-  ShufflingBuildProps,
-  computeEpochShuffling,
-  computeEpochShufflingAsync,
-} from "@lodestar/state-transition";
+import {CachedBeaconStateAllForks, EpochShuffling} from "@lodestar/state-transition";
 import {Epoch, RootHex} from "@lodestar/types";
 import {LodestarError, Logger, MapDef, pruneSetToMax} from "@lodestar/utils";
 import {Metrics} from "../metrics/metrics.js";
@@ -53,7 +46,7 @@ export type ShufflingCacheOpts = {
  * - if a shuffling is not available (which does not happen with default chain option of maxSkipSlots = 32), track a promise to make sure we don't compute the same shuffling twice
  * - skip computing shuffling when loading state bytes from disk
  */
-export class ShufflingCache implements IShufflingCache {
+export class ShufflingCache {
   /** LRU cache implemented as a map, pruned every time we add an item */
   private readonly itemsByDecisionRootByEpoch: MapDef<Epoch, Map<RootHex, CacheItem>> = new MapDef(
     () => new Map<RootHex, CacheItem>()
@@ -136,60 +129,22 @@ export class ShufflingCache implements IShufflingCache {
   }
 
   /**
-   * Gets a cached shuffling via the epoch and decision root.  If the shuffling is not
-   * available it will build it synchronously and return the shuffling.
-   *
-   * NOTE: If a shuffling is already queued and not calculated it will build and resolve
-   * the promise but the already queued build will happen at some later time
+   * Process a state to extract and cache all shufflings (previous, current, next).
+   * Uses the stored decision roots from epochCtx.
    */
-  getSync<T extends ShufflingBuildProps | undefined>(
-    epoch: Epoch,
-    decisionRoot: RootHex,
-    buildProps?: T
-  ): T extends ShufflingBuildProps ? EpochShuffling : EpochShuffling | null {
-    const cacheItem = this.itemsByDecisionRootByEpoch.getOrDefault(epoch).get(decisionRoot);
-    if (!cacheItem) {
-      this.metrics?.shufflingCache.miss.inc();
-    } else if (isShufflingCacheItem(cacheItem)) {
-      this.metrics?.shufflingCache.hit.inc();
-      return cacheItem.shuffling;
-    } else if (buildProps) {
-      // TODO: (@matthewkeil) This should possible log a warning??
-      this.metrics?.shufflingCache.shufflingPromiseNotResolvedAndThrownAway.inc();
-    } else {
-      this.metrics?.shufflingCache.shufflingPromiseNotResolved.inc();
-    }
+  processState(state: CachedBeaconStateAllForks): void {
+    const {epochCtx} = state;
 
-    let shuffling: EpochShuffling | null = null;
-    if (buildProps) {
-      const timer = this.metrics?.shufflingCache.shufflingCalculationTime.startTimer({source: "getSync"});
-      shuffling = computeEpochShuffling(buildProps.state, buildProps.activeIndices, epoch);
-      timer?.();
-      this.set(shuffling, decisionRoot);
-    }
-    return shuffling as T extends ShufflingBuildProps ? EpochShuffling : EpochShuffling | null;
-  }
+    // Cache previous shuffling
+    this.set(epochCtx.previousShuffling, epochCtx.previousDecisionRoot);
 
-  /**
-   * Queue asynchronous build for an EpochShuffling, triggered from state-transition
-   */
-  build(epoch: number, decisionRoot: string, state: BeaconStateAllForks, activeIndices: Uint32Array): void {
-    this.insertPromise(epoch, decisionRoot);
-    /**
-     * TODO: (@matthewkeil) This will get replaced by a proper build queue and a worker to do calculations
-     * on a NICE thread
-     */
-    const timer = this.metrics?.shufflingCache.shufflingCalculationTime.startTimer({source: "build"});
-    computeEpochShufflingAsync(state, activeIndices, epoch)
-      .then((shuffling) => {
-        this.set(shuffling, decisionRoot);
-      })
-      .catch((err) =>
-        this.logger?.error(`error building shuffling for epoch ${epoch} at decisionRoot ${decisionRoot}`, {}, err)
-      )
-      .finally(() => {
-        timer?.();
-      });
+    // Cache current shuffling
+    this.set(epochCtx.currentShuffling, epochCtx.currentDecisionRoot);
+
+    // Cache next shuffling if it exists
+    if (epochCtx.nextShuffling !== null) {
+      this.set(epochCtx.nextShuffling, epochCtx.nextDecisionRoot);
+    }
   }
 
   /**
@@ -207,7 +162,8 @@ export class ShufflingCache implements IShufflingCache {
           (Date.now() - cacheItem.timeInsertedMs) / 1000
         );
       } else {
-        this.metrics?.shufflingCache.shufflingBuiltMultipleTimes.inc();
+        this.metrics?.shufflingCache.shufflingSetMultipleTimes.inc();
+        return;
       }
     }
     // set the shuffling

packages/beacon-node/src/metrics/metrics/lodestar.ts

@@ -1308,32 +1308,23 @@ export function createLodestarMetrics(
         name: "lodestar_shuffling_cache_miss_count",
         help: "Count of shuffling cache miss",
       }),
-      shufflingBuiltMultipleTimes: register.gauge({
-        name: "lodestar_shuffling_cache_recalculated_shuffling_count",
-        help: "Count of shuffling that were build multiple times",
-      }),
-      shufflingPromiseNotResolvedAndThrownAway: register.gauge({
-        name: "lodestar_shuffling_cache_promise_not_resolved_and_thrown_away_count",
-        help: "Count of shuffling cache promises that were discarded and the shuffling was built synchronously",
+      shufflingSetMultipleTimes: register.gauge({
+        name: "lodestar_shuffling_cache_set_multiple_times_count",
+        help: "Count of shuffling that were set multiple times",
       }),
       shufflingPromiseNotResolved: register.gauge({
         name: "lodestar_shuffling_cache_promise_not_resolved_count",
         help: "Count of shuffling cache promises that were requested before the promise was resolved",
       }),
-      nextShufflingNotOnEpochCache: register.gauge({
-        name: "lodestar_shuffling_cache_next_shuffling_not_on_epoch_cache",
-        help: "The next shuffling was not on the epoch cache before the epoch transition",
-      }),
       shufflingPromiseResolutionTime: register.histogram({
         name: "lodestar_shuffling_cache_promise_resolution_time_seconds",
         help: "Time from promise insertion until promise resolution when shuffling was ready in seconds",
         buckets: [0.5, 1, 1.5, 2],
       }),
-      shufflingCalculationTime: register.histogram<{source: "build" | "getSync"}>({
+      shufflingCalculationTime: register.histogram({
         name: "lodestar_shuffling_cache_shuffling_calculation_time_seconds",
         help: "Run time of shuffling calculation",
         buckets: [0.5, 0.75, 1, 1.25, 1.5],
-        labelNames: ["source"],
       }),
     },