update custody group count

ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/datacolumns/DataColumnSidecarCustodyImpl.java

@@ -15,6 +15,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -153,15 +154,12 @@ public void onSlot(final UInt64 slot) {
     if (updateEpoch(spec.computeEpochAtSlot(slot))) {
       int groupCount = custodyGroupCountManager.getCustodyGroupCount();
       final int oldGroupCount = totalCustodyGroupCount.getAndSet(groupCount);
-      // TODO-fulu: ignoring the case when it's less, let's skip pruning in early version, to
-      // implement in future invalidating current custody as number of required groups have
-      // increased (https://github.com/Consensys/teku/issues/9468)
-      if (groupCount > oldGroupCount) {
+      if (groupCount != oldGroupCount) {
         LOG.debug("Custody group count changed from {} to {}", oldGroupCount, groupCount);
-        final UInt64 minCustodyPeriodSlot =
-            minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(currentSlot);
-        db.setFirstCustodyIncompleteSlot(minCustodyPeriodSlot).ifExceptionGetsHereRaiseABug();
       }
+      final UInt64 minCustodyPeriodSlot =
+          minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(currentSlot);
+      db.setFirstCustodyIncompleteSlot(minCustodyPeriodSlot).ifExceptionGetsHereRaiseABug();
     }
   }
 
@@ -171,6 +169,11 @@ public void onNewFinalizedCheckpoint(
     advanceFirstIncompleteSlot(checkpoint.getEpoch()).ifExceptionGetsHereRaiseABug();
   }
 
+  @VisibleForTesting
+  public int getTotalCustodyGroupCount() {
+    return totalCustodyGroupCount.get();
+  }
+
   private synchronized boolean updateEpoch(final UInt64 epoch) {
     if (!lastEpoch.equals(epoch)) {
       lastEpoch = epoch;
@@ -189,10 +192,6 @@ private SafeFuture<Void> advanceFirstIncompleteSlot(final UInt64 finalizedEpoch)
                 maybeFirstIncompleteOrLastComplete
                     .map(
                         firstIncompleteOrLastComplete -> {
-                          // TODO-fulu: if we don't have finalization, we will not advance it and
-                          // it's an issue
-                          //  non-finalized epochs could be still not synced with up-to-date custody
-                          // (https://github.com/Consensys/teku/issues/9469)
                           if (firstIncompleteOrLastComplete.slot().equals(firstNonFinalizedSlot)) {
                             LOG.debug(
                                 "Custody group count synced to {}", totalCustodyGroupCount.get());

ethereum/statetransition/src/test/java/tech/pegasys/teku/statetransition/datacolumns/DataColumnSidecarCustodyImplTest.java

@@ -19,7 +19,9 @@
 
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import tech.pegasys.teku.infrastructure.async.SafeFuture;
 import tech.pegasys.teku.infrastructure.unsigned.UInt64;
@@ -50,29 +52,28 @@ public class DataColumnSidecarCustodyImplTest {
   final int groupCount = config.getNumberOfCustodyGroups();
 
   private final DataStructureUtil dataStructureUtil = new DataStructureUtil(0, spec);
+  private final MinCustodyPeriodSlotCalculator minCustodyPeriodSlotCalculator =
+      MinCustodyPeriodSlotCalculator.createFromSpec(spec);
 
-  private DataColumnSidecar createSidecar(final BeaconBlock block, final int column) {
-    return dataStructureUtil.randomDataColumnSidecar(createSigned(block), UInt64.valueOf(column));
-  }
-
-  private SignedBeaconBlockHeader createSigned(final BeaconBlock block) {
-    return dataStructureUtil.signedBlock(block).asHeader();
-  }
+  private DataColumnSidecarCustodyImpl custody;
 
-  @Test
-  void sanityTest() throws Throwable {
-    DataColumnSidecarCustodyImpl custody =
+  @BeforeEach
+  public void setup() {
+    custody =
         new DataColumnSidecarCustodyImpl(
             spec,
             blockResolver,
             dbAccessor,
-            MinCustodyPeriodSlotCalculator.createFromSpec(spec),
+            minCustodyPeriodSlotCalculator,
             custodyGroupCountManager,
             groupCount);
     when(custodyGroupCountManager.getCustodyColumnIndices())
         .thenReturn(
             List.of(UInt64.valueOf(0), UInt64.valueOf(1), UInt64.valueOf(2), UInt64.valueOf(3)));
+  }
 
+  @Test
+  void sanityTest() throws Throwable {
     BeaconBlock block = blockResolver.addBlock(10, true);
     DataColumnSidecar sidecar0 = createSidecar(block, 0);
     DataColumnSidecar sidecar1 = createSidecar(block, 1);
@@ -94,4 +95,44 @@ void sanityTest() throws Throwable {
     assertThat(fRet2_1.get().get()).isEqualTo(sidecar0);
     assertThat(fRet2_2.get().get()).isEqualTo(sidecar0);
   }
+
+  @Test
+  public void onSlot_shouldUpdateCustodyGroupCountWhenDifferent() {
+    assertThat(custody.getTotalCustodyGroupCount()).isEqualTo(groupCount);
+    // Group count decreases
+    when(custodyGroupCountManager.getCustodyGroupCount()).thenReturn(groupCount - 2);
+    custody.onSlot(UInt64.ZERO);
+    assertThat(custody.getTotalCustodyGroupCount()).isEqualTo(groupCount - 2);
+    // Group count increases
+    when(custodyGroupCountManager.getCustodyGroupCount()).thenReturn(groupCount + 3);
+    custody.onSlot(UInt64.valueOf(config.getSlotsPerEpoch()).increment());
+    assertThat(custody.getTotalCustodyGroupCount()).isEqualTo(groupCount + 3);
+  }
+
+  @Test
+  public void onSlot_shouldUpdateFirstCustodyIncompleteSlot()
+      throws ExecutionException, InterruptedException {
+    custody.onSlot(UInt64.ZERO);
+    final UInt64 minCustodyPeriodSlot =
+        minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(UInt64.ZERO);
+    assertThat(db.getFirstCustodyIncompleteSlot().get()).hasValue(minCustodyPeriodSlot);
+    final UInt64 firstSlotAfterMinEpochsForDataColumnSidecarsRequests =
+        UInt64.valueOf(config.getMinEpochsForDataColumnSidecarsRequests())
+            .increment()
+            .times(config.getSlotsPerEpoch());
+    custody.onSlot(firstSlotAfterMinEpochsForDataColumnSidecarsRequests);
+    final UInt64 advancedMinCustodyPeriodSlot =
+        minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(
+            firstSlotAfterMinEpochsForDataColumnSidecarsRequests);
+    assertThat(advancedMinCustodyPeriodSlot).isGreaterThan(minCustodyPeriodSlot);
+    assertThat(db.getFirstCustodyIncompleteSlot().get()).hasValue(advancedMinCustodyPeriodSlot);
+  }
+
+  private DataColumnSidecar createSidecar(final BeaconBlock block, final int column) {
+    return dataStructureUtil.randomDataColumnSidecar(createSigned(block), UInt64.valueOf(column));
+  }
+
+  private SignedBeaconBlockHeader createSigned(final BeaconBlock block) {
+    return dataStructureUtil.signedBlock(block).asHeader();
+  }
 }