From 5fe87f8e443e86511b8dae1acd0894226d6ae3fa Mon Sep 17 00:00:00 2001
From: twoeths <tuyen@chainsafe.io>
Date: Tue, 2 Jul 2024 00:37:35 +0700
Subject: [PATCH] fix: prepareNextEpoch metric (#6924)

* fix: prepareNextEpoch metric

* Apply suggestions from code review

---------

Co-authored-by: Cayman <caymannava@gmail.com>
---
 .../beacon-node/src/chain/prepareNextSlot.ts  | 48 +++++++++----------
 .../state-transition/src/stateTransition.ts   |  1 +
 2 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/packages/beacon-node/src/chain/prepareNextSlot.ts b/packages/beacon-node/src/chain/prepareNextSlot.ts
index 3f730df3bf1d..48724ab25b0b 100644
--- a/packages/beacon-node/src/chain/prepareNextSlot.ts
+++ b/packages/beacon-node/src/chain/prepareNextSlot.ts
@@ -115,26 +115,6 @@ export class PrepareNextSlotScheduler {
         RegenCaller.precomputeEpoch
       );
 
-      // assuming there is no reorg, it caches the checkpoint state & helps avoid doing a full state transition in the next slot
-      //  + when gossip block comes, we need to validate and run state transition
-      //  + if next slot is a skipped slot, it'd help getting target checkpoint state faster to validate attestations
-      if (isEpochTransition) {
-        this.metrics?.precomputeNextEpochTransition.count.inc({result: "success"}, 1);
-        const previousHits = this.chain.regen.updatePreComputedCheckpoint(headRoot, nextEpoch);
-        if (previousHits === 0) {
-          this.metrics?.precomputeNextEpochTransition.waste.inc();
-        }
-        this.metrics?.precomputeNextEpochTransition.hits.set(previousHits ?? 0);
-        this.logger.verbose("Completed PrepareNextSlotScheduler epoch transition", {
-          nextEpoch,
-          headSlot,
-          prepareSlot,
-          previousHits,
-        });
-
-        precomputeEpochTransitionTimer?.();
-      }
-
       if (isExecutionStateType(prepareState)) {
         const proposerIndex = prepareState.epochCtx.getBeaconProposer(prepareSlot);
         const feeRecipient = this.chain.beaconProposerCache.get(proposerIndex);
@@ -198,7 +178,7 @@ export class PrepareNextSlotScheduler {
           });
         }
 
-        this.computeStateHashTreeRoot(updatedPrepareState);
+        this.computeStateHashTreeRoot(updatedPrepareState, isEpochTransition);
 
         // If emitPayloadAttributes is true emit a SSE payloadAttributes event
         if (this.chain.opts.emitPayloadAttributes === true) {
@@ -213,7 +193,27 @@ export class PrepareNextSlotScheduler {
           this.chain.emitter.emit(routes.events.EventType.payloadAttributes, {data, version: fork});
         }
       } else {
-        this.computeStateHashTreeRoot(prepareState);
+        this.computeStateHashTreeRoot(prepareState, isEpochTransition);
+      }
+
+      // assuming there is no reorg, it caches the checkpoint state & helps avoid doing a full state transition in the next slot
+      //  + when gossip block comes, we need to validate and run state transition
+      //  + if next slot is a skipped slot, it'd help getting target checkpoint state faster to validate attestations
+      if (isEpochTransition) {
+        this.metrics?.precomputeNextEpochTransition.count.inc({result: "success"}, 1);
+        const previousHits = this.chain.regen.updatePreComputedCheckpoint(headRoot, nextEpoch);
+        if (previousHits === 0) {
+          this.metrics?.precomputeNextEpochTransition.waste.inc();
+        }
+        this.metrics?.precomputeNextEpochTransition.hits.set(previousHits ?? 0);
+        this.logger.verbose("Completed PrepareNextSlotScheduler epoch transition", {
+          nextEpoch,
+          headSlot,
+          prepareSlot,
+          previousHits,
+        });
+
+        precomputeEpochTransitionTimer?.();
       }
     } catch (e) {
       if (!isErrorAborted(e) && !isQueueErrorAborted(e)) {
@@ -223,11 +223,11 @@ export class PrepareNextSlotScheduler {
     }
   };
 
-  computeStateHashTreeRoot(state: CachedBeaconStateAllForks): void {
+  computeStateHashTreeRoot(state: CachedBeaconStateAllForks, isEpochTransition: boolean): void {
     // cache HashObjects for faster hashTreeRoot() later, especially for computeNewStateRoot() if we need to produce a block at slot 0 of epoch
     // see https://github.com/ChainSafe/lodestar/issues/6194
     const hashTreeRootTimer = this.metrics?.stateHashTreeRootTime.startTimer({
-      source: StateHashTreeRootSource.prepareNextSlot,
+      source: isEpochTransition ? StateHashTreeRootSource.prepareNextEpoch : StateHashTreeRootSource.prepareNextSlot,
     });
     state.hashTreeRoot();
     hashTreeRootTimer?.();
diff --git a/packages/state-transition/src/stateTransition.ts b/packages/state-transition/src/stateTransition.ts
index 3ecd24ea9813..7602f4d9acc2 100644
--- a/packages/state-transition/src/stateTransition.ts
+++ b/packages/state-transition/src/stateTransition.ts
@@ -51,6 +51,7 @@ export enum StateHashTreeRootSource {
   stateTransition = "state_transition",
   blockTransition = "block_transition",
   prepareNextSlot = "prepare_next_slot",
+  prepareNextEpoch = "prepare_next_epoch",
   computeNewStateRoot = "compute_new_state_root",
 }