Add pruner for finalized states

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

@@ -17,6 +17,8 @@
 import static tech.pegasys.teku.spec.config.Constants.STORAGE_QUERY_CHANNEL_PARALLELISM;
 
 import java.util.Optional;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import tech.pegasys.teku.ethereum.pow.api.Eth1EventsChannel;
 import tech.pegasys.teku.infrastructure.async.AsyncRunner;
 import tech.pegasys.teku.infrastructure.async.SafeFuture;
@@ -40,6 +42,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,8 +52,10 @@
   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;
+  private static final Logger LOG = LogManager.getLogger();
 
   public StorageService(
       final ServiceConfig serviceConfig,
@@ -61,6 +66,7 @@
     this.config = storageConfiguration;
     this.depositSnapshotStorageEnabled = depositSnapshotStorageEnabled;
     this.blobSidecarsStorageCountersEnabled = blobSidecarsStorageCountersEnabled;
+    LOG.info("Storage service created");
   }
 
   @Override
@@ -111,6 +117,25 @@
                             pruningTimingsLabelledGauge,
                             pruningActiveLabelledGauge));
               }
+              LOG.info(
+                  "Data storage mode: {}, Retained epochs {}",
+                  config.getDataStorageMode(),
+                  config.getRetainedEpochs());
+              if (config.getDataStorageMode().storesFinalizedStates()
+                  && config.getRetainedEpochs() > -1) {
+                statePruner =
+                    Optional.of(
+                        new StatePruner(
+                            config.getSpec(),
+                            database,
+                            storagePrunerAsyncRunner,
+                            config.getStatePruningInterval(),
+                            config.getRetainedEpochs(),
+                            config.getStatePruningLimit(),
+                            "state",
+                            pruningTimingsLabelledGauge,
+                            pruningActiveLabelledGauge));
+              }
               if (config.getSpec().isMilestoneSupported(SpecMilestone.DENEB)) {
                 blobsPruner =
                     Optional.of(
@@ -170,6 +195,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,7 @@ default Stream<SlotAndBlockRootAndBlobIndex> streamBlobSidecarKeys(final UInt64
    * @return actual last pruned slot
    */
   UInt64 pruneFinalizedBlocks(UInt64 lastSlotToPrune, int pruneLimit);
+
+  Optional<UInt64> pruneFinalizedStates(
+      Optional<UInt64> lastPrunedSlot, UInt64 lastSlotToPruneStateFor, long pruneLimit);
 }

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

@@ -38,6 +38,9 @@ public class StorageConfiguration {
   public static final Duration DEFAULT_BLOCK_PRUNING_INTERVAL = Duration.ofMinutes(15);
   public static final int DEFAULT_BLOCK_PRUNING_LIMIT = 5000;
   public static final Duration DEFAULT_BLOBS_PRUNING_INTERVAL = Duration.ofMinutes(1);
+  public static final Duration DEFAULT_STATE_PRUNING_INTERVAL = Duration.ofMinutes(10);
+  public static final long DEFAULT_STORAGE_RETAINED_EPOCHS = -1;
+  public static final int DEFAULT_STATE_PRUNING_LIMIT = 3;
 
   // 60/12 = 5 blocks per minute * 6 max blobs per block = 30 blobs per minute at maximum, 15 as
   // target. Let's configure 48 pruning per minute, so we have some room for catching up.
@@ -53,8 +56,11 @@ public class StorageConfiguration {
   private final int maxKnownNodeCacheSize;
   private final Duration blockPruningInterval;
   private final int blockPruningLimit;
+  private final Duration statePruningInterval;
   private final Duration blobsPruningInterval;
   private final int blobsPruningLimit;
+  private final long retainedEpochs;
+  private final int statePruningLimit;
 
   private final int stateRebuildTimeoutSeconds;
 
@@ -70,6 +76,9 @@ private StorageConfiguration(
       final Duration blobsPruningInterval,
       final int blobsPruningLimit,
       final int stateRebuildTimeoutSeconds,
+      final long retainedEpochs,
+      final Duration statePruningInterval,
+      final int statePruningLimit,
       final Spec spec) {
     this.eth1DepositContract = eth1DepositContract;
     this.dataStorageMode = dataStorageMode;
@@ -82,6 +91,9 @@ private StorageConfiguration(
     this.blobsPruningInterval = blobsPruningInterval;
     this.blobsPruningLimit = blobsPruningLimit;
     this.stateRebuildTimeoutSeconds = stateRebuildTimeoutSeconds;
+    this.retainedEpochs = retainedEpochs;
+    this.statePruningInterval = statePruningInterval;
+    this.statePruningLimit = statePruningLimit;
     this.spec = spec;
   }
 
@@ -133,6 +145,18 @@ public int getBlobsPruningLimit() {
     return blobsPruningLimit;
   }
 
+  public long getRetainedEpochs() {
+    return retainedEpochs;
+  }
+
+  public Duration getStatePruningInterval() {
+    return statePruningInterval;
+  }
+
+  public int getStatePruningLimit() {
+    return statePruningLimit;
+  }
+
   public Spec getSpec() {
     return spec;
   }
@@ -152,6 +176,9 @@ 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 Duration statePruningInterval = DEFAULT_STATE_PRUNING_INTERVAL;
+    private long retainedEpochs = DEFAULT_STORAGE_RETAINED_EPOCHS;
+    private int statePruningLimit = DEFAULT_STATE_PRUNING_LIMIT;
 
     private Builder() {}
 
@@ -244,6 +271,31 @@ public Builder blobsPruningLimit(final int blobsPruningLimit) {
       return this;
     }
 
+    public Builder retainedEpochs(final long retainedEpochs) {
+      if (retainedEpochs < -1) {
+        throw new InvalidConfigurationException("Invalid number of states to be retained");
+      }
+      this.retainedEpochs = retainedEpochs;
+      return this;
+    }
+
+    public Builder statePruningInterval(final Duration statePruningInterval) {
+      if (statePruningInterval.isNegative() || statePruningInterval.isZero()) {
+        throw new InvalidConfigurationException("Block pruning interval must be positive");
+      }
+      this.statePruningInterval = statePruningInterval;
+      return this;
+    }
+
+    public Builder statePruningLimit(final int statePruningLimit) {
+      if (statePruningLimit < 0) {
+        throw new InvalidConfigurationException(
+            String.format("Invalid statePruningLimit: %d", statePruningLimit));
+      }
+      this.statePruningLimit = statePruningLimit;
+      return this;
+    }
+
     public StorageConfiguration build() {
       determineDataStorageMode();
       return new StorageConfiguration(
@@ -258,6 +310,9 @@ public StorageConfiguration build() {
           blobsPruningInterval,
           blobsPruningLimit,
           stateRebuildTimeoutSeconds,
+          retainedEpochs,
+          statePruningInterval,
+          statePruningLimit,
           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,84 @@ private void deleteFinalizedBlocks(final List<Pair<UInt64, Bytes32>> blocksToPru
     }
   }
 
+  @Override
+  public Optional<UInt64> pruneFinalizedStates(
+      final Optional<UInt64> lastPrunedSlot, final UInt64 lastSlotToPrune, final long pruneLimit) {
+    final Optional<UInt64> earliestFinalizedStateSlot;
+
+    if (lastPrunedSlot.isEmpty()) {
+      earliestFinalizedStateSlot = dao.getEarliestFinalizedStateSlot();
+    } else {
+      earliestFinalizedStateSlot = lastPrunedSlot;
+    }
+
+    LOG.debug(
+        "Earliest finalized state stored is for slot {}",
+        () ->
+            earliestFinalizedStateSlot.isEmpty()
+                ? "EMPTY"
+                : earliestFinalizedStateSlot.get().toString());
+    return earliestFinalizedStateSlot
+        .map(uInt64 -> pruneFinalizedStateForSlots(uInt64, lastSlotToPrune, pruneLimit))
+        .or(() -> Optional.of(lastSlotToPrune));
+  }
+
+  private UInt64 pruneFinalizedStateForSlots(
+      final UInt64 earliestFinalizedStateSlot,
+      final UInt64 lastSlotToPrune,
+      final long pruneLimit) {
+    final List<Triple<UInt64, Bytes32, Bytes32>> slotsToPruneStateFor;
+
+    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 slots ({}) to delete finalized state for",
+            () -> 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();
+      } catch (Exception e) {
+        LOG.error("Failed to prune finalized states", e);
+      }
+    }
+  }
+
   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);