From 444381318cedb2f4b0bc509a804067c7123185fa Mon Sep 17 00:00:00 2001 From: Carlo Maria Proietti Date: Wed, 16 Sep 2026 11:29:33 +0200 Subject: [PATCH 1/3] Fix BufferingDoFnRunner.checkpointCompleted() leveraging polymorphism in BufferedElement --- .../streaming/stableinput/BufferedElement.java | 4 ++++ .../streaming/stableinput/BufferedElements.java | 10 ++++++++++ .../streaming/stableinput/BufferingDoFnRunner.java | 7 ++++++- 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java index bc0b04211a6f..66e091a62403 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java @@ -18,10 +18,14 @@ package org.apache.beam.runners.flink.translation.wrappers.streaming.stableinput; import org.apache.beam.runners.core.DoFnRunner; +import org.joda.time.Instant; /** An interface for elements buffered during a checkpoint when using @RequiresStableInput. */ public interface BufferedElement { /** Processes this element with the provided DoFnRunner. */ void processWith(DoFnRunner doFnRunner); + + /** Returns the timestamp associated with this buffered element */ + Instant getTimestamp(); } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java index d0a1b12e01e2..d8f048135959 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java @@ -52,6 +52,11 @@ public void processWith(DoFnRunner doFnRunner) { doFnRunner.processElement(element); } + @Override + public Instant getTimestamp() { + return this.element.getTimestamp(); + } + @Override public boolean equals(@Nullable Object o) { if (this == o) { @@ -115,6 +120,11 @@ public void processWith(DoFnRunner doFnRunner) { CausedByDrain.NORMAL); } + @Override + public Instant getTimestamp() { + return outputTimestamp; + } + @Override public boolean equals(@Nullable Object o) { if (this == o) { 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 6d1d4085f0bc..204848b00320 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 @@ -299,7 +299,12 @@ public void checkpointCompleted(long checkpointId) throws Exception { bufferingElementsHandler.clear(); } } - minBufferedElementTimestamp = Long.MAX_VALUE; + minBufferedElementTimestamp = + currentBufferingElementsHandler + .getElements() + .mapToLong(e -> e.getTimestamp().getMillis()) + .min() + .orElse(Long.MAX_VALUE); } public long getOutputWatermarkHold() { From 9b29363eb131b5817be8f96942eedae494a5f6d6 Mon Sep 17 00:00:00 2001 From: Carlo Maria Proietti Date: Thu, 17 Sep 2026 08:22:08 +0200 Subject: [PATCH 2/3] Add a test for the fix --- .../wrappers/streaming/DoFnOperatorTest.java | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) 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..af583e5d2eab 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,87 @@ 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)); + + Supplier> doFnOperatorSupplier = + () -> + 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()); + + DoFnOperator doFnOperator = doFnOperatorSupplier.get(); + 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)); + } + @Test(expected = IllegalStateException.class) public void testFailOnRequiresStableInputAndDisabledCheckpointing() { TupleTag> outputTag = new TupleTag<>("main-output"); From d226cb6e3276849e82b4076444c763e27634ed99 Mon Sep 17 00:00:00 2001 From: Carlo Maria Proietti Date: Thu, 8 Oct 2026 18:31:01 +0200 Subject: [PATCH 3/3] refine the logic introducing minTimestampPerBuffer and improve the test --- .../stableinput/BufferedElement.java | 4 -- .../stableinput/BufferedElements.java | 10 ---- .../stableinput/BufferingDoFnRunner.java | 28 ++++++----- .../wrappers/streaming/DoFnOperatorTest.java | 47 +++++++++++-------- 4 files changed, 45 insertions(+), 44 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java index 66e091a62403..bc0b04211a6f 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElement.java @@ -18,14 +18,10 @@ package org.apache.beam.runners.flink.translation.wrappers.streaming.stableinput; import org.apache.beam.runners.core.DoFnRunner; -import org.joda.time.Instant; /** An interface for elements buffered during a checkpoint when using @RequiresStableInput. */ public interface BufferedElement { /** Processes this element with the provided DoFnRunner. */ void processWith(DoFnRunner doFnRunner); - - /** Returns the timestamp associated with this buffered element */ - Instant getTimestamp(); } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java index d8f048135959..d0a1b12e01e2 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java @@ -52,11 +52,6 @@ public void processWith(DoFnRunner doFnRunner) { doFnRunner.processElement(element); } - @Override - public Instant getTimestamp() { - return this.element.getTimestamp(); - } - @Override public boolean equals(@Nullable Object o) { if (this == o) { @@ -120,11 +115,6 @@ public void processWith(DoFnRunner doFnRunner) { CausedByDrain.NORMAL); } - @Override - public Instant getTimestamp() { - return outputTimestamp; - } - @Override public boolean equals(@Nullable Object o) { if (this == o) { 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 204848b00320..377b52d3bcef 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 @@ -18,6 +18,7 @@ package org.apache.beam.runners.flink.translation.wrappers.streaming.stableinput; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.Comparator; import java.util.Iterator; @@ -118,6 +119,8 @@ public static BufferingDoFnRunner create( private final BufferingElementsHandlerFactory bufferingElementsHandlerFactory; /** 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; /** The current handler used for buffering. */ @@ -176,6 +179,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; @@ -217,8 +222,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())); @@ -238,8 +245,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); @@ -280,8 +289,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(); @@ -297,14 +307,10 @@ public void checkpointCompleted(long checkpointId) throws Exception { underlying.finishBundle(); } bufferingElementsHandler.clear(); + minTimestampPerBuffer[toBeAckedId] = Long.MAX_VALUE; } } - minBufferedElementTimestamp = - currentBufferingElementsHandler - .getElements() - .mapToLong(e -> e.getTimestamp().getMillis()) - .min() - .orElse(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 af583e5d2eab..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 @@ -2265,26 +2265,24 @@ public void processElement(ProcessContext context) { WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE), new SerializablePipelineOptions(options)); - Supplier> doFnOperatorSupplier = - () -> - 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()); + 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()); - DoFnOperator doFnOperator = doFnOperatorSupplier.get(); OneInputStreamOperatorTestHarness, WindowedValue> testHarness = new OneInputStreamOperatorTestHarness<>(doFnOperator); @@ -2313,6 +2311,17 @@ public void processElement(ProcessContext context) { "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)