From f91b5af808b38dad5c8b8f2c74fdaa664f40fc9e Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Sun, 10 May 2026 10:04:54 +0000 Subject: [PATCH 01/12] [Dataflow] Add Operation::finishKey() and move timers into it Moving the processTimers call from finish() to finishKey(). In upcoming changes there'll be multiple streaming work items in a single beam bundle. With multiple work items, we've to process elements and timers of each work item before moving to the next work items. finishKey() will be called by the NativeIterator classes after iterating through all elements from a work item. Batch processes timers in BatchModeUngroupingParDoFn and does not rely on the processTimers() in ParDoOperation::finish(). So removing the processTimers() call from ParDoOperation::finish() is safe. Batch also does not use the new finishKey() method. --- ...tCommit_Java_ValidatesRunner_Dataflow.json | 2 +- ...va_ValidatesRunner_Dataflow_Streaming.json | 2 +- .../runners/dataflow/worker/PubsubReader.java | 14 +---- .../worker/StreamingModeExecutionContext.java | 12 +++- .../worker/UngroupedWindmillReader.java | 14 +---- .../worker/WindmillReaderIteratorBase.java | 15 ++++- .../worker/WindowingWindmillReader.java | 5 ++ .../WorkerCustomSourceOperationExecutor.java | 3 + .../streaming/ComputationWorkExecutor.java | 2 +- .../util/common/worker/FlattenOperation.java | 3 + .../util/common/worker/MapTaskExecutor.java | 7 +++ .../worker/util/common/worker/Operation.java | 3 + .../util/common/worker/ParDoOperation.java | 7 ++- .../util/common/worker/ReadOperation.java | 3 + .../util/common/worker/WorkExecutor.java | 3 + .../util/common/worker/WriteOperation.java | 3 + .../worker/IntrinsicMapTaskExecutorTest.java | 12 ++++ .../worker/StreamingDataflowWorkerTest.java | 4 +- .../StreamingModeExecutionContextTest.java | 11 +++- .../WindmillReaderIteratorBaseTest.java | 60 ++++++++++++++++++- .../worker/WorkerCustomSourcesTest.java | 7 ++- .../util/common/worker/ExecutorTestUtils.java | 3 + .../common/worker/MapTaskExecutorTest.java | 12 ++++ .../common/worker/ParDoOperationTest.java | 2 + 24 files changed, 170 insertions(+), 39 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json index e9d869cc508c..8144784f5f02 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 1, + "modification": 2, } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json index e623d3373a93..50d17c108f2e 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 1, + "modification": 2, } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubReader.java index b60cb84415ff..0c80909a0f3b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubReader.java @@ -117,20 +117,12 @@ public NativeReader create( @Override public NativeReaderIterator> iterator() throws IOException { - return new PubsubReaderIterator(context.getWorkItem()); + return new PubsubReaderIterator(); } class PubsubReaderIterator extends WindmillReaderIteratorBase { - protected PubsubReaderIterator(Windmill.WorkItem work) { - super(work, skipUndecodableElements); - } - - @Override - public boolean advance() throws IOException { - if (context.workIsFailed()) { - return false; - } - return super.advance(); + protected PubsubReaderIterator() { + super(context, skipUndecodableElements); } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index f75d452b211b..0255d5c577ec 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -56,6 +56,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInput; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputState; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; +import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest; @@ -157,6 +158,8 @@ public class StreamingModeExecutionContext extends DataflowExecutionContext activeReader; + private @Nullable WorkExecutor workExecutor; + public StreamingModeExecutionContext( CounterFactory counterFactory, String computationId, @@ -240,9 +243,11 @@ public void start( Work work, WindmillStateReader stateReader, SideInputStateFetcher sideInputStateFetcher, - Windmill.WorkItemCommitRequest.Builder outputBuilder) { + Windmill.WorkItemCommitRequest.Builder outputBuilder, + WorkExecutor workExecutor) { this.key = key; this.work = work; + this.workExecutor = workExecutor; this.computationKey = WindmillComputationKey.create(computationId, work.getShardedKey()); this.sideInputStateFetcher = sideInputStateFetcher; StreamingGlobalConfig config = globalConfigHandle.getConfig(); @@ -270,6 +275,11 @@ public void start( } } + public void finishKey() throws Exception { + checkNotNull(workExecutor, "workExecutor must be set before calling finishKey()"); + workExecutor.finishKey(); + } + /** * Ensure that the processing time is greater than any fired processing time timers. Otherwise, a * trigger could ignore the timer and orphan the window. diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java index 8ef0bf80323a..f6924493f190 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java @@ -109,20 +109,12 @@ public NativeReader create( @Override public NativeReaderIterator> iterator() throws IOException { - return new UngroupedWindmillReaderIterator(context.getWorkItem()); + return new UngroupedWindmillReaderIterator(); } class UngroupedWindmillReaderIterator extends WindmillReaderIteratorBase { - UngroupedWindmillReaderIterator(Windmill.WorkItem work) { - super(work, skipUndecodableElements); - } - - @Override - public boolean advance() throws IOException { - if (context.workIsFailed()) { - return false; - } - return super.advance(); + UngroupedWindmillReaderIterator() { + super(context, skipUndecodableElements); } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java index 7e6508a4788c..b8232ca9f61a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java @@ -34,6 +34,7 @@ */ public abstract class WindmillReaderIteratorBase extends NativeReader.NativeReaderIterator> { + private final StreamingModeExecutionContext context; private final Windmill.WorkItem work; private int bundleIndex = 0; private int messageIndex = -1; @@ -42,9 +43,10 @@ public abstract class WindmillReaderIteratorBase private static final Logger LOG = LoggerFactory.getLogger(WindmillReaderIteratorBase.class); protected WindmillReaderIteratorBase( - Windmill.WorkItem work, ValueProvider skipUndecodableElements) { + StreamingModeExecutionContext context, ValueProvider skipUndecodableElements) { + this.context = context; this.skipUndecodableElements = skipUndecodableElements; - this.work = work; + this.work = context.getWorkItem(); } @Override @@ -54,9 +56,18 @@ public boolean start() throws IOException { @Override public boolean advance() throws IOException { + if (context.workIsFailed()) { + throw new WorkItemCancelledException(context.getWorkItem().getShardingKey()); + } + while (true) { if (bundleIndex >= work.getMessageBundlesCount()) { current = null; + try { + context.finishKey(); + } catch (Exception e) { + throw new RuntimeException(e); + } return false; } Windmill.InputMessageBundle bundle = work.getMessageBundles(bundleIndex); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java index 173b254f6395..7483f2d138a2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java @@ -182,6 +182,11 @@ public boolean start() throws IOException { @Override public boolean advance() throws IOException { current = null; + try { + context.finishKey(); + } catch (Exception e) { + throw new RuntimeException(e); + } return false; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourceOperationExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourceOperationExecutor.java index 31528e96e07f..a1321d57ebb6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourceOperationExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourceOperationExecutor.java @@ -89,6 +89,9 @@ public void execute() throws Exception { LOG.debug("Source operation execution complete"); } + @Override + public void finishKey() throws Exception {} + @Override public SourceOperationResponse getResponse() { return response; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationWorkExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationWorkExecutor.java index 8dc681fc640c..b4f3a22a7f52 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationWorkExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationWorkExecutor.java @@ -74,7 +74,7 @@ public final void executeWork( SideInputStateFetcher sideInputStateFetcher, Windmill.WorkItemCommitRequest.Builder outputBuilder) throws Exception { - context().start(key, work, stateReader, sideInputStateFetcher, outputBuilder); + context().start(key, work, stateReader, sideInputStateFetcher, outputBuilder, workExecutor()); workExecutor().execute(); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/FlattenOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/FlattenOperation.java index 4847e9f2ea9c..af1b2b9c48bd 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/FlattenOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/FlattenOperation.java @@ -43,6 +43,9 @@ public void process(Object elem) throws Exception { } } + @Override + public void finishKey() throws Exception {} + @Override public boolean supportsRestart() { return true; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java index 58b95f286d55..3c33e1904069 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java @@ -112,6 +112,13 @@ public void execute() throws Exception { // TODO: support for success / failure ports? } + @Override + public void finishKey() throws Exception { + for (Operation op : operations) { + op.finishKey(); + } + } + @Override public NativeReader.Progress getWorkerProgress() throws Exception { return getReadOperation().getProgress(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Operation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Operation.java index b7b4e255cfa5..b630da33cfad 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Operation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Operation.java @@ -137,6 +137,9 @@ public void finish() throws Exception { } } + /** Called when all elements for a specific key have been processed. */ + public abstract void finishKey() throws Exception; + /** Aborts this Operation's execution. */ public void abort() throws Exception { synchronized (initializationStateLock) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java index 27b6e9d1fb35..5aec82073366 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java @@ -46,12 +46,17 @@ public void process(Object elem) throws Exception { } @Override - public void finish() throws Exception { + // Batch mode does not use this method and instead relies on BatchModeUngroupingParDoFn + // to process timers per key. + public void finishKey() throws Exception { try (Closeable scope = context.enterProcessTimers()) { checkStarted(); fn.processTimers(); } + } + @Override + public void finish() throws Exception { try (Closeable scope = context.enterFinish()) { fn.finishBundle(); super.finish(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ReadOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ReadOperation.java index d6b020483d4c..fabc8d6af25b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ReadOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ReadOperation.java @@ -271,6 +271,9 @@ public void finish() throws Exception { } } + @Override + public void finishKey() throws Exception {} + @Override public void abort() throws Exception { if (readerIterator != null) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java index b7170c80ced9..1083fdbb9c42 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java @@ -34,6 +34,9 @@ public interface WorkExecutor extends AutoCloseable { /** Executes the task. */ public abstract void execute() throws Exception; + /** Called when all elements for a specific key have been processed. */ + void finishKey() throws Exception; + /** * Returns the worker's current progress. * diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java index 673140d58d89..d28e7f3e5d3d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java @@ -105,6 +105,9 @@ public void finish() throws Exception { } } + @Override + public void finishKey() throws Exception {} + @Override public void abort() throws Exception { if (writer == null) { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java index c519efd4172c..396c8db87e6b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java @@ -104,6 +104,9 @@ public void abort() throws Exception { aborted = true; super.abort(); } + + @Override + public void finishKey() throws Exception {} } // A mock ReadOperation fed to a MapTaskExecutor in test. @@ -312,6 +315,9 @@ public void start() throws Exception { Metrics.counter("TestMetric", "MetricCounter").inc(1L); } } + + @Override + public void finishKey() throws Exception {} }, new Operation(new OutputReceiver[] {}, context2) { @Override @@ -321,6 +327,9 @@ public void start() throws Exception { Metrics.counter("TestMetric", "MetricCounter").inc(2L); } } + + @Override + public void finishKey() throws Exception {} }, new Operation(new OutputReceiver[] {}, context3) { @Override @@ -330,6 +339,9 @@ public void start() throws Exception { Metrics.counter("TestMetric", "MetricCounter").inc(3L); } } + + @Override + public void finishKey() throws Exception {} }); try (IntrinsicMapTaskExecutor executor = diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index d8a1d1b90d47..ff82a1ab5c4c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java @@ -3657,8 +3657,8 @@ public void testActiveWorkFailure() throws Exception { server.waitForAndGetCommitsWithTimeout(1, Duration.standardSeconds(5)); assertEquals(1, commits.size()); - assertEquals(0, BlockingFn.teardownCounter.get()); - assertEquals(1, BlockingFn.setupCounter.get()); + assertEquals(1, BlockingFn.teardownCounter.get()); + assertEquals(2, BlockingFn.setupCounter.get()); worker.stop(); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 4bfa6efc8880..970a9a537c0d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -62,6 +62,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.config.FakeGlobalConfigHandle; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; +import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache; @@ -99,6 +100,7 @@ public class StreamingModeExecutionContextTest { @Rule public transient Timeout globalTimeout = Timeout.seconds(600); @Mock private SideInputStateFetcher sideInputStateFetcher; @Mock private WindmillStateReader stateReader; + @Mock private WorkExecutor workExecutor; private static final String COMPUTATION_ID = "computationId"; @@ -168,7 +170,8 @@ public void testTimerInternalsSetTimer() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder); + outputBuilder, + workExecutor); TimerInternals timerInternals = stepContext.timerInternals(); @@ -218,7 +221,8 @@ public void testTimerInternalsProcessingTimeSkew() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder); + outputBuilder, + workExecutor); TimerInternals timerInternals = stepContext.timerInternals(); assertTrue(timerTimestamp.isBefore(timerInternals.currentProcessingTime())); } @@ -427,7 +431,8 @@ public void testStateTagEncodingBasedOnConfig() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder); + outputBuilder, + workExecutor); assertEquals(expectedEncoding, executionContext.getWindmillTagEncoding().getClass()); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java index 61e2f4250d06..539c38eeb1da 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java @@ -19,6 +19,11 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import java.io.IOException; import java.util.ArrayList; @@ -40,8 +45,8 @@ public class WindmillReaderIteratorBaseTest { private static class TestWindmillReaderIterator extends WindmillReaderIteratorBase { protected TestWindmillReaderIterator( - Windmill.WorkItem work, ValueProvider skipUndecodableElements) { - super(work, skipUndecodableElements); + StreamingModeExecutionContext context, ValueProvider skipUndecodableElements) { + super(context, skipUndecodableElements); } @Override @@ -81,6 +86,51 @@ public void testSkipErrors() throws IOException { testForMessageBundleCounts(true, 0, 0, 1, 3, 0, 1, 0, 0, 0, 0); } + @Test + public void testWorkItemCancelledException() throws IOException { + StreamingModeExecutionContext mockContext = mock(StreamingModeExecutionContext.class); + when(mockContext.workIsFailed()).thenReturn(true); + Windmill.WorkItem workItem = + Windmill.WorkItem.newBuilder().setKey(ByteString.EMPTY).setWorkToken(0L).build(); + when(mockContext.getWorkItem()).thenReturn(workItem); + + try (TestWindmillReaderIterator iter = + new TestWindmillReaderIterator(mockContext, ValueProvider.StaticValueProvider.of(false))) { + iter.start(); + fail("Expected WorkItemCancelledException"); + } catch (WorkItemCancelledException e) { + // Expected + } + } + + @Test + public void testFinishKeyCalled() throws Exception { + StreamingModeExecutionContext mockContext = mock(StreamingModeExecutionContext.class); + when(mockContext.workIsFailed()).thenReturn(false); + Windmill.WorkItem workItem = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.EMPTY) + .setWorkToken(0L) + .addMessageBundles( + Windmill.InputMessageBundle.newBuilder() + .setSourceComputationId("foo") + .addMessages( + Windmill.Message.newBuilder() + .setTimestamp(0) + .setData(ByteString.EMPTY) + .build()) + .build()) + .build(); + when(mockContext.getWorkItem()).thenReturn(workItem); + + try (TestWindmillReaderIterator iter = + new TestWindmillReaderIterator(mockContext, ValueProvider.StaticValueProvider.of(false))) { + assertTrue(iter.start()); + assertFalse(iter.advance()); // This should trigger finishKey + verify(mockContext).finishKey(); + } + } + private void testForMessageBundleCounts(int... messageBundleCounts) throws IOException { testForMessageBundleCounts(false, messageBundleCounts); } @@ -111,9 +161,13 @@ private void testForMessageBundleCounts(boolean skipErrors, int... messageBundle .setWorkToken(0L) .addAllMessageBundles(bundles) .build(); + + StreamingModeExecutionContext mockContext = mock(StreamingModeExecutionContext.class); + when(mockContext.getWorkItem()).thenReturn(workItem); + try (TestWindmillReaderIterator iter = new TestWindmillReaderIterator( - workItem, ValueProvider.StaticValueProvider.of(skipErrors))) { + mockContext, ValueProvider.StaticValueProvider.of(skipErrors))) { List actual = ReaderTestUtils.windowedValuesToValues( ReaderUtils.readRemainingFromIterator(iter, false)); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java index ce8ad32f71aa..5dfd65b46123 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java @@ -97,6 +97,7 @@ import org.apache.beam.runners.dataflow.worker.testing.TestCountingSource; import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader.NativeReaderIterator; +import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache; @@ -643,7 +644,8 @@ public void testReadUnboundedReader() throws Exception { Watermarks.builder().setInputDataWatermark(new Instant(0)).build()), mock(WindmillStateReader.class), mock(SideInputStateFetcher.class), - Windmill.WorkItemCommitRequest.newBuilder()); + Windmill.WorkItemCommitRequest.newBuilder(), + mock(WorkExecutor.class)); @SuppressWarnings({"unchecked", "rawtypes"}) NativeReader>>> reader = @@ -1023,7 +1025,8 @@ public void testFailedWorkItemsAbort() throws Exception { dummyWork, mock(WindmillStateReader.class), mock(SideInputStateFetcher.class), - Windmill.WorkItemCommitRequest.newBuilder()); + Windmill.WorkItemCommitRequest.newBuilder(), + mock(WorkExecutor.class)); @SuppressWarnings({"unchecked", "rawtypes"}) NativeReader>>> reader = diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java index 2c35f4bf99db..d5e3b9c87139 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java @@ -59,6 +59,9 @@ private static OutputReceiver[] createOutputReceivers(int numOutputs, CounterSet } return receivers; } + + @Override + public void finishKey() throws Exception {} } /** A {@code Reader} that yields a specified set of values. */ diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java index 188466a50572..5d8f8eebb6f6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java @@ -100,6 +100,9 @@ public void abort() throws Exception { aborted = true; super.abort(); } + + @Override + public void finishKey() throws Exception {} } // A mock ReadOperation fed to a MapTaskExecutor in test. @@ -309,6 +312,9 @@ public void start() throws Exception { Metrics.counter("TestMetric", "MetricCounter").inc(1L); } } + + @Override + public void finishKey() throws Exception {} }, new Operation(new OutputReceiver[] {}, context2) { @Override @@ -318,6 +324,9 @@ public void start() throws Exception { Metrics.counter("TestMetric", "MetricCounter").inc(2L); } } + + @Override + public void finishKey() throws Exception {} }, new Operation(new OutputReceiver[] {}, context3) { @Override @@ -327,6 +336,9 @@ public void start() throws Exception { Metrics.counter("TestMetric", "MetricCounter").inc(3L); } } + + @Override + public void finishKey() throws Exception {} }); assertEquals(TimeUnit.MINUTES.toMillis(10), stateTracker.getNextBundleLullDurationReportMs()); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java index ba327f92cc44..5d058b1968cb 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java @@ -104,6 +104,7 @@ public void testRunParDoOperation() throws Exception { parDoOperation.process(""); parDoOperation.process("bob"); + parDoOperation.finishKey(); parDoOperation.finish(); parDoOperation.abort(); @@ -147,6 +148,7 @@ public void testParDoOperationContext() throws Exception { operation.start(); operation.process("hello"); + operation.finishKey(); operation.finish(); InOrder inOrder = From a8e4b0eef083d3a7125947f9bd89be4a9f35b11d Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 11 May 2026 05:48:08 +0000 Subject: [PATCH 02/12] Update empty work item iterator --- .../runners/dataflow/worker/WindowingWindmillReader.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java index 7483f2d138a2..7108b4b7992d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java @@ -156,6 +156,11 @@ public NativeReaderIterator>> iterator() throw return new NativeReaderIterator>>() { @Override public boolean start() throws IOException { + try { + context.finishKey(); + } catch (Exception e) { + throw new RuntimeException(e); + } return false; } From 986b90331df73c0b29f41b75bb3b3381d3c956f7 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 11 May 2026 06:28:43 +0000 Subject: [PATCH 03/12] Check if finishKey is called before flushState --- .../dataflow/worker/StreamingModeExecutionContext.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index 0255d5c577ec..b5ced4400848 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -159,6 +159,7 @@ public class StreamingModeExecutionContext extends DataflowExecutionContext activeReader; private @Nullable WorkExecutor workExecutor; + private boolean finishKeyCalled = false; public StreamingModeExecutionContext( CounterFactory counterFactory, @@ -248,6 +249,7 @@ public void start( this.key = key; this.work = work; this.workExecutor = workExecutor; + this.finishKeyCalled = false; this.computationKey = WindmillComputationKey.create(computationId, work.getShardedKey()); this.sideInputStateFetcher = sideInputStateFetcher; StreamingGlobalConfig config = globalConfigHandle.getConfig(); @@ -276,8 +278,10 @@ public void start( } public void finishKey() throws Exception { + checkState(!finishKeyCalled, "finishKey was already called"); checkNotNull(workExecutor, "workExecutor must be set before calling finishKey()"); workExecutor.finishKey(); + this.finishKeyCalled = true; } /** @@ -461,6 +465,7 @@ public void invalidateCache() { } public Map> flushState() { + checkState(finishKeyCalled, "finishKey must be called before flushState"); Map> callbacks = new HashMap<>(); for (StepContext stepContext : getAllStepContexts()) { From 107517b169c5c4aeb36a6376a03d33a96b8b2c41 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 11 May 2026 07:25:35 +0000 Subject: [PATCH 04/12] Fix tests --- .../dataflow/worker/StreamingModeExecutionContextTest.java | 3 ++- .../beam/runners/dataflow/worker/WorkerCustomSourcesTest.java | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 970a9a537c0d..850c0988ac8a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -154,7 +154,7 @@ COMPUTATION_ID, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.cla } @Test - public void testTimerInternalsSetTimer() { + public void testTimerInternalsSetTimer() throws Exception { Windmill.WorkItemCommitRequest.Builder outputBuilder = Windmill.WorkItemCommitRequest.newBuilder(); NameContext nameContext = NameContextsForTests.nameContextForTest(); @@ -182,6 +182,7 @@ public void testTimerInternalsSetTimer() { new Instant(5000), TimeDomain.EVENT_TIME, CausedByDrain.NORMAL)); + executionContext.finishKey(); executionContext.flushState(); Windmill.Timer timer = outputBuilder.buildPartial().getOutputTimers(0); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java index 5dfd65b46123..bbc8bed7187e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java @@ -680,6 +680,7 @@ public void testReadUnboundedReader() throws Exception { numReadOnThisIteration, lessThanOrEqualTo(debugOptions.getUnboundedReaderMaxElements())); // Extract and verify state modifications. + context.finishKey(); context.flushState(); state = context.getOutputBuilder().getSourceStateUpdates().getState(); // CountingSource's watermark is the last record + 1. i is now one past the last record, From 9e83f33a4a620d4c947e5a1459d8e7c819654d3c Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 11 May 2026 08:01:57 +0000 Subject: [PATCH 05/12] Call finishKey from WorkerCustomSources --- .../worker/StreamingModeExecutionContext.java | 8 ++++++-- .../dataflow/worker/WindmillReaderIteratorBase.java | 6 +----- .../dataflow/worker/WindowingWindmillReader.java | 12 ++---------- .../runners/dataflow/worker/WorkerCustomSources.java | 9 +++++++-- 4 files changed, 16 insertions(+), 19 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index b5ced4400848..1ff2c1bc4a1c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -277,10 +277,14 @@ public void start( } } - public void finishKey() throws Exception { + public void finishKey() { checkState(!finishKeyCalled, "finishKey was already called"); checkNotNull(workExecutor, "workExecutor must be set before calling finishKey()"); - workExecutor.finishKey(); + try { + workExecutor.finishKey(); + } catch (Exception e) { + throw new RuntimeException(e); + } this.finishKeyCalled = true; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java index b8232ca9f61a..075a1a8a4250 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBase.java @@ -63,11 +63,7 @@ public boolean advance() throws IOException { while (true) { if (bundleIndex >= work.getMessageBundlesCount()) { current = null; - try { - context.finishKey(); - } catch (Exception e) { - throw new RuntimeException(e); - } + context.finishKey(); return false; } Windmill.InputMessageBundle bundle = work.getMessageBundles(bundleIndex); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java index 7108b4b7992d..488684769bd9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java @@ -156,11 +156,7 @@ public NativeReaderIterator>> iterator() throw return new NativeReaderIterator>>() { @Override public boolean start() throws IOException { - try { - context.finishKey(); - } catch (Exception e) { - throw new RuntimeException(e); - } + context.finishKey(); return false; } @@ -187,11 +183,7 @@ public boolean start() throws IOException { @Override public boolean advance() throws IOException { current = null; - try { - context.finishKey(); - } catch (Exception e) { - throw new RuntimeException(e); - } + context.finishKey(); return false; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java index c00edffeaf95..29d5fb3561a1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java @@ -824,6 +824,7 @@ public boolean start() throws IOException { } try { if (!reader.start()) { + context.finishKey(); return false; } } catch (Exception e) { @@ -841,10 +842,13 @@ public boolean advance() throws IOException { // that there are regular checkpoints and that state does not become too large. BackOff backoff = backoffFactory.backoff(); while (true) { + if (context.workIsFailed()) { + throw new WorkItemCancelledException(context.getWorkItem().getShardingKey()); + } if (elemsRead >= maxElems || Instant.now().isAfter(endTime) - || context.isSinkFullHintSet() - || context.workIsFailed()) { + || context.isSinkFullHintSet()) { + context.finishKey(); return false; } try { @@ -857,6 +861,7 @@ public boolean advance() throws IOException { } long nextBackoff = backoff.nextBackOffMillis(); if (nextBackoff == BackOff.STOP) { + context.finishKey(); return false; } Uninterruptibles.sleepUninterruptibly(nextBackoff, TimeUnit.MILLISECONDS); From d5ef9e5909d2991d878d7b778deb0c01b1aba560 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 11 May 2026 19:39:02 +0000 Subject: [PATCH 06/12] Fix tests --- .../worker/WorkerCustomSourcesTest.java | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java index bbc8bed7187e..d5cf2948d928 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java @@ -680,7 +680,6 @@ public void testReadUnboundedReader() throws Exception { numReadOnThisIteration, lessThanOrEqualTo(debugOptions.getUnboundedReaderMaxElements())); // Extract and verify state modifications. - context.finishKey(); context.flushState(); state = context.getOutputBuilder().getSourceStateUpdates().getState(); // CountingSource's watermark is the last record + 1. i is now one past the last record, @@ -1042,14 +1041,19 @@ public void testFailedWorkItemsAbort() throws Exception { NativeReaderIterator>>> readerIterator = reader.iterator(); int numReads = 0; - while ((numReads == 0) ? readerIterator.start() : readerIterator.advance()) { - WindowedValue>> value = readerIterator.getCurrent(); - assertEquals(KV.of(0, numReads), value.getValue().getValue()); - numReads++; - // Fail the work item after reading two elements. - if (numReads == 2) { - dummyWork.setFailed(); + try { + while ((numReads == 0) ? readerIterator.start() : readerIterator.advance()) { + WindowedValue>> value = readerIterator.getCurrent(); + assertEquals(KV.of(0, numReads), value.getValue().getValue()); + numReads++; + // Fail the work item after reading two elements. + if (numReads == 2) { + dummyWork.setFailed(); + } } + fail("Expected WorkItemCancelledException"); + } catch (WorkItemCancelledException e) { + // Expected } assertThat(numReads, equalTo(2)); } From 5085977a4e6ef5e78f0630fd619be660ab541bcd Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Tue, 12 May 2026 05:20:10 +0000 Subject: [PATCH 07/12] run test --- ...beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json index 50d17c108f2e..e1b083e439cc 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 2, + "modification": 3, } From 35e621cbec92c4d0ebeaef8d3bde67ec3da9da68 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Wed, 27 May 2026 00:25:26 +0000 Subject: [PATCH 08/12] fix comments --- .../dataflow/worker/StreamingModeExecutionContext.java | 4 ++-- .../dataflow/worker/util/common/worker/ParDoOperation.java | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index 1ff2c1bc4a1c..f4d2a44dbd66 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -17,7 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.services.dataflow.model.CounterUpdate; @@ -279,7 +279,7 @@ public void start( public void finishKey() { checkState(!finishKeyCalled, "finishKey was already called"); - checkNotNull(workExecutor, "workExecutor must be set before calling finishKey()"); + checkStateNotNull(workExecutor, "workExecutor must be set before calling finishKey()"); try { workExecutor.finishKey(); } catch (Exception e) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java index 5aec82073366..68f5fbe688da 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java @@ -45,9 +45,9 @@ public void process(Object elem) throws Exception { } } - @Override // Batch mode does not use this method and instead relies on BatchModeUngroupingParDoFn // to process timers per key. + @Override public void finishKey() throws Exception { try (Closeable scope = context.enterProcessTimers()) { checkStarted(); From c397c9d5b720335b007fdf2bb89ccb0fec903279 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Wed, 27 May 2026 00:39:44 +0000 Subject: [PATCH 09/12] fix comments --- .../runners/dataflow/worker/StreamingModeExecutionContext.java | 1 + 1 file changed, 1 insertion(+) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index f4d2a44dbd66..874a9b07b939 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.dataflow.worker; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.services.dataflow.model.CounterUpdate; From 0d5384f26f493056c266e1d219e49ef9bef55c7e Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Wed, 27 May 2026 01:31:59 +0000 Subject: [PATCH 10/12] fix tests --- .../dataflow/worker/StreamingModeExecutionContextTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index bc15840599ce..7eda169d0228 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -455,7 +455,8 @@ public void testSetBacklogBytes() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder); + outputBuilder, + workExecutor); stepContext.setBacklogBytes(1234.0); executionContext.flushState(); From 5514964e42baea2a9b1d02cc72eb7855603b0491 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Wed, 27 May 2026 06:33:47 +0000 Subject: [PATCH 11/12] fix tests --- .../dataflow/worker/StreamingModeExecutionContextTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 7eda169d0228..216ca5386675 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -459,6 +459,7 @@ public void testSetBacklogBytes() { workExecutor); stepContext.setBacklogBytes(1234.0); + executionContext.finishKey(); executionContext.flushState(); assertEquals(1234, outputBuilder.getSourceBacklogBytes()); From 01c31a2f25852ec56fae55a190f0ae4514657eea Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Wed, 27 May 2026 06:34:23 +0000 Subject: [PATCH 12/12] fix tests --- .../beam_PostCommit_Java_ValidatesRunner_Dataflow.json | 2 +- ...beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json index 8144784f5f02..e9d869cc508c 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 2, + "modification": 1, } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json index e1b083e439cc..e623d3373a93 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 3, + "modification": 1, }