This is an automated email from the ASF dual-hosted git repository.
MartijnVisser pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new ebfb1b0c945 [FLINK-40499][runtime] Emit ACTIVE before final
MAX_WATERMARK at drain
ebfb1b0c945 is described below
commit ebfb1b0c945d3a769e5c17184f198854b75df0f3
Author: Martijn Visser <[email protected]>
AuthorDate: Fri Aug 28 14:26:22 2026 +0200
[FLINK-40499][runtime] Emit ACTIVE before final MAX_WATERMARK at drain
At drain, advanceToEndOfEventTime() emitted Watermark.MAX_WATERMARK
directly into the chain output, but RecordWriterOutput and ChainingOutput
drop watermarks while the last announced WatermarkStatus is IDLE, so a
source that went idle before finishing never delivered the final watermark
and downstream event-time timers and windows did not fire. The drain path
now emits WatermarkStatus.ACTIVE first (deduplicated downstream when
already active). SourceOperatorStreamTask and SourceStreamTask skip the
status for tasks deployed as finished, whose
FinishedOnRestoreMainOperatorOutput rejects status events;
MultipleInputStreamTask needs no such guard because its chained source
outputs start ACTIVE and deduplicate the redundant status.
Generated-by: Claude Code (Fable 5)
Generated-by: Claude Code (Opus 5.5)
---
.../runtime/tasks/MultipleInputStreamTask.java | 7 +
.../runtime/tasks/SourceOperatorStreamTask.java | 9 +-
.../streaming/runtime/tasks/SourceStreamTask.java | 8 ++
.../tasks/ChainingOutputIdleMaxWatermarkTest.java | 139 ++++++++++++++++++++
.../MultipleInputStreamTaskIdleDrainTest.java | 88 +++++++++++++
.../SourceOperatorStreamTaskIdleDrainTest.java | 142 +++++++++++++++++++++
.../tasks/SourceStreamTaskIdleDrainTest.java | 71 +++++++++++
7 files changed, 463 insertions(+), 1 deletion(-)
diff --git
a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTask.java
b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTask.java
index e27876eccd7..1dc1df5e3b6 100644
---
a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTask.java
+++
b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTask.java
@@ -45,6 +45,7 @@ import
org.apache.flink.streaming.runtime.metrics.MinWatermarkGauge;
import org.apache.flink.streaming.runtime.metrics.WatermarkGauge;
import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
import org.apache.flink.util.concurrent.FutureUtils;
import javax.annotation.Nullable;
@@ -338,6 +339,12 @@ public class MultipleInputStreamTask<OUT>
@Override
protected void advanceToEndOfEventTime() throws Exception {
for (Output<StreamRecord<?>> sourceOutput :
operatorChain.getChainedSourceOutputs()) {
+ // Chained source outputs drop watermarks while the announced
status is IDLE, so
+ // re-activate first to ensure the final MAX_WATERMARK is
delivered. A redundant
+ // ACTIVE status is deduplicated by the output; this also holds
for tasks deployed
+ // as finished, whose chained sources never ran and whose outputs
therefore never
+ // left the initial ACTIVE status.
+ sourceOutput.emitWatermarkStatus(WatermarkStatus.ACTIVE);
sourceOutput.emitWatermark(Watermark.MAX_WATERMARK);
}
}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTask.java
b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTask.java
index 8c60a02e8e4..a635676ade0 100644
---
a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTask.java
+++
b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTask.java
@@ -214,7 +214,14 @@ public class SourceOperatorStreamTask<T> extends
StreamTask<T, SourceOperator<T,
}
@Override
- protected void advanceToEndOfEventTime() {
+ protected void advanceToEndOfEventTime() throws Exception {
+ // Downstream outputs drop watermarks while the announced status is
IDLE, so re-activate
+ // first to ensure the final MAX_WATERMARK is delivered. A redundant
ACTIVE status is
+ // deduplicated downstream. Tasks deployed as finished never ran the
source, so their
+ // output was never idle and rejects status events.
+ if (!operatorChain.isTaskDeployedAsFinished()) {
+ output.emitWatermarkStatus(WatermarkStatus.ACTIVE);
+ }
output.emitWatermark(Watermark.MAX_WATERMARK);
}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTask.java
b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTask.java
index 8768ded79c6..9234c0ffbc8 100644
---
a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTask.java
+++
b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTask.java
@@ -37,6 +37,7 @@ import
org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
import org.apache.flink.streaming.api.operators.StreamSource;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.tasks.mailbox.MailboxDefaultAction;
+import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.FatalExitExceptionHandler;
import org.apache.flink.util.FlinkException;
@@ -181,6 +182,13 @@ public class SourceStreamTask<
@Override
protected void advanceToEndOfEventTime() throws Exception {
+ // Downstream outputs drop watermarks while the announced status is
IDLE, so re-activate
+ // first to ensure the final MAX_WATERMARK is delivered. A redundant
ACTIVE status is
+ // deduplicated downstream. Tasks deployed as finished never ran the
source, so their
+ // output was never idle and rejects status events.
+ if (!operatorChain.isTaskDeployedAsFinished()) {
+
operatorChain.getMainOperatorOutput().emitWatermarkStatus(WatermarkStatus.ACTIVE);
+ }
operatorChain.getMainOperatorOutput().emitWatermark(Watermark.MAX_WATERMARK);
}
diff --git
a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/ChainingOutputIdleMaxWatermarkTest.java
b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/ChainingOutputIdleMaxWatermarkTest.java
new file mode 100644
index 00000000000..212e3f7cc2f
--- /dev/null
+++
b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/ChainingOutputIdleMaxWatermarkTest.java
@@ -0,0 +1,139 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.tasks;
+
+import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
+import org.apache.flink.streaming.api.operators.Input;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.io.RecordWriterOutput;
+import org.apache.flink.streaming.runtime.streamrecord.LatencyMarker;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Tests for FLINK-40499: {@link ChainingOutput#emitWatermark(Watermark)}
drops watermarks while the
+ * announced status is IDLE, so tasks must emit {@link WatermarkStatus#ACTIVE}
before the final
+ * {@link Watermark#MAX_WATERMARK} at end-of-job drain to make sure it is
delivered.
+ */
+class ChainingOutputIdleMaxWatermarkTest {
+
+ @Test
+ void activeStatusFollowedByMaxWatermarkIsDeliveredAfterIdle() {
+ final CollectingInput input = new CollectingInput();
+ final ChainingOutput<String> chainingOutput =
createChainingOutput(input);
+
+ chainingOutput.emitWatermarkStatus(WatermarkStatus.IDLE);
+ // this is what the drain path (advanceToEndOfEventTime) does:
re-activate, then emit MAX
+ chainingOutput.emitWatermarkStatus(WatermarkStatus.ACTIVE);
+ chainingOutput.emitWatermark(Watermark.MAX_WATERMARK);
+
+ assertThat(input.events)
+ .as(
+ "Re-activating the output before the final
MAX_WATERMARK must deliver "
+ + "both events downstream")
+ .containsExactly(
+ WatermarkStatus.IDLE, WatermarkStatus.ACTIVE,
Watermark.MAX_WATERMARK);
+ }
+
+ /**
+ * Characterization of the idle gate: a watermark emitted while the
announced status is IDLE is
+ * dropped. This is why the task-level drain path must re-activate the
output first instead of
+ * relying on the watermark passing through.
+ */
+ @Test
+ void maxWatermarkAloneIsDroppedWhileIdle() {
+ final CollectingInput input = new CollectingInput();
+ final ChainingOutput<String> chainingOutput =
createChainingOutput(input);
+
+ chainingOutput.emitWatermarkStatus(WatermarkStatus.IDLE);
+ chainingOutput.emitWatermark(Watermark.MAX_WATERMARK);
+
+ assertThat(input.events).containsExactly(WatermarkStatus.IDLE);
+ }
+
+ /**
+ * Documents why {@code MultipleInputStreamTask#advanceToEndOfEventTime()}
may emit ACTIVE
+ * unconditionally: on a task deployed as finished, the chained source
output never left its
+ * initial ACTIVE status, so the redundant ACTIVE is deduplicated and
never reaches the
+ * status-rejecting {@link FinishedOnRestoreInput} underneath. The input
is genuinely reachable
+ * (a differing status does throw) — it is the deduplication that makes
ACTIVE safe.
+ */
+ @Test
+ void redundantActiveIsDeduplicatedBeforeFinishedOnRestoreInput() {
+ final ChainingOutput<String> chainingOutput =
+ new ChainingOutput<>(
+ new FinishedOnRestoreInput<>(new
RecordWriterOutput[0], 1),
+ null,
+
UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup(),
+ null);
+
+ assertThatCode(() ->
chainingOutput.emitWatermarkStatus(WatermarkStatus.ACTIVE))
+ .as("A redundant ACTIVE must be deduplicated and never reach
the input")
+ .doesNotThrowAnyException();
+
+ assertThatThrownBy(() ->
chainingOutput.emitWatermarkStatus(WatermarkStatus.IDLE))
+ .as("A status change does reach FinishedOnRestoreInput, which
rejects it")
+ .isInstanceOf(ExceptionInChainedOperatorException.class)
+ .hasCauseInstanceOf(IllegalStateException.class);
+ }
+
+ private static ChainingOutput<String> createChainingOutput(CollectingInput
input) {
+ return new ChainingOutput<>(
+ input,
+ null,
+
UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup(),
+ null);
+ }
+
+ private static final class CollectingInput implements Input<String> {
+ final List<Object> events = new ArrayList<>();
+
+ @Override
+ public void processElement(StreamRecord<String> element) {
+ events.add(element);
+ }
+
+ @Override
+ public void processWatermark(Watermark mark) {
+ events.add(mark);
+ }
+
+ @Override
+ public void processWatermarkStatus(WatermarkStatus watermarkStatus) {
+ events.add(watermarkStatus);
+ }
+
+ @Override
+ public void processLatencyMarker(LatencyMarker latencyMarker) {
+ events.add(latencyMarker);
+ }
+
+ @Override
+ public void setKeyContextElement(StreamRecord<String> record) {}
+ }
+}
diff --git
a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTaskIdleDrainTest.java
b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTaskIdleDrainTest.java
new file mode 100644
index 00000000000..4de80975472
--- /dev/null
+++
b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTaskIdleDrainTest.java
@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.tasks;
+
+import org.apache.flink.api.common.eventtime.WatermarkStrategy;
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.api.connector.source.mocks.MockSource;
+import org.apache.flink.streaming.api.operators.SourceOperatorFactory;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Test for FLINK-40499: a {@link MultipleInputStreamTask} whose chained
source went IDLE must still
+ * deliver {@link Watermark#MAX_WATERMARK} through the chained source output
when the task is
+ * drained, so that downstream event-time timers and windows fire at end of
job.
+ *
+ * <p>This is the multi-input counterpart of {@code
SourceOperatorStreamTaskIdleDrainTest}: {@code
+ * ChainingOutput#emitWatermark} drops watermarks while the announced status
is IDLE, so {@code
+ * MultipleInputStreamTask#advanceToEndOfEventTime()} must emit {@link
WatermarkStatus#ACTIVE}
+ * before the MAX watermark.
+ */
+class MultipleInputStreamTaskIdleDrainTest {
+
+ @Test
+ void maxWatermarkReachesOutputAtDrainWhenChainedSourceWasIdle() throws
Exception {
+ try (StreamTaskMailboxTestHarness<String> testHarness =
+ new StreamTaskMailboxTestHarnessBuilder<>(
+ MultipleInputStreamTask::new,
BasicTypeInfo.STRING_TYPE_INFO)
+ .addInput(BasicTypeInfo.STRING_TYPE_INFO, 1)
+ .addSourceInput(
+ new SourceOperatorFactory<>(
+ new MockSource(
+
Boundedness.CONTINUOUS_UNBOUNDED, 2, true, true),
+ WatermarkStrategy.noWatermarks()),
+ BasicTypeInfo.INT_TYPE_INFO)
+ .addInput(BasicTypeInfo.DOUBLE_TYPE_INFO, 1)
+ .setupOutputForSingletonOperatorChain(
+ new MultipleInputStreamTaskTest
+
.MapToStringMultipleInputOperatorFactory(3))
+ .build()) {
+
+ // the source reader has no splits and marks the chained source
output IDLE
+ testHarness.processAll();
+
+ // make both network inputs idle so the chained source's watermark
is the deciding one
+ testHarness.processElement(WatermarkStatus.IDLE, 0, 0);
+ testHarness.processElement(WatermarkStatus.IDLE, 1, 0);
+ testHarness.processAll();
+
+ assertThat(testHarness.getOutput())
+ .as("Precondition: every input (incl. the chained source)
went idle")
+ .contains(WatermarkStatus.IDLE);
+
+ // this is what StreamTask.endData(DRAIN) invokes at drain
+ testHarness.getStreamTask().advanceToEndOfEventTime();
+ testHarness.processAll();
+
+ assertThat(testHarness.getOutput())
+ .as(
+ "At drain, WatermarkStatus.ACTIVE followed by
MAX_WATERMARK must be "
+ + "emitted even though the chained
source's last announced "
+ + "status was IDLE")
+ .containsSubsequence(
+ WatermarkStatus.IDLE, WatermarkStatus.ACTIVE,
Watermark.MAX_WATERMARK);
+ }
+ }
+}
diff --git
a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTaskIdleDrainTest.java
b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTaskIdleDrainTest.java
new file mode 100644
index 00000000000..f182f00ec0e
--- /dev/null
+++
b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTaskIdleDrainTest.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.tasks;
+
+import org.apache.flink.api.common.eventtime.WatermarkStrategy;
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.api.connector.source.ReaderOutput;
+import org.apache.flink.api.connector.source.SourceReader;
+import org.apache.flink.api.connector.source.SourceReaderContext;
+import org.apache.flink.api.connector.source.mocks.MockSource;
+import org.apache.flink.api.connector.source.mocks.MockSourceSplit;
+import org.apache.flink.core.io.InputStatus;
+import org.apache.flink.runtime.io.network.api.EndOfData;
+import org.apache.flink.runtime.io.network.api.StopMode;
+import org.apache.flink.streaming.api.operators.SourceOperatorFactory;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Test for FLINK-40499: a {@link SourceOperatorStreamTask} whose source went
IDLE before finishing
+ * must still emit {@link Watermark#MAX_WATERMARK} when the task is drained,
so that downstream
+ * event-time timers and windows fire at end of job.
+ *
+ * <p>This is the task-level counterpart of {@code
ChainingOutputIdleMaxWatermarkTest}: {@code
+ * RecordWriterOutput#emitWatermark} drops watermarks while the announced
status is IDLE, so {@code
+ * SourceOperatorStreamTask#advanceToEndOfEventTime()} must emit {@link
WatermarkStatus#ACTIVE}
+ * before the MAX watermark.
+ */
+class SourceOperatorStreamTaskIdleDrainTest {
+
+ @Test
+ void maxWatermarkIsEmittedAtDrainEvenIfSourceWentIdle() throws Exception {
+ SourceOperatorFactory<Integer> sourceOperatorFactory =
+ new SourceOperatorFactory<>(
+ new IdleBeforeFinishingSource(),
WatermarkStrategy.noWatermarks());
+
+ try (StreamTaskMailboxTestHarness<Integer> testHarness =
+ new StreamTaskMailboxTestHarnessBuilder<>(
+ SourceOperatorStreamTask::new,
BasicTypeInfo.INT_TYPE_INFO)
+ .setCollectNetworkEvents()
+
.setupOutputForSingletonOperatorChain(sourceOperatorFactory)
+ .build()) {
+
+ testHarness.processAll();
+ testHarness.finishProcessing();
+
+ // sanity check: the source really announced idleness before
finishing
+ assertThat(testHarness.getOutput())
+ .as("Precondition: the source should have announced IDLE
before finishing")
+ .contains(WatermarkStatus.IDLE);
+
+ // the drain path must re-activate the (idle) output before the
final MAX_WATERMARK,
+ // otherwise the watermark is dropped by the idle gate in the
output
+ assertThat(testHarness.getOutput())
+ .as(
+ "At drain, WatermarkStatus.ACTIVE followed by
MAX_WATERMARK must be "
+ + "emitted before EndOfData even if the
source's last "
+ + "announced status was IDLE")
+ .containsExactly(
+ WatermarkStatus.IDLE,
+ WatermarkStatus.ACTIVE,
+ Watermark.MAX_WATERMARK,
+ new EndOfData(StopMode.DRAIN));
+ }
+ }
+
+ /** A bounded source whose reader marks itself idle and then immediately
finishes. */
+ private static class IdleBeforeFinishingSource extends MockSource {
+ private static final long serialVersionUID = 1L;
+
+ IdleBeforeFinishingSource() {
+ super(Boundedness.BOUNDED, 1);
+ }
+
+ @Override
+ public SourceReader<Integer, MockSourceSplit> createReader(
+ SourceReaderContext readerContext) {
+ return new IdleThenFinishSourceReader();
+ }
+ }
+
+ /**
+ * A reader that marks the output IDLE (e.g. like a reader with no
assigned work, or an idleness
+ * timeout would) and then reaches end of input.
+ */
+ private static class IdleThenFinishSourceReader
+ implements SourceReader<Integer, MockSourceSplit> {
+
+ @Override
+ public InputStatus pollNext(ReaderOutput<Integer> output) {
+ output.markIdle();
+ return InputStatus.END_OF_INPUT;
+ }
+
+ @Override
+ public void start() {}
+
+ @Override
+ public List<MockSourceSplit> snapshotState(long checkpointId) {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public CompletableFuture<Void> isAvailable() {
+ return CompletableFuture.completedFuture(null);
+ }
+
+ @Override
+ public void addSplits(List<MockSourceSplit> splits) {}
+
+ @Override
+ public void notifyNoMoreSplits() {}
+
+ @Override
+ public void close() {}
+ }
+}
diff --git
a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTaskIdleDrainTest.java
b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTaskIdleDrainTest.java
new file mode 100644
index 00000000000..2a0a4b68ffa
--- /dev/null
+++
b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTaskIdleDrainTest.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.tasks;
+
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.runtime.io.network.api.EndOfData;
+import org.apache.flink.runtime.io.network.api.StopMode;
+import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
+import org.apache.flink.streaming.api.operators.StreamSource;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Test for FLINK-40499: a {@link SourceStreamTask} whose legacy source went
IDLE before finishing
+ * must still emit {@link Watermark#MAX_WATERMARK} when the task is drained.
+ */
+class SourceStreamTaskIdleDrainTest {
+
+ @Test
+ void maxWatermarkIsEmittedAtDrainEvenIfSourceWentIdle() throws Exception {
+ try (StreamTaskMailboxTestHarness<Integer> testHarness =
+ new StreamTaskMailboxTestHarnessBuilder<>(
+ SourceStreamTask::new,
BasicTypeInfo.INT_TYPE_INFO)
+ .setCollectNetworkEvents()
+ .setupOutputForSingletonOperatorChain(
+ new StreamSource<>(new
IdleBeforeFinishingSource()))
+ .build()) {
+
+ final EndOfData endOfData = new EndOfData(StopMode.DRAIN);
+ testHarness.processUntil(() ->
testHarness.getOutput().contains(endOfData));
+
+ assertThat(testHarness.getOutput())
+ .containsExactly(
+ WatermarkStatus.IDLE,
+ WatermarkStatus.ACTIVE,
+ Watermark.MAX_WATERMARK,
+ endOfData);
+ }
+ }
+
+ private static class IdleBeforeFinishingSource implements
SourceFunction<Integer> {
+
+ @Override
+ public void run(SourceContext<Integer> ctx) {
+ ctx.markAsTemporarilyIdle();
+ }
+
+ @Override
+ public void cancel() {}
+ }
+}
