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;
@@ -150,25 +151,45 @@ public SafeFuture<Boolean> hasCustodyDataColumnSidecar(
   @Override
   public void onSlot(final UInt64 slot) {
     currentSlot = 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) {
-        LOG.debug("Custody group count changed from {} to {}", oldGroupCount, groupCount);
-        final UInt64 minCustodyPeriodSlot =
-            minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(currentSlot);
-        db.setFirstCustodyIncompleteSlot(minCustodyPeriodSlot).ifExceptionGetsHereRaiseABug();
-      }
+    if (!updateEpoch(spec.computeEpochAtSlot(slot))) {
+      return;
+    }
+    final int newCustodyGroupCount = custodyGroupCountManager.getCustodyGroupCount();
+    final int oldCustodyGroupCount = totalCustodyGroupCount.getAndSet(newCustodyGroupCount);
+    if (newCustodyGroupCount == oldCustodyGroupCount) {
+      return;
+    }
+    LOG.debug(
+        "Custody group count changed from {} to {}", oldCustodyGroupCount, newCustodyGroupCount);
+    if (newCustodyGroupCount > oldCustodyGroupCount) {
+      final UInt64 minCustodyPeriodSlot =
+          minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(currentSlot);
+      db.setFirstCustodyIncompleteSlot(minCustodyPeriodSlot)
+          .finish(
+              error ->
+                  LOG.error(
+                      "Unexpected error while updating first custody incomplete slot with a new value: {}.",
+                      minCustodyPeriodSlot,
+                      error));
     }
   }
 
   @Override
   public void onNewFinalizedCheckpoint(
       final Checkpoint checkpoint, final boolean fromOptimisticBlock) {
-    advanceFirstIncompleteSlot(checkpoint.getEpoch()).ifExceptionGetsHereRaiseABug();
+    advanceFirstIncompleteSlot(checkpoint.getEpoch())
+        .finish(
+            error ->
+                LOG.error(
+                    "Unexpected error while advancing first custody incomplete slot for checkpoint {}{}",
+                    checkpoint,
+                    fromOptimisticBlock ? " (from optimistic block)" : "",
+                    error));
+  }
+
+  @VisibleForTesting
+  public int getTotalCustodyGroupCount() {
+    return totalCustodyGroupCount.get();
   }
 
   private synchronized boolean updateEpoch(final UInt64 epoch) {
@@ -189,10 +210,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

@@ -14,12 +14,16 @@
 package tech.pegasys.teku.statetransition.datacolumns;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 import java.util.List;
 import java.util.Optional;
 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 +54,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 +97,67 @@ void sanityTest() throws Throwable {
     assertThat(fRet2_1.get().get()).isEqualTo(sidecar0);
     assertThat(fRet2_2.get().get()).isEqualTo(sidecar0);
   }
+
+  @Test
+  public void onSlot_shouldUpdateCustodyGroupCount() {
+    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_shouldUpdateFirstCustodyIncompleteSlotOnlyWhenCustodyGroupCountIncreases() {
+    final DataColumnSidecarDbAccessor dbAccessorMock = mock(DataColumnSidecarDbAccessor.class);
+    when(dbAccessorMock.setFirstCustodyIncompleteSlot(any())).thenReturn(SafeFuture.COMPLETE);
+    // initial cgc is 0
+    custody =
+        new DataColumnSidecarCustodyImpl(
+            spec,
+            blockResolver,
+            dbAccessorMock,
+            minCustodyPeriodSlotCalculator,
+            custodyGroupCountManager,
+            0);
+    custody.onSlot(UInt64.ZERO);
+    // next epoch
+    final UInt64 firstEpochSlot = UInt64.valueOf(config.getSlotsPerEpoch()).increment();
+    // cgc increases to groupCount
+    when(custodyGroupCountManager.getCustodyGroupCount()).thenReturn(groupCount);
+    custody.onSlot(firstEpochSlot);
+    final UInt64 minCustodyPeriodSlotForFirstEpoch =
+        minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(firstEpochSlot);
+    verify(dbAccessorMock).setFirstCustodyIncompleteSlot(minCustodyPeriodSlotForFirstEpoch);
+
+    // cgc stays the same
+    when(custodyGroupCountManager.getCustodyGroupCount()).thenReturn(groupCount);
+    // next epoch slot
+    final UInt64 secondEpochSlot = firstEpochSlot.plus(config.getSlotsPerEpoch()).increment();
+    custody.onSlot(secondEpochSlot);
+    // min custody slot is not updated when cgc is the same
+    verifyNoMoreInteractions(dbAccessorMock);
+
+    // cgc decreases
+    when(custodyGroupCountManager.getCustodyGroupCount()).thenReturn(groupCount - 10);
+    // next epoch slot
+    final UInt64 thirdEpocSlot = secondEpochSlot.plus(config.getSlotsPerEpoch()).increment();
+    custody.onSlot(thirdEpocSlot);
+    // min custody slot is not updated when cgc decreases
+    final UInt64 minCustodyPeriodSlotForThirdEpoch =
+        minCustodyPeriodSlotCalculator.getMinCustodyPeriodSlot(thirdEpocSlot);
+    verify(dbAccessorMock).setFirstCustodyIncompleteSlot(minCustodyPeriodSlotForThirdEpoch);
+  }
+
+  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();
+  }
 }