From 550c7007bb3bcf9ea20d8c3c13a79fcb56844d46 Mon Sep 17 00:00:00 2001
From: Tuyen Nguyen <vutuyen2636@gmail.com>
Date: Thu, 17 Oct 2024 12:22:29 +0700
Subject: [PATCH] fix: balancesTreeCache to store one balances tree per epoch

---
 .../src/chain/balancesTreeCache.ts            | 27 ++++++++++++++-----
 1 file changed, 20 insertions(+), 7 deletions(-)

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<UintNumberType>[] = [];
+  private readonly unusedBalancesTrees: Map<Epoch, ListBasicTreeViewDU<UintNumberType>> = 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<UintNumberType> | 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;
   }
 }