From c72879101dc5668e6830d619f36f5f17196b2df0 Mon Sep 17 00:00:00 2001
From: Nazar Hussain <nazarhussain@gmail.com>
Date: Fri, 1 Nov 2024 10:50:26 +0100
Subject: [PATCH] Simplify the interface for archiver strategy

---
 .../src/chain/archiver/archiver.ts            | 22 +++++++++++++------
 .../src/chain/archiver/interface.ts           |  3 +--
 .../strategies/diffStateArchiveStrategy.ts    | 10 ++++-----
 .../frequencyStateArchiveStrategy.ts          |  8 +++----
 4 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/packages/beacon-node/src/chain/archiver/archiver.ts b/packages/beacon-node/src/chain/archiver/archiver.ts
index a6011918ca71..7bd9a95e8862 100644
--- a/packages/beacon-node/src/chain/archiver/archiver.ts
+++ b/packages/beacon-node/src/chain/archiver/archiver.ts
@@ -1,4 +1,4 @@
-import {Logger} from "@lodestar/utils";
+import {fromHex, Logger} from "@lodestar/utils";
 import {CheckpointWithHex} from "@lodestar/fork-choice";
 import {IBeaconDb} from "../../db/index.js";
 import {JobItemQueue} from "../../util/queue/index.js";
@@ -9,6 +9,7 @@ import {FrequencyStateArchiveStrategy} from "./strategies/frequencyStateArchiveS
 import {archiveBlocks} from "./archiveBlocks.js";
 import {StateArchiveMode, ArchiverOpts, StateArchiveStrategy} from "./interface.js";
 import {DifferentialStateArchiveStrategy} from "./strategies/diffStateArchiveStrategy.js";
+import {computeEpochAtSlot} from "@lodestar/state-transition";
 
 export const DEFAULT_STATE_ARCHIVE_MODE = StateArchiveMode.Frequency;
 
@@ -87,16 +88,23 @@ export class Archiver {
   }
 
   private onCheckpoint(): void {
-    const headStateRoot = this.chain.forkChoice.getHead().stateRoot;
+    const {stateRoot, slot} = this.chain.forkChoice.getHead();
+
     this.chain.regen.pruneOnCheckpoint(
       this.chain.forkChoice.getFinalizedCheckpoint().epoch,
       this.chain.forkChoice.getJustifiedCheckpoint().epoch,
-      headStateRoot
+      stateRoot
     );
 
-    this.statesArchiverStrategy.onCheckpoint(headStateRoot, this.metrics).catch((err) => {
-      this.logger.error("Error during state archive", {stateArchiveMode: this.stateArchiveMode}, err);
-    });
+    this.statesArchiverStrategy
+      .onCheckpoint(
+        {root: fromHex(stateRoot), rootHex: stateRoot, epoch: computeEpochAtSlot(slot)},
+        false,
+        this.metrics
+      )
+      .catch((err) => {
+        this.logger.error("Error during state archive", {stateArchiveMode: this.stateArchiveMode}, err);
+      });
   }
 
   private processFinalizedCheckpoint = async (finalized: CheckpointWithHex): Promise<void> => {
@@ -115,7 +123,7 @@ export class Archiver {
       );
       this.prevFinalized = finalized;
 
-      await this.statesArchiverStrategy.onFinalizedCheckpoint(finalized, this.metrics);
+      await this.statesArchiverStrategy.onCheckpoint(finalized, true, this.metrics);
 
       this.chain.regen.pruneOnFinalized(finalizedEpoch);
 
diff --git a/packages/beacon-node/src/chain/archiver/interface.ts b/packages/beacon-node/src/chain/archiver/interface.ts
index 311d0e5465e1..0533eecdda78 100644
--- a/packages/beacon-node/src/chain/archiver/interface.ts
+++ b/packages/beacon-node/src/chain/archiver/interface.ts
@@ -39,7 +39,6 @@ export type FinalizedStats = {
 };
 
 export interface StateArchiveStrategy {
-  onCheckpoint(stateRoot: RootHex, metrics?: Metrics | null): Promise<void>;
-  onFinalizedCheckpoint(finalized: CheckpointWithHex, metrics?: Metrics | null): Promise<void>;
+  onCheckpoint(checkpoint: CheckpointWithHex, finalized: boolean, metrics?: Metrics | null): Promise<void>;
   maybeArchiveState(finalized: CheckpointWithHex, metrics?: Metrics | null): Promise<void>;
 }
diff --git a/packages/beacon-node/src/chain/archiver/strategies/diffStateArchiveStrategy.ts b/packages/beacon-node/src/chain/archiver/strategies/diffStateArchiveStrategy.ts
index 7c1057ccc05b..9b8ac04c712f 100644
--- a/packages/beacon-node/src/chain/archiver/strategies/diffStateArchiveStrategy.ts
+++ b/packages/beacon-node/src/chain/archiver/strategies/diffStateArchiveStrategy.ts
@@ -16,12 +16,10 @@ export class DifferentialStateArchiveStrategy implements StateArchiveStrategy {
     }
   ) {}
 
-  onCheckpoint(_stateRoot: RootHex, _metrics?: Metrics | null): Promise<void> {
-    throw new Error("Method not implemented.");
-  }
-
-  async onFinalizedCheckpoint(finalized: CheckpointWithHex, _metrics?: Metrics | null): Promise<void> {
-    await this.maybeArchiveState(finalized);
+  async onCheckpoint(checkpoint: CheckpointWithHex, finalized: boolean, _metrics?: Metrics | null): Promise<void> {
+    if (finalized) {
+      await this.maybeArchiveState(checkpoint);
+    }
   }
 
   async maybeArchiveState(finalized: CheckpointWithHex): Promise<void> {
diff --git a/packages/beacon-node/src/chain/archiver/strategies/frequencyStateArchiveStrategy.ts b/packages/beacon-node/src/chain/archiver/strategies/frequencyStateArchiveStrategy.ts
index f8ea80b0260c..567c4c2be39a 100644
--- a/packages/beacon-node/src/chain/archiver/strategies/frequencyStateArchiveStrategy.ts
+++ b/packages/beacon-node/src/chain/archiver/strategies/frequencyStateArchiveStrategy.ts
@@ -28,12 +28,12 @@ export class FrequencyStateArchiveStrategy implements StateArchiveStrategy {
     protected readonly opts: StatesArchiverOpts
   ) {}
 
-  async onFinalizedCheckpoint(finalized: CheckpointWithHex, metrics?: Metrics | null): Promise<void> {
-    await this.maybeArchiveState(finalized, metrics);
+  async onCheckpoint(checkpoint: CheckpointWithHex, finalized: boolean, metrics?: Metrics | null): Promise<void> {
+    if (finalized) {
+      await this.maybeArchiveState(checkpoint, metrics);
+    }
   }
 
-  async onCheckpoint(_stateRoot: RootHex, _metrics?: Metrics | null): Promise<void> {}
-
   /**
    * Persist states every some epochs to
    * - Minimize disk space, storing the least states possible