Added a unit test

apache · Jun 21, 2024 · bcc1550 · bcc1550
1 parent f87a08b
commit bcc1550
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 2 deletions.
diff --git a/...k/src/main/java/org/apache/iceberg/flink/source/enumerator/AbstractIcebergEnumerator.java b/...k/src/main/java/org/apache/iceberg/flink/source/enumerator/AbstractIcebergEnumerator.java
@@ -105,8 +105,12 @@ public void handleSourceEvent(int subTaskId, int attemptNumber, SourceEvent sour
         LOG.info(
             "Received request split event from subtask {} attempt {}", subTaskId, attemptNumber);
         enumeratorContext.assignSplit(splitIdAssignedBySubtask.get(subTaskId), subTaskId);
-      } else {
-        LOG.info("Received request split event from subtask {}", subTaskId);
+      } else if (attemptNumber == 0) {
+        LOG.warn(
+            "Received request split event unexpectedly from subtask {} for speculative execution",
+            subTaskId);
+      } else if (!splitIdAssignedBySubtask.containsKey(subTaskId)) {
+        LOG.info("Looks like the subtask {} failed already and got cleaned up", subTaskId);
       }
     } else {
       throw new IllegalArgumentException(

diff --git a/...test/java/org/apache/iceberg/flink/source/enumerator/TestContinuousIcebergEnumerator.java b/...test/java/org/apache/iceberg/flink/source/enumerator/TestContinuousIcebergEnumerator.java
@@ -147,6 +147,41 @@ public void testRequestingReaderUnavailableWhenSplitDiscovered() throws Exceptio
         .contains(splits.get(0));
   }
 
+  @Test
+  public void testSpeculativeExecutionAttempt() throws Exception {
+    ManualContinuousSplitPlanner splitPlanner = new ManualContinuousSplitPlanner();
+    TestingSplitEnumeratorContext<IcebergSourceSplit> enumeratorContext =
+        new TestingSplitEnumeratorContext<>(4);
+    ScanContext config =
+        ScanContext.builder()
+            .streaming(true)
+            .startingStrategy(StreamingStartingStrategy.TABLE_SCAN_THEN_INCREMENTAL)
+            .build();
+    ContinuousIcebergEnumerator enumerator =
+        createEnumerator(enumeratorContext, config, splitPlanner);
+
+    // register one reader, and let it request a split
+    enumeratorContext.registerReader(2, "localhost");
+    enumerator.addReader(2);
+    enumerator.handleSourceEvent(2, new SplitRequestEvent());
+
+    // make one split available and trigger the periodic discovery
+    List<IcebergSourceSplit> splits =
+        SplitHelpers.createSplitsFromTransientHadoopTable(TEMPORARY_FOLDER, 1, 1);
+    splitPlanner.addSplits(splits, IcebergEnumeratorPosition.of(1L, 1L));
+    enumeratorContext.triggerAllActions();
+
+    List<IcebergSourceSplit> assignedSplits =
+        enumeratorContext.getSplitAssignments().get(2).getAssignedSplits();
+
+    enumerator.handleSourceEvent(2, 2, new SplitRequestEvent());
+
+    List<IcebergSourceSplit> speculatedAssignedSplits =
+        enumeratorContext.getSplitAssignments().get(2).getAssignedSplits();
+
+    Assert.assertEquals(assignedSplits, speculatedAssignedSplits);
+  }
+
   private static ContinuousIcebergEnumerator createEnumerator(
       SplitEnumeratorContext<IcebergSourceSplit> context,
       ScanContext scanContext,