[FLINK-40499][runtime] Emit ACTIVE before final MAX_WATERMARK at drain

flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/MultipleInputStreamTask.java

@@ -45,6 +45,7 @@
 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 void abortCheckpointOnBarrier(long checkpointId, CheckpointException caus
     @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);
         }
     }

flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceOperatorStreamTask.java

@@ -214,7 +214,14 @@ private CompletableFuture<Boolean> triggerStopWithSavepointAsync(
     }
 
     @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);
     }
 

flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SourceStreamTask.java

@@ -37,6 +37,7 @@
 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 void triggerCheckpoint(long checkpointId) throws FlinkException {
 
     @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);
     }
 

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) {}
+    }
+}

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);
+        }
+    }
+}