Add pruner for finalized states

services/chainstorage/src/main/java/tech/pegasys/teku/services/chainstorage/StorageService.java

@@ -40,6 +40,7 @@
 import tech.pegasys.teku.storage.server.VersionedDatabaseFactory;
 import tech.pegasys.teku.storage.server.pruner.BlobSidecarPruner;
 import tech.pegasys.teku.storage.server.pruner.BlockPruner;
+import tech.pegasys.teku.storage.server.pruner.StatePruner;
 
 public class StorageService extends Service implements StorageServiceFacade {
   private final StorageConfiguration config;
@@ -49,6 +50,7 @@
   private volatile BatchingVoteUpdateChannel batchingVoteUpdateChannel;
   private volatile Optional<BlockPruner> blockPruner = Optional.empty();
   private volatile Optional<BlobSidecarPruner> blobsPruner = Optional.empty();
+  private volatile Optional<StatePruner> statePruner = Optional.empty();
   private final boolean depositSnapshotStorageEnabled;
   private final boolean blobSidecarsStorageCountersEnabled;
 
@@ -111,6 +113,20 @@
                             pruningTimingsLabelledGauge,
                             pruningActiveLabelledGauge));
               }
+              if (config.getDataStorageMode().storesFinalizedStates()
+                  && config.getRetainedEpochs() > -1) {
+                statePruner =
+                    Optional.of(
+                        new StatePruner(
+                            config.getSpec(),
+                            database,
+                            storagePrunerAsyncRunner,
+                            config.getBlockPruningInterval(),
+                            config.getRetainedEpochs(),
+                            "state",
+                            pruningTimingsLabelledGauge,
+                            pruningActiveLabelledGauge));
+              }
               if (config.getSpec().isMilestoneSupported(SpecMilestone.DENEB)) {
                 blobsPruner =
                     Optional.of(
@@ -170,6 +186,11 @@
             __ ->
                 blobsPruner
                     .map(BlobSidecarPruner::start)
+                    .orElseGet(() -> SafeFuture.completedFuture(null)))
+        .thenCompose(
+            __ ->
+                statePruner
+                    .map(StatePruner::start)
                     .orElseGet(() -> SafeFuture.completedFuture(null)));
   }
 

storage/src/main/java/tech/pegasys/teku/storage/server/Database.java

@@ -236,4 +236,6 @@ default Stream<SlotAndBlockRootAndBlobIndex> streamBlobSidecarKeys(final UInt64
    * @return actual last pruned slot
    */
   UInt64 pruneFinalizedBlocks(UInt64 lastSlotToPrune, int pruneLimit);
+
+  UInt64 pruneFinalizedStates(UInt64 lastSlotToPruneStateFor, long pruneLimit);
 }

storage/src/main/java/tech/pegasys/teku/storage/server/StorageConfiguration.java

@@ -37,6 +37,7 @@ public class StorageConfiguration {
   public static final boolean DEFAULT_STORE_NON_CANONICAL_BLOCKS_ENABLED = false;
   public static final int DEFAULT_STATE_REBUILD_TIMEOUT_SECONDS = 120;
   public static final long DEFAULT_STORAGE_FREQUENCY = 2048L;
+  public static final long DEFAULT_STORAGE_RETAINED_EPOCHS = 1024L;
   public static final int DEFAULT_MAX_KNOWN_NODE_CACHE_SIZE = 100_000;
   public static final Duration DEFAULT_BLOCK_PRUNING_INTERVAL = Duration.ofMinutes(15);
   public static final int DEFAULT_BLOCK_PRUNING_LIMIT = 5000;
@@ -58,6 +59,7 @@ public class StorageConfiguration {
   private final int blockPruningLimit;
   private final Duration blobsPruningInterval;
   private final int blobsPruningLimit;
+  private final long retainedEpochs;
 
   private final int stateRebuildTimeoutSeconds;
 
@@ -73,6 +75,7 @@ private StorageConfiguration(
       final Duration blobsPruningInterval,
       final int blobsPruningLimit,
       final int stateRebuildTimeoutSeconds,
+      final long retainedEpochs,
       final Spec spec) {
     this.eth1DepositContract = eth1DepositContract;
     this.dataStorageMode = dataStorageMode;
@@ -85,6 +88,7 @@ private StorageConfiguration(
     this.blobsPruningInterval = blobsPruningInterval;
     this.blobsPruningLimit = blobsPruningLimit;
     this.stateRebuildTimeoutSeconds = stateRebuildTimeoutSeconds;
+    this.retainedEpochs = retainedEpochs;
     this.spec = spec;
   }
 
@@ -136,6 +140,10 @@ public int getBlobsPruningLimit() {
     return blobsPruningLimit;
   }
 
+  public long getRetainedEpochs() {
+    return retainedEpochs;
+  }
+
   public Spec getSpec() {
     return spec;
   }
@@ -155,6 +163,7 @@ public static final class Builder {
     private Duration blobsPruningInterval = DEFAULT_BLOBS_PRUNING_INTERVAL;
     private int blobsPruningLimit = DEFAULT_BLOBS_PRUNING_LIMIT;
     private int stateRebuildTimeoutSeconds = DEFAULT_STATE_REBUILD_TIMEOUT_SECONDS;
+    private long retainedEpochs = DEFAULT_STORAGE_RETAINED_EPOCHS;
 
     private Builder() {}
 
@@ -247,6 +256,14 @@ public Builder blobsPruningLimit(final int blobsPruningLimit) {
       return this;
     }
 
+    public Builder retainedEpochs(final long retainedEpochs) {
+      if (retainedEpochs < 0) {
+        throw new InvalidConfigurationException("Invalid number of states to be retained");
+      }
+      this.retainedEpochs = retainedEpochs;
+      return this;
+    }
+
     public StorageConfiguration build() {
       determineDataStorageMode();
       return new StorageConfiguration(
@@ -261,6 +278,7 @@ public StorageConfiguration build() {
           blobsPruningInterval,
           blobsPruningLimit,
           stateRebuildTimeoutSeconds,
+          retainedEpochs,
           spec);
     }
 

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/KvStoreDatabase.java

@@ -37,6 +37,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.commons.lang3.tuple.Pair;
+import org.apache.commons.lang3.tuple.Triple;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.tuweni.bytes.Bytes;
@@ -432,6 +433,76 @@ private void deleteFinalizedBlocks(final List<Pair<UInt64, Bytes32>> blocksToPru
     }
   }
 
+  @Override
+  public UInt64 pruneFinalizedStates(final UInt64 lastSlotToPrune, final long pruneLimit) {
+    final Optional<UInt64> earliestFinalizedStateSlot = dao.getEarliestFinalizedStateSlot();
+    LOG.debug(
+        "Earliest finalized state stored is for slot {}",
+        () ->
+            earliestFinalizedStateSlot.isEmpty()
+                ? "EMPTY"
+                : earliestFinalizedStateSlot.get().toString());
+    if (earliestFinalizedStateSlot.isEmpty()) {
+      return lastSlotToPrune;
+    }
+    return pruneFinalizedStateForSlots(
+        earliestFinalizedStateSlot.get(), lastSlotToPrune, pruneLimit);
+  }
+
+  private UInt64 pruneFinalizedStateForSlots(
+      final UInt64 earliestFinalizedStateSlot,
+      final UInt64 lastSlotToPrune,
+      final long pruneLimit) {
+    final List<Triple<UInt64, Bytes32, Bytes32>> slotsToPruneStateFor;
+    LOG.debug("Pruning finalized states to slot {} (included)", lastSlotToPrune);
+    try (final Stream<SignedBeaconBlock> stream =
+        dao.streamFinalizedBlocks(earliestFinalizedStateSlot, lastSlotToPrune)) {
+      slotsToPruneStateFor =
+          // fetch slot for log purposes, stateroot and blockroot (keys used to store state)
+          stream
+              .limit(pruneLimit)
+              .map(block -> Triple.of(block.getSlot(), block.getRoot(), block.getStateRoot()))
+              .toList();
+    }
+    if (slotsToPruneStateFor.isEmpty()) {
+      LOG.debug("No finalized state to prune up to {} slot", lastSlotToPrune);
+      return lastSlotToPrune;
+    }
+
+    final UInt64 lastPrunedBlockSlot =
+        slotsToPruneStateFor.get(slotsToPruneStateFor.size() - 1).getLeft();
+    LOG.debug(
+        "Pruning {} finalized state, last block slot is {}",
+        slotsToPruneStateFor.size(),
+        lastPrunedBlockSlot);
+    deleteFinalizedStatesForSlot(slotsToPruneStateFor);
+
+    return slotsToPruneStateFor.size() < pruneLimit ? lastSlotToPrune : lastPrunedBlockSlot;
+  }
+
+  private void deleteFinalizedStatesForSlot(
+      final List<Triple<UInt64, Bytes32, Bytes32>> slotsToPruneStateFor) {
+    if (!slotsToPruneStateFor.isEmpty()) {
+      if (slotsToPruneStateFor.size() < 20) {
+        LOG.debug(
+            "Received blocks ({}) to delete",
+            () -> slotsToPruneStateFor.stream().map(Triple::getLeft).toList());
+      } else {
+        LOG.debug("Received {} finalized blocks to delete", slotsToPruneStateFor.size());
+      }
+
+      try (final FinalizedUpdater updater = finalizedUpdater()) {
+        slotsToPruneStateFor.forEach(
+            triple -> {
+              updater.deleteFinalizedState(triple.getLeft());
+              updater.deleteFinalizedStateRoot(triple.getRight());
+            });
+
+        updater.commit();
+      }
+    }
+  }
+
   protected void updateHotBlocks(
       final HotUpdater updater,
       final Map<Bytes32, BlockAndCheckpoints> addedBlocks,

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/dataaccess/CombinedKvStoreDao.java

@@ -275,6 +275,11 @@ public Optional<SignedBeaconBlock> getEarliestFinalizedBlock() {
     return db.getFirstEntry(schema.getColumnFinalizedBlocksBySlot()).map(ColumnEntry::getValue);
   }
 
+  @Override
+  public Optional<UInt64> getEarliestFinalizedStateSlot() {
+    return stateStorageLogic.getEarliestAvailableFinalizedStateSlot(db, schema);
+  }
+
   @Override
   public Optional<SignedBeaconBlock> getLatestFinalizedBlockAtSlot(final UInt64 slot) {
     return db.getFloorEntry(schema.getColumnFinalizedBlocksBySlot(), slot)
@@ -711,6 +716,11 @@ public void addFinalizedState(final Bytes32 blockRoot, final BeaconState state)
       stateStorageUpdater.addFinalizedState(db, transaction, schema, state);
     }
 
+    @Override
+    public void deleteFinalizedState(final UInt64 slot) {
+      stateStorageUpdater.deleteFinalizedState(transaction, schema, slot);
+    }
+
     @Override
     public void addReconstructedFinalizedState(final Bytes32 blockRoot, final BeaconState state) {
       stateStorageUpdater.addReconstructedFinalizedState(db, transaction, schema, state);
@@ -721,6 +731,11 @@ public void addFinalizedStateRoot(final Bytes32 stateRoot, final UInt64 slot) {
       transaction.put(schema.getColumnSlotsByFinalizedStateRoot(), stateRoot, slot);
     }
 
+    @Override
+    public void deleteFinalizedStateRoot(final Bytes32 stateRoot) {
+      transaction.delete(schema.getColumnSlotsByFinalizedStateRoot(), stateRoot);
+    }
+
     @Override
     public void setOptimisticTransitionBlockSlot(final Optional<UInt64> transitionBlockSlot) {
       if (transitionBlockSlot.isPresent()) {

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/dataaccess/KvStoreCombinedDao.java

@@ -67,6 +67,8 @@ public interface KvStoreCombinedDao extends AutoCloseable {
 
   Optional<SignedBeaconBlock> getEarliestFinalizedBlock();
 
+  Optional<UInt64> getEarliestFinalizedStateSlot();
+
   Optional<SignedBeaconBlock> getLatestFinalizedBlockAtSlot(UInt64 slot);
 
   List<SignedBeaconBlock> getNonCanonicalBlocksAtSlot(UInt64 slot);
@@ -234,10 +236,14 @@ interface FinalizedUpdater extends AutoCloseable {
 
     void addFinalizedState(final Bytes32 blockRoot, final BeaconState state);
 
+    void deleteFinalizedState(final UInt64 slot);
+
     void addReconstructedFinalizedState(final Bytes32 blockRoot, final BeaconState state);
 
     void addFinalizedStateRoot(final Bytes32 stateRoot, final UInt64 slot);
 
+    void deleteFinalizedStateRoot(final Bytes32 stateRoot);
+
     void setOptimisticTransitionBlockSlot(final Optional<UInt64> transitionBlockSlot);
 
     void addNonCanonicalRootAtSlot(final UInt64 slot, final Set<Bytes32> blockRoots);

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/dataaccess/KvStoreCombinedDaoAdapter.java

@@ -167,6 +167,11 @@ public Optional<SignedBeaconBlock> getEarliestFinalizedBlock() {
     return finalizedDao.getEarliestFinalizedBlock();
   }
 
+  @Override
+  public Optional<UInt64> getEarliestFinalizedStateSlot() {
+    return finalizedDao.getEarliestFinalizedStateSlot();
+  }
+
   @Override
   public Optional<SignedBeaconBlock> getLatestFinalizedBlockAtSlot(final UInt64 slot) {
     return finalizedDao.getLatestFinalizedBlockAtSlot(slot);
@@ -553,6 +558,11 @@ public void addFinalizedState(final Bytes32 blockRoot, final BeaconState state)
       finalizedUpdater.addFinalizedState(blockRoot, state);
     }
 
+    @Override
+    public void deleteFinalizedState(final UInt64 slot) {
+      finalizedUpdater.deleteFinalizedState(slot);
+    }
+
     @Override
     public void addReconstructedFinalizedState(final Bytes32 blockRoot, final BeaconState state) {
       finalizedUpdater.addReconstructedFinalizedState(blockRoot, state);
@@ -563,6 +573,11 @@ public void addFinalizedStateRoot(final Bytes32 stateRoot, final UInt64 slot) {
       finalizedUpdater.addFinalizedStateRoot(stateRoot, slot);
     }
 
+    @Override
+    public void deleteFinalizedStateRoot(final Bytes32 stateRoot) {
+      finalizedUpdater.deleteFinalizedStateRoot(stateRoot);
+    }
+
     @Override
     public void setOptimisticTransitionBlockSlot(final Optional<UInt64> transitionBlockSlot) {
       finalizedUpdater.setOptimisticTransitionBlockSlot(transitionBlockSlot);

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/dataaccess/V4FinalizedKvStoreDao.java

@@ -68,6 +68,10 @@ public Optional<UInt64> getEarliestFinalizedBlockSlot() {
     return db.getFirstEntry(schema.getColumnFinalizedBlocksBySlot()).map(ColumnEntry::getKey);
   }
 
+  public Optional<UInt64> getEarliestFinalizedStateSlot() {
+    return stateStorageLogic.getEarliestAvailableFinalizedStateSlot(db, schema);
+  }
+
   public Optional<SignedBeaconBlock> getEarliestFinalizedBlock() {
     return db.getFirstEntry(schema.getColumnFinalizedBlocksBySlot()).map(ColumnEntry::getValue);
   }
@@ -337,6 +341,11 @@ public void addFinalizedState(final Bytes32 blockRoot, final BeaconState state)
       stateStorageUpdater.addFinalizedState(db, transaction, schema, state);
     }
 
+    @Override
+    public void deleteFinalizedState(final UInt64 slot) {
+      transaction.delete(schema.getColumnFinalizedStatesBySlot(), slot);
+    }
+
     @Override
     public void addReconstructedFinalizedState(final Bytes32 blockRoot, final BeaconState state) {
       stateStorageUpdater.addReconstructedFinalizedState(db, transaction, schema, state);
@@ -347,6 +356,11 @@ public void addFinalizedStateRoot(final Bytes32 stateRoot, final UInt64 slot) {
       transaction.put(schema.getColumnSlotsByFinalizedStateRoot(), stateRoot, slot);
     }
 
+    @Override
+    public void deleteFinalizedStateRoot(final Bytes32 stateRoot) {
+      transaction.delete(schema.getColumnSlotsByFinalizedStateRoot(), stateRoot);
+    }
+
     @Override
     public void setOptimisticTransitionBlockSlot(final Optional<UInt64> transitionBlockSlot) {
       if (transitionBlockSlot.isPresent()) {

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/dataaccess/V4FinalizedStateSnapshotStorageLogic.java

@@ -39,6 +39,12 @@ public Optional<BeaconState> getLatestAvailableFinalizedState(
         .map(ColumnEntry::getValue);
   }
 
+  @Override
+  public Optional<UInt64> getEarliestAvailableFinalizedStateSlot(
+      final KvStoreAccessor db, final S schema) {
+    return db.getFirstEntry(schema.getColumnFinalizedStatesBySlot()).map(ColumnEntry::getKey);
+  }
+
   @Override
   public FinalizedStateUpdater<S> updater() {
     return new FinalizedStateSnapshotUpdater<>(stateStorageFrequency);

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/dataaccess/V4FinalizedStateStorageLogic.java

@@ -25,6 +25,8 @@ public interface V4FinalizedStateStorageLogic<S> {
   Optional<BeaconState> getLatestAvailableFinalizedState(
       KvStoreAccessor db, S schema, UInt64 maxSlot);
 
+  Optional<UInt64> getEarliestAvailableFinalizedStateSlot(KvStoreAccessor db, S schema);
+
   FinalizedStateUpdater<S> updater();
 
   @MustBeClosed

storage/src/main/java/tech/pegasys/teku/storage/server/kvstore/dataaccess/V4FinalizedStateTreeStorageLogic.java

@@ -79,6 +79,12 @@ public Optional<BeaconState> getLatestAvailableFinalizedState(
                         GIndexUtil.SELF_G_INDEX));
   }
 
+  @Override
+  public Optional<UInt64> getEarliestAvailableFinalizedStateSlot(
+      final KvStoreAccessor db, final SchemaCombinedTreeState dbSchema) {
+    return db.getFirstEntry(dbSchema.getColumnFinalizedStateRootsBySlot()).map(ColumnEntry::getKey);
+  }
+
   @Override
   public FinalizedStateUpdater<SchemaCombinedTreeState> updater() {
     return new StateTreeUpdater(

storage/src/main/java/tech/pegasys/teku/storage/server/noop/NoOpDatabase.java

@@ -229,6 +229,11 @@ public UInt64 pruneFinalizedBlocks(final UInt64 lastSlotToPrune, final int prune
     return lastSlotToPrune;
   }
 
+  @Override
+  public UInt64 pruneFinalizedStates(final UInt64 lastSlotToPruneStateFor, final long pruneLimit) {
+    return lastSlotToPruneStateFor;
+  }
+
   @Override
   public void addMinGenesisTimeBlock(final MinGenesisTimeBlockEvent event) {}