[FLINK-40505][runtime] Prune failed subtasks from watermark alignment

flink-runtime/src/main/java/org/apache/flink/runtime/source/coordinator/SourceCoordinator.java

@@ -380,6 +380,13 @@ public void subtaskReset(int subtaskId, long checkpointId) {
 
                     context.subtaskReset(subtaskId);
 
+                    // Remove the last watermark reported by the failed subtask so that a stale
+                    // value does not keep constraining the watermark alignment group. The
+                    // restarted attempt re-registers itself with its next ReportedWatermarkEvent.
+                    combinedWatermark
+                            .remove(subtaskId)
+                            .ifPresent(this::updateAggregatedWatermarkOfGroup);
+
                     final List<SplitT> splitsToAddBack =
                             context.getAndRemoveUncheckpointedAssignment(subtaskId, checkpointId);
                     LOG.debug(
@@ -726,17 +733,18 @@ private void handleReportedWatermark(int subtask, Watermark watermark) throws Fl
 
         combinedWatermark
                 .aggregate(subtask, watermark)
-                .ifPresent(
-                        newCombinedWatermark ->
-                                coordinatorStore.computeIfPresent(
-                                        watermarkAlignmentParams.getWatermarkGroup(),
-                                        (key, oldValue) -> {
-                                            WatermarkAggregator<String> watermarkAggregator =
-                                                    (WatermarkAggregator<String>) oldValue;
-                                            watermarkAggregator.aggregate(
-                                                    operatorName, newCombinedWatermark);
-                                            return watermarkAggregator;
-                                        }));
+                .ifPresent(this::updateAggregatedWatermarkOfGroup);
+    }
+
+    private void updateAggregatedWatermarkOfGroup(Watermark newCombinedWatermark) {
+        coordinatorStore.computeIfPresent(
+                watermarkAlignmentParams.getWatermarkGroup(),
+                (key, oldValue) -> {
+                    WatermarkAggregator<String> watermarkAggregator =
+                            (WatermarkAggregator<String>) oldValue;
+                    watermarkAggregator.aggregate(operatorName, newCombinedWatermark);
+                    return watermarkAggregator;
+                });
     }
 
     private void ensureStarted() {
@@ -864,6 +872,27 @@ public Optional<Watermark> aggregate(T key, Watermark watermark) {
             return Optional.of(newAggregatedWatermark);
         }
 
+        /**
+         * Removes the {@link Watermark} for the given {@code key}.
+         *
+         * @return the new updated combined {@link Watermark} if the value has changed. {@code
+         *     Optional.empty()} otherwise.
+         */
+        public Optional<Watermark> remove(T key) {
+            Watermark oldAggregatedWatermark = getAggregatedWatermark();
+
+            WatermarkElement removedWatermarkElement = watermarks.remove(key);
+            if (removedWatermarkElement != null) {
+                orderedWatermarks.remove(removedWatermarkElement);
+            }
+
+            Watermark newAggregatedWatermark = getAggregatedWatermark();
+            if (newAggregatedWatermark.equals(oldAggregatedWatermark)) {
+                return Optional.empty();
+            }
+            return Optional.of(newAggregatedWatermark);
+        }
+
         public Set<T> keySet() {
             return watermarks.keySet();
         }

flink-runtime/src/test/java/org/apache/flink/runtime/source/coordinator/SourceCoordinatorAlignmentTest.java

@@ -116,6 +116,70 @@ void testWatermarkAlignmentWithIdleness() throws Exception {
         }
     }
 
+    @Test
+    void testWatermarkAlignmentStatePrunedAfterSubtaskFailureAndReset() throws Exception {
+        try (AutoCloseableRegistry closeableRegistry = new AutoCloseableRegistry()) {
+            SourceCoordinator<?, ?> sourceCoordinator1 =
+                    getAndStartNewSourceCoordinator(
+                            new WatermarkAlignmentParams(1000L, "group1", Long.MAX_VALUE),
+                            closeableRegistry);
+
+            int subtask0 = 0;
+            int subtask1 = 1;
+
+            // baseline sanity: the group is constrained by the lowest reported watermark
+            reportWatermarkEvent(sourceCoordinator1, subtask0, 100);
+            assertLatestWatermarkAlignmentEvent(subtask0, 1100);
+
+            reportWatermarkEvent(sourceCoordinator1, subtask1, 200);
+            assertLatestWatermarkAlignmentEvent(subtask0, 1100);
+            assertLatestWatermarkAlignmentEvent(subtask1, 1100);
+
+            // the subtask holding the group back fails and is reset
+            sourceCoordinator1.executionAttemptFailed(
+                    subtask0, 0, new RuntimeException("Artificial failure for subtask 0"));
+            sourceCoordinator1.subtaskReset(subtask0, 1L);
+            CoordinatorTestUtils.waitForCoordinatorToProcessActions(
+                    sourceCoordinator1.getContext());
+
+            sourceCoordinator1.announceCombinedWatermark();
+
+            // CORRECT expectation: with the failed subtask's state pruned, the group should now
+            // only be constrained by the surviving subtask (200 + 1000 drift). The failed subtask
+            // will re-report its watermark after restart anyway.
+            assertLatestWatermarkAlignmentEvent(subtask1, 1200);
+        }
+    }
+
+    @Test
+    void testSubtaskThatNeverReportedReceivesNoAlignmentEvents() throws Exception {
+        try (AutoCloseableRegistry closeableRegistry = new AutoCloseableRegistry()) {
+            SourceCoordinator<?, ?> sourceCoordinator1 =
+                    getAndStartNewSourceCoordinator(
+                            new WatermarkAlignmentParams(1000L, "group1", Long.MAX_VALUE),
+                            closeableRegistry);
+
+            int subtask0 = 0;
+            int subtask2 = 2;
+
+            // subtask0 reports; subtask2 is ready (gateway registered) but never reports a
+            // watermark (e.g. idle from birth)
+            reportWatermarkEvent(sourceCoordinator1, subtask0, 100);
+            assertLatestWatermarkAlignmentEvent(subtask0, 1100);
+
+            sourceCoordinator1.announceCombinedWatermark();
+            sourceCoordinator1.announceCombinedWatermark();
+
+            // characterization: announceCombinedWatermark only iterates over subtasks that have
+            // reported (combinedWatermark.keySet()), so a never-reporting subtask never receives
+            // any WatermarkAlignmentEvent
+            List<OperatorEvent> eventsForSilentSubtask =
+                    receivingTasks.getSentEventsForSubtask(subtask2);
+            assertThat(eventsForSilentSubtask)
+                    .noneMatch(event -> event instanceof WatermarkAlignmentEvent);
+        }
+    }
+
     @Test
     void testWatermarkAlignmentWithTwoGroups() throws Exception {
         try (AutoCloseableRegistry closeableRegistry = new AutoCloseableRegistry()) {