This is an automated email from the ASF dual-hosted git repository.
Abacn 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 dd896e2b239 [Dataflow Streaming] [Multi Key] Drop failed work in
BoundedQueueExecutor::pollWork (#38920)
dd896e2b239 is described below
commit dd896e2b2395937201b7721251d9f66936f9340c
Author: Arun Pandian <[email protected]>
AuthorDate: Fri Aug 14 17:37:18 2026 -0700
[Dataflow Streaming] [Multi Key] Drop failed work in
BoundedQueueExecutor::pollWork (#38920)
* Drop failed work in BoundedQueueExecutor::pollWork
* address comment
* Revert "address comment"
This reverts commit 3670a036aa83a3c18e18f3b9435214f3cbb0ad13.
* Revert "Drop failed work in BoundedQueueExecutor::pollWork"
This reverts commit 2c1278636a603fb5e06a5f5bf6f179b0fbf41d7c.
* [Dataflow Streaming] Remove finalizeCommits from processWork
* Plumb ComputationState to ProcessingContext
* Drop failed workitems during pollwork
* Improve tests
* address comments
* address comments
* address comments
---
.../worker/StreamingModeExecutionContext.java | 17 ++-
.../worker/streaming/ComputationWorkExecutor.java | 6 +-
.../worker/streaming/FailedWorkHandler.java | 24 ++++
.../dataflow/worker/util/BoundedQueueExecutor.java | 26 +++-
.../work/processing/StreamingWorkScheduler.java | 22 +--
.../processing/failures/WorkFailureProcessor.java | 6 +-
.../worker/StreamingDataflowWorkerTest.java | 154 ++++++++++++++++++++-
.../worker/StreamingModeExecutionContextTest.java | 95 +++++++++++--
.../dataflow/worker/WorkerCustomSourcesTest.java | 9 +-
.../worker/util/BoundedQueueExecutorTest.java | 129 ++++++++++++++++-
10 files changed, 442 insertions(+), 46 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 68dbd61f15f..365ebbdc1f9 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
@@ -53,6 +53,7 @@ import
org.apache.beam.runners.dataflow.worker.counters.NameContext;
import
org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
import
org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
@@ -172,10 +173,7 @@ public class StreamingModeExecutionContext
private @Nullable WorkExecutor workExecutor;
private boolean finishKeyCalled = false;
- @SuppressWarnings("UnusedVariable")
private @Nullable BoundedQueueExecutor workQueueExecutor;
-
- @SuppressWarnings("UnusedVariable")
private @Nullable BoundedQueueExecutorWorkHandle budgetHandle;
private final HotKeyLogger hotKeyLogger;
@@ -192,8 +190,8 @@ public class StreamingModeExecutionContext
void onKeyTransition(@Nullable Work oldWork, Work newWork);
}
- @SuppressWarnings("UnusedVariable")
private @Nullable KeyTransitionListener keyTransitionListener;
+ private @Nullable FailedWorkHandler onFailedWorkHandler;
private List<Work> executedWorks = Collections.emptyList();
private List<Windmill.WorkItemCommitRequest.Builder> outputBuilders =
Collections.emptyList();
@@ -335,6 +333,7 @@ public class StreamingModeExecutionContext
this.workQueueExecutor = null;
this.budgetHandle = null;
this.keyTransitionListener = null;
+ this.onFailedWorkHandler = null;
this.work = null;
this.key = null;
this.outputBuilder = null;
@@ -350,7 +349,8 @@ public class StreamingModeExecutionContext
BoundedQueueExecutor workQueueExecutor,
BoundedQueueExecutorWorkHandle budgetHandle,
@Nullable Coder<?> keyCoder,
- KeyTransitionListener keyTransitionListener)
+ KeyTransitionListener keyTransitionListener,
+ FailedWorkHandler onFailedWorkHandler)
throws CoderException {
reset();
this.executedWorks = new ArrayList<>();
@@ -361,6 +361,7 @@ public class StreamingModeExecutionContext
this.workQueueExecutor = workQueueExecutor;
this.budgetHandle = budgetHandle;
this.keyTransitionListener = keyTransitionListener;
+ this.onFailedWorkHandler = checkStateNotNull(onFailedWorkHandler);
this.workItemsPolled = 1;
this.bundleStartTimeNanos = System.nanoTime();
@@ -779,7 +780,11 @@ public class StreamingModeExecutionContext
@Nullable
ExecutableWork additionalWork =
- executor.pollWork(computationId, activeWork.getKeyGroup(), handle);
+ executor.pollWork(
+ computationId,
+ activeWork.getKeyGroup(),
+ handle,
+ checkStateNotNull(onFailedWorkHandler));
if (additionalWork != null) {
flushStateInternal();
Work newWork = additionalWork.work();
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 9391b842f03..dabf72ba4ea 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
@@ -65,7 +65,8 @@ public abstract class ComputationWorkExecutor {
Work work,
BoundedQueueExecutor workQueueExecutor,
BoundedQueueExecutorWorkHandle budgetHandle,
- KeyTransitionListener keyTransitionListener)
+ KeyTransitionListener keyTransitionListener,
+ FailedWorkHandler onFailedWorkHandler)
throws Exception {
context()
.start(
@@ -74,7 +75,8 @@ public abstract class ComputationWorkExecutor {
workQueueExecutor,
budgetHandle,
keyCoder().orElse(null),
- keyTransitionListener);
+ keyTransitionListener,
+ onFailedWorkHandler);
workExecutor().execute();
return context();
}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/FailedWorkHandler.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/FailedWorkHandler.java
new file mode 100644
index 00000000000..683ecf5d660
--- /dev/null
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/FailedWorkHandler.java
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming;
+
+/** Handler for failed {@link Work}. */
+@FunctionalInterface
+public interface FailedWorkHandler {
+ void onFailedWork(Work work);
+}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java
index 9eb9a37b1b7..2dd0f971168 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java
@@ -18,6 +18,7 @@
package org.apache.beam.runners.dataflow.worker.util;
import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+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.checkArgument;
import java.util.ArrayList;
@@ -31,6 +32,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.concurrent.GuardedBy;
import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
@@ -387,20 +389,32 @@ public class BoundedQueueExecutor {
}
public @Nullable ExecutableWork pollWork(
- String computationId, Work.KeyGroup keyGroup,
BoundedQueueExecutorWorkHandle handle) {
+ String computationId,
+ Work.KeyGroup keyGroup,
+ BoundedQueueExecutorWorkHandle handle,
+ FailedWorkHandler onFailedWorkHandler) {
checkArgument(
computationId != null && keyGroup != null &&
!keyGroup.equals(Work.KeyGroup.DEFAULT));
checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl);
+ checkStateNotNull(onFailedWorkHandler);
BoundedQueueExecutorWorkHandleImpl internalHandle =
(BoundedQueueExecutorWorkHandleImpl) handle;
if (keyGroupWorkQueue == null) {
return null;
}
- @Nullable QueuedWork queuedWork =
keyGroupWorkQueue.pollWork(computationId, keyGroup);
- if (queuedWork == null) {
- return null;
+ while (true) {
+ @Nullable QueuedWork queuedWork =
keyGroupWorkQueue.pollWork(computationId, keyGroup);
+ if (queuedWork == null) {
+ return null;
+ }
+ Work work = queuedWork.getWork().work();
+ if (work.isFailed()) {
+ queuedWork.getHandle().close();
+ onFailedWorkHandler.onFailedWork(work);
+ continue;
+ }
+ internalHandle.merge(queuedWork.getHandle());
+ return queuedWork.getWork();
}
- internalHandle.merge(queuedWork.getHandle());
- return queuedWork.getWork();
}
private void decrementCounters(int elements, long bytes) {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
index 27952569e0e..63cfad5a9a6 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
@@ -45,6 +45,7 @@ import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWor
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
import
org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
@@ -321,8 +322,11 @@ public class StreamingWorkScheduler {
try {
StreamingModeExecutionContext context =
computationWorkExecutor.context();
+ FailedWorkHandler onFailedWorkHandler =
getFailedWorkHandler(computationState);
+
// Blocks while executing work.
- computationWorkExecutor.executeWork(work, workExecutor, handle,
keyTransitionListener);
+ computationWorkExecutor.executeWork(
+ work, workExecutor, handle, keyTransitionListener,
onFailedWorkHandler);
List<Work> workBatch;
List<Windmill.WorkItemCommitRequest> workItemCommits;
@@ -465,14 +469,10 @@ public class StreamingWorkScheduler {
ExecutableWork.create(w, (retry, h) ->
processWork(computationState, retry, h)));
}
+ FailedWorkHandler onFailedWorkHandler =
getFailedWorkHandler(computationState);
+
workFailureProcessor.logAndProcessFailureBatch(
- computationId,
- systemName,
- executableWorks,
- t,
- invalidWork ->
- computationState.completeWorkAndScheduleNextWorkForKey(
- invalidWork.getShardedKey(), invalidWork.id()));
+ computationId, systemName, executableWorks, t, onFailedWorkHandler);
} catch (OutOfMemoryError oom) {
throw oom;
} catch (Throwable t2) {
@@ -481,6 +481,12 @@ public class StreamingWorkScheduler {
}
}
+ private static FailedWorkHandler getFailedWorkHandler(ComputationState
computationState) {
+ return failedWork ->
+ computationState.completeWorkAndScheduleNextWorkForKey(
+ failedWork.getShardedKey(), failedWork.id());
+ }
+
private void recordProcessingTime(
StageInfo stageInfo, List<Work> workBatch, long
processingStartTimeNanos) {
long processingTimeMsecs =
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java
index de33d3d3961..d23c870178e 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java
@@ -19,12 +19,12 @@ package
org.apache.beam.runners.dataflow.worker.windmill.work.processing.failure
import java.util.List;
import java.util.concurrent.TimeUnit;
-import java.util.function.Consumer;
import java.util.function.Supplier;
import javax.annotation.Nullable;
import javax.annotation.concurrent.ThreadSafe;
import
org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
import org.apache.beam.sdk.annotations.Internal;
@@ -103,7 +103,7 @@ public final class WorkFailureProcessor {
String systemName,
List<ExecutableWork> executableWorks,
Throwable t,
- Consumer<Work> onInvalidWork)
+ FailedWorkHandler onFailedWorkHandler)
throws Throwable {
List<ExecutableWork> worksToRetryLocally = new java.util.ArrayList<>();
@@ -112,7 +112,7 @@ public final class WorkFailureProcessor {
case DO_NOT_RETRY:
// Consider the item invalid. It will eventually be retried by
Windmill if it still needs
// to be processed.
- onInvalidWork.accept(executableWork.work());
+ onFailedWorkHandler.onFailedWork(executableWork.work());
break;
case RETRY_LOCALLY:
// Try again after some delay and at the end of the queue to avoid a
tight loop.
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 2d35da51a79..c48b30ecf64 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
@@ -1522,13 +1522,146 @@ public class StreamingDataflowWorkerTest {
worker.stop();
}
+ @Test
+ public void
testMultiKeyCommit_queuedWorkItemFailsAndSubsequentWorkItemPickedUp()
+ throws Exception {
+ if (!streamingEngine) {
+ return;
+ }
+ BlockingKvDoFn.reset();
+ StreamingDataflowWorker worker = makeMultiKeyEnabledWorker(new
BlockingKvDoFn());
+ worker.start();
+
+ String batchInputText1 =
+ "work {"
+ + " computation_id: \""
+ + DEFAULT_COMPUTATION_ID
+ + "\""
+ + " input_data_watermark: 0"
+ + " work {"
+ + " key: \"key1\""
+ + " sharding_key: 1"
+ + " work_token: 1"
+ + " cache_token: 2"
+ + " key_group { high: 0 low: 1 }"
+ + " message_bundles {"
+ + " source_computation_id: \""
+ + DEFAULT_SOURCE_COMPUTATION_ID
+ + "\""
+ + " messages {"
+ + " timestamp: 0"
+ + " data: \"data1\""
+ + " }"
+ + " }"
+ + " }"
+ + " work {"
+ + " key: \"key2\""
+ + " sharding_key: 2"
+ + " work_token: 2"
+ + " cache_token: 3"
+ + " key_group { high: 0 low: 1 }"
+ + " message_bundles {"
+ + " source_computation_id: \""
+ + DEFAULT_SOURCE_COMPUTATION_ID
+ + "\""
+ + " messages {"
+ + " timestamp: 0"
+ + " data: \"data2\""
+ + " }"
+ + " }"
+ + " }"
+ + "}";
+
+ String batchInputText2 =
+ "work {"
+ + " computation_id: \""
+ + DEFAULT_COMPUTATION_ID
+ + "\""
+ + " input_data_watermark: 0"
+ + " work {"
+ + " key: \"key2\""
+ + " sharding_key: 2"
+ + " work_token: 3"
+ + " cache_token: 4"
+ + " key_group { high: 0 low: 1 }"
+ + " message_bundles {"
+ + " source_computation_id: \""
+ + DEFAULT_SOURCE_COMPUTATION_ID
+ + "\""
+ + " messages {"
+ + " timestamp: 0"
+ + " data: \"data3\""
+ + " }"
+ + " }"
+ + " }"
+ + "}";
+ Windmill.GetWorkResponse batchInput1 =
+ buildInput(
+ batchInputText1,
+ CoderUtils.encodeToByteArray(
+ CollectionCoder.of(IntervalWindow.getCoder()),
+ Collections.singletonList(DEFAULT_WINDOW)));
+ Windmill.GetWorkResponse batchInput2 =
+ buildInput(
+ batchInputText2,
+ CoderUtils.encodeToByteArray(
+ CollectionCoder.of(IntervalWindow.getCoder()),
+ Collections.singletonList(DEFAULT_WINDOW)));
+
+
server.whenGetDataCalled().answerByDefault(StreamingDataflowWorkerTest::emptyDataResponder);
+
+ server.whenGetWorkCalled().thenReturn(batchInput1).thenReturn(batchInput2);
+ server.waitForEmptyWorkQueue();
+
+ // Wait for key1 to start processing and block on BlockingKvDoFn.
+ BlockingKvDoFn.counter.get().acquire(1);
+
+ // Fail key2 (work token 2) via failed heartbeat while key1 is still
processing.
+ ComputationHeartbeatResponse.Builder failedHeartbeat =
+ ComputationHeartbeatResponse.newBuilder();
+ failedHeartbeat
+ .setComputationId(DEFAULT_COMPUTATION_ID)
+ .addHeartbeatResponsesBuilder()
+ .setCacheToken(3)
+ .setWorkToken(2)
+ .setShardingKey(2)
+ .setFailed(true);
+
+ // Fake server processes heartbeat responses are processed synchronously
in sendFailedHeartbeats
+
server.sendFailedHeartbeats(Collections.singletonList(failedHeartbeat.build()));
+
+ // Unblock key1 to allow bundle to poll key2 (token 2 -> failed, skipped)
and key2 (token 3).
+ BlockingKvDoFn.blocker.get().countDown();
+
+ Map<Long, Windmill.WorkItemCommitRequest> result =
server.waitForAndGetCommits(2);
+
+ assertTrue(result.containsKey(1L));
+ assertTrue(result.containsKey(3L));
+ assertFalse(result.containsKey(2L));
+
+ List<Windmill.MultiKeyWorkItemCommitRequest> multiKeyCommits =
+ server.getMultiKeyCommitsReceived();
+ assertEquals(1, multiKeyCommits.size());
+ Windmill.MultiKeyWorkItemCommitRequest multiKeyCommit =
multiKeyCommits.get(0);
+ assertEquals(2, multiKeyCommit.getRequestsCount());
+ assertEquals(1, multiKeyCommit.getRequests(0).getWorkToken());
+ assertEquals(3, multiKeyCommit.getRequests(1).getWorkToken());
+
+ worker.stop();
+ }
+
private StreamingDataflowWorker makeMultiKeyEnabledWorker() {
+ return makeMultiKeyEnabledWorker(new WorkDoFn());
+ }
+
+ private StreamingDataflowWorker makeMultiKeyEnabledWorker(
+ DoFn<KV<String, String>, KV<String, String>> doFn) {
KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(),
StringUtf8Coder.of());
List<ParallelInstruction> instructions =
Arrays.asList(
makeSourceInstruction(kvCoder),
- makeDoFnInstruction(new WorkDoFn(), 0, kvCoder),
+ makeDoFnInstruction(doFn, 0, kvCoder),
makeSinkInstruction(kvCoder, 1));
StreamingDataflowWorker worker =
@@ -4825,6 +4958,25 @@ public class StreamingDataflowWorkerTest {
}
}
+ static class BlockingKvDoFn extends DoFn<KV<String, String>, KV<String,
String>> {
+ public static final AtomicReference<CountDownLatch> blocker =
+ new AtomicReference<>(new CountDownLatch(1));
+ public static final AtomicReference<Semaphore> counter =
+ new AtomicReference<>(new Semaphore(0));
+
+ @ProcessElement
+ public void processElement(ProcessContext c) throws InterruptedException {
+ counter.get().release();
+ blocker.get().await();
+ c.output(c.element());
+ }
+
+ public static void reset() {
+ blocker.set(new CountDownLatch(1));
+ counter.set(new Semaphore(0));
+ }
+ }
+
static class LargeCommitFn extends DoFn<KV<String, String>, KV<String,
String>> {
@ProcessElement
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 eb6bb51e420..5aceb0ca956 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
@@ -54,6 +54,7 @@ import
org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState
import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions;
import
org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowExecutionStateTracker;
import
org.apache.beam.runners.dataflow.worker.MetricsToCounterUpdateConverter.Kind;
+import
org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.KeyTransitionListener;
import
org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionState;
import
org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionStateRegistry;
import org.apache.beam.runners.dataflow.worker.counters.CounterSet;
@@ -62,6 +63,7 @@ import
org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.NoopProfi
import
org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
import
org.apache.beam.runners.dataflow.worker.streaming.config.FakeGlobalConfigHandle;
@@ -96,6 +98,7 @@ import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.joda.time.Instant;
+import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -109,6 +112,15 @@ import org.mockito.MockitoAnnotations;
@RunWith(JUnit4.class)
public class StreamingModeExecutionContextTest {
+ private static final FailedWorkHandler FAILING_FAILED_WORK_HANDLER =
+ ignored -> {
+ Assert.fail();
+ };
+
+ private static final KeyTransitionListener FAILING_KEY_TRANSISITON =
+ (oldWork, newWork) -> {
+ Assert.fail();
+ };
@Rule public transient Timeout globalTimeout = Timeout.seconds(600);
@Mock private WorkExecutor workExecutor;
@@ -198,10 +210,11 @@ public class StreamingModeExecutionContextTest {
context.start(
work,
workExecutor,
- /* workQueueExecutor= */ null,
- /* budgetHandle= */ null,
+ /* workQueueExecutor= */ mock(BoundedQueueExecutor.class),
+ /* budgetHandle= */ mock(BoundedQueueExecutorWorkHandle.class),
keyCoder,
- /* keyTransitionListener= */ (k, c) -> {});
+ FAILING_KEY_TRANSISITON,
+ /* onFailedWorkHandler= */ FAILING_FAILED_WORK_HANDLER);
} catch (CoderException e) {
throw new RuntimeException(e);
}
@@ -546,13 +559,20 @@ public class StreamingModeExecutionContextTest {
workItem2,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
ExecutableWork executableWork2 = ExecutableWork.create(work2, (w, h) ->
{});
- when(mockExecutor.pollWork(eq(COMPUTATION_ID), eq(work1.getKeyGroup()),
eq(mockHandle)))
+ when(mockExecutor.pollWork(eq(COMPUTATION_ID), eq(work1.getKeyGroup()),
eq(mockHandle), any()))
.thenReturn(executableWork2)
.thenReturn(null);
StreamingModeExecutionContext.KeyTransitionListener mockListener =
mock(StreamingModeExecutionContext.KeyTransitionListener.class);
- executionContext.start(work1, workExecutor, mockExecutor, mockHandle,
null, mockListener);
+ executionContext.start(
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ mockListener,
+ FAILING_FAILED_WORK_HANDLER);
assertTrue(executionContext.advance());
assertEquals("key2", executionContext.getSerializedKey().toStringUtf8());
@@ -577,11 +597,17 @@ public class StreamingModeExecutionContextTest {
createMockWork(
workItem1,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
- when(mockExecutor.pollWork(eq(COMPUTATION_ID), eq(work1.getKeyGroup()),
eq(mockHandle)))
+ when(mockExecutor.pollWork(eq(COMPUTATION_ID), eq(work1.getKeyGroup()),
eq(mockHandle), any()))
.thenReturn(null);
executionContext.start(
- work1, workExecutor, mockExecutor, mockHandle, null, (oldWork,
newWork) -> {});
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ FAILING_KEY_TRANSISITON,
+ FAILING_FAILED_WORK_HANDLER);
assertFalse(executionContext.advance());
}
@@ -611,7 +637,14 @@ public class StreamingModeExecutionContextTest {
createMockWork(
workItem1,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
- context.start(work1, workExecutor, mockExecutor, mockHandle, null,
(oldWork, newWork) -> {});
+ context.start(
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ FAILING_KEY_TRANSISITON,
+ FAILING_FAILED_WORK_HANDLER);
assertFalse(context.advance());
verifyNoInteractions(mockExecutor);
@@ -642,7 +675,14 @@ public class StreamingModeExecutionContextTest {
createMockWork(
workItem1,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
- context.start(work1, workExecutor, mockExecutor, mockHandle, null,
(oldWork, newWork) -> {});
+ context.start(
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ FAILING_KEY_TRANSISITON,
+ FAILING_FAILED_WORK_HANDLER);
assertFalse(context.advance());
verifyNoInteractions(mockExecutor);
@@ -666,7 +706,13 @@ public class StreamingModeExecutionContextTest {
workItem1,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
executionContext.start(
- work1, workExecutor, mockExecutor, mockHandle, null, (oldWork,
newWork) -> {});
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ FAILING_KEY_TRANSISITON,
+ FAILING_FAILED_WORK_HANDLER);
work1.setFailed();
@@ -689,7 +735,13 @@ public class StreamingModeExecutionContextTest {
workItem1,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
executionContext.start(
- work1, workExecutor, mockExecutor, mockHandle, null, (oldWork,
newWork) -> {});
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ FAILING_KEY_TRANSISITON,
+ FAILING_FAILED_WORK_HANDLER);
assertFalse(executionContext.advance());
verifyNoInteractions(mockExecutor);
@@ -717,7 +769,14 @@ public class StreamingModeExecutionContextTest {
createMockWork(
workItem1,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
- context.start(work1, workExecutor, mockExecutor, mockHandle, null,
(oldWork, newWork) -> {});
+ context.start(
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ FAILING_KEY_TRANSISITON,
+ FAILING_FAILED_WORK_HANDLER);
assertFalse(context.advance());
verifyNoInteractions(mockExecutor);
@@ -750,11 +809,19 @@ public class StreamingModeExecutionContextTest {
createMockWork(
workItem1,
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build());
- context.start(work1, workExecutor, mockExecutor, mockHandle, null,
(oldWork, newWork) -> {});
+ context.start(
+ work1,
+ workExecutor,
+ mockExecutor,
+ mockHandle,
+ null,
+ FAILING_KEY_TRANSISITON,
+ FAILING_FAILED_WORK_HANDLER);
context.reportBytesSinked(50);
assertFalse(context.advance());
- verify(mockExecutor).pollWork(COMPUTATION_ID, work1.getKeyGroup(),
mockHandle);
+ verify(mockExecutor)
+ .pollWork(eq(COMPUTATION_ID), eq(work1.getKeyGroup()), eq(mockHandle),
any());
reset(mockExecutor);
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 679227a11dc..2a741da7fd8 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
@@ -89,6 +89,7 @@ import
org.apache.beam.runners.dataflow.worker.WorkerCustomSources.SplittableOnl
import org.apache.beam.runners.dataflow.worker.counters.CounterSet;
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
import
org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.NoopProfileScope;
+import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
import
org.apache.beam.runners.dataflow.worker.streaming.config.FixedGlobalConfigHandle;
@@ -97,6 +98,7 @@ import
org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalC
import
org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters;
import
org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory;
import org.apache.beam.runners.dataflow.worker.testing.TestCountingSource;
+import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
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;
@@ -218,10 +220,11 @@ public class WorkerCustomSourcesTest {
context.start(
work,
mock(WorkExecutor.class),
- /* workQueueExecutor= */ null,
- /* budgetHandle= */ null,
+ /* workQueueExecutor= */ mock(BoundedQueueExecutor.class),
+ /* budgetHandle= */ mock(BoundedQueueExecutorWorkHandle.class),
/* keyCoder= */ null,
- /* keyTransitionListener= */ mock(KeyTransitionListener.class));
+ /* keyTransitionListener= */ mock(KeyTransitionListener.class),
+ /* onFailedWorkHandler= */ ignored -> {});
} catch (CoderException e) {
throw new RuntimeException(e);
}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java
index 0e75fa01f4f..75230594fc1 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java
@@ -25,6 +25,7 @@ import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
import java.util.Arrays;
import java.util.Collection;
@@ -35,6 +36,7 @@ import java.util.function.BiConsumer;
import java.util.function.Consumer;
import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
import
org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor.BoundedQueueExecutorWorkHandleImpl;
@@ -46,6 +48,7 @@ import
org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.joda.time.Instant;
+import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -67,6 +70,10 @@ public class BoundedQueueExecutorTest {
return Arrays.asList(new Object[][] {{false}, {true}});
}
+ private static final FailedWorkHandler FAILING_FAILED_WORK_HANDLER =
+ ignored -> {
+ Assert.fail();
+ };
private static final long MAXIMUM_BYTES_OUTSTANDING = 10000000;
private static final int DEFAULT_MAX_THREADS = 2;
private static final int DEFAULT_THREAD_EXPIRATION_SEC = 60;
@@ -90,6 +97,33 @@ public class BoundedQueueExecutorTest {
computationId, keyGroup, (work, handle) -> executeWorkFn.accept(work));
}
+ private static ExecutableWork createWorkWithCompIdAndKeyGroupAndWorkToken(
+ String computationId, Work.KeyGroup keyGroup, long workToken,
Consumer<Work> executeWorkFn) {
+ WorkItem workItem =
+ WorkItem.newBuilder()
+ .setKey(ByteString.EMPTY)
+ .setShardingKey(1)
+ .setWorkToken(workToken)
+ .setCacheToken(1)
+ .setKeyGroup(
+ Windmill.Uint128Proto.newBuilder()
+ .setHigh(keyGroup.high())
+ .setLow(keyGroup.low())
+ .build())
+ .build();
+ return ExecutableWork.create(
+ Work.create(
+ workItem,
+ workItem.getSerializedSize(),
+ Watermarks.builder().setInputDataWatermark(Instant.now()).build(),
+ Work.createProcessingContext(
+ computationId, new FakeGetDataClient(), ignored -> {},
mock(HeartbeatSender.class)),
+ false,
+ Instant::now,
+ ImmutableList.of()),
+ (work, handle) -> executeWorkFn.accept(work));
+ }
+
private static ExecutableWork createWorkWithHandle(
String computationId,
Work.KeyGroup keyGroup,
@@ -511,7 +545,7 @@ public class BoundedQueueExecutorTest {
assertEquals(3, testExecutor.elementsOutstanding());
// Steal work2 using pollWork with compA and keyGroup2
- ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup2,
stealHandle);
+ ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup2,
stealHandle, ignored -> {});
assertNotNull(stolen);
assertEquals(work2, stolen);
@@ -520,7 +554,8 @@ public class BoundedQueueExecutorTest {
targetStart.await();
// Steal work1 using pollWork with compA and keyGroup1
- ExecutableWork stolen1 = testExecutor.pollWork("compA", keyGroup1,
stealHandle);
+ ExecutableWork stolen1 =
+ testExecutor.pollWork("compA", keyGroup1, stealHandle,
FAILING_FAILED_WORK_HANDLER);
assertNotNull(stolen1);
assertEquals(work1, stolen1);
@@ -569,10 +604,98 @@ public class BoundedQueueExecutorTest {
ExecutableWork work = createWorkWithCompIdAndKeyGroup("compA", keyGroup,
ignored -> {});
testExecutor.execute(work, 100);
- ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup,
stealHandle);
+ ExecutableWork stolen =
+ testExecutor.pollWork("compA", keyGroup, stealHandle,
FAILING_FAILED_WORK_HANDLER);
assertNull(stolen);
blockerStop.countDown();
testExecutor.shutdown();
}
+
+ @Test
+ public void testPollWork_skipsFailedWorkAndCallsOnFailedWorkHandler() throws
Exception {
+ BoundedQueueExecutor testExecutor =
+ new BoundedQueueExecutor(
+ 1,
+ 60,
+ TimeUnit.SECONDS,
+ 100,
+ 10000000,
+ new
ThreadFactoryBuilder().setNameFormat("testPollWork-%d").setDaemon(true).build(),
+ useFairMonitor,
+ /* useKeyGroupWorkQueue= */ true);
+
+ CountDownLatch blockerStart = new CountDownLatch(1);
+ CountDownLatch blockerStop = new CountDownLatch(1);
+ AtomicReference<BoundedQueueExecutorWorkHandle> blockerHandleRef = new
AtomicReference<>();
+ ExecutableWork blockerWork =
+ createWorkWithHandle(
+ "compA",
+ DEFAULT_KEY_GROUP,
+ (work, handle) -> {
+ blockerHandleRef.set(handle);
+ blockerStart.countDown();
+ try {
+ blockerStop.await();
+ } catch (InterruptedException e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ testExecutor.execute(blockerWork, 10);
+ blockerStart.await();
+ BoundedQueueExecutorWorkHandleImpl stealHandle =
+ (BoundedQueueExecutorWorkHandleImpl) blockerHandleRef.get();
+ assertNotNull(stealHandle);
+
+ Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1);
+ FailedWorkHandler onFailedWorkHandler = mock(FailedWorkHandler.class);
+
+ ExecutableWork work1 =
+ createWorkWithCompIdAndKeyGroupAndWorkToken("compA", keyGroup, 101,
ignored -> {});
+ ExecutableWork work2 =
+ createWorkWithCompIdAndKeyGroupAndWorkToken("compA", keyGroup, 102,
ignored -> {});
+ ExecutableWork work3 =
+ createWorkWithCompIdAndKeyGroupAndWorkToken("compA", keyGroup, 103,
ignored -> {});
+
+ // Enqueue both tasks (they will wait in the queue because the thread is
blocked).
+ testExecutor.execute(work1, 100);
+ testExecutor.execute(work2, 150);
+ testExecutor.execute(work3, 200);
+
+ assertEquals(4, testExecutor.elementsOutstanding());
+ assertEquals(460, testExecutor.bytesOutstanding());
+
+ // Mark work1, work3 as failed while waiting in the queue.
+ work1.work().setFailed();
+ work3.work().setFailed();
+
+ // pollWork should skip work1, close work1's handle, invoke
+ // onFailedWorkHandler callback, and return work2.
+ ExecutableWork stolen =
+ testExecutor.pollWork("compA", keyGroup, stealHandle,
onFailedWorkHandler);
+ assertNotNull(stolen);
+ assertEquals(work2, stolen);
+
+ verify(onFailedWorkHandler).onFailedWork(work1.work());
+
+ // Verify stealHandle merged (blockerWork: 10 bytes, work2: 150 bytes).
+ assertEquals(160, stealHandle.bytes());
+
+ stolen = testExecutor.pollWork("compA", keyGroup, stealHandle,
onFailedWorkHandler);
+ assertNull(stolen);
+ verify(onFailedWorkHandler).onFailedWork(work3.work());
+
+ // Still 160, nothing should be merged in.
+ assertEquals(160, stealHandle.bytes());
+
+ // Polling again should return null since no more tasks exist for keyGroup.
+ assertNull(testExecutor.pollWork("compA", keyGroup, stealHandle,
onFailedWorkHandler));
+
+ blockerStop.countDown();
+ stealHandle.close();
+ assertEquals(0, testExecutor.elementsOutstanding());
+ assertEquals(0, testExecutor.bytesOutstanding());
+ testExecutor.shutdown();
+ }
}
