diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java index 4d44d4633a2a..428ab09f5987 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java +++ b/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 BufferingDoFnRunner 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 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 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 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 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() { diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java index cce8f808c242..6a000f308907 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java +++ b/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 outputTag = new TupleTag<>("main-output"); + + WindowedValues.FullWindowedValueCoder windowedValueCoder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + + DoFn doFn = + new DoFn() { + @ProcessElement + // Use RequiresStableInput to force buffering elements + @RequiresStableInput + public void processElement(ProcessContext context) { + context.output(context.element()); + } + }; + + DoFnOperator.MultiOutputOutputManagerFactory outputManagerFactory = + new DoFnOperator.MultiOutputOutputManagerFactory<>( + outputTag, + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE), + new SerializablePipelineOptions(options)); + + DoFnOperator 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> 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> outputTag = new TupleTag<>("main-output");