diff --git a/packages/beacon-node/src/chain/balancesTreeCache.ts b/packages/beacon-node/src/chain/balancesTreeCache.ts index 462ae860809..6658ce66cf6 100644 --- a/packages/beacon-node/src/chain/balancesTreeCache.ts +++ b/packages/beacon-node/src/chain/balancesTreeCache.ts @@ -1,16 +1,21 @@ import {ListBasicTreeViewDU, UintNumberType} from "@chainsafe/ssz"; import {IBalancesTreeCache, CachedBeaconStateAllForks} from "@lodestar/state-transition"; import {Metrics} from "../metrics/index.js"; +import {Epoch} from "@lodestar/types"; const MAX_ITEMS = 2; +/** + * A cached of unused balances tree + * States in the same epoch share the same balances tree so we only want to cache max once per epoch + */ export class BalancesTreeCache implements IBalancesTreeCache { - private readonly unusedBalancesTrees: ListBasicTreeViewDU[] = []; + private readonly unusedBalancesTrees: Map> = new Map(); constructor(private readonly metrics: Metrics | null = null) { if (metrics) { metrics.balancesTreeCache.size.addCollect(() => { - metrics.balancesTreeCache.size.set(this.unusedBalancesTrees.length); + metrics.balancesTreeCache.size.set(this.unusedBalancesTrees.size); }); } } @@ -19,20 +24,28 @@ export class BalancesTreeCache implements IBalancesTreeCache { if (state === undefined) { return; } + const stateEpoch = state.epochCtx.epoch; + if (this.unusedBalancesTrees.has(stateEpoch)) { + return; + } - this.unusedBalancesTrees.push(state.balances); - while (this.unusedBalancesTrees.length > MAX_ITEMS) { - this.unusedBalancesTrees.shift(); + this.unusedBalancesTrees.set(stateEpoch, state.balances); + while (this.unusedBalancesTrees.size > MAX_ITEMS) { + const firstEpoch = Array.from(this.unusedBalancesTrees.keys())[0]; + this.unusedBalancesTrees.delete(firstEpoch); } } getUnusedBalances(): ListBasicTreeViewDU | undefined { - if (this.unusedBalancesTrees.length === 0) { + if (this.unusedBalancesTrees.size === 0) { this.metrics?.balancesTreeCache.miss.inc(); return undefined; } this.metrics?.balancesTreeCache.hit.inc(); - return this.unusedBalancesTrees.shift(); + const firstEpoch = Array.from(this.unusedBalancesTrees.keys())[0]; + const unusedBalances = this.unusedBalancesTrees.get(firstEpoch); + this.unusedBalancesTrees.delete(firstEpoch); + return unusedBalances; } }