Improve DasSamplerBasic check as recentChainData.contains doesn't mean we should drop in Gloas

Author: zilm13

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

@@ -255,8 +255,19 @@ public void onSlot(final UInt64 slot) {
         .values()
         .removeIf(
             tracker -> {
-              if (tracker.slot().isLessThan(firstNonFinalizedSlot)
-                  || recentChainData.containsBlock(tracker.blockRoot())) {
+              final boolean isCleanupDue;
+              if (tracker.slot().isLessThan(firstNonFinalizedSlot)) {
+                isCleanupDue = true;
+              } else {
+                final boolean blockImported = recentChainData.containsBlock(tracker.blockRoot());
+                final boolean isDataAvailabilityDeferred =
+                    spec.atSlot(tracker.slot())
+                        .getForkChoiceUtil()
+                        .isDataAvailabilityCheckDeferredToExecutionPayloadEnvelope();
+                isCleanupDue = blockImported && !isDataAvailabilityDeferred;
+              }
+
+              if (isCleanupDue) {
                 // Outdated
                 if (!tracker.completionFuture().isDone()) {
                   // make sure the future releases any pending waiters

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

@@ -563,6 +563,36 @@ void onBlockImported_shouldRetainTrackerWhenDataAvailabilityDeferredToExecutionP
         .doesNotContainKey(blockWithBlobs.getRoot());
   }
 
+  @Test
+  void onSlot_shouldRetainGloasImportedBlockTrackerUntilExecutionPayloadImported() {
+    final DasSamplerBasic gloasSampler = createGloasSampler();
+    final SignedBeaconBlock blockWithBlobs =
+        dataStructureUtil.randomSignedBeaconBlockWithCommitments(UInt64.ONE, 1);
+    when(rpcFetchDelayProvider.calculate(blockWithBlobs.getSlot())).thenReturn(Duration.ZERO);
+    when(recentChainData.getFinalizedEpoch()).thenReturn(UInt64.ZERO);
+    when(recentChainData.containsBlock(blockWithBlobs.getRoot())).thenReturn(true);
+
+    gloasSampler.onNewBlock(blockWithBlobs, Optional.of(RemoteOrigin.GOSSIP));
+    final SafeFuture<List<UInt64>> completionFuture =
+        gloasSampler
+            .getRecentlySampledColumnsByRoot()
+            .get(blockWithBlobs.getRoot())
+            .completionFuture();
+
+    gloasSampler.onSlot(UInt64.valueOf(2));
+
+    assertThat(gloasSampler.getRecentlySampledColumnsByRoot())
+        .containsKey(blockWithBlobs.getRoot());
+    assertThat(completionFuture).isNotDone();
+
+    // Data availability is deferred to the payload, so the tracker is pruned now.
+    gloasSampler.onExecutionPayloadImported(blockWithBlobs.getSlotAndBlockRoot());
+
+    assertThat(gloasSampler.getRecentlySampledColumnsByRoot())
+        .doesNotContainKey(blockWithBlobs.getRoot());
+    assertThat(completionFuture).isCancelled();
+  }
+
   @Test
   void enableBlockImportOnCompletion_shouldImportOnlyOnceWhenCalledMultipleTimes() {
     final SignedBeaconBlock blockWithBlobs =