This is an automated email from the ASF dual-hosted git repository.
scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new d553b78ce2f [Dataflow Streaming] Add Operation::finishKey() and move
processTimers into it (#38430)
d553b78ce2f is described below
commit d553b78ce2fd2ed01abf71d21aa7f065dacc72df
Author: Arun Pandian <[email protected]>
AuthorDate: Wed May 27 01:17:07 2026 -0700
[Dataflow Streaming] Add Operation::finishKey() and move processTimers into
it (#38430)
* [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.
---
.../beam/runners/dataflow/worker/PubsubReader.java | 14 ++---
.../worker/StreamingModeExecutionContext.java | 22 +++++++-
.../dataflow/worker/UngroupedWindmillReader.java | 14 ++---
.../worker/WindmillReaderIteratorBase.java | 11 +++-
.../dataflow/worker/WindowingWindmillReader.java | 2 +
.../WorkerCustomSourceOperationExecutor.java | 3 ++
.../dataflow/worker/WorkerCustomSources.java | 9 +++-
.../worker/streaming/ComputationWorkExecutor.java | 2 +-
.../util/common/worker/FlattenOperation.java | 3 ++
.../worker/util/common/worker/MapTaskExecutor.java | 7 +++
.../worker/util/common/worker/Operation.java | 3 ++
.../worker/util/common/worker/ParDoOperation.java | 7 ++-
.../worker/util/common/worker/ReadOperation.java | 3 ++
.../worker/util/common/worker/WorkExecutor.java | 3 ++
.../worker/util/common/worker/WriteOperation.java | 3 ++
.../worker/IntrinsicMapTaskExecutorTest.java | 12 +++++
.../worker/StreamingDataflowWorkerTest.java | 4 +-
.../worker/StreamingModeExecutionContextTest.java | 18 +++++--
.../worker/WindmillReaderIteratorBaseTest.java | 60 ++++++++++++++++++++--
.../dataflow/worker/WorkerCustomSourcesTest.java | 26 ++++++----
.../util/common/worker/ExecutorTestUtils.java | 3 ++
.../util/common/worker/MapTaskExecutorTest.java | 12 +++++
.../util/common/worker/ParDoOperationTest.java | 2 +
23 files changed, 195 insertions(+), 48 deletions(-)
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 b60cb84415f..0c80909a0f3 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 @@ class PubsubReader<T> extends
NativeReader<WindowedValue<T>> {
@Override
public NativeReaderIterator<WindowedValue<T>> iterator() throws IOException {
- return new PubsubReaderIterator(context.getWorkItem());
+ return new PubsubReaderIterator();
}
class PubsubReaderIterator extends WindmillReaderIteratorBase<T> {
- 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 e1f1b21e135..25ce299adf7 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,6 +17,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;
@@ -56,6 +57,7 @@ import
org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalC
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 +159,9 @@ public class StreamingModeExecutionContext extends
DataflowExecutionContext<Step
*/
private @Nullable UnboundedReader<?> activeReader;
+ private @Nullable WorkExecutor workExecutor;
+ private boolean finishKeyCalled = false;
+
public StreamingModeExecutionContext(
CounterFactory counterFactory,
String computationId,
@@ -240,9 +245,12 @@ public class StreamingModeExecutionContext extends
DataflowExecutionContext<Step
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.finishKeyCalled = false;
this.computationKey = WindmillComputationKey.create(computationId,
work.getShardedKey());
this.sideInputStateFetcher = sideInputStateFetcher;
StreamingGlobalConfig config = globalConfigHandle.getConfig();
@@ -271,6 +279,17 @@ public class StreamingModeExecutionContext extends
DataflowExecutionContext<Step
}
}
+ public void finishKey() {
+ checkState(!finishKeyCalled, "finishKey was already called");
+ checkStateNotNull(workExecutor, "workExecutor must be set before calling
finishKey()");
+ try {
+ workExecutor.finishKey();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ this.finishKeyCalled = true;
+ }
+
/**
* Ensure that the processing time is greater than any fired processing time
timers. Otherwise, a
* trigger could ignore the timer and orphan the window.
@@ -452,6 +471,7 @@ public class StreamingModeExecutionContext extends
DataflowExecutionContext<Step
}
public Map<Long, Pair<Instant, Runnable>> flushState() {
+ checkState(finishKeyCalled, "finishKey must be called before flushState");
Map<Long, Pair<Instant, Runnable>> callbacks = new HashMap<>();
for (StepContext stepContext : getAllStepContexts()) {
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 1bcf8cf179c..bd68adeddfb 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
@@ -110,20 +110,12 @@ class UngroupedWindmillReader<T> extends
NativeReader<WindowedValue<T>> {
@Override
public NativeReaderIterator<WindowedValue<T>> iterator() throws IOException {
- return new UngroupedWindmillReaderIterator(context.getWorkItem());
+ return new UngroupedWindmillReaderIterator();
}
class UngroupedWindmillReaderIterator extends WindmillReaderIteratorBase<T> {
- 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 7e6508a4788..075a1a8a425 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 @@ import org.slf4j.LoggerFactory;
*/
public abstract class WindmillReaderIteratorBase<T>
extends NativeReader.NativeReaderIterator<WindowedValue<T>> {
+ 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<T>
private static final Logger LOG =
LoggerFactory.getLogger(WindmillReaderIteratorBase.class);
protected WindmillReaderIteratorBase(
- Windmill.WorkItem work, ValueProvider<Boolean> skipUndecodableElements) {
+ StreamingModeExecutionContext context, ValueProvider<Boolean>
skipUndecodableElements) {
+ this.context = context;
this.skipUndecodableElements = skipUndecodableElements;
- this.work = work;
+ this.work = context.getWorkItem();
}
@Override
@@ -54,9 +56,14 @@ public abstract class WindmillReaderIteratorBase<T>
@Override
public boolean advance() throws IOException {
+ if (context.workIsFailed()) {
+ throw new
WorkItemCancelledException(context.getWorkItem().getShardingKey());
+ }
+
while (true) {
if (bundleIndex >= work.getMessageBundlesCount()) {
current = null;
+ 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 173b254f639..488684769bd 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,7 @@ class WindowingWindmillReader<K, T> extends
NativeReader<WindowedValue<KeyedWork
return new NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>>() {
@Override
public boolean start() throws IOException {
+ context.finishKey();
return false;
}
@@ -182,6 +183,7 @@ class WindowingWindmillReader<K, T> extends
NativeReader<WindowedValue<KeyedWork
@Override
public boolean advance() throws IOException {
current = null;
+ context.finishKey();
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 31528e96e07..a1321d57ebb 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 class WorkerCustomSourceOperationExecutor implements
SourceOperationExecu
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/WorkerCustomSources.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java
index c00edffeaf9..29d5fb3561a 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 class WorkerCustomSources {
}
try {
if (!reader.start()) {
+ context.finishKey();
return false;
}
} catch (Exception e) {
@@ -841,10 +842,13 @@ public class WorkerCustomSources {
// 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 class WorkerCustomSources {
}
long nextBackoff = backoff.nextBackOffMillis();
if (nextBackoff == BackOff.STOP) {
+ context.finishKey();
return false;
}
Uninterruptibles.sleepUninterruptibly(nextBackoff,
TimeUnit.MILLISECONDS);
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 8dc681fc640..b4f3a22a7f5 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 abstract class ComputationWorkExecutor {
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 4847e9f2ea9..af1b2b9c48b 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 class FlattenOperation extends ReceivingOperation {
}
}
+ @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 58b95f286d5..3c33e190406 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 class MapTaskExecutor implements WorkExecutor {
// 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 b7b4e255cfa..b630da33cfa 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 abstract class Operation {
}
}
+ /** 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 27b6e9d1fb3..68f5fbe688d 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,13 +45,18 @@ public class ParDoOperation extends ReceivingOperation {
}
}
+ // Batch mode does not use this method and instead relies on
BatchModeUngroupingParDoFn
+ // to process timers per key.
@Override
- public void finish() throws Exception {
+ 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 d6b020483d4..fabc8d6af25 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 class ReadOperation extends Operation {
}
}
+ @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 b7170c80ced..1083fdbb9c4 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 673140d58d8..d28e7f3e5d3 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 class WriteOperation extends ReceivingOperation {
}
}
+ @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 c519efd4172..396c8db87e6 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 class IntrinsicMapTaskExecutorTest {
aborted = true;
super.abort();
}
+
+ @Override
+ public void finishKey() throws Exception {}
}
// A mock ReadOperation fed to a MapTaskExecutor in test.
@@ -312,6 +315,9 @@ public class IntrinsicMapTaskExecutorTest {
Metrics.counter("TestMetric", "MetricCounter").inc(1L);
}
}
+
+ @Override
+ public void finishKey() throws Exception {}
},
new Operation(new OutputReceiver[] {}, context2) {
@Override
@@ -321,6 +327,9 @@ public class IntrinsicMapTaskExecutorTest {
Metrics.counter("TestMetric", "MetricCounter").inc(2L);
}
}
+
+ @Override
+ public void finishKey() throws Exception {}
},
new Operation(new OutputReceiver[] {}, context3) {
@Override
@@ -330,6 +339,9 @@ public class IntrinsicMapTaskExecutorTest {
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 d8a1d1b90d4..ff82a1ab5c4 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 class StreamingDataflowWorkerTest {
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 a1c7609e5af..216ca538667 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.Work;
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";
@@ -152,7 +154,7 @@ public class StreamingModeExecutionContextTest {
}
@Test
- public void testTimerInternalsSetTimer() {
+ public void testTimerInternalsSetTimer() throws Exception {
Windmill.WorkItemCommitRequest.Builder outputBuilder =
Windmill.WorkItemCommitRequest.newBuilder();
NameContext nameContext = NameContextsForTests.nameContextForTest();
@@ -168,7 +170,8 @@ public class StreamingModeExecutionContextTest {
Watermarks.builder().setInputDataWatermark(new
Instant(1000)).build()),
stateReader,
sideInputStateFetcher,
- outputBuilder);
+ outputBuilder,
+ workExecutor);
TimerInternals timerInternals = stepContext.timerInternals();
@@ -179,6 +182,7 @@ public class StreamingModeExecutionContextTest {
new Instant(5000),
TimeDomain.EVENT_TIME,
CausedByDrain.NORMAL));
+ executionContext.finishKey();
executionContext.flushState();
Windmill.Timer timer = outputBuilder.buildPartial().getOutputTimers(0);
@@ -218,7 +222,8 @@ public class StreamingModeExecutionContextTest {
Watermarks.builder().setInputDataWatermark(new
Instant(1000)).build()),
stateReader,
sideInputStateFetcher,
- outputBuilder);
+ outputBuilder,
+ workExecutor);
TimerInternals timerInternals = stepContext.timerInternals();
assertTrue(timerTimestamp.isBefore(timerInternals.currentProcessingTime()));
}
@@ -427,7 +432,8 @@ public class StreamingModeExecutionContextTest {
Watermarks.builder().setInputDataWatermark(new
Instant(1000)).build()),
stateReader,
sideInputStateFetcher,
- outputBuilder);
+ outputBuilder,
+ workExecutor);
assertEquals(expectedEncoding,
executionContext.getWindmillTagEncoding().getClass());
}
}
@@ -449,9 +455,11 @@ public class StreamingModeExecutionContextTest {
Watermarks.builder().setInputDataWatermark(new
Instant(1000)).build()),
stateReader,
sideInputStateFetcher,
- outputBuilder);
+ outputBuilder,
+ workExecutor);
stepContext.setBacklogBytes(1234.0);
+ executionContext.finishKey();
executionContext.flushState();
assertEquals(1234, outputBuilder.getSourceBacklogBytes());
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 61e2f4250d0..539c38eeb1d 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 @@ package org.apache.beam.runners.dataflow.worker;
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 @@ import org.junit.runners.JUnit4;
public class WindmillReaderIteratorBaseTest {
private static class TestWindmillReaderIterator extends
WindmillReaderIteratorBase<Long> {
protected TestWindmillReaderIterator(
- Windmill.WorkItem work, ValueProvider<Boolean>
skipUndecodableElements) {
- super(work, skipUndecodableElements);
+ StreamingModeExecutionContext context, ValueProvider<Boolean>
skipUndecodableElements) {
+ super(context, skipUndecodableElements);
}
@Override
@@ -81,6 +86,51 @@ public class WindmillReaderIteratorBaseTest {
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 @@ public class WindmillReaderIteratorBaseTest {
.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<Long> 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 ce8ad32f71a..d5cf2948d92 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.streaming.sideinput.SideInputStat
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 class WorkerCustomSourcesTest {
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<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>>
reader =
@@ -1023,7 +1025,8 @@ public class WorkerCustomSourcesTest {
dummyWork,
mock(WindmillStateReader.class),
mock(SideInputStateFetcher.class),
- Windmill.WorkItemCommitRequest.newBuilder());
+ Windmill.WorkItemCommitRequest.newBuilder(),
+ mock(WorkExecutor.class));
@SuppressWarnings({"unchecked", "rawtypes"})
NativeReader<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>>
reader =
@@ -1038,14 +1041,19 @@ public class WorkerCustomSourcesTest {
NativeReaderIterator<WindowedValue<ValueWithRecordId<KV<Integer,
Integer>>>> readerIterator =
reader.iterator();
int numReads = 0;
- while ((numReads == 0) ? readerIterator.start() :
readerIterator.advance()) {
- WindowedValue<ValueWithRecordId<KV<Integer, Integer>>> 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<ValueWithRecordId<KV<Integer, Integer>>> 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));
}
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 2c35f4bf99d..d5e3b9c8713 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 @@ public class ExecutorTestUtils {
}
return receivers;
}
+
+ @Override
+ public void finishKey() throws Exception {}
}
/** A {@code Reader<String>} 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 188466a5057..5d8f8eebb6f 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 class MapTaskExecutorTest {
aborted = true;
super.abort();
}
+
+ @Override
+ public void finishKey() throws Exception {}
}
// A mock ReadOperation fed to a MapTaskExecutor in test.
@@ -309,6 +312,9 @@ public class MapTaskExecutorTest {
Metrics.counter("TestMetric", "MetricCounter").inc(1L);
}
}
+
+ @Override
+ public void finishKey() throws Exception {}
},
new Operation(new OutputReceiver[] {}, context2) {
@Override
@@ -318,6 +324,9 @@ public class MapTaskExecutorTest {
Metrics.counter("TestMetric", "MetricCounter").inc(2L);
}
}
+
+ @Override
+ public void finishKey() throws Exception {}
},
new Operation(new OutputReceiver[] {}, context3) {
@Override
@@ -327,6 +336,9 @@ public class MapTaskExecutorTest {
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 ba327f92cc4..5d058b1968c 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 class ParDoOperationTest {
parDoOperation.process("");
parDoOperation.process("bob");
+ parDoOperation.finishKey();
parDoOperation.finish();
parDoOperation.abort();
@@ -147,6 +148,7 @@ public class ParDoOperationTest {
operation.start();
operation.process("hello");
+ operation.finishKey();
operation.finish();
InOrder inOrder =
