From c3a7efa39daea8b1b9098ae7102b8df6c4db3def Mon Sep 17 00:00:00 2001 From: Martijn Visser <2989614+MartijnVisser@users.noreply.github.com> Date: Fri, 28 Aug 2026 14:27:12 +0200 Subject: [PATCH] [FLINK-40505][runtime] Prune failed subtasks from watermark alignment The per-subtask WatermarkAggregator in SourceCoordinator never removed entries, so a failed subtask's last reported watermark kept constraining maxAllowedWatermark for the whole alignment group after failover. Remove the subtask's watermark in subtaskReset (the per-subtask callback that only fires once no execution attempt is alive) and propagate a changed aggregate to the group-level aggregator; the restarted attempt re-registers itself with its next ReportedWatermarkEvent. The updated maxAllowedWatermark reaches subtasks via the existing periodic announceCombinedWatermark, so no immediate announcement is needed. Generated-by: Claude Code (Fable 5) --- .../source/coordinator/SourceCoordinator.java | 51 +++++++++++---- .../SourceCoordinatorAlignmentTest.java | 64 +++++++++++++++++++ 2 files changed, 104 insertions(+), 11 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/source/coordinator/SourceCoordinator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/source/coordinator/SourceCoordinator.java index f75c9be0cea466..18e9a82a0efcc2 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/source/coordinator/SourceCoordinator.java +++ b/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 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 watermarkAggregator = - (WatermarkAggregator) oldValue; - watermarkAggregator.aggregate( - operatorName, newCombinedWatermark); - return watermarkAggregator; - })); + .ifPresent(this::updateAggregatedWatermarkOfGroup); + } + + private void updateAggregatedWatermarkOfGroup(Watermark newCombinedWatermark) { + coordinatorStore.computeIfPresent( + watermarkAlignmentParams.getWatermarkGroup(), + (key, oldValue) -> { + WatermarkAggregator watermarkAggregator = + (WatermarkAggregator) oldValue; + watermarkAggregator.aggregate(operatorName, newCombinedWatermark); + return watermarkAggregator; + }); } private void ensureStarted() { @@ -864,6 +872,27 @@ public Optional 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 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 keySet() { return watermarks.keySet(); } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/source/coordinator/SourceCoordinatorAlignmentTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/source/coordinator/SourceCoordinatorAlignmentTest.java index f0cba2b43f4334..5182ce545025c4 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/source/coordinator/SourceCoordinatorAlignmentTest.java +++ b/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 eventsForSilentSubtask = + receivingTasks.getSentEventsForSubtask(subtask2); + assertThat(eventsForSilentSubtask) + .noneMatch(event -> event instanceof WatermarkAlignmentEvent); + } + } + @Test void testWatermarkAlignmentWithTwoGroups() throws Exception { try (AutoCloseableRegistry closeableRegistry = new AutoCloseableRegistry()) {