From f94ff92dbe054cc6a1f3e20d41a1682dfd1fa2dc Mon Sep 17 00:00:00 2001
From: Tuyen Nguyen <vutuyen2636@gmail.com>
Date: Thu, 17 Oct 2024 13:16:29 +0700
Subject: [PATCH] chore: track source of cached balances tree

---
 packages/beacon-node/src/chain/balancesTreeCache.ts  | 9 ++++++++-
 packages/beacon-node/src/chain/blocks/importBlock.ts | 3 ++-
 packages/beacon-node/src/chain/chain.ts              | 4 ++--
 packages/beacon-node/src/metrics/metrics/lodestar.ts | 6 ++++++
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/packages/beacon-node/src/chain/balancesTreeCache.ts b/packages/beacon-node/src/chain/balancesTreeCache.ts
index 6658ce66cf6..2c70264b15d 100644
--- a/packages/beacon-node/src/chain/balancesTreeCache.ts
+++ b/packages/beacon-node/src/chain/balancesTreeCache.ts
@@ -5,6 +5,11 @@ import {Epoch} from "@lodestar/types";
 
 const MAX_ITEMS = 2;
 
+export enum BalancesTreeSource {
+  PRUNE_ON_FINALIZED = "pruned_on_finalized",
+  IMPORT_BLOCK = "import_block",
+};
+
 /**
  * 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
@@ -20,7 +25,7 @@ export class BalancesTreeCache implements IBalancesTreeCache {
     }
   }
 
-  processUnusedState(state: CachedBeaconStateAllForks | undefined): void {
+  processUnusedState(state: CachedBeaconStateAllForks | undefined, source: BalancesTreeSource): void {
     if (state === undefined) {
       return;
     }
@@ -29,6 +34,8 @@ export class BalancesTreeCache implements IBalancesTreeCache {
       return;
     }
 
+    this.metrics?.balancesTreeCache.total.inc({source});
+
     this.unusedBalancesTrees.set(stateEpoch, state.balances);
     while (this.unusedBalancesTrees.size > MAX_ITEMS) {
       const firstEpoch = Array.from(this.unusedBalancesTrees.keys())[0];
diff --git a/packages/beacon-node/src/chain/blocks/importBlock.ts b/packages/beacon-node/src/chain/blocks/importBlock.ts
index 1c9f4e12a49..324c837b1ed 100644
--- a/packages/beacon-node/src/chain/blocks/importBlock.ts
+++ b/packages/beacon-node/src/chain/blocks/importBlock.ts
@@ -23,6 +23,7 @@ import {ForkchoiceCaller} from "../forkChoice/index.js";
 import {FullyVerifiedBlock, ImportBlockOpts, AttestationImportOpt, BlockInputType} from "./types.js";
 import {getCheckpointFromState} from "./utils/checkpoint.js";
 import {writeBlockInputToDb} from "./writeBlockInputToDb.js";
+import {BalancesTreeSource} from "../balancesTreeCache.js";
 
 /**
  * Fork-choice allows to import attestations from current (0) or past (1) epoch.
@@ -104,7 +105,7 @@ export async function importBlock(
       if (isCurrentSlot && prunedStates) {
         for (const states of prunedStates.values()) {
           // cp states on the same epoch shares the same balances seed tree so only need one of them
-          this.balancesTreeCache.processUnusedState(states[0]);
+          this.balancesTreeCache.processUnusedState(states[0], BalancesTreeSource.IMPORT_BLOCK);
         }
       }
     })
diff --git a/packages/beacon-node/src/chain/chain.ts b/packages/beacon-node/src/chain/chain.ts
index 87ed50b786d..81e0affd661 100644
--- a/packages/beacon-node/src/chain/chain.ts
+++ b/packages/beacon-node/src/chain/chain.ts
@@ -101,7 +101,7 @@ import {DbCPStateDatastore} from "./stateCache/datastore/db.js";
 import {FileCPStateDatastore} from "./stateCache/datastore/file.js";
 import {SyncCommitteeRewards, computeSyncCommitteeRewards} from "./rewards/syncCommitteeRewards.js";
 import {AttestationsRewards, computeAttestationsRewards} from "./rewards/attestationsRewards.js";
-import {BalancesTreeCache} from "./balancesTreeCache.js";
+import {BalancesTreeCache, BalancesTreeSource} from "./balancesTreeCache.js";
 
 /**
  * Arbitrary constants, blobs and payloads should be consumed immediately in the same slot
@@ -889,7 +889,7 @@ export class BeaconChain implements IBeaconChain {
     if (this.clock.currentEpoch - finalizedEpoch <= SAFE_FINALIZED_EPOCH_TO_CURRENT_EPOCH_DIFF && prunedStates) {
       // cp states on the same epoch shares the same balances seed tree so only need one of them
       for (const states of prunedStates.values()) {
-        this.balancesTreeCache.processUnusedState(states[0]);
+        this.balancesTreeCache.processUnusedState(states[0], BalancesTreeSource.PRUNE_ON_FINALIZED);
       }
     }
   }
diff --git a/packages/beacon-node/src/metrics/metrics/lodestar.ts b/packages/beacon-node/src/metrics/metrics/lodestar.ts
index be9ffd79ef1..2664c33b8f4 100644
--- a/packages/beacon-node/src/metrics/metrics/lodestar.ts
+++ b/packages/beacon-node/src/metrics/metrics/lodestar.ts
@@ -19,6 +19,7 @@ import {RegistryMetricCreator} from "../utils/registryMetricCreator.js";
 import {OpSource} from "../validatorMonitor.js";
 import {CacheItemType} from "../../chain/stateCache/types.js";
 import {AllocSource} from "../../util/bufferPool.js";
+import {BalancesTreeSource} from "../../chain/balancesTreeCache.js";
 
 export type LodestarMetrics = ReturnType<typeof createLodestarMetrics>;
 
@@ -1337,6 +1338,11 @@ export function createLodestarMetrics(
         name: "lodestar_balances_tree_cache_size",
         help: "Balances tree cache size",
       }),
+      total: register.gauge<{source: BalancesTreeSource}>({
+        name: "lodestar_balances_tree_cache_total",
+        help: "Total number of balances tree cache",
+        labelNames: ["source"],
+      }),
       hit: register.gauge({
         name: "lodestar_balances_tree_cache_hit_total",
         help: "Total number of balances tree cache hits",