Simplify the interface for archiver strategy

ChainSafe · Nov 1, 2024 · c728791 · c728791
1 parent b3b1698
commit c728791
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 19 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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