Flink dropping watermarks hold fix

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java

@@ -19,6 +19,7 @@
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.Iterator;
@@ -129,6 +130,9 @@ public static <InputT, OutputT> BufferingDoFnRunner<InputT, OutputT> create(
   /** The maximum number of buffers for data of not yet acknowledged checkpoints. */
   final int numCheckpointBuffers;
 
+  /** Minimum timestamp for each checkpoint buffer. */
+  private final long[] minTimestampPerBuffer;
+
   /** The current active state id which, on checkpoint, is linked to a checkpoint id. */
   int currentStateIndex;
 
@@ -193,6 +197,8 @@ private BufferingDoFnRunner(
           }
         };
     this.numCheckpointBuffers = initializeState(maxConcurrentCheckpoints);
+    this.minTimestampPerBuffer = new long[numCheckpointBuffers];
+    Arrays.fill(this.minTimestampPerBuffer, Long.MAX_VALUE);
     this.currentBufferingElementsHandler =
         bufferingElementsHandlerFactory.get(rotateAndGetStateIndex());
     this.keyedStateBackend = keyedStateBackend;
@@ -234,8 +240,10 @@ public void startBundle() {
 
   @Override
   public void processElement(WindowedValue<InputT> elem) {
-    minBufferedElementTimestamp =
-        Math.min(elem.getTimestamp().getMillis(), minBufferedElementTimestamp);
+    final long elemTimestamp = elem.getTimestamp().getMillis();
+    minTimestampPerBuffer[currentStateIndex] =
+        Math.min(minTimestampPerBuffer[currentStateIndex], elemTimestamp);
+    minBufferedElementTimestamp = Math.min(elemTimestamp, minBufferedElementTimestamp);
     try (Locker lock = locker != null ? locker.get() : null) {
       if (keySelector != null) {
         keyedStateBackend.setCurrentKey(keySelector.apply(elem.getValue()));
@@ -255,8 +263,10 @@ public <KeyT> void onTimer(
       TimeDomain timeDomain,
       CausedByDrain causedByDrain) {
 
-    minBufferedElementTimestamp =
-        Math.min(outputTimestamp.getMillis(), minBufferedElementTimestamp);
+    final long outputTimestampInMillis = outputTimestamp.getMillis();
+    minTimestampPerBuffer[currentStateIndex] =
+        Math.min(minTimestampPerBuffer[currentStateIndex], outputTimestampInMillis);
+    minBufferedElementTimestamp = Math.min(outputTimestampInMillis, minBufferedElementTimestamp);
     try (Locker lock = locker != null ? locker.get() : null) {
       if (keySelector != null) {
         keyedStateBackend.setCurrentKey(key);
@@ -297,8 +307,9 @@ public void checkpoint(long checkpointId) throws Exception {
   public void checkpointCompleted(long checkpointId) throws Exception {
     List<CheckpointIdentifier> allToAck = gatherToBeAcknowledgedCheckpoints(checkpointId);
     for (CheckpointIdentifier toBeAcked : allToAck) {
+      final int toBeAckedId = toBeAcked.internalId;
       BufferingElementsHandler bufferingElementsHandler =
-          bufferingElementsHandlerFactory.get(toBeAcked.internalId);
+          bufferingElementsHandlerFactory.get(toBeAckedId);
       try (Locker lock = locker != null ? locker.get() : null) {
         final Iterator<BufferedElement> iterator =
             bufferingElementsHandler.getElements().iterator();
@@ -314,9 +325,10 @@ public void checkpointCompleted(long checkpointId) throws Exception {
           underlying.finishBundle();
         }
         bufferingElementsHandler.clear();
+        minTimestampPerBuffer[toBeAckedId] = Long.MAX_VALUE;
       }
     }
-    minBufferedElementTimestamp = Long.MAX_VALUE;
+    minBufferedElementTimestamp = Arrays.stream(minTimestampPerBuffer).min().orElse(Long.MAX_VALUE);
   }
 
   public long getOutputWatermarkHold() {

runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java

@@ -99,6 +99,7 @@
 import org.apache.flink.api.java.functions.KeySelector;
 import org.apache.flink.api.java.typeutils.ValueTypeInfo;
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.streaming.api.watermark.Watermark;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
 import org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness;
@@ -2233,6 +2234,96 @@ public void finishBundle(FinishBundleContext context) {
             WindowedValues.valueInGlobalWindow(KV.of("key3", "finishBundle"))));
   }
 
+  /** Ensures that output watermark does not pass buffered elements */
+  @Test
+  public void testExactlyOnceBufferingOutputWatermarkCorrectness() throws Exception {
+    FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
+    options.setStreaming(true);
+    options.setMaxBundleSize(1L);
+    options.setCheckpointingInterval(1L);
+    options.setCheckpointingMode("EXACTLY_ONCE");
+    options.setNumConcurrentCheckpoints(1);
+
+    TupleTag<String> outputTag = new TupleTag<>("main-output");
+
+    WindowedValues.FullWindowedValueCoder<String> windowedValueCoder =
+        WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE);
+
+    DoFn<String, String> doFn =
+        new DoFn<String, String>() {
+          @ProcessElement
+          // Use RequiresStableInput to force buffering elements
+          @RequiresStableInput
+          public void processElement(ProcessContext context) {
+            context.output(context.element());
+          }
+        };
+
+    DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
+        new DoFnOperator.MultiOutputOutputManagerFactory<>(
+            outputTag,
+            WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE),
+            new SerializablePipelineOptions(options));
+
+    DoFnOperator<String, String, String> doFnOperator =
+        new DoFnOperator<>(
+            doFn,
+            "stepName",
+            windowedValueCoder,
+            Collections.emptyMap(),
+            outputTag,
+            Collections.emptyList(),
+            outputManagerFactory,
+            WindowingStrategy.globalDefault(),
+            new HashMap<>(), /* side-input mapping */
+            Collections.emptyList(), /* side inputs */
+            options,
+            null,
+            null,
+            DoFnSchemaInformation.create(),
+            Collections.emptyMap());
+
+    OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
+        new OneInputStreamOperatorTestHarness<>(doFnOperator);
+
+    testHarness.open();
+
+    testHarness.processElement(
+        new StreamRecord<>(WindowedValues.timestampedValueInGlobalWindow("A", new Instant(10))));
+    testHarness.snapshot(1L, 0L);
+    testHarness.processElement(
+        new StreamRecord<>(WindowedValues.timestampedValueInGlobalWindow("B", new Instant(20))));
+    testHarness.processWatermark(new Watermark(100));
+    assertThat(
+        "A@10 has not yet completed the checkpoint; output watermark must be held at 10 until checkpoint completion",
+        doFnOperator.getCurrentOutputWatermark(),
+        is(10L));
+
+    doFnOperator.notifyCheckpointComplete(1L);
+    assertThat(
+        "A@10 has been checked but B@20 is still buffered; output watermark must not pass it",
+        doFnOperator.getCurrentOutputWatermark(),
+        is(20L));
+
+    testHarness.snapshot(2L, 0L);
+    doFnOperator.notifyCheckpointComplete(2L);
+    assertThat(
+        "B@20 has been checked and the buffer is empty; output watermark must advance to input watermark",
+        doFnOperator.getCurrentOutputWatermark(),
+        is(100L));
+
+    assertThat(
+        testHarness.getOutput(),
+        contains(
+            new Watermark(10),
+            new StreamRecord<>(WindowedValues.timestampedValueInGlobalWindow("A", new Instant(10))),
+            new Watermark(20),
+            new StreamRecord<>(WindowedValues.timestampedValueInGlobalWindow("B", new Instant(20))),
+            new Watermark(100)));
+
+    testHarness.close();
+  }
+
   @Test(expected = IllegalStateException.class)
   public void testFailOnRequiresStableInputAndDisabledCheckpointing() {
     TupleTag<KV<String, String>> outputTag = new TupleTag<>("main-output");