scwhittle commented on code in PR #38920:
URL: https://github.com/apache/beam/pull/38920#discussion_r3756490940
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java:
##########
@@ -361,6 +361,7 @@ public void start(
this.workQueueExecutor = workQueueExecutor;
this.budgetHandle = budgetHandle;
this.keyTransitionListener = keyTransitionListener;
+ this.onFailedWorkHandler = onFailedWorkHandler;
Review Comment:
if multikeybundleoptions say it is enabled, verify that failedworkhandler is
non-null here?
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java:
##########
@@ -577,11 +579,17 @@ public void testAdvance_noMoreWork() throws Exception {
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,
+ (oldWork, newWork) -> {},
+ ignored -> {});
Review Comment:
change to one that fails if ever called throughout this test? similar with
transition?
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java:
##########
@@ -1524,13 +1524,144 @@ public void
testCompleteCommit_retryableFailureTriggersReExecution() throws Exce
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);
+
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();
Review Comment:
is this racy with receiving the heartbeats sent above? will the test be
flaky?
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java:
##########
@@ -511,7 +540,7 @@ public void testPollWork() throws Exception {
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 -> {});
Review Comment:
instead of ingnoring, have a failure handler that we can assert is not called
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java:
##########
@@ -569,10 +598,83 @@ public void testPollWorkWithLinkedBlockingQueue() throws
Exception {
ExecutableWork work = createWorkWithCompIdAndKeyGroup("compA", keyGroup,
ignored -> {});
testExecutor.execute(work, 100);
- ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup,
stealHandle);
+ ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup,
stealHandle, ignored -> {});
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 -> {});
+
+ // Enqueue both tasks (they will wait in the queue because the thread is
blocked).
+ testExecutor.execute(work1, 100);
+ testExecutor.execute(work2, 150);
+
+ assertEquals(3, testExecutor.elementsOutstanding());
+ assertEquals(260, testExecutor.bytesOutstanding());
+
+ // Mark work1 as failed while waiting in the queue.
+ work1.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());
+
+ // Polling again should return null since no more tasks exist for keyGroup.
Review Comment:
how about a failed work3 to cover the case where we clear the failed stuff
and there is not another item available?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
