From dbe2188a37135b2ec34b0322459cdd012395dcc9 Mon Sep 17 00:00:00 2001 From: Matthew Keil Date: Wed, 4 Dec 2024 01:50:26 -0500 Subject: [PATCH] feat: debug too many shuffling promises (#7251) * feat: add asyncShufflingCalculation to StateTransitionOpts * feat: add asyncShufflingCalculation to all regen / processSlots consumers * fix: default to false for async shuffling and remove unnecessary props * fix: remove unnecessary flags from stateTransition * feat: implement conditional build of shuffling for prepareNextSlot * fix: spec test bug where shufflingCache is present from BeaconChain constructor * feat: sync build next shuffling if not queued async * fix: use getSync to pull next shuffling correctly * docs: add comment to prepareNextSlot * refactor: rename StateCloneOpts to StateRegenerationOpts * feat: pass asyncShufflingCalculation through to afterProcessEpoch and refactor conditional to run purely sync * docs: add issue number to comment * chore: lint --- .../beacon-node/src/chain/prepareNextSlot.ts | 7 ++- .../beacon-node/src/chain/regen/interface.ts | 18 ++++-- .../beacon-node/src/chain/regen/queued.ts | 18 ++++-- packages/beacon-node/src/chain/regen/regen.ts | 14 ++--- .../chain/stateCache/blockStateCacheImpl.ts | 4 +- .../chain/stateCache/fifoBlockStateCache.ts | 4 +- .../stateCache/inMemoryCheckpointsCache.ts | 10 ++-- .../stateCache/persistentCheckpointsCache.ts | 12 ++-- .../beacon-node/src/chain/stateCache/types.ts | 12 ++-- .../state-transition/src/cache/epochCache.ts | 59 ++++++++++--------- .../src/cache/epochTransitionCache.ts | 17 +++++- 11 files changed, 105 insertions(+), 70 deletions(-) diff --git a/packages/beacon-node/src/chain/prepareNextSlot.ts b/packages/beacon-node/src/chain/prepareNextSlot.ts index 40cadb3774c7..f78c1842bd78 100644 --- a/packages/beacon-node/src/chain/prepareNextSlot.ts +++ b/packages/beacon-node/src/chain/prepareNextSlot.ts @@ -114,7 +114,12 @@ 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 - {dontTransferCache: !isEpochTransition}, + // + // 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}, RegenCaller.precomputeEpoch ); diff --git a/packages/beacon-node/src/chain/regen/interface.ts b/packages/beacon-node/src/chain/regen/interface.ts index b70fbc059875..b9a4e38b5b68 100644 --- a/packages/beacon-node/src/chain/regen/interface.ts +++ b/packages/beacon-node/src/chain/regen/interface.ts @@ -28,8 +28,12 @@ export enum RegenFnName { getCheckpointState = "getCheckpointState", } -export type StateCloneOpts = { +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 { @@ -56,7 +60,11 @@ export interface IStateRegeneratorInternal { * Return a valid pre-state for a beacon block * This will always return a state in the latest viable epoch */ - getPreState(block: BeaconBlock, opts: StateCloneOpts, rCaller: RegenCaller): Promise; + getPreState( + block: BeaconBlock, + opts: StateRegenerationOpts, + rCaller: RegenCaller + ): Promise; /** * Return a valid checkpoint state @@ -64,7 +72,7 @@ export interface IStateRegeneratorInternal { */ getCheckpointState( cp: phase0.Checkpoint, - opts: StateCloneOpts, + opts: StateRegenerationOpts, rCaller: RegenCaller ): Promise; @@ -74,12 +82,12 @@ export interface IStateRegeneratorInternal { getBlockSlotState( blockRoot: RootHex, slot: Slot, - opts: StateCloneOpts, + opts: StateRegenerationOpts, rCaller: RegenCaller ): Promise; /** * Return the exact state with `stateRoot` */ - getState(stateRoot: RootHex, rCaller: RegenCaller, opts?: StateCloneOpts): Promise; + getState(stateRoot: RootHex, rCaller: RegenCaller, opts?: StateRegenerationOpts): Promise; } diff --git a/packages/beacon-node/src/chain/regen/queued.ts b/packages/beacon-node/src/chain/regen/queued.ts index b5084d593356..9069b384fd58 100644 --- a/packages/beacon-node/src/chain/regen/queued.ts +++ b/packages/beacon-node/src/chain/regen/queued.ts @@ -8,7 +8,13 @@ import {JobItemQueue} from "../../util/queue/index.js"; import {CheckpointHex, toCheckpointHex} from "../stateCache/index.js"; import {BlockStateCache, CheckpointStateCache} from "../stateCache/types.js"; import {RegenError, RegenErrorCode} from "./errors.js"; -import {IStateRegenerator, IStateRegeneratorInternal, RegenCaller, RegenFnName, StateCloneOpts} from "./interface.js"; +import { + IStateRegenerator, + IStateRegeneratorInternal, + RegenCaller, + RegenFnName, + StateRegenerationOpts, +} from "./interface.js"; import {RegenModules, StateRegenerator} from "./regen.js"; const REGEN_QUEUE_MAX_LEN = 256; @@ -86,7 +92,7 @@ export class QueuedStateRegenerator implements IStateRegenerator { */ getPreStateSync( block: BeaconBlock, - opts: StateCloneOpts = {dontTransferCache: true} + opts: StateRegenerationOpts = {dontTransferCache: true} ): CachedBeaconStateAllForks | null { const parentRoot = toRootHex(block.parentRoot); const parentBlock = this.forkChoice.getBlockHex(parentRoot); @@ -212,7 +218,7 @@ export class QueuedStateRegenerator implements IStateRegenerator { */ async getPreState( block: BeaconBlock, - opts: StateCloneOpts, + opts: StateRegenerationOpts, rCaller: RegenCaller ): Promise { this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getPreState}); @@ -231,7 +237,7 @@ export class QueuedStateRegenerator implements IStateRegenerator { async getCheckpointState( cp: phase0.Checkpoint, - opts: StateCloneOpts, + opts: StateRegenerationOpts, rCaller: RegenCaller ): Promise { this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getCheckpointState}); @@ -256,7 +262,7 @@ export class QueuedStateRegenerator implements IStateRegenerator { async getBlockSlotState( blockRoot: RootHex, slot: Slot, - opts: StateCloneOpts, + opts: StateRegenerationOpts, rCaller: RegenCaller ): Promise { this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getBlockSlotState}); @@ -268,7 +274,7 @@ export class QueuedStateRegenerator implements IStateRegenerator { async getState( stateRoot: RootHex, rCaller: RegenCaller, - opts: StateCloneOpts = {dontTransferCache: true} + opts: StateRegenerationOpts = {dontTransferCache: true} ): Promise { this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getState}); diff --git a/packages/beacon-node/src/chain/regen/regen.ts b/packages/beacon-node/src/chain/regen/regen.ts index 073556d8162f..06d1cee71332 100644 --- a/packages/beacon-node/src/chain/regen/regen.ts +++ b/packages/beacon-node/src/chain/regen/regen.ts @@ -20,7 +20,7 @@ import {getCheckpointFromState} from "../blocks/utils/checkpoint.js"; import {ChainEvent, ChainEventEmitter} from "../emitter.js"; import {BlockStateCache, CheckpointStateCache} from "../stateCache/types.js"; import {RegenError, RegenErrorCode} from "./errors.js"; -import {IStateRegeneratorInternal, RegenCaller, StateCloneOpts} from "./interface.js"; +import {IStateRegeneratorInternal, RegenCaller, StateRegenerationOpts} from "./interface.js"; export type RegenModules = { db: IBeaconDb; @@ -51,7 +51,7 @@ export class StateRegenerator implements IStateRegeneratorInternal { */ async getPreState( block: BeaconBlock, - opts: StateCloneOpts, + opts: StateRegenerationOpts, regenCaller: RegenCaller ): Promise { const parentBlock = this.modules.forkChoice.getBlock(block.parentRoot); @@ -84,7 +84,7 @@ export class StateRegenerator implements IStateRegeneratorInternal { */ async getCheckpointState( cp: phase0.Checkpoint, - opts: StateCloneOpts, + opts: StateRegenerationOpts, regenCaller: RegenCaller, allowDiskReload = false ): Promise { @@ -99,7 +99,7 @@ export class StateRegenerator implements IStateRegeneratorInternal { async getBlockSlotState( blockRoot: RootHex, slot: Slot, - opts: StateCloneOpts, + opts: StateRegenerationOpts, regenCaller: RegenCaller, allowDiskReload = false ): Promise { @@ -146,7 +146,7 @@ export class StateRegenerator implements IStateRegeneratorInternal { async getState( stateRoot: RootHex, caller: RegenCaller, - opts?: StateCloneOpts, + opts?: StateRegenerationOpts, // internal option, don't want to expose to external caller allowDiskReload = false ): Promise { @@ -322,7 +322,7 @@ async function processSlotsByCheckpoint( preState: CachedBeaconStateAllForks, slot: Slot, regenCaller: RegenCaller, - opts: StateCloneOpts + opts: StateRegenerationOpts ): Promise { let postState = await processSlotsToNearestCheckpoint(modules, preState, slot, regenCaller, opts); if (postState.slot < slot) { @@ -343,7 +343,7 @@ async function processSlotsToNearestCheckpoint( preState: CachedBeaconStateAllForks, slot: Slot, regenCaller: RegenCaller, - opts: StateCloneOpts + opts: StateRegenerationOpts ): Promise { const preSlot = preState.slot; const postSlot = slot; diff --git a/packages/beacon-node/src/chain/stateCache/blockStateCacheImpl.ts b/packages/beacon-node/src/chain/stateCache/blockStateCacheImpl.ts index f57c9a411923..7d87675b7bbc 100644 --- a/packages/beacon-node/src/chain/stateCache/blockStateCacheImpl.ts +++ b/packages/beacon-node/src/chain/stateCache/blockStateCacheImpl.ts @@ -3,7 +3,7 @@ import {CachedBeaconStateAllForks} from "@lodestar/state-transition"; import {Epoch, RootHex} from "@lodestar/types"; import {toRootHex} from "@lodestar/utils"; import {Metrics} from "../../metrics/index.js"; -import {StateCloneOpts} from "../regen/interface.js"; +import {StateRegenerationOpts} from "../regen/interface.js"; import {MapTracker} from "./mapMetrics.js"; import {BlockStateCache} from "./types.js"; @@ -39,7 +39,7 @@ export class BlockStateCacheImpl implements BlockStateCache { } } - get(rootHex: RootHex, opts?: StateCloneOpts): CachedBeaconStateAllForks | null { + get(rootHex: RootHex, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null { this.metrics?.lookups.inc(); const item = this.head?.stateRoot === rootHex ? this.head.state : this.cache.get(rootHex); if (!item) { diff --git a/packages/beacon-node/src/chain/stateCache/fifoBlockStateCache.ts b/packages/beacon-node/src/chain/stateCache/fifoBlockStateCache.ts index eec1fce5d6c2..a119efe66887 100644 --- a/packages/beacon-node/src/chain/stateCache/fifoBlockStateCache.ts +++ b/packages/beacon-node/src/chain/stateCache/fifoBlockStateCache.ts @@ -4,7 +4,7 @@ import {RootHex} from "@lodestar/types"; import {toRootHex} from "@lodestar/utils"; import {Metrics} from "../../metrics/index.js"; import {LinkedList} from "../../util/array.js"; -import {StateCloneOpts} from "../regen/interface.js"; +import {StateRegenerationOpts} from "../regen/interface.js"; import {MapTracker} from "./mapMetrics.js"; import {BlockStateCache} from "./types.js"; @@ -93,7 +93,7 @@ export class FIFOBlockStateCache implements BlockStateCache { /** * Get a state from this cache given a state root hex. */ - get(rootHex: RootHex, opts?: StateCloneOpts): CachedBeaconStateAllForks | null { + get(rootHex: RootHex, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null { this.metrics?.lookups.inc(); const item = this.cache.get(rootHex); if (!item) { diff --git a/packages/beacon-node/src/chain/stateCache/inMemoryCheckpointsCache.ts b/packages/beacon-node/src/chain/stateCache/inMemoryCheckpointsCache.ts index 4caa6779f697..81562d669365 100644 --- a/packages/beacon-node/src/chain/stateCache/inMemoryCheckpointsCache.ts +++ b/packages/beacon-node/src/chain/stateCache/inMemoryCheckpointsCache.ts @@ -3,7 +3,7 @@ import {CachedBeaconStateAllForks} from "@lodestar/state-transition"; import {Epoch, RootHex, phase0} from "@lodestar/types"; import {MapDef, toRootHex} from "@lodestar/utils"; import {Metrics} from "../../metrics/index.js"; -import {StateCloneOpts} from "../regen/interface.js"; +import {StateRegenerationOpts} from "../regen/interface.js"; import {MapTracker} from "./mapMetrics.js"; import {CacheItemType, CheckpointStateCache} from "./types.js"; @@ -42,7 +42,7 @@ export class InMemoryCheckpointStateCache implements CheckpointStateCache { this.maxEpochs = maxEpochs; } - async getOrReload(cp: CheckpointHex, opts?: StateCloneOpts): Promise { + async getOrReload(cp: CheckpointHex, opts?: StateRegenerationOpts): Promise { return this.get(cp, opts); } @@ -54,7 +54,7 @@ export class InMemoryCheckpointStateCache implements CheckpointStateCache { async getOrReloadLatest( rootHex: string, maxEpoch: number, - opts?: StateCloneOpts + opts?: StateRegenerationOpts ): Promise { return this.getLatest(rootHex, maxEpoch, opts); } @@ -64,7 +64,7 @@ export class InMemoryCheckpointStateCache implements CheckpointStateCache { return 0; } - get(cp: CheckpointHex, opts?: StateCloneOpts): CachedBeaconStateAllForks | null { + get(cp: CheckpointHex, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null { this.metrics?.lookups.inc(); const cpKey = toCheckpointKey(cp); const item = this.cache.get(cpKey); @@ -98,7 +98,7 @@ export class InMemoryCheckpointStateCache implements CheckpointStateCache { /** * Searches for the latest cached state with a `root`, starting with `epoch` and descending */ - getLatest(rootHex: RootHex, maxEpoch: Epoch, opts?: StateCloneOpts): CachedBeaconStateAllForks | null { + getLatest(rootHex: RootHex, maxEpoch: Epoch, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null { // sort epochs in descending order, only consider epochs lte `epoch` const epochs = Array.from(this.epochIndex.keys()) .sort((a, b) => b - a) diff --git a/packages/beacon-node/src/chain/stateCache/persistentCheckpointsCache.ts b/packages/beacon-node/src/chain/stateCache/persistentCheckpointsCache.ts index a7e0a7cdecfe..9a08aaa75461 100644 --- a/packages/beacon-node/src/chain/stateCache/persistentCheckpointsCache.ts +++ b/packages/beacon-node/src/chain/stateCache/persistentCheckpointsCache.ts @@ -7,7 +7,7 @@ import {Logger, MapDef, fromHex, sleep, toHex, toRootHex} from "@lodestar/utils" import {Metrics} from "../../metrics/index.js"; import {AllocSource, BufferPool, BufferWithKey} from "../../util/bufferPool.js"; import {IClock} from "../../util/clock.js"; -import {StateCloneOpts} from "../regen/interface.js"; +import {StateRegenerationOpts} from "../regen/interface.js"; import {serializeState} from "../serializeState.js"; import {CPStateDatastore, DatastoreKey} from "./datastore/index.js"; import {MapTracker} from "./mapMetrics.js"; @@ -168,7 +168,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache { * - Get block for processing * - Regen head state */ - async getOrReload(cp: CheckpointHex, opts?: StateCloneOpts): Promise { + async getOrReload(cp: CheckpointHex, opts?: StateRegenerationOpts): Promise { const stateOrStateBytesData = await this.getStateOrLoadDb(cp, opts); if (stateOrStateBytesData === null || isCachedBeaconState(stateOrStateBytesData)) { return stateOrStateBytesData?.clone(opts?.dontTransferCache) ?? null; @@ -240,7 +240,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache { */ async getStateOrLoadDb( cp: CheckpointHex, - opts?: StateCloneOpts + opts?: StateRegenerationOpts ): Promise { const cpKey = toCacheKey(cp); const inMemoryState = this.get(cpKey, opts); @@ -272,7 +272,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache { /** * Similar to get() api without reloading from disk */ - get(cpOrKey: CheckpointHex | string, opts?: StateCloneOpts): CachedBeaconStateAllForks | null { + get(cpOrKey: CheckpointHex | string, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null { this.metrics?.cpStateCache.lookups.inc(); const cpKey = typeof cpOrKey === "string" ? cpOrKey : toCacheKey(cpOrKey); const cacheItem = this.cache.get(cpKey); @@ -323,7 +323,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache { /** * Searches in-memory state for the latest cached state with a `root` without reload, starting with `epoch` and descending */ - getLatest(rootHex: RootHex, maxEpoch: Epoch, opts?: StateCloneOpts): CachedBeaconStateAllForks | null { + getLatest(rootHex: RootHex, maxEpoch: Epoch, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null { // sort epochs in descending order, only consider epochs lte `epoch` const epochs = Array.from(this.epochIndex.keys()) .sort((a, b) => b - a) @@ -349,7 +349,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache { async getOrReloadLatest( rootHex: RootHex, maxEpoch: Epoch, - opts?: StateCloneOpts + opts?: StateRegenerationOpts ): Promise { // sort epochs in descending order, only consider epochs lte `epoch` const epochs = Array.from(this.epochIndex.keys()) diff --git a/packages/beacon-node/src/chain/stateCache/types.ts b/packages/beacon-node/src/chain/stateCache/types.ts index 403b469dd352..19f05c23ee35 100644 --- a/packages/beacon-node/src/chain/stateCache/types.ts +++ b/packages/beacon-node/src/chain/stateCache/types.ts @@ -1,7 +1,7 @@ import {routes} from "@lodestar/api"; import {CachedBeaconStateAllForks} from "@lodestar/state-transition"; import {Epoch, RootHex, phase0} from "@lodestar/types"; -import {StateCloneOpts} from "../regen/interface.js"; +import {StateRegenerationOpts} from "../regen/interface.js"; export type CheckpointHex = {epoch: Epoch; rootHex: RootHex}; @@ -21,7 +21,7 @@ export type CheckpointHex = {epoch: Epoch; rootHex: RootHex}; * The cache key is state root */ export interface BlockStateCache { - get(rootHex: RootHex, opts?: StateCloneOpts): CachedBeaconStateAllForks | null; + get(rootHex: RootHex, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null; add(item: CachedBeaconStateAllForks): void; setHeadState(item: CachedBeaconStateAllForks | null): void; /** @@ -60,15 +60,15 @@ export interface BlockStateCache { */ export interface CheckpointStateCache { init?: () => Promise; - getOrReload(cp: CheckpointHex, opts?: StateCloneOpts): Promise; + getOrReload(cp: CheckpointHex, opts?: StateRegenerationOpts): Promise; getStateOrBytes(cp: CheckpointHex): Promise; - get(cpOrKey: CheckpointHex | string, opts?: StateCloneOpts): CachedBeaconStateAllForks | null; + get(cpOrKey: CheckpointHex | string, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null; add(cp: phase0.Checkpoint, state: CachedBeaconStateAllForks): void; - getLatest(rootHex: RootHex, maxEpoch: Epoch, opts?: StateCloneOpts): CachedBeaconStateAllForks | null; + getLatest(rootHex: RootHex, maxEpoch: Epoch, opts?: StateRegenerationOpts): CachedBeaconStateAllForks | null; getOrReloadLatest( rootHex: RootHex, maxEpoch: Epoch, - opts?: StateCloneOpts + opts?: StateRegenerationOpts ): Promise; updatePreComputedCheckpoint(rootHex: RootHex, epoch: Epoch): number | null; prune(finalizedEpoch: Epoch, justifiedEpoch: Epoch): void; diff --git a/packages/state-transition/src/cache/epochCache.ts b/packages/state-transition/src/cache/epochCache.ts index 86e63c672024..af9e79bc831c 100644 --- a/packages/state-transition/src/cache/epochCache.ts +++ b/packages/state-transition/src/cache/epochCache.ts @@ -49,6 +49,7 @@ import { import {computeBaseRewardPerIncrement, computeSyncParticipantReward} from "../util/syncCommittee.js"; import {sumTargetUnslashedBalanceIncrements} from "../util/targetUnslashedBalance.js"; import {EffectiveBalanceIncrements, getEffectiveBalanceIncrementsWithLen} from "./effectiveBalanceIncrements.js"; +import {EpochTransitionCache} from "./epochTransitionCache.js"; import {Index2PubkeyCache, syncPubkeys} from "./pubkeyCache.js"; import {CachedBeaconStateAllForks} from "./stateCache.js"; import { @@ -605,14 +606,7 @@ export class EpochCache { * Steps for afterProcessEpoch * 1) update previous/current/next values of cached items */ - afterProcessEpoch( - state: CachedBeaconStateAllForks, - epochTransitionCache: { - nextShufflingDecisionRoot: RootHex; - nextShufflingActiveIndices: Uint32Array; - nextEpochTotalActiveBalanceByIncrement: number; - } - ): void { + afterProcessEpoch(state: CachedBeaconStateAllForks, epochTransitionCache: EpochTransitionCache): void { // Because the slot was incremented before entering this function the "next epoch" is actually the "current epoch" // in this context but that is not actually true because the state transition happens in the last 4 seconds of the // epoch. For the context of this function "upcoming epoch" is used to denote the epoch that will begin after this @@ -657,28 +651,35 @@ export class EpochCache { this.nextDecisionRoot = epochTransitionCache.nextShufflingDecisionRoot; this.nextActiveIndices = epochTransitionCache.nextShufflingActiveIndices; if (this.shufflingCache) { - this.nextShuffling = null; - // This promise will resolve immediately after the synchronous code of the state-transition runs. Until - // the build is done on a worker thread it will be calculated immediately after the epoch transition - // completes. Once the work is done concurrently it should be ready by time this get runs so the promise - // will resolve directly on the next spin of the event loop because the epoch transition and shuffling take - // about the same time to calculate so theoretically its ready now. Do not await here though in case it - // is not ready yet as the transition must not be asynchronous. - this.shufflingCache - .get(epochAfterUpcoming, this.nextDecisionRoot) - .then((shuffling) => { - if (!shuffling) { - throw new Error("EpochShuffling not returned from get in afterProcessEpoch"); - } - this.nextShuffling = shuffling; - }) - .catch((err) => { - this.shufflingCache?.logger?.error( - "EPOCH_CONTEXT_SHUFFLING_BUILD_ERROR", - {epoch: epochAfterUpcoming, decisionRoot: epochTransitionCache.nextShufflingDecisionRoot}, - err - ); + if (!epochTransitionCache.asyncShufflingCalculation) { + this.nextShuffling = this.shufflingCache.getSync(epochAfterUpcoming, this.nextDecisionRoot, { + state, + activeIndices: this.nextActiveIndices, }); + } else { + this.nextShuffling = null; + // This promise will resolve immediately after the synchronous code of the state-transition runs. Until + // the build is done on a worker thread it will be calculated immediately after the epoch transition + // completes. Once the work is done concurrently it should be ready by time this get runs so the promise + // will resolve directly on the next spin of the event loop because the epoch transition and shuffling take + // about the same time to calculate so theoretically its ready now. Do not await here though in case it + // is not ready yet as the transition must not be asynchronous. + this.shufflingCache + .get(epochAfterUpcoming, this.nextDecisionRoot) + .then((shuffling) => { + if (!shuffling) { + throw new Error("EpochShuffling not returned from get in afterProcessEpoch"); + } + this.nextShuffling = shuffling; + }) + .catch((err) => { + this.shufflingCache?.logger?.error( + "EPOCH_CONTEXT_SHUFFLING_BUILD_ERROR", + {epoch: epochAfterUpcoming, decisionRoot: epochTransitionCache.nextShufflingDecisionRoot}, + err + ); + }); + } } else { // Only for testing. shufflingCache should always be available in prod this.nextShuffling = computeEpochShuffling(state, this.nextActiveIndices, epochAfterUpcoming); diff --git a/packages/state-transition/src/cache/epochTransitionCache.ts b/packages/state-transition/src/cache/epochTransitionCache.ts index b21f940c28d5..d159ca5f3763 100644 --- a/packages/state-transition/src/cache/epochTransitionCache.ts +++ b/packages/state-transition/src/cache/epochTransitionCache.ts @@ -33,6 +33,10 @@ export type EpochTransitionCacheOpts = { * Assert progressive balances the same to EpochTransitionCache */ assertCorrectProgressiveBalances?: boolean; + /** + * Do not queue shuffling calculation async. Forces sync JIT calculation in afterProcessEpoch + */ + asyncShufflingCalculation?: boolean; }; /** @@ -176,6 +180,12 @@ export interface EpochTransitionCache { */ nextEpochTotalActiveBalanceByIncrement: number; + /** + * Compute the shuffling sync or async. Defaults to synchronous. Need to pass `true` with the + * `EpochTransitionCacheOpts` + */ + asyncShufflingCalculation: boolean; + /** * Track by validator index if it's active in the prev epoch. * Used in metrics @@ -387,7 +397,11 @@ export function beforeProcessEpoch( for (let i = 0; i < nextEpochShufflingActiveIndicesLength; i++) { nextShufflingActiveIndices[i] = nextEpochShufflingActiveValidatorIndices[i]; } - state.epochCtx.shufflingCache?.build(epochAfterNext, nextShufflingDecisionRoot, state, nextShufflingActiveIndices); + + const asyncShufflingCalculation = opts?.asyncShufflingCalculation ?? false; + if (asyncShufflingCalculation) { + state.epochCtx.shufflingCache?.build(epochAfterNext, nextShufflingDecisionRoot, state, nextShufflingActiveIndices); + } if (totalActiveStakeByIncrement < 1) { totalActiveStakeByIncrement = 1; @@ -514,6 +528,7 @@ export function beforeProcessEpoch( indicesToEject, nextShufflingDecisionRoot, nextShufflingActiveIndices, + asyncShufflingCalculation, // to be updated in processEffectiveBalanceUpdates nextEpochTotalActiveBalanceByIncrement: 0, isActivePrevEpoch,