Skip to content

Commit

Permalink
Add getPreStateSync()
Browse files Browse the repository at this point in the history
  • Loading branch information
ensi321 committed Jan 30, 2024
1 parent 47e9650 commit 719ab1d
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 30 deletions.
10 changes: 7 additions & 3 deletions packages/beacon-node/src/chain/chain.ts
Original file line number Diff line number Diff line change
Expand Up @@ -994,9 +994,13 @@ export class BeaconChain implements IBeaconChain {
}

async getBlockRewards(block: allForks.FullOrBlindedBeaconBlock): Promise<BlockRewards> {
const preState = (await this.regen.getPreState(block, {dontTransferCache: true}, RegenCaller.restApi)).clone();
const preState = this.regen.getPreStateSync(block);
const postState = this.regen.getStateSync(toHexString(block.stateRoot)) ?? undefined;
const result = computeBlockRewards(block, preState, postState);
return result;

if (preState === null) {
throw Error(`Pre-state is unavailable given block's parent root ${toHexString(block.parentRoot)}`);
}

return computeBlockRewards(block, preState, postState);
}
}
1 change: 1 addition & 0 deletions packages/beacon-node/src/chain/regen/interface.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ export interface IStateRegenerator extends IStateRegeneratorInternal {
dropCache(): void;
dumpCacheSummary(): routes.lodestar.StateCacheItem[];
getStateSync(stateRoot: RootHex): CachedBeaconStateAllForks | null;
getPreStateSync(block: allForks.BeaconBlock): CachedBeaconStateAllForks | null;
getCheckpointStateSync(cp: CheckpointHex): CachedBeaconStateAllForks | null;
getClosestHeadState(head: ProtoBlock): CachedBeaconStateAllForks | null;
pruneOnCheckpoint(finalizedEpoch: Epoch, justifiedEpoch: Epoch, headStateRoot: RootHex): void;
Expand Down
64 changes: 37 additions & 27 deletions packages/beacon-node/src/chain/regen/queued.ts
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,40 @@ export class QueuedStateRegenerator implements IStateRegenerator {
return this.stateCache.get(stateRoot);
}

getPreStateSync(block: allForks.BeaconBlock): CachedBeaconStateAllForks | null {
const parentRoot = toHexString(block.parentRoot);
const parentBlock = this.forkChoice.getBlockHex(parentRoot);
if (!parentBlock) {
throw new RegenError({
code: RegenErrorCode.BLOCK_NOT_IN_FORKCHOICE,
blockRoot: block.parentRoot,
});
}

const parentEpoch = computeEpochAtSlot(parentBlock.slot);
const blockEpoch = computeEpochAtSlot(block.slot);

// Check the checkpoint cache (if the pre-state is a checkpoint state)
if (parentEpoch < blockEpoch) {
const checkpointState = this.checkpointStateCache.getLatest(parentRoot, blockEpoch);
if (checkpointState && computeEpochAtSlot(checkpointState.slot) === blockEpoch) {
return checkpointState;
}
}

// Check the state cache, only if the state doesn't need to go through an epoch transition.
// Otherwise the state transition may not be cached and wasted. Queue for regen since the
// work required will still be significant.
if (parentEpoch === blockEpoch) {
const state = this.stateCache.get(parentBlock.stateRoot);
if (state) {
return state;
}
}

return null;
}

getCheckpointStateSync(cp: CheckpointHex): CachedBeaconStateAllForks | null {
return this.checkpointStateCache.get(cp);
}
Expand Down Expand Up @@ -137,34 +171,10 @@ export class QueuedStateRegenerator implements IStateRegenerator {
this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getPreState});

// First attempt to fetch the state from caches before queueing
const parentRoot = toHexString(block.parentRoot);
const parentBlock = this.forkChoice.getBlockHex(parentRoot);
if (!parentBlock) {
throw new RegenError({
code: RegenErrorCode.BLOCK_NOT_IN_FORKCHOICE,
blockRoot: block.parentRoot,
});
}
const cachedState = this.getPreStateSync(block);

const parentEpoch = computeEpochAtSlot(parentBlock.slot);
const blockEpoch = computeEpochAtSlot(block.slot);

// Check the checkpoint cache (if the pre-state is a checkpoint state)
if (parentEpoch < blockEpoch) {
const checkpointState = this.checkpointStateCache.getLatest(parentRoot, blockEpoch);
if (checkpointState && computeEpochAtSlot(checkpointState.slot) === blockEpoch) {
return checkpointState;
}
}

// Check the state cache, only if the state doesn't need to go through an epoch transition.
// Otherwise the state transition may not be cached and wasted. Queue for regen since the
// work required will still be significant.
if (parentEpoch === blockEpoch) {
const state = this.stateCache.get(parentBlock.stateRoot);
if (state) {
return state;
}
if (cachedState !== null) {
return cachedState;
}

// The state is not immediately available in the caches, enqueue the job
Expand Down

0 comments on commit 719ab1d

Please sign in to comment.