Add archiveSidecarsProofs to Database: persist KZG proofs and drop extension columns

storage/src/integration-test/java/tech/pegasys/teku/storage/server/kvstore/DatabaseTest.java

@@ -46,6 +46,7 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.tuweni.bytes.Bytes32;
@@ -55,6 +56,7 @@
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.parallel.Execution;
 import org.junit.jupiter.api.parallel.ExecutionMode;
+import tech.pegasys.infrastructure.logging.LogCaptor;
 import tech.pegasys.teku.bls.BLSKeyGenerator;
 import tech.pegasys.teku.bls.BLSKeyPair;
 import tech.pegasys.teku.dataproviders.lookup.BlockProvider;
@@ -3071,6 +3073,142 @@ public void dataColumnSidecarsProofs_lifecycle(final DatabaseContext context) th
     assertThat(database.getDataColumnSidecarsProofs(header.getMessage().getSlot())).isEmpty();
   }
 
+  @TestTemplate
+  public void archiveSidecarsProofs_dropsExtensionColumnsAndRetainsProofs(
+      final DatabaseContext context) throws IOException {
+    setupWithSpec(TestSpecFactory.createMinimalFulu());
+    initialize(context);
+
+    final int numberOfColumns = spec.getNumberOfDataColumns().orElseThrow();
+    final int halfColumns = numberOfColumns / 2;
+    final UInt64 slot = UInt64.valueOf(3);
+
+    final List<List<KZGProof>> expectedExtensionProofs = storeFullColumnSet(slot, numberOfColumns);
+
+    database.archiveSidecarsProofs(ZERO, slot);
+
+    // only the reconstructable first half of the columns is retained
+    assertThat(getStoredColumnIndices(slot))
+        .containsExactlyElementsOf(
+            Stream.iterate(ZERO, UInt64::increment).limit(halfColumns).toList());
+    // the proofs of the dropped extension columns are archived for reconstruction
+    assertThat(database.getDataColumnSidecarsProofs(slot)).contains(expectedExtensionProofs);
+    assertThat(database.getLastDataColumnSidecarsProofsSlot()).contains(slot);
+  }
+
+  @TestTemplate
+  public void archiveSidecarsProofs_skipsSlotWithIncompleteExtensionColumns(
+      final DatabaseContext context) throws IOException {
+    setupWithSpec(TestSpecFactory.createMinimalFulu());
+    initialize(context);
+
+    final int numberOfColumns = spec.getNumberOfDataColumns().orElseThrow();
+    final UInt64 slot = UInt64.valueOf(3);
+
+    final SignedBeaconBlockHeader header = dataStructureUtil.randomSignedBeaconBlockHeader(slot);
+    final SszList<SszKZGCommitment> kzgCommitments = randomFuluKzgCommitments();
+    // store every column except the last extension one, so the extension half is incomplete
+    Stream.iterate(ZERO, UInt64::increment)
+        .limit(numberOfColumns - 1L)
+        .map(index -> dataStructureUtil.randomDataColumnSidecar(header, kzgCommitments, index))
+        .forEach(database::addSidecar);
+
+    database.archiveSidecarsProofs(ZERO, slot);
+
+    // nothing archived and no column dropped: reconstruction from a partial half is impossible
+    assertThat(database.getDataColumnSidecarsProofs(slot)).isEmpty();
+    assertThat(database.getLastDataColumnSidecarsProofsSlot()).isEmpty();
+    assertThat(getStoredColumnIndices(slot)).hasSize(numberOfColumns - 1);
+  }
+
+  @TestTemplate
+  public void archiveSidecarsProofs_skipsSlotWithMissingFirstHalfColumns(
+      final DatabaseContext context) throws IOException {
+    setupWithSpec(TestSpecFactory.createMinimalFulu());
+    initialize(context);
+
+    final int numberOfColumns = spec.getNumberOfDataColumns().orElseThrow();
+    final UInt64 slot = UInt64.valueOf(3);
+
+    final SignedBeaconBlockHeader header = dataStructureUtil.randomSignedBeaconBlockHeader(slot);
+    final SszList<SszKZGCommitment> kzgCommitments = randomFuluKzgCommitments();
+    // store all columns except index 0 (one first-half gap) — a single gap must prevent archiving
+    Stream.iterate(UInt64.ONE, UInt64::increment)
+        .limit(numberOfColumns - 1L)
+        .map(index -> dataStructureUtil.randomDataColumnSidecar(header, kzgCommitments, index))
+        .forEach(database::addSidecar);
+
+    database.archiveSidecarsProofs(ZERO, slot);
+
+    // archiving must be skipped: a missing first-half column makes extension data irrecoverable
+    assertThat(database.getDataColumnSidecarsProofs(slot)).isEmpty();
+    assertThat(database.getLastDataColumnSidecarsProofsSlot()).isEmpty();
+    assertThat(getStoredColumnIndices(slot)).hasSize(numberOfColumns - 1);
+  }
+
+  @TestTemplate
+  public void pruneAllSidecars_alsoRemovesArchivedProofs(final DatabaseContext context)
+      throws IOException {
+    setupWithSpec(TestSpecFactory.createMinimalFulu());
+    initialize(context);
+
+    final int numberOfColumns = spec.getNumberOfDataColumns().orElseThrow();
+    final UInt64 slot = UInt64.valueOf(3);
+
+    try (final LogCaptor logCaptor = LogCaptor.forClass(KvStoreDatabase.class, Level.DEBUG)) {
+      // archive the extension columns down to proofs: slot now holds first-half sidecars + proofs
+      storeFullColumnSet(slot, numberOfColumns);
+      database.archiveSidecarsProofs(ZERO, slot);
+      assertThat(database.getDataColumnSidecarsProofs(slot)).isPresent();
+      assertThat(database.getLastDataColumnSidecarsProofsSlot()).contains(slot);
+
+      // pruning the slot must drop the retained proofs alongside the sidecars
+      database.pruneAllSidecars(slot, 10);
+
+      assertThat(getStoredColumnIndices(slot)).isEmpty();
+      assertThat(database.getDataColumnSidecarsProofs(slot)).isEmpty();
+      assertThat(database.getLastDataColumnSidecarsProofsSlot()).isEmpty();
+
+      // both the archiving and the prune-time removal are observable in the logs at debug level
+      assertThat(logCaptor.getDebugLogs())
+          .anyMatch(log -> log.contains("Archiving data column sidecars to proofs"))
+          .anyMatch(log -> log.contains("Removing archived data column sidecar proofs"));
+    }
+  }
+
+  private List<List<KZGProof>> storeFullColumnSet(final UInt64 slot, final int numberOfColumns) {
+    final SignedBeaconBlockHeader header = dataStructureUtil.randomSignedBeaconBlockHeader(slot);
+    final SszList<SszKZGCommitment> kzgCommitments = randomFuluKzgCommitments();
+    final List<DataColumnSidecar> sidecars =
+        Stream.iterate(ZERO, UInt64::increment)
+            .limit(numberOfColumns)
+            .map(index -> dataStructureUtil.randomDataColumnSidecar(header, kzgCommitments, index))
+            .toList();
+    sidecars.forEach(database::addSidecar);
+    return sidecars.stream()
+        .skip(numberOfColumns / 2)
+        .map(sidecar -> sidecar.getKzgProofs().stream().map(SszKZGProof::getKZGProof).toList())
+        .toList();
+  }
+
+  private SszList<SszKZGCommitment> randomFuluKzgCommitments() {
+    return SchemaDefinitionsFulu.required(
+            spec.forMilestone(SpecMilestone.FULU).getSchemaDefinitions())
+        .getDataColumnSidecarSchema()
+        .getKzgCommitmentsSchema()
+        .createFromElements(
+            dataStructureUtil.randomKZGCommitments(14).stream()
+                .map(SszKZGCommitment::new)
+                .toList());
+  }
+
+  private List<UInt64> getStoredColumnIndices(final UInt64 slot) {
+    try (final Stream<DataColumnSlotAndIdentifier> identifiers =
+        database.streamDataColumnIdentifiers(slot, slot)) {
+      return identifiers.map(DataColumnSlotAndIdentifier::columnIndex).sorted().toList();
+    }
+  }
+
   private List<Map.Entry<Bytes32, UInt64>> getFinalizedStateRootsList() {
     try (final Stream<Map.Entry<Bytes32, UInt64>> roots = database.getFinalizedStateRoots()) {
       return roots.map(entry -> Map.entry(entry.getKey(), entry.getValue())).collect(toList());

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

@@ -331,4 +331,17 @@ default Stream<DataColumnSlotAndIdentifier> streamNonCanonicalDataColumnIdentifi
   // Triggers a full, blocking compaction of the underlying storage to physically reclaim the disk
   // space left behind by pruning. Expensive and I/O-heavy; intended for offline/CLI use.
   void compactStorage();
+
+  /**
+   * Archives the reconstructable extension data column sidecars (column indices >=
+   * NUMBER_OF_COLUMNS / 2) in [startSlot, tillSlotInclusive]: for each fully populated slot it
+   * persists their KZG proofs and drops the sidecars themselves, retaining only enough data to
+   * reconstruct them on demand.
+   *
+   * <p><b>Callers must submit small ranges.</b> There is no internal limit; the entire range is
+   * scanned in a single pass. Use {@link
+   * tech.pegasys.teku.storage.server.pruner.DataColumnSidecarPruner} which breaks the work into
+   * fixed-size chunks.
+   */
+  void archiveSidecarsProofs(UInt64 startSlot, UInt64 tillSlotInclusive);
 }

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

@@ -69,6 +69,7 @@
 import tech.pegasys.teku.spec.datastructures.state.AnchorPoint;
 import tech.pegasys.teku.spec.datastructures.state.Checkpoint;
 import tech.pegasys.teku.spec.datastructures.state.beaconstate.BeaconState;
+import tech.pegasys.teku.spec.datastructures.type.SszKZGProof;
 import tech.pegasys.teku.spec.datastructures.util.DataColumnSlotAndIdentifier;
 import tech.pegasys.teku.spec.datastructures.util.SlotAndBlockRootAndBlobIndex;
 import tech.pegasys.teku.storage.api.GloasForkChoiceRebuildData;
@@ -1389,6 +1390,97 @@ public void pruneAllSidecars(final UInt64 tillSlotInclusive, final int pruneLimi
         "Data column sidecars pruning completed in {} ms", System.currentTimeMillis() - startTime);
   }
 
+  @Override
+  public void archiveSidecarsProofs(final UInt64 startSlot, final UInt64 tillSlotInclusive) {
+    // Extension columns (indices >= NUMBER_OF_COLUMNS / 2) are the reconstructable half whose
+    // proofs we retain while dropping the columns themselves.
+    final int halfColumns =
+        spec.getNumberOfDataColumns()
+                .orElseThrow(
+                    () ->
+                        new IllegalStateException(
+                            "Cannot archive data column sidecar proofs before the Fulu milestone"))
+            / 2;
+    try (final Stream<DataColumnSlotAndIdentifier> dataColumnSidecars =
+        streamDataColumnIdentifiers(startSlot, tillSlotInclusive)) {
+
+      int archivedSlots = 0;
+
+      final Map<UInt64, List<DataColumnSlotAndIdentifier>> archiveMap = new HashMap<>();
+
+      dataColumnSidecars.forEach(
+          item -> archiveMap.computeIfAbsent(item.slot(), k -> new ArrayList<>()).add(item));
+
+      final List<UInt64> slots = archiveMap.keySet().stream().sorted().toList();
+
+      if (!slots.isEmpty()) {
+        LOG.debug(
+            "Archiving data column sidecars to proofs from slots {} to {}",
+            slots.getFirst(),
+            slots.getLast());
+        try (final FinalizedUpdater updater = finalizedUpdater()) {
+          for (final UInt64 slot : slots) {
+            final List<DataColumnSlotAndIdentifier> extensionKeys =
+                archiveMap.get(slot).stream()
+                    .filter(id -> id.columnIndex().isGreaterThanOrEqualTo(halfColumns))
+                    .sorted()
+                    .toList();
+            final long firstHalfCount =
+                archiveMap.get(slot).stream()
+                    .filter(id -> id.columnIndex().isLessThan(halfColumns))
+                    .count();
+
+            // Skip slots where either half is incomplete: first-half columns are required for
+            // reconstruction, and archiving extension columns without them would make the archived
+            // data irrecoverable.
+            if (extensionKeys.size() != halfColumns || firstHalfCount != halfColumns) {
+              LOG.trace(
+                  "Skipping archival for slot {}: have {}/{} extension and {}/{} first-half columns",
+                  slot,
+                  extensionKeys.size(),
+                  halfColumns,
+                  firstHalfCount,
+                  halfColumns);
+              continue;
+            }
+
+            final List<DataColumnSidecar> sidecars = new ArrayList<>();
+
+            for (final DataColumnSlotAndIdentifier key : extensionKeys) {
+              final Optional<Bytes> sidecar = dao.getSidecar(key);
+              if (sidecar.isEmpty()) {
+                break;
+              }
+              sidecars.add(spec.deserializeSidecar(sidecar.get(), key.slot()));
+            }
+
+            if (sidecars.size() == halfColumns) {
+              final List<List<KZGProof>> proofs =
+                  sidecars.stream()
+                      .map(
+                          sidecar ->
+                              sidecar.getKzgProofs().stream()
+                                  .map(SszKZGProof::getKZGProof)
+                                  .toList())
+                      .toList();
+              updater.addDataColumnSidecarsProofs(slot, proofs);
+              for (final DataColumnSlotAndIdentifier key : extensionKeys) {
+                updater.removeSidecar(key);
+              }
+              ++archivedSlots;
+              LOG.trace(
+                  "Pruned {} extension data column sidecars, keeping their proofs, at slot {}",
+                  extensionKeys.size(),
+                  slot);
+            }
+          }
+          updater.commit();
+        }
+        LOG.debug("Archived data column sidecars to proofs across {} slots", archivedSlots);
+      }
+    }
+  }
+
   /**
    * Prunes data column sidecars oldest-first: each run removes the oldest (up to) {@code
    * pruneSlotLimit} distinct populated slots at or before the cutoff, in a single committed
@@ -1448,6 +1540,14 @@ void pruneDataColumnSidecars(
       // supported slot.
       if (sidecarType == DataColumnSidecarType.CANONICAL) {
         updater.setLastDataColumnSidecarPrunedSlot(toPrune.keys().getLast().slot());
+        final List<UInt64> prunedSlots =
+            toPrune.keys().stream().map(DataColumnSlotAndIdentifier::slot).distinct().toList();
+        LOG.debug(
+            "Removing archived data column sidecar proofs for {} pruned canonical slots ({}..{})",
+            prunedSlots.size(),
+            prunedSlots.getFirst(),
+            prunedSlots.getLast());
+        prunedSlots.forEach(updater::removeDataColumnSidecarsProofs);
       }
       updater.commit();
     }

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

@@ -448,6 +448,9 @@ public void pruneAllSidecars(final UInt64 tillSlotInclusive, final int pruneLimi
   @Override
   public void compactStorage() {}
 
+  @Override
+  public void archiveSidecarsProofs(final UInt64 startSlot, final UInt64 tillSlotInclusive) {}
+
   @Override
   public void close() {}
 }

storage/src/property-test/java/tech/pegasys/teku/storage/server/kvstore/serialization/DataColumnSidecarsProofsSerializerTest.java

@@ -14,6 +14,7 @@
 package tech.pegasys.teku.storage.server.kvstore.serialization;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import java.util.List;
 import java.util.stream.Stream;
@@ -73,4 +74,25 @@ final void testRoundTrip() {
     final byte[] serialize = serializer.serialize(expectedProofs);
     assertThat(serializer.deserialize(serialize)).isEqualTo(expectedProofs);
   }
+
+  @Test
+  final void testSingleColumnSingleProofRoundTrip() {
+    final List<List<KZGProof>> proofs = List.of(List.of(dataStructureUtil.randomKZGProof()));
+    final byte[] serialize = serializer.serialize(proofs);
+    assertThat(serializer.deserialize(serialize)).isEqualTo(proofs);
+  }
+
+  @Test
+  final void testRaggedColumnsFailToDeserialize() {
+    // columns of unequal length leave a trailing partial column that cannot be regrouped by the
+    // uniform blob size written into the header
+    final List<List<KZGProof>> ragged =
+        List.of(
+            List.of(dataStructureUtil.randomKZGProof(), dataStructureUtil.randomKZGProof()),
+            List.of(dataStructureUtil.randomKZGProof()));
+    final byte[] serialize = serializer.serialize(ragged);
+    assertThatThrownBy(() -> serializer.deserialize(serialize))
+        .isInstanceOf(RuntimeException.class)
+        .hasMessageContaining("Unexpected proofs found");
+  }
 }

storage/src/test/java/tech/pegasys/teku/storage/client/CombinedChainDataClientTest.java

@@ -31,6 +31,7 @@
 import tech.pegasys.teku.infrastructure.async.SafeFuture;
 import tech.pegasys.teku.infrastructure.async.SafeFutureAssert;
 import tech.pegasys.teku.infrastructure.unsigned.UInt64;
+import tech.pegasys.teku.kzg.KZGProof;
 import tech.pegasys.teku.spec.Spec;
 import tech.pegasys.teku.spec.TestSpecFactory;
 import tech.pegasys.teku.spec.datastructures.blobs.DataColumnSidecar;
@@ -252,6 +253,29 @@ void getsDataColumnSidecarsBySlotAndBlockRoot() {
     verify(historicalChainData, never()).getSidecar(otherBlockIdentifier);
   }
 
+  @Test
+  void getDataColumnSidecarProofs_returnsStoredProofs() {
+    final UInt64 slot = UInt64.valueOf(42);
+    final List<List<KZGProof>> proofs =
+        List.of(
+            List.of(dataStructureUtil.randomKZGProof(), dataStructureUtil.randomKZGProof()),
+            List.of(dataStructureUtil.randomKZGProof()));
+    when(historicalChainData.getDataColumnSidecarsProofs(slot))
+        .thenReturn(SafeFuture.completedFuture(Optional.of(proofs)));
+
+    assertThat(SafeFutureAssert.safeJoin(client.getDataColumnSidecarProofs(slot)))
+        .isEqualTo(proofs);
+  }
+
+  @Test
+  void getDataColumnSidecarProofs_returnsEmptyListWhenAbsent() {
+    final UInt64 slot = UInt64.valueOf(42);
+    when(historicalChainData.getDataColumnSidecarsProofs(slot))
+        .thenReturn(SafeFuture.completedFuture(Optional.empty()));
+
+    assertThat(SafeFutureAssert.safeJoin(client.getDataColumnSidecarProofs(slot))).isEmpty();
+  }
+
   @Test
   void getBestFinalizedState_fetchesFinalizedState()
       throws ExecutionException, InterruptedException {