From 4753a8bea6de1ad5db9a0c5262ad0c9de01a86c4 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 22 May 2026 00:26:00 +0000 Subject: [PATCH 01/11] [Dataflow Streaming] Prepare BoundedQueueExecutor for MultiKey bundles --- .../worker/streaming/ExecutableWork.java | 18 ++-- .../dataflow/worker/streaming/WorkResult.java | 32 +++++++ .../worker/util/BoundedQueueExecutor.java | 59 ++++++++---- .../dataflow/worker/util/ExceptionUtils.java | 41 +++++++++ .../processing/StreamingCommitFinalizer.java | 2 +- .../processing/StreamingWorkScheduler.java | 17 ++-- .../worker/StreamingDataflowWorkerTest.java | 6 +- .../worker/streaming/ActiveWorkStateTest.java | 4 +- .../streaming/ComputationStateCacheTest.java | 2 +- .../worker/util/BoundedQueueExecutorTest.java | 90 +++++++++++++------ .../failures/WorkFailureProcessorTest.java | 6 +- .../work/refresh/ActiveWorkRefresherTest.java | 6 +- 12 files changed, 222 insertions(+), 61 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkResult.java create mode 100644 runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ExceptionUtils.java diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java index db279f066630..9f29a6496f77 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java @@ -18,24 +18,28 @@ package org.apache.beam.runners.dataflow.worker.streaming; import com.google.auto.value.AutoValue; -import java.util.function.Consumer; +import java.util.function.Function; +import org.apache.beam.runners.dataflow.worker.util.ExceptionUtils; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; /** {@link Work} instance and a processing function used to process the work. */ @AutoValue -public abstract class ExecutableWork implements Runnable { +public abstract class ExecutableWork { - public static ExecutableWork create(Work work, Consumer executeWorkFn) { + public static ExecutableWork create(Work work, Function executeWorkFn) { return new AutoValue_ExecutableWork(work, executeWorkFn); } public abstract Work work(); - public abstract Consumer executeWorkFn(); + public abstract Function executeWorkFn(); - @Override - public void run() { - executeWorkFn().accept(work()); + public WorkResult run() { + try { + return executeWorkFn().apply(work()); + } catch (Throwable t) { + throw ExceptionUtils.propagate(t); + } } public final WorkId id() { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkResult.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkResult.java new file mode 100644 index 000000000000..07db0e98aca4 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkResult.java @@ -0,0 +1,32 @@ +/* + * 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; + +import com.google.auto.value.AutoValue; + +/** The result of executing an {@link ExecutableWork}. */ +@AutoValue +public abstract class WorkResult { + public static WorkResult create(int itemsProcessed, long bytesProcessed) { + return new AutoValue_WorkResult(itemsProcessed, bytesProcessed); + } + + public abstract int itemsProcessed(); + + public abstract long bytesProcessed(); +} 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 9079c3cc69b8..5a2a7584bc9b 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 @@ -24,6 +24,8 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import javax.annotation.concurrent.GuardedBy; +import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; +import org.apache.beam.runners.dataflow.worker.streaming.WorkResult; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor.Guard; @@ -38,7 +40,7 @@ public class BoundedQueueExecutor { // Used to guard elementsOutstanding and bytesOutstanding. private final Monitor monitor; - private final ConcurrentLinkedQueue decrementQueue = new ConcurrentLinkedQueue<>(); + private final ConcurrentLinkedQueue decrementQueue = new ConcurrentLinkedQueue<>(); private final Object decrementQueueDrainLock = new Object(); private final AtomicBoolean isDecrementBatchPending = new AtomicBoolean(false); private int elementsOutstanding = 0; @@ -106,7 +108,7 @@ protected void afterExecute(Runnable r, Throwable t) { // Before adding a Work to the queue, check that there are enough bytes of space or no other // outstanding elements of work. - public void execute(Runnable work, long workBytes) { + public void execute(ExecutableWork work, long workBytes) { monitor.enterWhenUninterruptibly( new Guard(monitor) { @Override @@ -119,12 +121,17 @@ public boolean isSatisfied() { executeMonitorHeld(work, workBytes); } - // Forcibly add something to the queue, ignoring the length limit. - public void forceExecute(Runnable work, long workBytes) { + public void forceExecute(ExecutableWork work, long workBytes) { monitor.enter(); executeMonitorHeld(work, workBytes); } + /** Forcibly execute a Runnable callback with 0 bytes of size. */ + public void forceExecute(Runnable work) { + monitor.enter(); + executeMonitorHeld(work); + } + // Set the maximum/core pool size of the executor. public synchronized void setMaximumPoolSize(int maximumPoolSize, int maximumElementsOutstanding) { // For ThreadPoolExecutor, the maximum pool size should always greater than or equal to core @@ -221,32 +228,54 @@ public String summaryHtml() { } } - private void executeMonitorHeld(Runnable work, long workBytes) { + private void executeMonitorHeld(ExecutableWork work, long workBytes) { bytesOutstanding += workBytes; ++elementsOutstanding; monitor.leave(); + executor.execute( + () -> { + // Any execution exception thrown by work.run() propagates uncaught, triggering + // the default JVM UncaughtExceptionHandler which immediately crashes/terminates + // the JVM. Since the process exits immediately, reclaiming resource budgets in + // this JVM is unnecessary. Furthermore, since a failed execution does not return + // a WorkResult, we do not have a good/accurate fallback value to decrement. + WorkResult result = work.run(); + decrementCounters(result); + }); + } + + private void executeMonitorHeld(Runnable work) { + ++elementsOutstanding; + monitor.leave(); + try { executor.execute( () -> { try { work.run(); } finally { - decrementCounters(workBytes); + // Commit finalizer callbacks catch and swallow all exceptions downstream + // to keep the worker alive (so the JVM does not crash). Therefore, to + // prevent elements outstanding capacity leaks under swallowed failures, + // we must guarantee decrementing element counts in the finally block. + decrementCounters(WorkResult.create(1, 0L)); } }); - } catch (RuntimeException e) { - // If the execute() call threw an exception, decrement counters here. - decrementCounters(workBytes); - throw e; + } catch (Throwable e) { + // Since finalizer rejections are caught and swallowed downstream, we must + // decrement elements outstanding immediately on task submission failure to + // prevent permanent capacity leaks in the running JVM. + decrementCounters(WorkResult.create(1, 0L)); + throw ExceptionUtils.propagate(e); } } - private void decrementCounters(long workBytes) { + private void decrementCounters(WorkResult result) { // All threads queue decrements and one thread grabs the monitor and updates // counters. We do this to reduce contention on monitor which is locked by // GetWork thread - decrementQueue.add(workBytes); + decrementQueue.add(result); boolean submittedToExistingBatch = isDecrementBatchPending.getAndSet(true); if (submittedToExistingBatch) { // There is already a thread about to drain the decrement queue @@ -265,12 +294,12 @@ private void decrementCounters(long workBytes) { long bytesToDecrement = 0; int elementsToDecrement = 0; while (true) { - Long pollResult = decrementQueue.poll(); + WorkResult pollResult = decrementQueue.poll(); if (pollResult == null) { break; } - bytesToDecrement += pollResult; - ++elementsToDecrement; + bytesToDecrement += pollResult.bytesProcessed(); + elementsToDecrement += pollResult.itemsProcessed(); } if (elementsToDecrement == 0) { return; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ExceptionUtils.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ExceptionUtils.java new file mode 100644 index 000000000000..4bbdbb3216ca --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ExceptionUtils.java @@ -0,0 +1,41 @@ +/* + * 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.util; + +import org.apache.beam.sdk.annotations.Internal; + +/** Utility methods for simplifying work with exceptions and throwables. */ +@Internal +public final class ExceptionUtils { + + private ExceptionUtils() {} + + /** + * Propagates {@code throwable} as-is if it is an instance of {@link RuntimeException} or {@link + * Error}, or else as a last resort wraps it in a {@code RuntimeException} and then propagates. + */ + public static RuntimeException propagate(Throwable throwable) { + if (throwable instanceof RuntimeException) { + throw (RuntimeException) throwable; + } else if (throwable instanceof Error) { + throw (Error) throwable; + } else { + throw new RuntimeException(throwable); + } + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java index 5a66545ab335..22573bf1ced2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java @@ -156,7 +156,7 @@ public void finalizeCommits(Iterable finalizeIds) { } for (Runnable callback : callbacksToExecute) { try { - finalizationExecutor.forceExecute(callback, 0); + finalizationExecutor.forceExecute(callback); } catch (OutOfMemoryError oom) { throw oom; } catch (Throwable t) { 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 1428037d9ca0..65a1325cd40e 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 @@ -42,11 +42,13 @@ 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; +import org.apache.beam.runners.dataflow.worker.streaming.WorkResult; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle; import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; +import org.apache.beam.runners.dataflow.worker.util.ExceptionUtils; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution; import org.apache.beam.runners.dataflow.worker.windmill.client.commits.Commit; @@ -233,15 +235,15 @@ public void queueAppliedFinalizeIds(ImmutableList appliedFinalizeIds) { * * @implNote This will block the calling thread during execution of user DoFns. */ - private void processWork( + private WorkResult processWork( ComputationState computationState, Work work, ImmutableList getWorkStreamLatencies) { work.recordGetWorkStreamLatencies(getWorkStreamLatencies); - processWork(computationState, work); + return processWork(computationState, work); } - private void processWork(ComputationState computationState, Work work) { + private WorkResult processWork(ComputationState computationState, Work work) { Windmill.WorkItem workItem = work.getWorkItem(); String computationId = computationState.getComputationId(); ByteString key = workItem.getKey(); @@ -258,7 +260,7 @@ private void processWork(ComputationState computationState, Work work) { outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true)); work.setState(Work.State.COMMIT_QUEUED); work.queueCommit(outputBuilder.build(), computationState); - return; + return WorkResult.create(1, work.getSerializedWorkItemSize()); } long processingStartTimeNanos = System.nanoTime(); @@ -284,6 +286,7 @@ private void processWork(ComputationState computationState, Work work) { work.queueCommit(validatedCommitRequest, computationState); recordProcessingStats(commitRequest, workItem, executeWorkResult); LOG.debug("Processing done for work token: {}", workItem.getWorkToken()); + return WorkResult.create(1, work.getSerializedWorkItemSize()); } catch (Throwable t) { // OutOfMemoryError that are caught will be rethrown and trigger jvm termination. try { @@ -294,10 +297,14 @@ private void processWork(ComputationState computationState, Work work) { invalidWork -> computationState.completeWorkAndScheduleNextWorkForKey( invalidWork.getShardedKey(), invalidWork.id())); + // Failure successfully processed/invalidated/rescheduled. Return failure WorkResult to + // release budget cleanly. + return WorkResult.create(1, work.getSerializedWorkItemSize()); } catch (OutOfMemoryError oom) { throw oom; } catch (Throwable t2) { - throw new RuntimeException(t2); + LOG.warn("Failed to process work failure safely for work {}", work.id(), t2); + throw ExceptionUtils.propagate(t2); } } finally { // Update total processing time counters. Updating in finally clause ensures that 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..593883343bd4 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 @@ -106,6 +106,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey; 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.WorkResult; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandleImpl; import org.apache.beam.runners.dataflow.worker.streaming.harness.FanOutStreamingEngineWorkerHarness; @@ -373,7 +374,10 @@ private static ExecutableWork createMockWork( computationId, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), false, Instant::now), - processWorkFn); + work -> { + processWorkFn.accept(work); + return WorkResult.create(1, work.getSerializedWorkItemSize()); + }); } private byte[] intervalWindowBytes(IntervalWindow window) throws Exception { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java index 865ae2612803..4942e9610fd6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java @@ -73,7 +73,7 @@ private static ExecutableWork createWork(Windmill.WorkItem workItem) { createWorkProcessingContext(), false, Instant::now), - ignored -> {}); + work -> WorkResult.create(1, work.getSerializedWorkItemSize())); } private static ExecutableWork expiredWork(Windmill.WorkItem workItem) { @@ -85,7 +85,7 @@ private static ExecutableWork expiredWork(Windmill.WorkItem workItem) { createWorkProcessingContext(), false, () -> Instant.EPOCH), - ignored -> {}); + work -> WorkResult.create(1, work.getSerializedWorkItemSize())); } private static Work.ProcessingContext createWorkProcessingContext() { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java index 1c8b8fca131d..d9ad1157ee4c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java @@ -77,7 +77,7 @@ private static ExecutableWork createWork(ShardedKey shardedKey, long workToken, mock(HeartbeatSender.class)), false, Instant::now), - ignored -> {}); + work -> WorkResult.create(1, work.getSerializedWorkItemSize())); } @Before 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 d7ea039bb809..bdadcd2b9e8c 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 @@ -32,6 +32,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; 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.WorkResult; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; @@ -85,18 +86,22 @@ private static ExecutableWork createWork(Consumer executeWorkFn) { mock(HeartbeatSender.class)), false, Instant::now), - executeWorkFn); + work -> { + executeWorkFn.accept(work); + return WorkResult.create(1, work.getSerializedWorkItemSize()); + }); } - private Runnable createSleepProcessWorkFn(CountDownLatch start, CountDownLatch stop) { - return () -> { - start.countDown(); - try { - stop.await(); - } catch (Exception e) { - throw new RuntimeException(e); - } - }; + private ExecutableWork createSleepProcessWork(CountDownLatch start, CountDownLatch stop) { + return createWork( + ignored -> { + start.countDown(); + try { + stop.await(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); } @Before @@ -123,9 +128,9 @@ public void testScheduleWorkWhenExceedMaximumPoolSize() throws Exception { CountDownLatch processStop2 = new CountDownLatch(1); CountDownLatch processStart3 = new CountDownLatch(1); CountDownLatch processStop3 = new CountDownLatch(1); - Runnable m1 = createSleepProcessWorkFn(processStart1, processStop1); - Runnable m2 = createSleepProcessWorkFn(processStart2, processStop2); - Runnable m3 = createSleepProcessWorkFn(processStart3, processStop3); + ExecutableWork m1 = createSleepProcessWork(processStart1, processStop1); + ExecutableWork m2 = createSleepProcessWork(processStart2, processStop2); + ExecutableWork m3 = createSleepProcessWork(processStart3, processStop3); executor.execute(m1, 1); processStart1.await(); @@ -152,8 +157,8 @@ public void testScheduleWorkWhenExceedMaximumBytesOutstanding() throws Exception CountDownLatch processStop1 = new CountDownLatch(1); CountDownLatch processStart2 = new CountDownLatch(1); CountDownLatch processStop2 = new CountDownLatch(1); - Runnable m1 = createSleepProcessWorkFn(processStart1, processStop1); - Runnable m2 = createSleepProcessWorkFn(processStart2, processStop2); + ExecutableWork m1 = createSleepProcessWork(processStart1, processStop1); + ExecutableWork m2 = createSleepProcessWork(processStart2, processStop2); executor.execute(m1, 10000000); processStart1.await(); @@ -187,9 +192,9 @@ public void testOverrideMaximumPoolSize() throws Exception { CountDownLatch processStart2 = new CountDownLatch(1); CountDownLatch processStart3 = new CountDownLatch(1); CountDownLatch stop = new CountDownLatch(1); - Runnable m1 = createSleepProcessWorkFn(processStart1, stop); - Runnable m2 = createSleepProcessWorkFn(processStart2, stop); - Runnable m3 = createSleepProcessWorkFn(processStart3, stop); + ExecutableWork m1 = createSleepProcessWork(processStart1, stop); + ExecutableWork m2 = createSleepProcessWork(processStart2, stop); + ExecutableWork m3 = createSleepProcessWork(processStart3, stop); // Initial state. assertEquals(0, executor.activeCount()); @@ -225,9 +230,9 @@ public void testRecordTotalTimeMaxActiveThreadsUsed() throws Exception { CountDownLatch processStart2 = new CountDownLatch(1); CountDownLatch processStart3 = new CountDownLatch(1); CountDownLatch stop = new CountDownLatch(1); - Runnable m1 = createSleepProcessWorkFn(processStart1, stop); - Runnable m2 = createSleepProcessWorkFn(processStart2, stop); - Runnable m3 = createSleepProcessWorkFn(processStart3, stop); + ExecutableWork m1 = createSleepProcessWork(processStart1, stop); + ExecutableWork m2 = createSleepProcessWork(processStart2, stop); + ExecutableWork m3 = createSleepProcessWork(processStart3, stop); // Initial state. assertEquals(0, executor.activeCount()); @@ -264,9 +269,9 @@ public void testRecordTotalTimeMaxActiveThreadsUsedWhenMaximumPoolSizeIsIncrease CountDownLatch processStart2 = new CountDownLatch(1); CountDownLatch processStart3 = new CountDownLatch(1); CountDownLatch stop = new CountDownLatch(1); - Runnable m1 = createSleepProcessWorkFn(processStart1, stop); - Runnable m2 = createSleepProcessWorkFn(processStart2, stop); - Runnable m3 = createSleepProcessWorkFn(processStart3, stop); + ExecutableWork m1 = createSleepProcessWork(processStart1, stop); + ExecutableWork m2 = createSleepProcessWork(processStart2, stop); + ExecutableWork m3 = createSleepProcessWork(processStart3, stop); // Initial state. assertEquals(0, executor.activeCount()); @@ -308,9 +313,9 @@ public void testRecordTotalTimeMaxActiveThreadsUsedWhenMaximumPoolSizeIsReduced( CountDownLatch processStop2 = new CountDownLatch(1); CountDownLatch processStart3 = new CountDownLatch(1); CountDownLatch processStop3 = new CountDownLatch(1); - Runnable m1 = createSleepProcessWorkFn(processStart1, processStop1); - Runnable m2 = createSleepProcessWorkFn(processStart2, processStop2); - Runnable m3 = createSleepProcessWorkFn(processStart3, processStop3); + ExecutableWork m1 = createSleepProcessWork(processStart1, processStop1); + ExecutableWork m2 = createSleepProcessWork(processStart2, processStop2); + ExecutableWork m3 = createSleepProcessWork(processStart3, processStop3); // Initial state. assertEquals(0, executor.activeCount()); @@ -351,6 +356,37 @@ public void testRecordTotalTimeMaxActiveThreadsUsedWhenMaximumPoolSizeIsReduced( executor.shutdown(); } + @Test + public void testRunnableExceptionPropagationDecrementsCounters() throws Exception { + CountDownLatch processStart = new CountDownLatch(1); + CountDownLatch processStop = new CountDownLatch(1); + + Runnable work = + () -> { + processStart.countDown(); + try { + processStop.await(); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + throw new RuntimeException("Simulated finalizer processing exception"); + }; + + executor.forceExecute(work); + processStart.await(); + + assertEquals(1, executor.elementsOutstanding()); + + processStop.countDown(); + + // Wait until outstanding elements are released + while (executor.elementsOutstanding() != 0) { + Thread.sleep(10); + } + + assertEquals(0, executor.elementsOutstanding()); + } + @Test public void testRenderSummaryHtml() { String expectedSummaryHtml = diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java index 68a11895fa12..291be6f0bf9b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java @@ -33,6 +33,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; 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.WorkResult; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; @@ -98,7 +99,10 @@ private static ExecutableWork createWork(Supplier clock, Consumer mock(HeartbeatSender.class)), false, clock), - processWorkFn); + work -> { + processWorkFn.accept(work); + return WorkResult.create(1, work.getSerializedWorkItemSize()); + }); } private static ExecutableWork createWork(Consumer processWorkFn) { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java index 054db878c869..e9dc42de8aa6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java @@ -46,6 +46,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey; 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.WorkResult; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; @@ -137,7 +138,10 @@ private ExecutableWork createOldWork( "computationId", new FakeGetDataClient(), ignored -> {}, heartbeatSender), false, ActiveWorkRefresherTest::aLongTimeAgo), - processWork); + work -> { + processWork.accept(work); + return WorkResult.create(1, work.getSerializedWorkItemSize()); + }); } @Test From 0726da4e690b5e284b619812eeb1b4af0fe2d020 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 22 May 2026 22:32:05 +0000 Subject: [PATCH 02/11] [Dataflow Streaming] Refactor BoundedQueueExecutor with handles --- ...va => BoundedQueueExecutorWorkHandle.java} | 18 +- .../worker/streaming/ExecutableWork.java | 11 +- .../worker/util/BoundedQueueExecutor.java | 158 ++++++++++++++---- .../processing/StreamingWorkScheduler.java | 23 ++- .../worker/StreamingDataflowWorkerTest.java | 4 +- .../worker/streaming/ActiveWorkStateTest.java | 4 +- .../streaming/ComputationStateCacheTest.java | 2 +- .../worker/util/BoundedQueueExecutorTest.java | 145 +++++++++++++++- .../failures/WorkFailureProcessorTest.java | 4 +- .../work/refresh/ActiveWorkRefresherTest.java | 4 +- 10 files changed, 295 insertions(+), 78 deletions(-) rename runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/{WorkResult.java => BoundedQueueExecutorWorkHandle.java} (69%) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkResult.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java similarity index 69% rename from runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkResult.java rename to runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java index 07db0e98aca4..1ca534966947 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkResult.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java @@ -17,16 +17,8 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; -import com.google.auto.value.AutoValue; - -/** The result of executing an {@link ExecutableWork}. */ -@AutoValue -public abstract class WorkResult { - public static WorkResult create(int itemsProcessed, long bytesProcessed) { - return new AutoValue_WorkResult(itemsProcessed, bytesProcessed); - } - - public abstract int itemsProcessed(); - - public abstract long bytesProcessed(); -} +/** + * A handle to use when requesting pulling more work from @BoundedQueueExecutor + * via @BoundedQueueExecutor.pollWork + */ +public interface BoundedQueueExecutorWorkHandle {} diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java index 9f29a6496f77..161c16106373 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java @@ -18,7 +18,7 @@ package org.apache.beam.runners.dataflow.worker.streaming; import com.google.auto.value.AutoValue; -import java.util.function.Function; +import java.util.function.BiConsumer; import org.apache.beam.runners.dataflow.worker.util.ExceptionUtils; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; @@ -26,17 +26,18 @@ @AutoValue public abstract class ExecutableWork { - public static ExecutableWork create(Work work, Function executeWorkFn) { + public static ExecutableWork create( + Work work, BiConsumer executeWorkFn) { return new AutoValue_ExecutableWork(work, executeWorkFn); } public abstract Work work(); - public abstract Function executeWorkFn(); + public abstract BiConsumer executeWorkFn(); - public WorkResult run() { + public void run(BoundedQueueExecutorWorkHandle handle) { try { - return executeWorkFn().apply(work()); + executeWorkFn().accept(work(), handle); } catch (Throwable t) { throw ExceptionUtils.propagate(t); } 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 5a2a7584bc9b..386b9e0a436a 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 @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.util; +import java.util.Optional; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadFactory; @@ -24,8 +25,10 @@ import java.util.concurrent.TimeUnit; 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.WorkResult; +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; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor.Guard; @@ -40,7 +43,19 @@ public class BoundedQueueExecutor { // Used to guard elementsOutstanding and bytesOutstanding. private final Monitor monitor; - private final ConcurrentLinkedQueue decrementQueue = new ConcurrentLinkedQueue<>(); + + private static class Budget { + + final int elements; + final long bytes; + + Budget(int elements, long bytes) { + this.elements = elements; + this.bytes = bytes; + } + } + + private final ConcurrentLinkedQueue decrementQueue = new ConcurrentLinkedQueue<>(); private final Object decrementQueueDrainLock = new Object(); private final AtomicBoolean isDecrementBatchPending = new AtomicBoolean(false); private int elementsOutstanding = 0; @@ -121,15 +136,16 @@ public boolean isSatisfied() { executeMonitorHeld(work, workBytes); } + // Forcibly add ExecutableWork to the queue, ignoring the limits. public void forceExecute(ExecutableWork work, long workBytes) { monitor.enter(); executeMonitorHeld(work, workBytes); } /** Forcibly execute a Runnable callback with 0 bytes of size. */ - public void forceExecute(Runnable work) { + public void forceExecute(Runnable runnable) { monitor.enter(); - executeMonitorHeld(work); + executeMonitorHeld(runnable); } // Set the maximum/core pool size of the executor. @@ -228,21 +244,84 @@ public String summaryHtml() { } } + class BoundedQueueExecutorWorkHandleImpl + implements BoundedQueueExecutorWorkHandle, AutoCloseable { + + private int elements; + private long bytes; + private boolean closed = false; + + private BoundedQueueExecutorWorkHandleImpl(int elements, long bytes) { + this.elements = elements; + this.bytes = bytes; + } + + public synchronized void addBudget(int elements, long bytes) { + Preconditions.checkState(!closed, "Cannot add budget to a closed WorkBudgetHandle"); + this.elements += elements; + this.bytes += bytes; + } + + public synchronized void cancel() { + this.closed = true; + } + + @Override + public synchronized void close() { + Preconditions.checkArgument(!closed); + closed = true; + decrementCounters(this.elements, this.bytes); + } + } + + private static class QueuedWork implements Runnable { + + private final ExecutableWork work; + private final BoundedQueueExecutorWorkHandleImpl handle; + private final long workBytes; + + public QueuedWork( + ExecutableWork work, BoundedQueueExecutorWorkHandleImpl handle, long workBytes) { + this.work = work; + this.handle = handle; + this.workBytes = workBytes; + } + + public void cancelHandle() { + handle.cancel(); + } + + public ExecutableWork getWork() { + return work; + } + + public long getWorkBytes() { + return workBytes; + } + + @Override + public void run() { + Preconditions.checkArgument(!handle.closed); + try { + work.run(handle); + } finally { + handle.close(); + } + } + } + private void executeMonitorHeld(ExecutableWork work, long workBytes) { - bytesOutstanding += workBytes; ++elementsOutstanding; + bytesOutstanding += workBytes; monitor.leave(); - - executor.execute( - () -> { - // Any execution exception thrown by work.run() propagates uncaught, triggering - // the default JVM UncaughtExceptionHandler which immediately crashes/terminates - // the JVM. Since the process exits immediately, reclaiming resource budgets in - // this JVM is unnecessary. Furthermore, since a failed execution does not return - // a WorkResult, we do not have a good/accurate fallback value to decrement. - WorkResult result = work.run(); - decrementCounters(result); - }); + BoundedQueueExecutorWorkHandleImpl handle = + new BoundedQueueExecutorWorkHandleImpl(1, workBytes); + try { + executor.execute(new QueuedWork(work, handle, workBytes)); + } catch (Throwable e) { + handle.close(); + throw ExceptionUtils.propagate(e); + } } private void executeMonitorHeld(Runnable work) { @@ -255,27 +334,40 @@ private void executeMonitorHeld(Runnable work) { try { work.run(); } finally { - // Commit finalizer callbacks catch and swallow all exceptions downstream - // to keep the worker alive (so the JVM does not crash). Therefore, to - // prevent elements outstanding capacity leaks under swallowed failures, - // we must guarantee decrementing element counts in the finally block. - decrementCounters(WorkResult.create(1, 0L)); + decrementCounters(1, 0L); } }); } catch (Throwable e) { - // Since finalizer rejections are caught and swallowed downstream, we must - // decrement elements outstanding immediately on task submission failure to - // prevent permanent capacity leaks in the running JVM. - decrementCounters(WorkResult.create(1, 0L)); + decrementCounters(1, 0L); throw ExceptionUtils.propagate(e); } } - private void decrementCounters(WorkResult result) { - // All threads queue decrements and one thread grabs the monitor and updates - // counters. We do this to reduce contention on monitor which is locked by - // GetWork thread - decrementQueue.add(result); + @VisibleForTesting + BoundedQueueExecutorWorkHandleImpl createEmptyBudgetHandle() { + return new BoundedQueueExecutorWorkHandleImpl(0, 0L); + } + + public Optional pollWork(BoundedQueueExecutorWorkHandle handle) { + BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; + while (true) { + Runnable runnable = executor.getQueue().poll(); + if (runnable == null) { + return Optional.empty(); + } + if (runnable instanceof QueuedWork) { + QueuedWork queuedWork = (QueuedWork) runnable; + queuedWork.cancelHandle(); + internalHandle.addBudget(1, queuedWork.getWorkBytes()); + return Optional.of(queuedWork.getWork()); + } + // Pop and execute standard callbacks immediately on the calling thread to drain the queue + runnable.run(); + } + } + + private void decrementCounters(int elements, long bytes) { + decrementQueue.add(new Budget(elements, bytes)); boolean submittedToExistingBatch = isDecrementBatchPending.getAndSet(true); if (submittedToExistingBatch) { // There is already a thread about to drain the decrement queue @@ -294,12 +386,12 @@ private void decrementCounters(WorkResult result) { long bytesToDecrement = 0; int elementsToDecrement = 0; while (true) { - WorkResult pollResult = decrementQueue.poll(); + Budget pollResult = decrementQueue.poll(); if (pollResult == null) { break; } - bytesToDecrement += pollResult.bytesProcessed(); - elementsToDecrement += pollResult.itemsProcessed(); + bytesToDecrement += pollResult.bytes; + elementsToDecrement += pollResult.elements; } if (elementsToDecrement == 0) { return; 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 65a1325cd40e..0c3289d4cf58 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 @@ -35,6 +35,7 @@ import org.apache.beam.runners.dataflow.worker.ReaderCache; import org.apache.beam.runners.dataflow.worker.WorkItemCancelledException; import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC; +import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle; 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; @@ -42,7 +43,6 @@ 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; -import org.apache.beam.runners.dataflow.worker.streaming.WorkResult; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle; import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; @@ -220,7 +220,7 @@ public void scheduleWork( ExecutableWork.create( Work.create( workItem, serializedWorkItemSize, watermarks, processingContext, drainMode, clock), - work -> processWork(computationState, work, getWorkStreamLatencies))); + (work, handle) -> processWork(computationState, work, getWorkStreamLatencies, handle))); } /** Adds any applied finalize ids to the commit finalizer to have their callbacks executed. */ @@ -234,16 +234,19 @@ public void queueAppliedFinalizeIds(ImmutableList appliedFinalizeIds) { * internally if processing fails due to uncaught {@link Exception}(s). * * @implNote This will block the calling thread during execution of user DoFns. + * @param handle handled to pass to BoundedQueueExecutor.pollWork, currently unused */ - private WorkResult processWork( + private void processWork( ComputationState computationState, Work work, - ImmutableList getWorkStreamLatencies) { + ImmutableList getWorkStreamLatencies, + BoundedQueueExecutorWorkHandle handle) { work.recordGetWorkStreamLatencies(getWorkStreamLatencies); - return processWork(computationState, work); + processWork(computationState, work, handle); } - private WorkResult processWork(ComputationState computationState, Work work) { + private void processWork( + ComputationState computationState, Work work, BoundedQueueExecutorWorkHandle unusedHandle) { Windmill.WorkItem workItem = work.getWorkItem(); String computationId = computationState.getComputationId(); ByteString key = workItem.getKey(); @@ -260,7 +263,7 @@ private WorkResult processWork(ComputationState computationState, Work work) { outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true)); work.setState(Work.State.COMMIT_QUEUED); work.queueCommit(outputBuilder.build(), computationState); - return WorkResult.create(1, work.getSerializedWorkItemSize()); + return; } long processingStartTimeNanos = System.nanoTime(); @@ -286,20 +289,16 @@ private WorkResult processWork(ComputationState computationState, Work work) { work.queueCommit(validatedCommitRequest, computationState); recordProcessingStats(commitRequest, workItem, executeWorkResult); LOG.debug("Processing done for work token: {}", workItem.getWorkToken()); - return WorkResult.create(1, work.getSerializedWorkItemSize()); } catch (Throwable t) { // OutOfMemoryError that are caught will be rethrown and trigger jvm termination. try { workFailureProcessor.logAndProcessFailure( computationId, - ExecutableWork.create(work, retry -> processWork(computationState, retry)), + ExecutableWork.create(work, (retry, h) -> processWork(computationState, retry, h)), t, invalidWork -> computationState.completeWorkAndScheduleNextWorkForKey( invalidWork.getShardedKey(), invalidWork.id())); - // Failure successfully processed/invalidated/rescheduled. Return failure WorkResult to - // release budget cleanly. - return WorkResult.create(1, work.getSerializedWorkItemSize()); } catch (OutOfMemoryError oom) { throw oom; } catch (Throwable t2) { 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 593883343bd4..4d39e5d83f66 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 @@ -106,7 +106,6 @@ import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey; 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.WorkResult; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandleImpl; import org.apache.beam.runners.dataflow.worker.streaming.harness.FanOutStreamingEngineWorkerHarness; @@ -374,9 +373,8 @@ private static ExecutableWork createMockWork( computationId, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), false, Instant::now), - work -> { + (work, handle) -> { processWorkFn.accept(work); - return WorkResult.create(1, work.getSerializedWorkItemSize()); }); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java index 4942e9610fd6..0f14efdd0c0b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java @@ -73,7 +73,7 @@ private static ExecutableWork createWork(Windmill.WorkItem workItem) { createWorkProcessingContext(), false, Instant::now), - work -> WorkResult.create(1, work.getSerializedWorkItemSize())); + (work, handle) -> {}); } private static ExecutableWork expiredWork(Windmill.WorkItem workItem) { @@ -85,7 +85,7 @@ private static ExecutableWork expiredWork(Windmill.WorkItem workItem) { createWorkProcessingContext(), false, () -> Instant.EPOCH), - work -> WorkResult.create(1, work.getSerializedWorkItemSize())); + (work, handle) -> {}); } private static Work.ProcessingContext createWorkProcessingContext() { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java index d9ad1157ee4c..30ad97140e1e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java @@ -77,7 +77,7 @@ private static ExecutableWork createWork(ShardedKey shardedKey, long workToken, mock(HeartbeatSender.class)), false, Instant::now), - work -> WorkResult.create(1, work.getSerializedWorkItemSize())); + (work, handle) -> {}); } @Before 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 bdadcd2b9e8c..577d72da9e37 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 @@ -26,13 +26,14 @@ import java.util.Arrays; import java.util.Collection; +import java.util.Optional; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; 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.WorkResult; +import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor.BoundedQueueExecutorWorkHandleImpl; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; @@ -86,9 +87,8 @@ private static ExecutableWork createWork(Consumer executeWorkFn) { mock(HeartbeatSender.class)), false, Instant::now), - work -> { + (work, handle) -> { executeWorkFn.accept(work); - return WorkResult.create(1, work.getSerializedWorkItemSize()); }); } @@ -387,6 +387,145 @@ public void testRunnableExceptionPropagationDecrementsCounters() throws Exceptio assertEquals(0, executor.elementsOutstanding()); } + @Test + public void testPollWorkAndInlineBatchExecution() throws Exception { + BoundedQueueExecutor testExecutor = + new BoundedQueueExecutor( + 1, + DEFAULT_THREAD_EXPIRATION_SEC, + TimeUnit.SECONDS, + 10, + MAXIMUM_BYTES_OUTSTANDING, + new ThreadFactoryBuilder() + .setNameFormat("testPollWorkAndInlineBatchExecution-%d") + .setDaemon(true) + .build(), + useFairMonitor); + + CountDownLatch blockerStart = new CountDownLatch(1); + CountDownLatch blockerStop = new CountDownLatch(1); + ExecutableWork blockerWork = createSleepProcessWork(blockerStart, blockerStop); + + CountDownLatch start1 = new CountDownLatch(1); + CountDownLatch stop1 = new CountDownLatch(1); + ExecutableWork m1 = createSleepProcessWork(start1, stop1); + + CountDownLatch start2 = new CountDownLatch(1); + CountDownLatch stop2 = new CountDownLatch(1); + ExecutableWork m2 = createSleepProcessWork(start2, stop2); + + // 1. Occupy the single worker thread with blocker work so subsequent tasks remain queued. + testExecutor.execute(blockerWork, 0); + blockerStart.await(); + assertEquals(1, testExecutor.elementsOutstanding()); + assertEquals(0, testExecutor.bytesOutstanding()); + + // 2. Enqueue tasks to stay in the queue. + testExecutor.execute(m1, 1000); + testExecutor.execute(m2, 2000); + + assertEquals(3, testExecutor.elementsOutstanding()); + assertEquals(3000, testExecutor.bytesOutstanding()); + + // 3. Create the batch handle. + try (BoundedQueueExecutorWorkHandleImpl batchHandle = testExecutor.createEmptyBudgetHandle()) { + // 4. Poll tasks inline. + Optional polled1 = testExecutor.pollWork(batchHandle); + assertTrue(polled1.isPresent()); + assertEquals(m1, polled1.get()); + + Optional polled2 = testExecutor.pollWork(batchHandle); + assertTrue(polled2.isPresent()); + assertEquals(m2, polled2.get()); + + // Queue should now be empty. + Optional polled3 = testExecutor.pollWork(batchHandle); + assertFalse(polled3.isPresent()); + + // 5. Run polled tasks inline. + start1.countDown(); + stop1.countDown(); + polled1.get().run(batchHandle); + + start2.countDown(); + stop2.countDown(); + polled2.get().run(batchHandle); + + // Outstanding counts should NOT yet be decremented. + assertEquals(3, testExecutor.elementsOutstanding()); + assertEquals(3000, testExecutor.bytesOutstanding()); + } + + // 6. Upon close, outstanding counts should immediately reflect the batch decrement in one shot. + // Only the blocker task (0 bytes, 1 element) should remain outstanding. + while (testExecutor.elementsOutstanding() != 1) { + Thread.sleep(10); + } + assertEquals(1, testExecutor.elementsOutstanding()); + assertEquals(0, testExecutor.bytesOutstanding()); + + // Clean up blocker. + blockerStop.countDown(); + testExecutor.shutdown(); + } + + @Test + public void testPollWorkAndInlineBatchExecutionWithException() throws Exception { + BoundedQueueExecutor testExecutor = + new BoundedQueueExecutor( + 1, + DEFAULT_THREAD_EXPIRATION_SEC, + TimeUnit.SECONDS, + 10, + MAXIMUM_BYTES_OUTSTANDING, + new ThreadFactoryBuilder() + .setNameFormat("testPollWorkAndInlineBatchExecutionWithException-%d") + .setDaemon(true) + .build(), + useFairMonitor); + + CountDownLatch blockerStart = new CountDownLatch(1); + CountDownLatch blockerStop = new CountDownLatch(1); + ExecutableWork blockerWork = createSleepProcessWork(blockerStart, blockerStop); + + // Occupy all worker threads + testExecutor.execute(blockerWork, 0); + blockerStart.await(); + + ExecutableWork inlineWork1 = + createWork( + ignored -> { + throw new RuntimeException("Simulated inline execution exception"); + }); + + long size1 = inlineWork1.work().getSerializedWorkItemSize(); + testExecutor.execute(inlineWork1, size1); + + long outstandingBytesBefore = testExecutor.bytesOutstanding(); + int outstandingElementsBefore = testExecutor.elementsOutstanding(); + + try { + try (BoundedQueueExecutorWorkHandleImpl batchHandle = + testExecutor.createEmptyBudgetHandle()) { + Optional polled1 = testExecutor.pollWork(batchHandle); + assertTrue(polled1.isPresent()); + polled1.get().run(batchHandle); + } + } catch (RuntimeException e) { + assertEquals("Simulated inline execution exception", e.getMessage()); + } + + // Outstanding elements must still be released cleanly by try-with-resources close! + while (testExecutor.elementsOutstanding() != outstandingElementsBefore - 1) { + Thread.sleep(10); + } + assertEquals(outstandingElementsBefore - 1, testExecutor.elementsOutstanding()); + assertEquals(outstandingBytesBefore - size1, testExecutor.bytesOutstanding()); + + blockerStop.countDown(); + testExecutor.shutdown(); + } + @Test public void testRenderSummaryHtml() { String expectedSummaryHtml = diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java index 291be6f0bf9b..51bd4816b031 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java @@ -33,7 +33,6 @@ import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; 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.WorkResult; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; @@ -99,9 +98,8 @@ private static ExecutableWork createWork(Supplier clock, Consumer mock(HeartbeatSender.class)), false, clock), - work -> { + (work, handle) -> { processWorkFn.accept(work); - return WorkResult.create(1, work.getSerializedWorkItemSize()); }); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java index e9dc42de8aa6..88a82c6f76b6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java @@ -46,7 +46,6 @@ import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey; 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.WorkResult; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; @@ -138,9 +137,8 @@ private ExecutableWork createOldWork( "computationId", new FakeGetDataClient(), ignored -> {}, heartbeatSender), false, ActiveWorkRefresherTest::aLongTimeAgo), - work -> { + (work, handle) -> { processWork.accept(work); - return WorkResult.create(1, work.getSerializedWorkItemSize()); }); } From 32fa605bcf57c63ec91c434afae75396555566aa Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 22 May 2026 22:44:31 +0000 Subject: [PATCH 03/11] doc fix --- .../runners/dataflow/worker/util/BoundedQueueExecutor.java | 7 +++++++ 1 file changed, 7 insertions(+) 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 386b9e0a436a..f22d1d8b570e 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 @@ -348,6 +348,13 @@ BoundedQueueExecutorWorkHandleImpl createEmptyBudgetHandle() { return new BoundedQueueExecutorWorkHandleImpl(0, 0L); } + /** + * Poll additional work to be executed inline inside with the current execute(ExecutableWork work, + * long workBytes) call. It is the responsibility of the caller to execute or discard the returned + * ExecutableWork. Budget for the returned work is released when the execute() call finishes. + * + * @param handle the handle that was passed to ExecutableWork.executeWorkFn + */ public Optional pollWork(BoundedQueueExecutorWorkHandle handle) { BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; while (true) { From d6e6ff3feb18825ecde82a81a97e992db6061c94 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 22 May 2026 22:52:58 +0000 Subject: [PATCH 04/11] fix comments --- .../runners/dataflow/worker/util/BoundedQueueExecutor.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 f22d1d8b570e..c64ec9bbd568 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 @@ -268,7 +268,7 @@ public synchronized void cancel() { @Override public synchronized void close() { - Preconditions.checkArgument(!closed); + if (closed) return; closed = true; decrementCounters(this.elements, this.bytes); } @@ -356,6 +356,7 @@ BoundedQueueExecutorWorkHandleImpl createEmptyBudgetHandle() { * @param handle the handle that was passed to ExecutableWork.executeWorkFn */ public Optional pollWork(BoundedQueueExecutorWorkHandle handle) { + Preconditions.checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl); BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; while (true) { Runnable runnable = executor.getQueue().poll(); From 091949451061589a6b9fa5a81297203bce5a35cc Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Sun, 31 May 2026 10:22:59 +0000 Subject: [PATCH 05/11] Address comments --- .../worker/streaming/ExecutableWork.java | 36 ++++- .../worker/util/BoundedQueueExecutor.java | 110 ++++++------- .../worker/util/BoundedQueueExecutorTest.java | 146 ++---------------- 3 files changed, 100 insertions(+), 192 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java index 161c16106373..5586bc0217b0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java @@ -17,23 +17,42 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; -import com.google.auto.value.AutoValue; +import java.util.Objects; import java.util.function.BiConsumer; import org.apache.beam.runners.dataflow.worker.util.ExceptionUtils; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; /** {@link Work} instance and a processing function used to process the work. */ -@AutoValue -public abstract class ExecutableWork { +public final class ExecutableWork { + private final Work work; + private final BiConsumer executeWorkFn; + + private ExecutableWork( + Work work, BiConsumer executeWorkFn) { + this.work = Objects.requireNonNull(work); + this.executeWorkFn = Objects.requireNonNull(executeWorkFn); + } + + /** + * Creates an {@link ExecutableWork} instance. + * + * @param executeWorkFn The function executing the work. It'll be called along with a + * BoundedQueueExecutorWorkHandle. The handle needs to be passed to BoundedQueueExecutor when + * requesting more work to process inline. + */ public static ExecutableWork create( Work work, BiConsumer executeWorkFn) { - return new AutoValue_ExecutableWork(work, executeWorkFn); + return new ExecutableWork(work, executeWorkFn); } - public abstract Work work(); + public Work work() { + return work; + } - public abstract BiConsumer executeWorkFn(); + public BiConsumer executeWorkFn() { + return executeWorkFn; + } public void run(BoundedQueueExecutorWorkHandle handle) { try { @@ -50,4 +69,9 @@ public final WorkId id() { public final Windmill.WorkItem getWorkItem() { return work().getWorkItem(); } + + @Override + public String toString() { + return "ExecutableWork{" + id() + "}"; + } } 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 c64ec9bbd568..fe936b3c8da0 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 @@ -17,7 +17,9 @@ */ package org.apache.beam.runners.dataflow.worker.util; -import java.util.Optional; +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadFactory; @@ -244,26 +246,62 @@ public String summaryHtml() { } } - class BoundedQueueExecutorWorkHandleImpl + /** + * A handle to use when requesting pulling more work from @BoundedQueueExecutor + * via @BoundedQueueExecutor.pollWork. A single handle aggregates all budgets from work pulled for + * inline execution and releases the budget after the multi work bundle is complete. + */ + final class BoundedQueueExecutorWorkHandleImpl implements BoundedQueueExecutorWorkHandle, AutoCloseable { + @GuardedBy("this") private int elements; + + @GuardedBy("this") private long bytes; + + @GuardedBy("this") private boolean closed = false; private BoundedQueueExecutorWorkHandleImpl(int elements, long bytes) { + checkArgument(elements >= 0 && bytes >= 0); this.elements = elements; this.bytes = bytes; } - public synchronized void addBudget(int elements, long bytes) { - Preconditions.checkState(!closed, "Cannot add budget to a closed WorkBudgetHandle"); - this.elements += elements; - this.bytes += bytes; + /** + * Merges the budget from another handle into this handle. + * + *

This transfers the budget (elements and bytes) from the {@code other} handle to this + * handle, and marks the {@code other} handle as closed to prevent it from releasing the budget + * again if it is closed. + */ + public void merge(BoundedQueueExecutorWorkHandleImpl other) { + synchronized (this) { + Preconditions.checkState(!closed, "Cannot merge into a closed handle"); + synchronized (checkArgumentNotNull(other)) { + Preconditions.checkState(!other.closed, "Cannot merge a closed handle"); + this.elements += other.elements; + this.bytes += other.bytes; + other.closed = true; + other.elements = 0; + other.bytes = 0; + } + } + } + + public synchronized boolean isClosed() { + return closed; + } + + @VisibleForTesting + synchronized int elements() { + return elements; } - public synchronized void cancel() { - this.closed = true; + @VisibleForTesting + synchronized long bytes() { + return bytes; } @Override @@ -274,38 +312,29 @@ public synchronized void close() { } } - private static class QueuedWork implements Runnable { + private static final class QueuedWork implements Runnable { private final ExecutableWork work; private final BoundedQueueExecutorWorkHandleImpl handle; - private final long workBytes; - public QueuedWork( - ExecutableWork work, BoundedQueueExecutorWorkHandleImpl handle, long workBytes) { + public QueuedWork(ExecutableWork work, BoundedQueueExecutorWorkHandleImpl handle) { this.work = work; this.handle = handle; - this.workBytes = workBytes; - } - - public void cancelHandle() { - handle.cancel(); } public ExecutableWork getWork() { return work; } - public long getWorkBytes() { - return workBytes; + public BoundedQueueExecutorWorkHandleImpl getHandle() { + return handle; } @Override public void run() { - Preconditions.checkArgument(!handle.closed); - try { + checkArgument(!handle.isClosed()); + try (handle) { work.run(handle); - } finally { - handle.close(); } } } @@ -317,7 +346,7 @@ private void executeMonitorHeld(ExecutableWork work, long workBytes) { BoundedQueueExecutorWorkHandleImpl handle = new BoundedQueueExecutorWorkHandleImpl(1, workBytes); try { - executor.execute(new QueuedWork(work, handle, workBytes)); + executor.execute(new QueuedWork(work, handle)); } catch (Throwable e) { handle.close(); throw ExceptionUtils.propagate(e); @@ -344,37 +373,14 @@ private void executeMonitorHeld(Runnable work) { } @VisibleForTesting - BoundedQueueExecutorWorkHandleImpl createEmptyBudgetHandle() { - return new BoundedQueueExecutorWorkHandleImpl(0, 0L); - } - - /** - * Poll additional work to be executed inline inside with the current execute(ExecutableWork work, - * long workBytes) call. It is the responsibility of the caller to execute or discard the returned - * ExecutableWork. Budget for the returned work is released when the execute() call finishes. - * - * @param handle the handle that was passed to ExecutableWork.executeWorkFn - */ - public Optional pollWork(BoundedQueueExecutorWorkHandle handle) { - Preconditions.checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl); - BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; - while (true) { - Runnable runnable = executor.getQueue().poll(); - if (runnable == null) { - return Optional.empty(); - } - if (runnable instanceof QueuedWork) { - QueuedWork queuedWork = (QueuedWork) runnable; - queuedWork.cancelHandle(); - internalHandle.addBudget(1, queuedWork.getWorkBytes()); - return Optional.of(queuedWork.getWork()); - } - // Pop and execute standard callbacks immediately on the calling thread to drain the queue - runnable.run(); - } + BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) { + return new BoundedQueueExecutorWorkHandleImpl(elements, bytes); } private void decrementCounters(int elements, long bytes) { + // All threads queue decrements and one thread grabs the monitor and updates + // counters. We do this to reduce contention on monitor which is locked by + // GetWork thread decrementQueue.add(new Budget(elements, bytes)); boolean submittedToExistingBatch = isDecrementBatchPending.getAndSet(true); if (submittedToExistingBatch) { 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 577d72da9e37..55fe82c7163c 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 @@ -26,7 +26,6 @@ import java.util.Arrays; import java.util.Collection; -import java.util.Optional; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; @@ -388,142 +387,21 @@ public void testRunnableExceptionPropagationDecrementsCounters() throws Exceptio } @Test - public void testPollWorkAndInlineBatchExecution() throws Exception { - BoundedQueueExecutor testExecutor = - new BoundedQueueExecutor( - 1, - DEFAULT_THREAD_EXPIRATION_SEC, - TimeUnit.SECONDS, - 10, - MAXIMUM_BYTES_OUTSTANDING, - new ThreadFactoryBuilder() - .setNameFormat("testPollWorkAndInlineBatchExecution-%d") - .setDaemon(true) - .build(), - useFairMonitor); - - CountDownLatch blockerStart = new CountDownLatch(1); - CountDownLatch blockerStop = new CountDownLatch(1); - ExecutableWork blockerWork = createSleepProcessWork(blockerStart, blockerStop); - - CountDownLatch start1 = new CountDownLatch(1); - CountDownLatch stop1 = new CountDownLatch(1); - ExecutableWork m1 = createSleepProcessWork(start1, stop1); - - CountDownLatch start2 = new CountDownLatch(1); - CountDownLatch stop2 = new CountDownLatch(1); - ExecutableWork m2 = createSleepProcessWork(start2, stop2); - - // 1. Occupy the single worker thread with blocker work so subsequent tasks remain queued. - testExecutor.execute(blockerWork, 0); - blockerStart.await(); - assertEquals(1, testExecutor.elementsOutstanding()); - assertEquals(0, testExecutor.bytesOutstanding()); - - // 2. Enqueue tasks to stay in the queue. - testExecutor.execute(m1, 1000); - testExecutor.execute(m2, 2000); - - assertEquals(3, testExecutor.elementsOutstanding()); - assertEquals(3000, testExecutor.bytesOutstanding()); - - // 3. Create the batch handle. - try (BoundedQueueExecutorWorkHandleImpl batchHandle = testExecutor.createEmptyBudgetHandle()) { - // 4. Poll tasks inline. - Optional polled1 = testExecutor.pollWork(batchHandle); - assertTrue(polled1.isPresent()); - assertEquals(m1, polled1.get()); - - Optional polled2 = testExecutor.pollWork(batchHandle); - assertTrue(polled2.isPresent()); - assertEquals(m2, polled2.get()); - - // Queue should now be empty. - Optional polled3 = testExecutor.pollWork(batchHandle); - assertFalse(polled3.isPresent()); - - // 5. Run polled tasks inline. - start1.countDown(); - stop1.countDown(); - polled1.get().run(batchHandle); - - start2.countDown(); - stop2.countDown(); - polled2.get().run(batchHandle); - - // Outstanding counts should NOT yet be decremented. - assertEquals(3, testExecutor.elementsOutstanding()); - assertEquals(3000, testExecutor.bytesOutstanding()); - } - - // 6. Upon close, outstanding counts should immediately reflect the batch decrement in one shot. - // Only the blocker task (0 bytes, 1 element) should remain outstanding. - while (testExecutor.elementsOutstanding() != 1) { - Thread.sleep(10); - } - assertEquals(1, testExecutor.elementsOutstanding()); - assertEquals(0, testExecutor.bytesOutstanding()); - - // Clean up blocker. - blockerStop.countDown(); - testExecutor.shutdown(); - } - - @Test - public void testPollWorkAndInlineBatchExecutionWithException() throws Exception { - BoundedQueueExecutor testExecutor = - new BoundedQueueExecutor( - 1, - DEFAULT_THREAD_EXPIRATION_SEC, - TimeUnit.SECONDS, - 10, - MAXIMUM_BYTES_OUTSTANDING, - new ThreadFactoryBuilder() - .setNameFormat("testPollWorkAndInlineBatchExecutionWithException-%d") - .setDaemon(true) - .build(), - useFairMonitor); - - CountDownLatch blockerStart = new CountDownLatch(1); - CountDownLatch blockerStop = new CountDownLatch(1); - ExecutableWork blockerWork = createSleepProcessWork(blockerStart, blockerStop); - - // Occupy all worker threads - testExecutor.execute(blockerWork, 0); - blockerStart.await(); - - ExecutableWork inlineWork1 = - createWork( - ignored -> { - throw new RuntimeException("Simulated inline execution exception"); - }); + public void testHandleMerge() throws Exception { + BoundedQueueExecutorWorkHandleImpl handle1 = executor.createBudgetHandle(1, 100L); + BoundedQueueExecutorWorkHandleImpl handle2 = executor.createBudgetHandle(2, 200L); - long size1 = inlineWork1.work().getSerializedWorkItemSize(); - testExecutor.execute(inlineWork1, size1); - - long outstandingBytesBefore = testExecutor.bytesOutstanding(); - int outstandingElementsBefore = testExecutor.elementsOutstanding(); - - try { - try (BoundedQueueExecutorWorkHandleImpl batchHandle = - testExecutor.createEmptyBudgetHandle()) { - Optional polled1 = testExecutor.pollWork(batchHandle); - assertTrue(polled1.isPresent()); - polled1.get().run(batchHandle); - } - } catch (RuntimeException e) { - assertEquals("Simulated inline execution exception", e.getMessage()); - } + handle1.merge(handle2); - // Outstanding elements must still be released cleanly by try-with-resources close! - while (testExecutor.elementsOutstanding() != outstandingElementsBefore - 1) { - Thread.sleep(10); - } - assertEquals(outstandingElementsBefore - 1, testExecutor.elementsOutstanding()); - assertEquals(outstandingBytesBefore - size1, testExecutor.bytesOutstanding()); + // Verify that handle2 has 0 budget and is closed. + assertEquals(0, handle2.elements()); + assertEquals(0, handle2.bytes()); + assertTrue(handle2.isClosed()); - blockerStop.countDown(); - testExecutor.shutdown(); + // Verify that handle1 has the combined budget and is not closed. + assertEquals(3, handle1.elements()); + assertEquals(300L, handle1.bytes()); + assertFalse(handle1.isClosed()); } @Test From d5a44cf3464c5a62bf71d3f7cee7b726997a1bdb Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Sun, 31 May 2026 09:42:37 +0000 Subject: [PATCH 06/11] multikey 5 --- .../BoundedQueueExecutorWorkHandle.java | 4 +- .../worker/streaming/ExecutableWork.java | 4 + .../dataflow/worker/streaming/Work.java | 4 + .../worker/util/BoundedQueueExecutor.java | 29 +- .../worker/util/ComputationWorkQueue.java | 394 +++++++++++++++++ .../processing/StreamingWorkScheduler.java | 2 +- .../worker/util/BoundedQueueExecutorTest.java | 312 ++++++++++++- .../worker/util/ComputationWorkQueueTest.java | 411 ++++++++++++++++++ 8 files changed, 1148 insertions(+), 12 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueue.java create mode 100644 runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java index 1ca534966947..fb96ea8ea151 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java @@ -18,7 +18,7 @@ package org.apache.beam.runners.dataflow.worker.streaming; /** - * A handle to use when requesting pulling more work from @BoundedQueueExecutor - * via @BoundedQueueExecutor.pollWork + * A handle to use when requesting pulling more work from {@link BoundedQueueExecutor} via {@link + * BoundedQueueExecutor#pollWork(String, BoundedQueueExecutorWorkHandle)}. */ public interface BoundedQueueExecutorWorkHandle {} diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java index 5586bc0217b0..134dae6577ab 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java @@ -74,4 +74,8 @@ public final Windmill.WorkItem getWorkItem() { public String toString() { return "ExecutableWork{" + id() + "}"; } + + public final String getComputationId() { + return work().getComputationId(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java index cb01e1e508ce..27c3f21bb73c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java @@ -184,6 +184,10 @@ public long getSerializedWorkItemSize() { return serializedWorkItemSize; } + public String getComputationId() { + return processingContext.computationId(); + } + @Override public ShardedKey getShardedKey() { return shardedKey; 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 fe936b3c8da0..066f6c75fa9c 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 @@ -21,7 +21,6 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadFactory; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; @@ -94,7 +93,7 @@ public BoundedQueueExecutor( initialMaximumPoolSize, keepAliveTime, unit, - new LinkedBlockingQueue<>(), + new ComputationWorkQueue(), threadFactory) { @Override protected void beforeExecute(Thread t, Runnable r) { @@ -277,9 +276,10 @@ private BoundedQueueExecutorWorkHandleImpl(int elements, long bytes) { * again if it is closed. */ public void merge(BoundedQueueExecutorWorkHandleImpl other) { + checkArgumentNotNull(other); synchronized (this) { Preconditions.checkState(!closed, "Cannot merge into a closed handle"); - synchronized (checkArgumentNotNull(other)) { + synchronized (other) { Preconditions.checkState(!other.closed, "Cannot merge a closed handle"); this.elements += other.elements; this.bytes += other.bytes; @@ -312,9 +312,9 @@ public synchronized void close() { } } - private static final class QueuedWork implements Runnable { + static final class QueuedWork implements Runnable { - private final ExecutableWork work; + private volatile ExecutableWork work; private final BoundedQueueExecutorWorkHandleImpl handle; public QueuedWork(ExecutableWork work, BoundedQueueExecutorWorkHandleImpl handle) { @@ -330,6 +330,12 @@ public BoundedQueueExecutorWorkHandleImpl getHandle() { return handle; } + public boolean isCancelled() { + synchronized (handle) { + return handle.closed; + } + } + @Override public void run() { checkArgument(!handle.isClosed()); @@ -377,6 +383,19 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) return new BoundedQueueExecutorWorkHandleImpl(elements, bytes); } + /** Poll work for a specific computationId. */ + public java.util.Optional pollWork( + String computationId, BoundedQueueExecutorWorkHandle handle) { + checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl); + BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; + QueuedWork queuedWork = ((ComputationWorkQueue) executor.getQueue()).pollWork(computationId); + if (queuedWork == null) { + return java.util.Optional.empty(); + } + internalHandle.merge(queuedWork.getHandle()); + return java.util.Optional.of(queuedWork.getWork()); + } + private void decrementCounters(int elements, long bytes) { // All threads queue decrements and one thread grabs the monitor and updates // counters. We do this to reduce contention on monitor which is locked by diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueue.java new file mode 100644 index 000000000000..efa25393fc70 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueue.java @@ -0,0 +1,394 @@ +/* + * 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.util; + +import java.util.AbstractQueue; +import java.util.Collection; +import java.util.HashMap; +import java.util.Iterator; +import java.util.Map; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; +import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor.QueuedWork; + +/** + * A custom, thread-safe doubly-linked BlockingQueue that groups pending tasks by computation ID. + * Achieves true O(1) targeted work stealing and FIFO queue mutations. + */ +@SuppressWarnings({ + "nullness", // Suppress Checker Framework nullness warnings for pointer operations + "initialization" // Suppress initialization and underinitialization warnings for sentinel pointer + // setup +}) +class ComputationWorkQueue extends AbstractQueue implements BlockingQueue { + + static class Node { + final Runnable task; + final String computationId; + + Node prevGlobal; + Node nextGlobal; + Node prevComp; + Node nextComp; + + boolean dequeued = false; + + Node(Runnable task) { + this.task = task; + if (task instanceof QueuedWork) { + this.computationId = ((QueuedWork) task).getWork().getComputationId(); + } else { + this.computationId = null; + } + } + } + + private static class ComputationList { + final Node head; + final Node tail; + + ComputationList() { + head = new Node(null); + tail = new Node(null); + head.nextComp = tail; + tail.prevComp = head; + } + + boolean isEmpty() { + return head.nextComp == tail; + } + + void append(Node node) { + Node last = tail.prevComp; + node.prevComp = last; + node.nextComp = tail; + last.nextComp = node; + tail.prevComp = node; + } + + void remove(Node node) { + if (node.prevComp != null && node.nextComp != null) { + node.prevComp.nextComp = node.nextComp; + node.nextComp.prevComp = node.prevComp; + node.prevComp = null; + node.nextComp = null; + } + } + } + + private final ReentrantLock lock = new ReentrantLock(); + private final Condition notEmpty = lock.newCondition(); + + // Sentinels for the global list + private final Node globalHead = new Node(null); + private final Node globalTail = new Node(null); + + // Map of active computation queues + private final Map compLists = new HashMap<>(); + + private int size = 0; + + public ComputationWorkQueue() { + globalHead.nextGlobal = globalTail; + globalTail.prevGlobal = globalHead; + } + + private void unlinkNode(Node node) { + if (node.dequeued) { + return; + } + node.dequeued = true; + + // 1. Unlink from global list + Node prevG = node.prevGlobal; + Node nextG = node.nextGlobal; + if (prevG != null && nextG != null) { + prevG.nextGlobal = nextG; + nextG.prevGlobal = prevG; + } + node.prevGlobal = null; + node.nextGlobal = null; + + // 2. Unlink from computation list + if (node.computationId != null) { + ComputationList compList = compLists.get(node.computationId); + if (compList != null) { + compList.remove(node); + if (compList.isEmpty()) { + compLists.remove(node.computationId); // Prevent memory leaks of empty keys + } + } + } + + size--; + } + + private Node removeFirstGlobal() { + Node first = globalHead.nextGlobal; + if (first == globalTail) { + return null; + } + unlinkNode(first); + return first; + } + + public QueuedWork pollWork(String computationId) { + if (computationId == null) { + return null; + } + lock.lock(); + try { + ComputationList compList = compLists.get(computationId); + if (compList == null || compList.isEmpty()) { + return null; + } + + // Retrieve the first pending task for this computation in O(1) + Node firstNode = compList.head.nextComp; + unlinkNode(firstNode); + + return (QueuedWork) firstNode.task; + } finally { + lock.unlock(); + } + } + + @Override + public boolean offer(Runnable e) { + if (e == null) throw new NullPointerException(); + lock.lock(); + try { + Node node = new Node(e); + + // Append to global list tail + Node lastG = globalTail.prevGlobal; + node.prevGlobal = lastG; + node.nextGlobal = globalTail; + lastG.nextGlobal = node; + globalTail.prevGlobal = node; + + // Append to computation list if applicable + if (node.computationId != null) { + ComputationList compList = + compLists.computeIfAbsent(node.computationId, k -> new ComputationList()); + compList.append(node); + } + + size++; + notEmpty.signal(); + return true; + } finally { + lock.unlock(); + } + } + + @Override + public void put(Runnable e) throws InterruptedException { + offer(e); // Unbounded queue + } + + @Override + public boolean offer(Runnable e, long timeout, TimeUnit unit) throws InterruptedException { + return offer(e); // Unbounded queue + } + + @Override + public Runnable poll() { + lock.lock(); + try { + Node node = removeFirstGlobal(); + return (node != null) ? node.task : null; + } finally { + lock.unlock(); + } + } + + @Override + public Runnable take() throws InterruptedException { + lock.lockInterruptibly(); + try { + while (size == 0) { + notEmpty.await(); + } + Node node = removeFirstGlobal(); + return node.task; + } finally { + lock.unlock(); + } + } + + @Override + public Runnable poll(long timeout, TimeUnit unit) throws InterruptedException { + long nanos = unit.toNanos(timeout); + lock.lockInterruptibly(); + try { + while (size == 0) { + if (nanos <= 0) { + return null; + } + nanos = notEmpty.awaitNanos(nanos); + } + Node node = removeFirstGlobal(); + return node.task; + } finally { + lock.unlock(); + } + } + + @Override + public Runnable peek() { + lock.lock(); + try { + Node first = globalHead.nextGlobal; + if (first == globalTail) { + return null; + } + return first.task; + } finally { + lock.unlock(); + } + } + + @Override + public int size() { + lock.lock(); + try { + return size; + } finally { + lock.unlock(); + } + } + + @Override + public boolean isEmpty() { + lock.lock(); + try { + return size == 0; + } finally { + lock.unlock(); + } + } + + @Override + public boolean remove(Object o) { + if (o == null) return false; + lock.lock(); + try { + // Walk the global queue in O(N) to find and unlink the node + Node curr = globalHead.nextGlobal; + while (curr != globalTail) { + if (curr.task.equals(o)) { + unlinkNode(curr); + return true; + } + curr = curr.nextGlobal; + } + return false; + } finally { + lock.unlock(); + } + } + + @Override + public boolean contains(Object o) { + if (o == null) return false; + lock.lock(); + try { + Node curr = globalHead.nextGlobal; + while (curr != globalTail) { + if (curr.task.equals(o)) { + return true; + } + curr = curr.nextGlobal; + } + return false; + } finally { + lock.unlock(); + } + } + + @Override + public int remainingCapacity() { + return Integer.MAX_VALUE; + } + + @Override + public int drainTo(Collection c) { + return drainTo(c, Integer.MAX_VALUE); + } + + @Override + public int drainTo(Collection c, int maxElements) { + if (c == null) throw new NullPointerException(); + if (c == this) throw new IllegalArgumentException(); + if (maxElements <= 0) return 0; + lock.lock(); + try { + int added = 0; + Node curr = globalHead.nextGlobal; + while (curr != globalTail && added < maxElements) { + Node next = curr.nextGlobal; + unlinkNode(curr); + c.add(curr.task); + added++; + curr = next; + } + return added; + } finally { + lock.unlock(); + } + } + + @Override + public void clear() { + // WARNING: Calling clear() directly on this queue will unlink tasks but will NOT + // close their respective work budget handles, leaking elementsOutstanding/bytesOutstanding + // in BoundedQueueExecutor. Call executor.remove(r) or let worker threads consume them. + lock.lock(); + try { + Node curr = globalHead.nextGlobal; + while (curr != globalTail) { + Node next = curr.nextGlobal; + unlinkNode(curr); + curr = next; + } + } finally { + lock.unlock(); + } + } + + @Override + public Iterator iterator() { + lock.lock(); + try { + java.util.List snapshot = new java.util.ArrayList<>(size); + Node curr = globalHead.nextGlobal; + while (curr != globalTail) { + if (curr.task != null) { + snapshot.add(curr.task); + } + curr = curr.nextGlobal; + } + return java.util.Collections.unmodifiableList(snapshot).iterator(); + } finally { + lock.unlock(); + } + } +} 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 0c3289d4cf58..dbcb0f09c0a9 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 @@ -234,7 +234,7 @@ public void queueAppliedFinalizeIds(ImmutableList appliedFinalizeIds) { * internally if processing fails due to uncaught {@link Exception}(s). * * @implNote This will block the calling thread during execution of user DoFns. - * @param handle handled to pass to BoundedQueueExecutor.pollWork, currently unused + * @param handle budget handle associated with the executed work */ private void processWork( ComputationState computationState, 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 55fe82c7163c..98e12a3c86f1 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 @@ -67,6 +67,11 @@ public static Collection useFairMonitor() { private BoundedQueueExecutor executor; private static ExecutableWork createWork(Consumer executeWorkFn) { + return createWorkWithCompId("computationId", executeWorkFn); + } + + private static ExecutableWork createWorkWithCompId( + String computationId, Consumer executeWorkFn) { WorkItem workItem = WorkItem.newBuilder() .setKey(ByteString.EMPTY) @@ -80,10 +85,7 @@ private static ExecutableWork createWork(Consumer executeWorkFn) { workItem.getSerializedSize(), Watermarks.builder().setInputDataWatermark(Instant.now()).build(), Work.createProcessingContext( - "computationId", - new FakeGetDataClient(), - ignored -> {}, - mock(HeartbeatSender.class)), + computationId, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), false, Instant::now), (work, handle) -> { @@ -404,6 +406,308 @@ public void testHandleMerge() throws Exception { assertFalse(handle1.isClosed()); } + @Test + public void testTombstoneExecutionAndQueueSize() throws Exception { + BoundedQueueExecutor testExecutor = + new BoundedQueueExecutor( + 1, // 1 thread to strictly control execution order + DEFAULT_THREAD_EXPIRATION_SEC, + TimeUnit.SECONDS, + 10, + MAXIMUM_BYTES_OUTSTANDING, + new ThreadFactoryBuilder() + .setNameFormat("testTombstoneExecutionAndQueueSize-%d") + .setDaemon(true) + .build(), + useFairMonitor); + + CountDownLatch blockerStart = new CountDownLatch(1); + CountDownLatch blockerStop = new CountDownLatch(1); + ExecutableWork blockerWork = + createWorkWithCompId( + "comp-1", + ignored -> { + blockerStart.countDown(); + try { + blockerStop.await(); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + + CountDownLatch targetStart = new CountDownLatch(1); + ExecutableWork targetWork = + createWorkWithCompId( + "comp-2", + ignored -> { + targetStart.countDown(); + }); + + // 1. Occupy the worker thread with blockerWork + testExecutor.execute(blockerWork, 0); + blockerStart.await(); + + // 2. Enqueue targetWork (goes into queue) + testExecutor.execute(targetWork, 1000); + + // Wait a moment to ensure targetWork is registered in the queue + assertEquals(2, testExecutor.elementsOutstanding()); + assertEquals(1, testExecutor.executorQueueIsEmpty() ? 0 : 1); // it is in the queue + + // 3. Now steal targetWork using targeted pollWork from a "stealer" context + try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { + java.util.Optional stolen = testExecutor.pollWork("comp-2", stealHandle); + assertTrue(stolen.isPresent()); + assertEquals(targetWork, stolen.get()); + + // Since it's stolen, the budget is transferred to stealHandle. + // The outstanding size must still be 2 (blocker + stolen in stealHandle). + assertEquals(2, testExecutor.elementsOutstanding()); + + // 4. Run the stolen work inline in the stealer thread. + stolen.get().run(stealHandle); + targetStart.await(); + } + + // After stealHandle is closed, the outstanding size should decrement by 1 (the stolen task). + // Only the blocker remains. + while (testExecutor.elementsOutstanding() != 1) { + Thread.sleep(10); + } + assertEquals(1, testExecutor.elementsOutstanding()); + + // 5. Unblock the blocker. The executor worker thread should now pop the targetWork tombstone. + // Since it is cancelled (stolen), the worker thread should skip it as a no-op. + // The elementsOutstanding should drop to 0 after the blocker completes. + blockerStop.countDown(); + + while (testExecutor.elementsOutstanding() != 0) { + Thread.sleep(10); + } + assertEquals(0, testExecutor.elementsOutstanding()); + testExecutor.shutdown(); + } + + @Test + public void testConcurrentStealingAndPolling() throws Exception { + final int numComputations = 10; + final int tasksPerComp = 100; + final int totalTasks = numComputations * tasksPerComp; + + // 4 executor worker threads to execute work + BoundedQueueExecutor testExecutor = + new BoundedQueueExecutor( + 4, + DEFAULT_THREAD_EXPIRATION_SEC, + TimeUnit.SECONDS, + totalTasks * 2, + MAXIMUM_BYTES_OUTSTANDING, + new ThreadFactoryBuilder() + .setNameFormat("testConcurrentStealingAndPolling-%d") + .setDaemon(true) + .build(), + useFairMonitor); + + // Latches to coordinate + CountDownLatch allDone = new CountDownLatch(totalTasks); + + // Enqueue work for all computations + for (int i = 0; i < numComputations; i++) { + final String compId = "comp-" + i; + for (int j = 0; j < tasksPerComp; j++) { + ExecutableWork work = + createWorkWithCompId( + compId, + ignored -> { + allDone.countDown(); + }); + testExecutor.execute(work, 100); + } + } + + // Launch active "stealers" that target specific computations concurrently + Thread[] stealers = new Thread[numComputations]; + for (int i = 0; i < numComputations; i++) { + final String compId = "comp-" + i; + stealers[i] = + new Thread( + () -> { + try { + // Attempt to steal tasks for this computation + int stolenCount = 0; + while (stolenCount < tasksPerComp) { + try (BoundedQueueExecutorWorkHandleImpl stealHandle = + testExecutor.createBudgetHandle(0, 0L)) { + java.util.Optional stolen = + testExecutor.pollWork(compId, stealHandle); + if (stolen.isPresent()) { + stolen.get().run(stealHandle); + stolenCount++; + } else { + // Yield if none available + Thread.sleep(1); + } + } catch (InterruptedException e) { + break; + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + stealers[i].start(); + } + + // Wait for all tasks to complete (either executed by thread pool or stolen) + assertTrue(allDone.await(30, TimeUnit.SECONDS)); + + for (Thread stealer : stealers) { + stealer.interrupt(); + stealer.join(); + } + + // Wait until all outstanding elements are released cleanly + while (testExecutor.elementsOutstanding() != 0) { + Thread.sleep(10); + } + assertEquals(0, testExecutor.elementsOutstanding()); + assertEquals(0, testExecutor.bytesOutstanding()); + testExecutor.shutdown(); + } + + @Test + public void testConcurrentStealingAndPollingNoWastedThreads() throws Exception { + final int numComputations = 10; + final int tasksPerComp = 100; + final int totalTasks = numComputations * tasksPerComp; + + // 4 executor worker threads to execute work + BoundedQueueExecutor testExecutor = + new BoundedQueueExecutor( + 4, + DEFAULT_THREAD_EXPIRATION_SEC, + TimeUnit.SECONDS, + totalTasks * 2, + MAXIMUM_BYTES_OUTSTANDING, + new ThreadFactoryBuilder() + .setNameFormat("testConcurrentStealingAndPollingNoWastedThreads-%d") + .setDaemon(true) + .build(), + useFairMonitor); + + // Thread-safe counters for tracking executions + java.util.concurrent.atomic.AtomicInteger totalExecuted = + new java.util.concurrent.atomic.AtomicInteger(0); + java.util.concurrent.atomic.AtomicInteger workerExecuted = + new java.util.concurrent.atomic.AtomicInteger(0); + java.util.concurrent.atomic.AtomicInteger poolWorkerExecuted = + new java.util.concurrent.atomic.AtomicInteger(0); + + // Map to track duplicate executions + java.util.concurrent.ConcurrentHashMap + executionCounts = new java.util.concurrent.ConcurrentHashMap<>(); + + // Latches to coordinate + CountDownLatch allDone = new CountDownLatch(totalTasks); + + // Enqueue work for all computations + for (int i = 0; i < numComputations; i++) { + final String compId = "comp-" + i; + for (int j = 0; j < tasksPerComp; j++) { + final String taskId = compId + "-task-" + j; + ExecutableWork work = + createWorkWithCompId( + compId, + ignored -> { + int count = + executionCounts + .computeIfAbsent( + taskId, k -> new java.util.concurrent.atomic.AtomicInteger(0)) + .incrementAndGet(); + totalExecuted.incrementAndGet(); + String threadName = Thread.currentThread().getName(); + if (threadName.startsWith("testConcurrentStealingAndPollingNoWastedThreads-")) { + poolWorkerExecuted.incrementAndGet(); + } else { + workerExecuted.incrementAndGet(); + } + allDone.countDown(); + }); + testExecutor.execute(work, 100); + } + } + + // Launch active "stealers" that target specific computations concurrently + Thread[] stealers = new Thread[numComputations]; + for (int i = 0; i < numComputations; i++) { + final String compId = "comp-" + i; + stealers[i] = + new Thread( + () -> { + try { + while (!Thread.currentThread().isInterrupted()) { + try (BoundedQueueExecutorWorkHandleImpl stealHandle = + testExecutor.createBudgetHandle(0, 0L)) { + java.util.Optional stolen = + testExecutor.pollWork(compId, stealHandle); + if (stolen.isPresent()) { + stolen.get().run(stealHandle); + } else { + Thread.sleep(1); + } + } catch (InterruptedException e) { + break; + } + } + } catch (Exception e) { + // Ignore or log + } + }); + stealers[i].start(); + } + + // Wait for all tasks to complete (either executed by thread pool or stolen) + assertTrue(allDone.await(60, TimeUnit.SECONDS)); + + // Stop stealers + for (Thread stealer : stealers) { + stealer.interrupt(); + stealer.join(); + } + + // Verify execution invariants + assertEquals("Total executions must match total tasks", totalTasks, totalExecuted.get()); + assertEquals("No duplicate executions allowed", totalTasks, executionCounts.size()); + for (java.util.Map.Entry entry : + executionCounts.entrySet()) { + assertEquals( + "Task " + entry.getKey() + " executed more than once", 1, entry.getValue().get()); + } + + // Verify that completed task count of the inner executor perfectly matches poolWorkerExecuted + java.lang.reflect.Field executorField = BoundedQueueExecutor.class.getDeclaredField("executor"); + executorField.setAccessible(true); + java.util.concurrent.ThreadPoolExecutor innerExecutor = + (java.util.concurrent.ThreadPoolExecutor) executorField.get(testExecutor); + + // Wait a tiny bit for any pending pool tasks to fully complete their afterExecute/bookkeeping + Thread.sleep(100); + + long completedTasks = innerExecutor.getCompletedTaskCount(); + assertEquals( + "Completed task count of the pool must perfectly match the number of worker-executed tasks (exactly 0 wasted activations)", + (long) poolWorkerExecuted.get(), + completedTasks); + + // Wait until all outstanding elements are released cleanly + while (testExecutor.elementsOutstanding() != 0) { + Thread.sleep(10); + } + assertEquals(0, testExecutor.elementsOutstanding()); + assertEquals(0, testExecutor.bytesOutstanding()); + testExecutor.shutdown(); + } + @Test public void testRenderSummaryHtml() { String expectedSummaryHtml = diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java new file mode 100644 index 000000000000..e5025d5a29ef --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java @@ -0,0 +1,411 @@ +/* + * 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.util; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; +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.QueuedWork; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; +import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; +import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +@SuppressWarnings({ + "ReturnValueIgnored", + "UnusedVariable", + "FutureReturnValueIgnored", + "CatchAndPrintStackTrace", + "ThreadPriorityCheck", + "nullness" +}) +public class ComputationWorkQueueTest { + + private BoundedQueueExecutor executor; + + @Before + public void setUp() { + executor = + new BoundedQueueExecutor( + 2, + 60, + TimeUnit.SECONDS, + 100, + 10000000, + new ThreadFactoryBuilder().setNameFormat("Test-%d").setDaemon(true).build(), + false); + } + + private QueuedWork createQueuedWork(String computationId, long workBytes) { + WorkItem workItem = + WorkItem.newBuilder() + .setKey(ByteString.EMPTY) + .setShardingKey(1) + .setWorkToken(33) + .setCacheToken(1) + .build(); + ExecutableWork work = + 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), + (w, h) -> {}); + return new QueuedWork(work, executor.createBudgetHandle(1, workBytes)); + } + + private static class MockRunnable implements Runnable { + final String id; + + MockRunnable(String id) { + this.id = id; + } + + @Override + public void run() {} + } + + @Test + public void testBasicOfferAndPoll() { + ComputationWorkQueue queue = new ComputationWorkQueue(); + assertTrue(queue.isEmpty()); + assertEquals(0, queue.size()); + + MockRunnable task1 = new MockRunnable("1"); + MockRunnable task2 = new MockRunnable("2"); + + assertTrue(queue.offer(task1)); + assertTrue(queue.offer(task2)); + assertEquals(2, queue.size()); + + assertEquals(task1, queue.poll()); + assertEquals(task2, queue.poll()); + assertNull(queue.poll()); + assertTrue(queue.isEmpty()); + } + + @Test + public void testRemove() { + ComputationWorkQueue queue = new ComputationWorkQueue(); + MockRunnable task1 = new MockRunnable("1"); + MockRunnable task2 = new MockRunnable("2"); + + queue.offer(task1); + queue.offer(task2); + + assertTrue(queue.remove(task1)); + assertEquals(1, queue.size()); + assertEquals(task2, queue.poll()); + assertFalse(queue.remove(task1)); // Already gone + } + + @Test + public void testDrainTo() { + ComputationWorkQueue queue = new ComputationWorkQueue(); + MockRunnable task1 = new MockRunnable("1"); + MockRunnable task2 = new MockRunnable("2"); + queue.offer(task1); + queue.offer(task2); + + List drained = new ArrayList<>(); + assertEquals(2, queue.drainTo(drained)); + assertEquals(2, drained.size()); + assertEquals(task1, drained.get(0)); + assertEquals(task2, drained.get(1)); + assertTrue(queue.isEmpty()); + } + + @Test + public void testIteratorSafeTraversalAndImmutable() { + ComputationWorkQueue queue = new ComputationWorkQueue(); + MockRunnable task1 = new MockRunnable("1"); + MockRunnable task2 = new MockRunnable("2"); + queue.offer(task1); + queue.offer(task2); + + java.util.Iterator it = queue.iterator(); + assertTrue(it.hasNext()); + assertEquals(task1, it.next()); + assertTrue(it.hasNext()); + assertEquals(task2, it.next()); + assertFalse(it.hasNext()); + + // Assert that mutating the iterator throws UnsupportedOperationException + it = queue.iterator(); + assertTrue(it.hasNext()); + it.next(); + try { + it.remove(); + org.junit.Assert.fail("Iterator must be immutable"); + } catch (UnsupportedOperationException e) { + // Expected + } + } + + @Test + public void testPollWorkTargeted() { + ComputationWorkQueue queue = new ComputationWorkQueue(); + + QueuedWork workA1 = createQueuedWork("compA", 100); + QueuedWork workB1 = createQueuedWork("compB", 200); + QueuedWork workA2 = createQueuedWork("compA", 150); + + queue.offer(workA1); + queue.offer(workB1); + queue.offer(workA2); + + assertEquals(3, queue.size()); + + // Targeted poll A + QueuedWork polledA1 = queue.pollWork("compA"); + assertNotNull(polledA1); + assertEquals("compA", polledA1.getWork().getComputationId()); + assertEquals(100, polledA1.getHandle().bytes()); + + // Verify size decremented + assertEquals(2, queue.size()); + + // Poll next should be B1 (since A1 was stolen, B1 is now first global) + assertEquals(workB1, queue.poll()); + assertEquals(1, queue.size()); + + // Last should be A2 + assertEquals(workA2, queue.poll()); + assertTrue(queue.isEmpty()); + } + + @Test + public void testMemoryPruningLeavesZeroLeaks() { + ComputationWorkQueue queue = new ComputationWorkQueue(); + QueuedWork workA1 = createQueuedWork("compA", 100); + queue.offer(workA1); + + // Steal A1 + QueuedWork polled = queue.pollWork("compA"); + assertNotNull(polled); + assertTrue(queue.isEmpty()); + + // Offering another work with different computation ID + QueuedWork workB1 = createQueuedWork("compB", 200); + queue.offer(workB1); + assertEquals(1, queue.size()); + + // Steal B1 + QueuedWork polledB = queue.pollWork("compB"); + assertNotNull(polledB); + assertTrue(queue.isEmpty()); + } + + @Test + public void testConcurrentStress() throws InterruptedException { + final ComputationWorkQueue queue = new ComputationWorkQueue(); + final int producerThreads = 4; + final int consumerThreads = 4; + final int tasksPerProducer = 1000; + final int totalTasks = producerThreads * tasksPerProducer; + + ExecutorService executorService = + Executors.newFixedThreadPool(producerThreads + consumerThreads); + final CountDownLatch startLatch = new CountDownLatch(1); + final CountDownLatch doneLatch = new CountDownLatch(producerThreads + consumerThreads); + final AtomicInteger consumedCount = new AtomicInteger(0); + + // Start producers + for (int i = 0; i < producerThreads; i++) { + executorService.submit( + () -> { + try { + startLatch.await(); + for (int j = 0; j < tasksPerProducer; j++) { + String compId = "comp-" + (j % 5); + queue.offer(createQueuedWork(compId, 10)); + } + } catch (Exception e) { + e.printStackTrace(); + } finally { + doneLatch.countDown(); + } + }); + } + + // Start consumers (mix of poll and pollWork) + for (int i = 0; i < consumerThreads; i++) { + final int consumerId = i; + executorService.submit( + () -> { + try { + startLatch.await(); + while (consumedCount.get() < totalTasks) { + Runnable task; + if (consumerId % 2 == 0) { + // Targeted poll + String compId = "comp-" + (consumedCount.get() % 5); + task = queue.pollWork(compId); + } else { + // Global poll + task = queue.poll(); + } + if (task != null) { + consumedCount.incrementAndGet(); + } else { + Thread.yield(); + } + } + } catch (Exception e) { + e.printStackTrace(); + } finally { + doneLatch.countDown(); + } + }); + } + + startLatch.countDown(); + assertTrue(doneLatch.await(10, TimeUnit.SECONDS)); + executorService.shutdown(); + assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS)); + + assertEquals(0, queue.size()); + assertTrue(queue.isEmpty()); + } + + @Test + public void testTakeBlocksAndWakesUp() throws InterruptedException { + final ComputationWorkQueue queue = new ComputationWorkQueue(); + final MockRunnable task = new MockRunnable("take-task"); + final java.util.concurrent.atomic.AtomicReference result = + new java.util.concurrent.atomic.AtomicReference<>(); + final CountDownLatch started = new CountDownLatch(1); + final CountDownLatch finished = new CountDownLatch(1); + + Thread t = + new Thread( + () -> { + started.countDown(); + try { + result.set(queue.take()); + } catch (InterruptedException e) { + // Ignore + } finally { + finished.countDown(); + } + }); + t.setDaemon(true); + t.start(); + + assertTrue(started.await(2, TimeUnit.SECONDS)); + // Give thread a moment to enter await() + Thread.sleep(100); + assertEquals(Thread.State.WAITING, t.getState()); + + queue.offer(task); + + assertTrue(finished.await(2, TimeUnit.SECONDS)); + assertEquals(task, result.get()); + } + + @Test + public void testPollWithTimeout() throws InterruptedException { + final ComputationWorkQueue queue = new ComputationWorkQueue(); + final MockRunnable task = new MockRunnable("poll-task"); + final java.util.concurrent.atomic.AtomicReference result = + new java.util.concurrent.atomic.AtomicReference<>(); + final CountDownLatch started = new CountDownLatch(1); + final CountDownLatch finished = new CountDownLatch(1); + + // 1. Verify timeout returns null + Thread t1 = + new Thread( + () -> { + started.countDown(); + try { + result.set(queue.poll(50, TimeUnit.MILLISECONDS)); + } catch (InterruptedException e) { + // Ignore + } finally { + finished.countDown(); + } + }); + t1.setDaemon(true); + t1.start(); + + assertTrue(started.await(2, TimeUnit.SECONDS)); + Thread.sleep(10); + assertEquals(Thread.State.TIMED_WAITING, t1.getState()); + + assertTrue(finished.await(2, TimeUnit.SECONDS)); + assertNull(result.get()); + + // 2. Verify timed poll receives task offered concurrently + final CountDownLatch started2 = new CountDownLatch(1); + final CountDownLatch finished2 = new CountDownLatch(1); + final java.util.concurrent.atomic.AtomicReference result2 = + new java.util.concurrent.atomic.AtomicReference<>(); + + Thread t2 = + new Thread( + () -> { + started2.countDown(); + try { + result2.set(queue.poll(2, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + // Ignore + } finally { + finished2.countDown(); + } + }); + t2.setDaemon(true); + t2.start(); + + assertTrue(started2.await(2, TimeUnit.SECONDS)); + Thread.sleep(50); + assertEquals(Thread.State.TIMED_WAITING, t2.getState()); + + queue.offer(task); + + assertTrue(finished2.await(2, TimeUnit.SECONDS)); + assertEquals(task, result2.get()); + } +} From 3f92d406fa295ae0d5aca6b85e90150e88cda21f Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Sun, 31 May 2026 11:06:38 +0000 Subject: [PATCH 07/11] Update windmill protos for multikey bundles --- .../windmill/src/main/proto/windmill.proto | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto b/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto index 1da7ef9be8bb..a7a99e2ca5a1 100644 --- a/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto +++ b/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto @@ -421,6 +421,11 @@ message WatermarkHold { optional string state_family = 4; } +message Uint128Proto { + required fixed64 high = 1; + required fixed64 low = 2; +} + // Proto describing a hot key detected on a given WorkItem. message HotKeyInfo { // The age of the hot key measured from when it was first detected. @@ -448,6 +453,8 @@ message WorkItem { // present, this field includes metadata associated with any hot key. optional HotKeyInfo hot_key_info = 11; + optional Uint128Proto key_group = 18; + reserved 12, 13, 14, 15, 16; } @@ -671,9 +678,24 @@ message WorkItemCommitRequest { reserved 6, 23; } +message MultiKeyWorkItemCommitRequest { + optional Uint128Proto key_group = 7; + + repeated WorkItemCommitRequest requests = 1; + + repeated OutputMessageBundle output_messages = 2; + + repeated PubSubMessageBundle pubsub_messages = 3; + + repeated int64 finalize_ids = 4 [packed = true]; + + reserved 6; +} + message ComputationCommitWorkRequest { required string computation_id = 1; repeated WorkItemCommitRequest requests = 2; + repeated MultiKeyWorkItemCommitRequest multi_key_requests = 3; } message CommitWorkRequest { @@ -899,6 +921,14 @@ message StreamingCommitRequestChunk { // before handing off to the WindmillHost for processing. optional int64 remaining_bytes_for_work_item = 4; optional bytes serialized_work_item_commit = 5; + + enum CommitType { + COMMIT_TYPE_UNSPECIFIED = 0; + COMMIT_TYPE_SINGLE_KEY = 1; + COMMIT_TYPE_MULTI_KEY = 2; + } + + optional CommitType commit_type = 7; } message StreamingCommitResponse { From 68fce830318d2b69e8d0186787060c1362606ebc Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Sun, 31 May 2026 11:39:04 +0000 Subject: [PATCH 08/11] introduce key group queue --- .../worker/StreamingDataflowWorker.java | 7 +- .../worker/streaming/ExecutableWork.java | 5 + .../dataflow/worker/streaming/Work.java | 56 +++ .../worker/util/BoundedQueueExecutor.java | 41 +- .../worker/util/ComputationWorkQueue.java | 80 +++- .../worker/util/BoundedQueueExecutorTest.java | 357 +++++------------- .../worker/util/ComputationWorkQueueTest.java | 51 ++- 7 files changed, 313 insertions(+), 284 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java index 4d070da995b3..8ab0abd6772d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java @@ -178,6 +178,8 @@ public final class StreamingDataflowWorker { // Experiment make the monitor within BoundedQueueExecutor fair public static final String BOUNDED_QUEUE_EXECUTOR_USE_FAIR_MONITOR_EXPERIMENT = "windmill_bounded_queue_executor_use_fair_monitor"; + public static final String ENABLE_KEY_GROUP_WORK_QUEUE_EXPERIMENT = + "unstable_enable_multi_key_bundle"; private final WindmillStateCache stateCache; private AtomicReference statusPages = new AtomicReference<>(); @@ -1017,6 +1019,8 @@ private static JobHeader createJobHeader(DataflowWorkerHarnessOptions options, l private static BoundedQueueExecutor createWorkUnitExecutor(DataflowWorkerHarnessOptions options) { boolean useFairMonitor = DataflowRunner.hasExperiment(options, BOUNDED_QUEUE_EXECUTOR_USE_FAIR_MONITOR_EXPERIMENT); + boolean useKeyGroupWorkQueue = + DataflowRunner.hasExperiment(options, ENABLE_KEY_GROUP_WORK_QUEUE_EXPERIMENT); return new BoundedQueueExecutor( chooseMaxThreads(options), THREAD_EXPIRATION_TIME_SEC, @@ -1024,7 +1028,8 @@ private static BoundedQueueExecutor createWorkUnitExecutor(DataflowWorkerHarness chooseMaxBundlesOutstanding(options), chooseMaxBytesOutstanding(options), new ThreadFactoryBuilder().setNameFormat("DataflowWorkUnits-%d").setDaemon(true).build(), - useFairMonitor); + useFairMonitor, + useKeyGroupWorkQueue); } public static void main(String[] args) throws Exception { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java index 134dae6577ab..f18cdf3b58e7 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.dataflow.worker.streaming; import java.util.Objects; +import java.util.Optional; import java.util.function.BiConsumer; import org.apache.beam.runners.dataflow.worker.util.ExceptionUtils; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; @@ -78,4 +79,8 @@ public String toString() { public final String getComputationId() { return work().getComputationId(); } + + public final Optional getKeyGroup() { + return work().getKeyGroup(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java index 27c3f21bb73c..b0fda0e0516b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java @@ -25,6 +25,7 @@ import java.util.IntSummaryStatistics; import java.util.Map; import java.util.Map.Entry; +import java.util.Objects; import java.util.Optional; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; @@ -52,6 +53,7 @@ import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; @@ -74,6 +76,7 @@ public final class Work implements RefreshableWork { private final Instant startTime; private final Map totalDurationPerState; private final WorkId id; + private final Optional keyGroup; private final String latencyTrackingId; private final long serializedWorkItemSize; private volatile TimedState currentState; @@ -101,6 +104,11 @@ private Work( // keyUniverse inside EnumMap every time. this.totalDurationPerState = new EnumMap<>(EMPTY_ENUM_MAP); this.id = WorkId.of(workItem); + this.keyGroup = + workItem.hasKeyGroup() + ? Optional.of( + KeyGroup.create(workItem.getKeyGroup().getHigh(), workItem.getKeyGroup().getLow())) + : Optional.empty(); this.latencyTrackingId = Long.toHexString(workItem.getShardingKey()) + '-' @@ -274,6 +282,10 @@ public WorkId id() { return id; } + public Optional getKeyGroup() { + return keyGroup; + } + public void recordGetWorkStreamLatencies( ImmutableList getWorkStreamLatencies) { for (LatencyAttribution latency : getWorkStreamLatencies) { @@ -420,4 +432,48 @@ private Optional fetchKeyedState(KeyedGetDataRequest reque return Optional.ofNullable(getDataClient().getStateData(computationId(), request)); } } + + public static final class KeyGroup { + private final long high; + private final long low; + + private KeyGroup(long high, long low) { + this.high = high; + this.low = low; + } + + public static KeyGroup create(long high, long low) { + return new KeyGroup(high, low); + } + + public long high() { + return high; + } + + public long low() { + return low; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof KeyGroup)) { + return false; + } + KeyGroup other = (KeyGroup) o; + return high == other.high && low == other.low; + } + + @Override + public int hashCode() { + return Objects.hash(high, low); + } + + @Override + public String toString() { + return "KeyGroup{" + "high=" + high + ", low=" + low + '}'; + } + } } 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 066f6c75fa9c..d1b9f6a3ae03 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 @@ -20,7 +20,9 @@ import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import java.util.Optional; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadFactory; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; @@ -28,6 +30,7 @@ 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.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; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor; @@ -85,6 +88,26 @@ public BoundedQueueExecutor( long maximumBytesOutstanding, ThreadFactory threadFactory, boolean useFairMonitor) { + this( + initialMaximumPoolSize, + keepAliveTime, + unit, + maximumElementsOutstanding, + maximumBytesOutstanding, + threadFactory, + useFairMonitor, + /* useKeyGroupWorkQueue= */ true); + } + + public BoundedQueueExecutor( + int initialMaximumPoolSize, + long keepAliveTime, + TimeUnit unit, + int maximumElementsOutstanding, + long maximumBytesOutstanding, + ThreadFactory threadFactory, + boolean useFairMonitor, + boolean useKeyGroupWorkQueue) { this.maximumPoolSize = initialMaximumPoolSize; monitor = new Monitor(useFairMonitor); executor = @@ -93,7 +116,7 @@ public BoundedQueueExecutor( initialMaximumPoolSize, keepAliveTime, unit, - new ComputationWorkQueue(), + useKeyGroupWorkQueue ? new ComputationWorkQueue() : new LinkedBlockingQueue<>(), threadFactory) { @Override protected void beforeExecute(Thread t, Runnable r) { @@ -383,17 +406,21 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) return new BoundedQueueExecutorWorkHandleImpl(elements, bytes); } - /** Poll work for a specific computationId. */ - public java.util.Optional pollWork( - String computationId, BoundedQueueExecutorWorkHandle handle) { + /** Poll work for a specific computationId and keyGroup. */ + public Optional pollWork( + String computationId, Work.KeyGroup keyGroup, BoundedQueueExecutorWorkHandle handle) { checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl); BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; - QueuedWork queuedWork = ((ComputationWorkQueue) executor.getQueue()).pollWork(computationId); + if (!(executor.getQueue() instanceof ComputationWorkQueue)) { + return Optional.empty(); + } + QueuedWork queuedWork = + ((ComputationWorkQueue) executor.getQueue()).pollWork(computationId, keyGroup); if (queuedWork == null) { - return java.util.Optional.empty(); + return Optional.empty(); } internalHandle.merge(queuedWork.getHandle()); - return java.util.Optional.of(queuedWork.getWork()); + return Optional.of(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/util/ComputationWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueue.java index efa25393fc70..e4b3c09e8323 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueue.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueue.java @@ -22,11 +22,14 @@ import java.util.HashMap; import java.util.Iterator; import java.util.Map; +import java.util.Objects; import java.util.concurrent.BlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; +import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor.QueuedWork; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A custom, thread-safe doubly-linked BlockingQueue that groups pending tasks by computation ID. @@ -42,6 +45,7 @@ class ComputationWorkQueue extends AbstractQueue implements BlockingQu static class Node { final Runnable task; final String computationId; + final Work.KeyGroup keyGroup; Node prevGlobal; Node nextGlobal; @@ -54,17 +58,19 @@ static class Node { this.task = task; if (task instanceof QueuedWork) { this.computationId = ((QueuedWork) task).getWork().getComputationId(); + this.keyGroup = ((QueuedWork) task).getWork().getKeyGroup().orElse(null); } else { this.computationId = null; + this.keyGroup = null; } } } - private static class ComputationList { + private static class KeyGroupWorkQueue { final Node head; final Node tail; - ComputationList() { + KeyGroupWorkQueue() { head = new Node(null); tail = new Node(null); head.nextComp = tail; @@ -101,7 +107,7 @@ void remove(Node node) { private final Node globalTail = new Node(null); // Map of active computation queues - private final Map compLists = new HashMap<>(); + private final Map compLists = new HashMap<>(); private int size = 0; @@ -128,11 +134,12 @@ private void unlinkNode(Node node) { // 2. Unlink from computation list if (node.computationId != null) { - ComputationList compList = compLists.get(node.computationId); + QueueKey key = QueueKey.create(node.computationId, node.keyGroup); + KeyGroupWorkQueue compList = compLists.get(key); if (compList != null) { compList.remove(node); if (compList.isEmpty()) { - compLists.remove(node.computationId); // Prevent memory leaks of empty keys + compLists.remove(key); // Prevent memory leaks of empty keys } } } @@ -149,18 +156,19 @@ private Node removeFirstGlobal() { return first; } - public QueuedWork pollWork(String computationId) { - if (computationId == null) { + public QueuedWork pollWork(String computationId, Work.KeyGroup keyGroup) { + if (computationId == null || keyGroup == null) { return null; } lock.lock(); try { - ComputationList compList = compLists.get(computationId); + QueueKey key = QueueKey.create(computationId, keyGroup); + KeyGroupWorkQueue compList = compLists.get(key); if (compList == null || compList.isEmpty()) { return null; } - // Retrieve the first pending task for this computation in O(1) + // Retrieve the first pending task for this computation and keyGroup in O(1) Node firstNode = compList.head.nextComp; unlinkNode(firstNode); @@ -186,8 +194,8 @@ public boolean offer(Runnable e) { // Append to computation list if applicable if (node.computationId != null) { - ComputationList compList = - compLists.computeIfAbsent(node.computationId, k -> new ComputationList()); + QueueKey key = QueueKey.create(node.computationId, node.keyGroup); + KeyGroupWorkQueue compList = compLists.computeIfAbsent(key, k -> new KeyGroupWorkQueue()); compList.append(node); } @@ -391,4 +399,54 @@ public Iterator iterator() { lock.unlock(); } } + + static final class QueueKey { + private final String computationId; + private final Work.@Nullable KeyGroup keyGroup; + + private QueueKey(String computationId, Work.@Nullable KeyGroup keyGroup) { + this.computationId = Objects.requireNonNull(computationId); + this.keyGroup = keyGroup; + } + + public static QueueKey create(String computationId, Work.@Nullable KeyGroup keyGroup) { + return new QueueKey(computationId, keyGroup); + } + + public String computationId() { + return computationId; + } + + public Work.@Nullable KeyGroup keyGroup() { + return keyGroup; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof QueueKey)) { + return false; + } + QueueKey other = (QueueKey) o; + return computationId.equals(other.computationId) && Objects.equals(keyGroup, other.keyGroup); + } + + @Override + public int hashCode() { + return Objects.hash(computationId, keyGroup); + } + + @Override + public String toString() { + return "QueueKey{" + + "computationId='" + + computationId + + '\'' + + ", keyGroup=" + + keyGroup + + '}'; + } + } } 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 98e12a3c86f1..e1f65c3bbdae 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 @@ -33,6 +33,7 @@ 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; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; @@ -66,18 +67,30 @@ public static Collection useFairMonitor() { @Rule public transient Timeout globalTimeout = Timeout.seconds(300); private BoundedQueueExecutor executor; + private static final Work.KeyGroup DEFAULT_KEY_GROUP = Work.KeyGroup.create(1, 2); + private static ExecutableWork createWork(Consumer executeWorkFn) { return createWorkWithCompId("computationId", executeWorkFn); } private static ExecutableWork createWorkWithCompId( String computationId, Consumer executeWorkFn) { + return createWorkWithCompIdAndKeyGroup(computationId, DEFAULT_KEY_GROUP, executeWorkFn); + } + + private static ExecutableWork createWorkWithCompIdAndKeyGroup( + String computationId, Work.KeyGroup keyGroup, Consumer executeWorkFn) { WorkItem workItem = WorkItem.newBuilder() .setKey(ByteString.EMPTY) .setShardingKey(1) .setWorkToken(33) .setCacheToken(1) + .setKeyGroup( + Windmill.Uint128Proto.newBuilder() + .setHigh(keyGroup.high()) + .setLow(keyGroup.low()) + .build()) .build(); return ExecutableWork.create( Work.create( @@ -407,25 +420,35 @@ public void testHandleMerge() throws Exception { } @Test - public void testTombstoneExecutionAndQueueSize() throws Exception { + public void testRenderSummaryHtml() { + String expectedSummaryHtml = + "Worker Threads: 0/2
/n" + + "Active Threads: 0
/n" + + "Work Queue Size: 0/102
/n" + + "Work Queue Bytes: 0/10000000
/n"; + assertEquals(expectedSummaryHtml, executor.summaryHtml()); + } + + @Test + public void testPollWork() throws Exception { + // Create separate BoundedQueueExecutor with 1 thread so we can block it easily BoundedQueueExecutor testExecutor = new BoundedQueueExecutor( - 1, // 1 thread to strictly control execution order - DEFAULT_THREAD_EXPIRATION_SEC, + 1, + 60, TimeUnit.SECONDS, - 10, - MAXIMUM_BYTES_OUTSTANDING, - new ThreadFactoryBuilder() - .setNameFormat("testTombstoneExecutionAndQueueSize-%d") - .setDaemon(true) - .build(), + 100, + 10000000, + new ThreadFactoryBuilder().setNameFormat("testStealing-%d").setDaemon(true).build(), useFairMonitor); + // 1. Create blocker task to occupy the worker thread CountDownLatch blockerStart = new CountDownLatch(1); CountDownLatch blockerStop = new CountDownLatch(1); ExecutableWork blockerWork = - createWorkWithCompId( - "comp-1", + createWorkWithCompIdAndKeyGroup( + "blockerComp", + DEFAULT_KEY_GROUP, ignored -> { blockerStart.countDown(); try { @@ -435,286 +458,98 @@ public void testTombstoneExecutionAndQueueSize() throws Exception { } }); + testExecutor.execute(blockerWork, 0); + blockerStart.await(); + + // 2. Create two distinct key groups + Work.KeyGroup keyGroup1 = Work.KeyGroup.create(1, 1); + Work.KeyGroup keyGroup2 = Work.KeyGroup.create(1, 2); + + // Create executable tasks CountDownLatch targetStart = new CountDownLatch(1); - ExecutableWork targetWork = - createWorkWithCompId( - "comp-2", + ExecutableWork work1 = createWorkWithCompIdAndKeyGroup("compA", keyGroup1, ignored -> {}); + ExecutableWork work2 = + createWorkWithCompIdAndKeyGroup( + "compA", + keyGroup2, ignored -> { targetStart.countDown(); }); - // 1. Occupy the worker thread with blockerWork - testExecutor.execute(blockerWork, 0); - blockerStart.await(); - - // 2. Enqueue targetWork (goes into queue) - testExecutor.execute(targetWork, 1000); + // Enqueue tasks (they will wait in the queue because the thread is blocked) + testExecutor.execute(work1, 100); + testExecutor.execute(work2, 150); - // Wait a moment to ensure targetWork is registered in the queue - assertEquals(2, testExecutor.elementsOutstanding()); - assertEquals(1, testExecutor.executorQueueIsEmpty() ? 0 : 1); // it is in the queue + // Total outstanding elements must be 3 (blocker + work1 + work2) + assertEquals(3, testExecutor.elementsOutstanding()); - // 3. Now steal targetWork using targeted pollWork from a "stealer" context + // Steal work2 using pollWork with compA and keyGroup2 try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { - java.util.Optional stolen = testExecutor.pollWork("comp-2", stealHandle); + java.util.Optional stolen = + testExecutor.pollWork("compA", keyGroup2, stealHandle); assertTrue(stolen.isPresent()); - assertEquals(targetWork, stolen.get()); + assertEquals(work2, stolen.get()); - // Since it's stolen, the budget is transferred to stealHandle. - // The outstanding size must still be 2 (blocker + stolen in stealHandle). - assertEquals(2, testExecutor.elementsOutstanding()); - - // 4. Run the stolen work inline in the stealer thread. + // Run the stolen task stolen.get().run(stealHandle); targetStart.await(); } - // After stealHandle is closed, the outstanding size should decrement by 1 (the stolen task). - // Only the blocker remains. - while (testExecutor.elementsOutstanding() != 1) { - Thread.sleep(10); + // Steal work1 using pollWork with compA and keyGroup1 + try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { + java.util.Optional stolen = + testExecutor.pollWork("compA", keyGroup1, stealHandle); + assertTrue(stolen.isPresent()); + assertEquals(work1, stolen.get()); } - assertEquals(1, testExecutor.elementsOutstanding()); - // 5. Unblock the blocker. The executor worker thread should now pop the targetWork tombstone. - // Since it is cancelled (stolen), the worker thread should skip it as a no-op. - // The elementsOutstanding should drop to 0 after the blocker completes. + // Unblock the blocker and shut down blockerStop.countDown(); - - while (testExecutor.elementsOutstanding() != 0) { - Thread.sleep(10); - } - assertEquals(0, testExecutor.elementsOutstanding()); - testExecutor.shutdown(); - } - - @Test - public void testConcurrentStealingAndPolling() throws Exception { - final int numComputations = 10; - final int tasksPerComp = 100; - final int totalTasks = numComputations * tasksPerComp; - - // 4 executor worker threads to execute work - BoundedQueueExecutor testExecutor = - new BoundedQueueExecutor( - 4, - DEFAULT_THREAD_EXPIRATION_SEC, - TimeUnit.SECONDS, - totalTasks * 2, - MAXIMUM_BYTES_OUTSTANDING, - new ThreadFactoryBuilder() - .setNameFormat("testConcurrentStealingAndPolling-%d") - .setDaemon(true) - .build(), - useFairMonitor); - - // Latches to coordinate - CountDownLatch allDone = new CountDownLatch(totalTasks); - - // Enqueue work for all computations - for (int i = 0; i < numComputations; i++) { - final String compId = "comp-" + i; - for (int j = 0; j < tasksPerComp; j++) { - ExecutableWork work = - createWorkWithCompId( - compId, - ignored -> { - allDone.countDown(); - }); - testExecutor.execute(work, 100); - } - } - - // Launch active "stealers" that target specific computations concurrently - Thread[] stealers = new Thread[numComputations]; - for (int i = 0; i < numComputations; i++) { - final String compId = "comp-" + i; - stealers[i] = - new Thread( - () -> { - try { - // Attempt to steal tasks for this computation - int stolenCount = 0; - while (stolenCount < tasksPerComp) { - try (BoundedQueueExecutorWorkHandleImpl stealHandle = - testExecutor.createBudgetHandle(0, 0L)) { - java.util.Optional stolen = - testExecutor.pollWork(compId, stealHandle); - if (stolen.isPresent()) { - stolen.get().run(stealHandle); - stolenCount++; - } else { - // Yield if none available - Thread.sleep(1); - } - } catch (InterruptedException e) { - break; - } - } - } catch (Exception e) { - throw new RuntimeException(e); - } - }); - stealers[i].start(); - } - - // Wait for all tasks to complete (either executed by thread pool or stolen) - assertTrue(allDone.await(30, TimeUnit.SECONDS)); - - for (Thread stealer : stealers) { - stealer.interrupt(); - stealer.join(); - } - - // Wait until all outstanding elements are released cleanly - while (testExecutor.elementsOutstanding() != 0) { - Thread.sleep(10); - } - assertEquals(0, testExecutor.elementsOutstanding()); - assertEquals(0, testExecutor.bytesOutstanding()); testExecutor.shutdown(); } @Test - public void testConcurrentStealingAndPollingNoWastedThreads() throws Exception { - final int numComputations = 10; - final int tasksPerComp = 100; - final int totalTasks = numComputations * tasksPerComp; - - // 4 executor worker threads to execute work + public void testPollWorkWithLinkedBlockingQueue() throws Exception { BoundedQueueExecutor testExecutor = new BoundedQueueExecutor( - 4, - DEFAULT_THREAD_EXPIRATION_SEC, + 1, + 60, TimeUnit.SECONDS, - totalTasks * 2, - MAXIMUM_BYTES_OUTSTANDING, - new ThreadFactoryBuilder() - .setNameFormat("testConcurrentStealingAndPollingNoWastedThreads-%d") - .setDaemon(true) - .build(), - useFairMonitor); - - // Thread-safe counters for tracking executions - java.util.concurrent.atomic.AtomicInteger totalExecuted = - new java.util.concurrent.atomic.AtomicInteger(0); - java.util.concurrent.atomic.AtomicInteger workerExecuted = - new java.util.concurrent.atomic.AtomicInteger(0); - java.util.concurrent.atomic.AtomicInteger poolWorkerExecuted = - new java.util.concurrent.atomic.AtomicInteger(0); - - // Map to track duplicate executions - java.util.concurrent.ConcurrentHashMap - executionCounts = new java.util.concurrent.ConcurrentHashMap<>(); - - // Latches to coordinate - CountDownLatch allDone = new CountDownLatch(totalTasks); - - // Enqueue work for all computations - for (int i = 0; i < numComputations; i++) { - final String compId = "comp-" + i; - for (int j = 0; j < tasksPerComp; j++) { - final String taskId = compId + "-task-" + j; - ExecutableWork work = - createWorkWithCompId( - compId, - ignored -> { - int count = - executionCounts - .computeIfAbsent( - taskId, k -> new java.util.concurrent.atomic.AtomicInteger(0)) - .incrementAndGet(); - totalExecuted.incrementAndGet(); - String threadName = Thread.currentThread().getName(); - if (threadName.startsWith("testConcurrentStealingAndPollingNoWastedThreads-")) { - poolWorkerExecuted.incrementAndGet(); - } else { - workerExecuted.incrementAndGet(); - } - allDone.countDown(); - }); - testExecutor.execute(work, 100); - } - } + 100, + 10000000, + new ThreadFactoryBuilder().setNameFormat("testLinkedQueue-%d").setDaemon(true).build(), + useFairMonitor, + /* useKeyGroupWorkQueue= */ false); - // Launch active "stealers" that target specific computations concurrently - Thread[] stealers = new Thread[numComputations]; - for (int i = 0; i < numComputations; i++) { - final String compId = "comp-" + i; - stealers[i] = - new Thread( - () -> { - try { - while (!Thread.currentThread().isInterrupted()) { - try (BoundedQueueExecutorWorkHandleImpl stealHandle = - testExecutor.createBudgetHandle(0, 0L)) { - java.util.Optional stolen = - testExecutor.pollWork(compId, stealHandle); - if (stolen.isPresent()) { - stolen.get().run(stealHandle); - } else { - Thread.sleep(1); - } - } catch (InterruptedException e) { - break; - } - } - } catch (Exception e) { - // Ignore or log - } - }); - stealers[i].start(); - } + CountDownLatch blockerStart = new CountDownLatch(1); + CountDownLatch blockerStop = new CountDownLatch(1); + ExecutableWork blockerWork = + createWorkWithCompIdAndKeyGroup( + "blockerComp", + DEFAULT_KEY_GROUP, + ignored -> { + blockerStart.countDown(); + try { + blockerStop.await(); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); - // Wait for all tasks to complete (either executed by thread pool or stolen) - assertTrue(allDone.await(60, TimeUnit.SECONDS)); + testExecutor.execute(blockerWork, 0); + blockerStart.await(); - // Stop stealers - for (Thread stealer : stealers) { - stealer.interrupt(); - stealer.join(); - } + Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1); + ExecutableWork work = createWorkWithCompIdAndKeyGroup("compA", keyGroup, ignored -> {}); + testExecutor.execute(work, 100); - // Verify execution invariants - assertEquals("Total executions must match total tasks", totalTasks, totalExecuted.get()); - assertEquals("No duplicate executions allowed", totalTasks, executionCounts.size()); - for (java.util.Map.Entry entry : - executionCounts.entrySet()) { - assertEquals( - "Task " + entry.getKey() + " executed more than once", 1, entry.getValue().get()); + try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { + java.util.Optional stolen = + testExecutor.pollWork("compA", keyGroup, stealHandle); + assertFalse(stolen.isPresent()); } - // Verify that completed task count of the inner executor perfectly matches poolWorkerExecuted - java.lang.reflect.Field executorField = BoundedQueueExecutor.class.getDeclaredField("executor"); - executorField.setAccessible(true); - java.util.concurrent.ThreadPoolExecutor innerExecutor = - (java.util.concurrent.ThreadPoolExecutor) executorField.get(testExecutor); - - // Wait a tiny bit for any pending pool tasks to fully complete their afterExecute/bookkeeping - Thread.sleep(100); - - long completedTasks = innerExecutor.getCompletedTaskCount(); - assertEquals( - "Completed task count of the pool must perfectly match the number of worker-executed tasks (exactly 0 wasted activations)", - (long) poolWorkerExecuted.get(), - completedTasks); - - // Wait until all outstanding elements are released cleanly - while (testExecutor.elementsOutstanding() != 0) { - Thread.sleep(10); - } - assertEquals(0, testExecutor.elementsOutstanding()); - assertEquals(0, testExecutor.bytesOutstanding()); + blockerStop.countDown(); testExecutor.shutdown(); } - - @Test - public void testRenderSummaryHtml() { - String expectedSummaryHtml = - "Worker Threads: 0/2
/n" - + "Active Threads: 0
/n" - + "Work Queue Size: 0/102
/n" - + "Work Queue Bytes: 0/10000000
/n"; - assertEquals(expectedSummaryHtml, executor.summaryHtml()); - } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java index e5025d5a29ef..3ab5d03101aa 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/ComputationWorkQueueTest.java @@ -72,13 +72,25 @@ public void setUp() { false); } + private static final Work.KeyGroup DEFAULT_KEY_GROUP = Work.KeyGroup.create(1, 2); + private QueuedWork createQueuedWork(String computationId, long workBytes) { + return createQueuedWork(computationId, DEFAULT_KEY_GROUP, workBytes); + } + + private QueuedWork createQueuedWork( + String computationId, Work.KeyGroup keyGroup, long workBytes) { WorkItem workItem = WorkItem.newBuilder() .setKey(ByteString.EMPTY) .setShardingKey(1) .setWorkToken(33) .setCacheToken(1) + .setKeyGroup( + org.apache.beam.runners.dataflow.worker.windmill.Windmill.Uint128Proto.newBuilder() + .setHigh(keyGroup.high()) + .setLow(keyGroup.low()) + .build()) .build(); ExecutableWork work = ExecutableWork.create( @@ -200,7 +212,7 @@ public void testPollWorkTargeted() { assertEquals(3, queue.size()); // Targeted poll A - QueuedWork polledA1 = queue.pollWork("compA"); + QueuedWork polledA1 = queue.pollWork("compA", DEFAULT_KEY_GROUP); assertNotNull(polledA1); assertEquals("compA", polledA1.getWork().getComputationId()); assertEquals(100, polledA1.getHandle().bytes()); @@ -224,7 +236,7 @@ public void testMemoryPruningLeavesZeroLeaks() { queue.offer(workA1); // Steal A1 - QueuedWork polled = queue.pollWork("compA"); + QueuedWork polled = queue.pollWork("compA", DEFAULT_KEY_GROUP); assertNotNull(polled); assertTrue(queue.isEmpty()); @@ -234,7 +246,7 @@ public void testMemoryPruningLeavesZeroLeaks() { assertEquals(1, queue.size()); // Steal B1 - QueuedWork polledB = queue.pollWork("compB"); + QueuedWork polledB = queue.pollWork("compB", DEFAULT_KEY_GROUP); assertNotNull(polledB); assertTrue(queue.isEmpty()); } @@ -283,7 +295,7 @@ public void testConcurrentStress() throws InterruptedException { if (consumerId % 2 == 0) { // Targeted poll String compId = "comp-" + (consumedCount.get() % 5); - task = queue.pollWork(compId); + task = queue.pollWork(compId, DEFAULT_KEY_GROUP); } else { // Global poll task = queue.poll(); @@ -408,4 +420,35 @@ public void testPollWithTimeout() throws InterruptedException { assertTrue(finished2.await(2, TimeUnit.SECONDS)); assertEquals(task, result2.get()); } + + @Test + public void testPollWorkWithKeyGroup() { + ComputationWorkQueue queue = new ComputationWorkQueue(); + + Work.KeyGroup keyGroup1 = Work.KeyGroup.create(1, 1); + Work.KeyGroup keyGroup2 = Work.KeyGroup.create(1, 2); + + QueuedWork workA1 = createQueuedWork("compA", keyGroup1, 100); + QueuedWork workA2 = createQueuedWork("compA", keyGroup2, 150); + + queue.offer(workA1); + queue.offer(workA2); + + assertEquals(2, queue.size()); + + // Poll with keyGroup2 first - should return workA2 + QueuedWork polledA2 = queue.pollWork("compA", keyGroup2); + assertNotNull(polledA2); + assertEquals(workA2, polledA2); + assertEquals(1, queue.size()); + + // Poll with keyGroup2 again - should return null + assertNull(queue.pollWork("compA", keyGroup2)); + + // Poll with keyGroup1 - should return workA1 + QueuedWork polledA1 = queue.pollWork("compA", keyGroup1); + assertNotNull(polledA1); + assertEquals(workA1, polledA1); + assertTrue(queue.isEmpty()); + } } From cad35a1fa95962aafb06d66eacd9c2cbd3b35aa1 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 1 Jun 2026 00:15:00 +0000 Subject: [PATCH 09/11] context update --- gradle.properties | 4 +- runners/direct-java/build.gradle | 6 +- .../worker/KeyTokenInvalidException.java | 14 + .../worker/StreamingModeExecutionContext.java | 372 ++++++++++++- .../worker/WindmillReaderIteratorBase.java | 19 +- .../worker/WindowingWindmillReader.java | 85 +-- .../worker/WorkItemCancelledException.java | 12 + .../dataflow/worker/WorkerCustomSources.java | 292 ++++++---- .../streaming/ComputationWorkExecutor.java | 34 +- .../dataflow/worker/streaming/Work.java | 18 +- .../windmill/client/WindmillStream.java | 7 + .../windmill/client/commits/Commit.java | 34 +- .../commits/StreamingEngineWorkCommitter.java | 61 ++- .../client/grpc/GrpcCommitWorkStream.java | 66 ++- .../windmill/state/WindmillStateReader.java | 3 +- .../processing/StreamingWorkScheduler.java | 300 ++++++++--- .../failures/WorkFailureProcessor.java | 100 +++- .../dataflow/worker/FakeWindmillServer.java | 45 ++ .../worker/StreamingDataflowWorkerTest.java | 6 +- .../StreamingModeExecutionContextTest.java | 502 ++++++++++++++++++ .../WindmillReaderIteratorBaseTest.java | 131 +++++ .../worker/WorkerCustomSourcesTest.java | 123 +++++ .../StreamingEngineWorkCommitterTest.java | 63 +++ .../failures/WorkFailureProcessorTest.java | 151 ++++++ sdks/java/core/build.gradle | 9 + 25 files changed, 2151 insertions(+), 306 deletions(-) diff --git a/gradle.properties b/gradle.properties index 95e50105a494..d091210113a1 100644 --- a/gradle.properties +++ b/gradle.properties @@ -15,8 +15,8 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################ -org.gradle.caching=true -org.gradle.parallel=true +org.gradle.caching=false +org.gradle.parallel=false org.gradle.daemon=false org.gradle.configureondemand=true org.gradle.jvmargs=-Xss10240k diff --git a/runners/direct-java/build.gradle b/runners/direct-java/build.gradle index 1ab702da3213..b429a28ac8ec 100644 --- a/runners/direct-java/build.gradle +++ b/runners/direct-java/build.gradle @@ -66,7 +66,11 @@ dependencies { dependOnProjectsAndConfigs.each { // For projects producing shadowjar, use the packaged jar as dependency to // handle redirected packages from it - implementation project(path: it.key, configuration: it.value) + if (it.value != null) { + implementation project(path: it.key, configuration: it.value) + } else { + implementation project(it.key) + } } shadow library.java.vendored_grpc_1_69_0 shadow library.java.joda_time diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/KeyTokenInvalidException.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/KeyTokenInvalidException.java index 29b16b71883f..c57220665010 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/KeyTokenInvalidException.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/KeyTokenInvalidException.java @@ -17,12 +17,26 @@ */ package org.apache.beam.runners.dataflow.worker; +import java.util.Optional; import javax.annotation.Nullable; +import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey; /** Indicates that the key token was invalid when data was attempted to be fetched. */ public class KeyTokenInvalidException extends RuntimeException { + private final @Nullable ShardedKey shardedKey; + public KeyTokenInvalidException(String key) { super("Unable to fetch data due to token mismatch for key " + key); + this.shardedKey = null; + } + + public KeyTokenInvalidException(ShardedKey shardedKey, String key) { + super("Unable to fetch data due to token mismatch for key " + key); + this.shardedKey = shardedKey; + } + + public Optional getShardedKey() { + return Optional.ofNullable(shardedKey); } /** Returns whether an exception was caused by a {@link KeyTokenInvalidException}. */ 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 25ce299adf7a..5bca60b5661d 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 @@ -25,6 +25,7 @@ import com.google.api.services.dataflow.model.SideInputInfo; import java.io.Closeable; import java.io.IOException; +import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; import java.util.Iterator; @@ -34,6 +35,7 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.concurrent.NotThreadSafe; import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; @@ -50,6 +52,8 @@ import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; 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.Watermarks; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; @@ -57,6 +61,12 @@ 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.BoundedQueueExecutor; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; 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; @@ -162,6 +172,39 @@ public class StreamingModeExecutionContext extends DataflowExecutionContext keyCoder; + + // Key switch listener to delegate MDC logging context, sampler metrics, and thread name updates + public interface KeySwitchListener { + void onKeySwitch(Work oldWork, Work newWork); + } + + private @Nullable KeySwitchListener keySwitchListener; + + // Configurable batch limits (defaults parsed from pipeline options/experiments) + private int maxKeyGroupBatchSize = 100; + private long maxKeyGroupBatchTimeNanos = TimeUnit.MILLISECONDS.toNanos(100); + private long maxKeyGroupBatchBytes = 10L * 1024 * 1024; // 10MB + + // Batch-tracking state metrics + private int additionalWorkItemsPolled = 0; + private long bundleStartTimeNanos = 0; + private long accumulatedCommitBytes = 0; + + // Centrally accumulated bundle metadata (flat-map for callbacks) + private final List executedWorks = new ArrayList<>(); + private final List outputBuilders = new ArrayList<>(); + private final Map> accumulatedCallbacks = new HashMap<>(); + private volatile @Nullable Work failedWork = null; + private @Nullable WindmillStateReader activeStateReader; + private long stateBytesRead = 0; + private @Nullable String sourceBytesProcessCounterName; + public StreamingModeExecutionContext( CounterFactory counterFactory, String computationId, @@ -208,7 +251,11 @@ public boolean throwExceptionsForLargeOutput() { } public boolean workIsFailed() { - return work != null && work.isFailed(); + return failedWork != null; + } + + public @Nullable Work getFailedWork() { + return failedWork; } public boolean getDrainMode() { @@ -247,9 +294,111 @@ public void start( SideInputStateFetcher sideInputStateFetcher, Windmill.WorkItemCommitRequest.Builder outputBuilder, WorkExecutor workExecutor) { + start( + key, + work, + stateReader, + sideInputStateFetcher, + outputBuilder, + workExecutor, + /* workQueueExecutor= */ null, + /* budgetHandle= */ null, + new HotKeyLogger(), + /* hotKeyLoggingEnabled= */ false, + /* stepName= */ "", + /* keyCoder= */ null, + /* maxKeyGroupBatchSize= */ 1, + /* maxKeyGroupBatchTimeNanos= */ 0L, + /* maxKeyGroupBatchBytes= */ 0L, + /* keySwitchListener= */ (k, c) -> {}, + /* sourceBytesProcessCounterName= */ null); + } + + public void start( + @Nullable Object key, + Work work, + WindmillStateReader stateReader, + SideInputStateFetcher sideInputStateFetcher, + Windmill.WorkItemCommitRequest.Builder outputBuilder, + WorkExecutor workExecutor, + BoundedQueueExecutor workQueueExecutor, + BoundedQueueExecutorWorkHandle budgetHandle, + HotKeyLogger hotKeyLogger, + boolean hotKeyLoggingEnabled, + String stepName, + @Nullable Coder keyCoder, + int maxKeyGroupBatchSize, + long maxKeyGroupBatchTimeNanos, + long maxKeyGroupBatchBytes, + KeySwitchListener keySwitchListener) { + start( + key, + work, + stateReader, + sideInputStateFetcher, + outputBuilder, + workExecutor, + workQueueExecutor, + budgetHandle, + hotKeyLogger, + hotKeyLoggingEnabled, + stepName, + keyCoder, + maxKeyGroupBatchSize, + maxKeyGroupBatchTimeNanos, + maxKeyGroupBatchBytes, + keySwitchListener, + /* sourceBytesProcessCounterName= */ null); + } + + public void start( + @Nullable Object key, + Work work, + WindmillStateReader stateReader, + SideInputStateFetcher sideInputStateFetcher, + Windmill.WorkItemCommitRequest.Builder outputBuilder, + WorkExecutor workExecutor, + BoundedQueueExecutor workQueueExecutor, + BoundedQueueExecutorWorkHandle budgetHandle, + HotKeyLogger hotKeyLogger, + boolean hotKeyLoggingEnabled, + String stepName, + @Nullable Coder keyCoder, + int maxKeyGroupBatchSize, + long maxKeyGroupBatchTimeNanos, + long maxKeyGroupBatchBytes, + KeySwitchListener keySwitchListener, + @Nullable String sourceBytesProcessCounterName) { this.key = key; this.work = work; this.workExecutor = workExecutor; + this.workQueueExecutor = workQueueExecutor; + this.budgetHandle = budgetHandle; + this.hotKeyLogger = hotKeyLogger; + this.hotKeyLoggingEnabled = hotKeyLoggingEnabled; + this.stepName = stepName; + this.keyCoder = keyCoder; + + this.maxKeyGroupBatchSize = maxKeyGroupBatchSize; + this.maxKeyGroupBatchTimeNanos = maxKeyGroupBatchTimeNanos; + this.maxKeyGroupBatchBytes = maxKeyGroupBatchBytes; + this.keySwitchListener = keySwitchListener; + + // Clear and initialize the accumulated metadata for the primary key + this.additionalWorkItemsPolled = 0; + this.bundleStartTimeNanos = System.nanoTime(); + this.accumulatedCommitBytes = 0; + + this.executedWorks.clear(); + this.outputBuilders.clear(); + this.accumulatedCallbacks.clear(); + this.failedWork = null; + + work.setOnFailureListener(() -> this.failedWork = work); + this.executedWorks.add(work); + this.outputBuilders.add(outputBuilder); + this.outputBuilder = outputBuilder; + this.finishKeyCalled = false; this.computationKey = WindmillComputationKey.create(computationId, work.getShardedKey()); this.sideInputStateFetcher = sideInputStateFetcher; @@ -260,11 +409,14 @@ public void start( config.enableStateTagEncodingV2() ? WindmillTagEncodingV2.instance() : WindmillTagEncodingV1.instance(); - this.outputBuilder = outputBuilder; this.sideInputCache.clear(); this.backlogBytes = UnboundedReader.BACKLOG_UNKNOWN; clearSinkFullHint(); + this.activeStateReader = stateReader; + this.stateBytesRead = 0; + this.sourceBytesProcessCounterName = sourceBytesProcessCounterName; + Instant processingTime = computeProcessingTime(work.getWorkItem().getTimers().getTimersList()); Collection stepContexts = getAllStepContexts(); @@ -280,7 +432,12 @@ public void start( } public void finishKey() { - checkState(!finishKeyCalled, "finishKey was already called"); + if (finishKeyCalled) { + return; + } + if (activeStateReader != null) { + this.stateBytesRead += activeStateReader.getBytesRead(); + } checkStateNotNull(workExecutor, "workExecutor must be set before calling finishKey()"); try { workExecutor.finishKey(); @@ -288,6 +445,13 @@ public void finishKey() { throw new RuntimeException(e); } this.finishKeyCalled = true; + + flushStateInternal(); + + // Accumulate commit request size to enforce the size limit threshold (10MB) + if (outputBuilder != null) { + this.accumulatedCommitBytes += outputBuilder.buildPartial().getSerializedSize(); + } } /** @@ -441,20 +605,22 @@ public void setActiveReader(UnboundedReader reader) { /** Invalidate the state and reader caches for this computation and key. */ public void invalidateCache() { - ByteString key = getSerializedKey(); - if (key != null) { - readerCache.invalidateReader(getComputationKey()); - if (activeReader != null) { - try { - activeReader.close(); - } catch (IOException e) { - LOG.warn( - "Failed to close reader for {}-{}", computationId, getWorkItem().getShardingKey(), e); - } + for (Work w : executedWorks) { + WindmillComputationKey compKey = + WindmillComputationKey.create(computationId, w.getShardedKey()); + readerCache.invalidateReader(compKey); + stateCache.invalidate(w.getShardedKey()); + } + if (activeReader != null) { + try { + activeReader.close(); + } catch (IOException e) { + LOG.warn( + "Failed to close reader for {}-{}", computationId, getWorkItem().getShardingKey(), e); } - activeReader = null; - stateCache.invalidate(key, getWorkItem().getShardingKey()); } + activeReader = null; + activeStateReader = null; } public UnboundedSource.@Nullable CheckpointMark getReaderCheckpoint( @@ -470,8 +636,7 @@ public void invalidateCache() { } } - public Map> flushState() { - checkState(finishKeyCalled, "finishKey must be called before flushState"); + private void flushStateInternal() { Map> callbacks = new HashMap<>(); for (StepContext stepContext : getAllStepContexts()) { @@ -555,7 +720,178 @@ public Map> flushState() { // RestrictionTracker.getProgress() or GetSize() are not defined. outputBuilder.setSourceBacklogBytes(backlogBytes); } - return callbacks; + + this.accumulatedCallbacks.putAll(callbacks); + + if (sourceBytesProcessCounterName != null && workExecutor instanceof MapTaskExecutor) { + MapTaskExecutor mapTaskExecutor = (MapTaskExecutor) workExecutor; + ReadOperation readOperation = mapTaskExecutor.getReadOperation(); + long sourceBytesProcessed = 0; + if (readOperation != null && readOperation.receivers != null) { + for (OutputReceiver receiver : readOperation.receivers) { + if (receiver != null && receiver.getOutputCounters() != null) { + ElementCounter elementCounter = + receiver.getOutputCounters().get(sourceBytesProcessCounterName); + if (elementCounter instanceof OutputObjectAndByteCounter) { + OutputObjectAndByteCounter byteCounter = (OutputObjectAndByteCounter) elementCounter; + if (byteCounter.getByteCount() != null) { + sourceBytesProcessed += byteCounter.getByteCount().getAndReset(); + } + } + } + } + } + outputBuilder.setSourceBytesProcessed(sourceBytesProcessed); + } + } + + public Map> flushState() { + if (!finishKeyCalled) { + finishKey(); + } + return accumulatedCallbacks; + } + + public boolean advance() { + if (workIsFailed()) { + throw new WorkItemCancelledException(failedWork.getWorkItem().getShardingKey()); + } + + if (workQueueExecutor == null || budgetHandle == null || work == null) { + finishKey(); + return false; + } + + // 1. Check if we hit the batch count limit (default 100 additional works) + if (additionalWorkItemsPolled >= maxKeyGroupBatchSize) { + finishKey(); // Finalize and flush the last key + return false; + } + + // 2. Check if we hit the batch time limit (default 100ms) + long elapsedNanos = System.nanoTime() - bundleStartTimeNanos; + if (elapsedNanos >= maxKeyGroupBatchTimeNanos) { + finishKey(); + return false; + } + + // 3. Check if we hit the batch size limit (default 10MB) + long currentKeyBytes = + outputBuilder != null ? outputBuilder.buildPartial().getSerializedSize() : 0; + if (accumulatedCommitBytes + currentKeyBytes >= maxKeyGroupBatchBytes) { + finishKey(); + return false; + } + + // 4. Poll next work item in the same key group + if (work.getKeyGroup().isPresent()) { + Optional additionalWorkOpt = + workQueueExecutor.pollWork(computationId, work.getKeyGroup().get(), budgetHandle); + if (additionalWorkOpt.isPresent()) { + Work oldWork = work; + Work newWork = additionalWorkOpt.get().work(); + additionalWorkItemsPolled++; + + // 5. Conclude active key's execution and flush state atomically + finishKey(); + + // 6. Trigger listener to align MDC logging context and metrics + if (keySwitchListener != null) { + keySwitchListener.onKeySwitch(oldWork, newWork); + } + + // 7. Bind the context to Key B internally + startForNewKey(newWork); + return true; + } + } + + // No more work available. Finalize and flush the last active key atomically + finishKey(); + return false; + } + + private void startForNewKey(Work newWork) { + Object newKey = null; + if (keyCoder != null) { + try { + newKey = keyCoder.decode(newWork.getWorkItem().getKey().newInput(), Coder.Context.OUTER); + } catch (IOException e) { + throw new RuntimeException("Failed to decode key during key switch", e); + } + } + + this.key = newKey; + this.work = newWork; + this.finishKeyCalled = false; + this.computationKey = WindmillComputationKey.create(computationId, newWork.getShardedKey()); + + // 1. Create a fresh output builder for the new key and append it to outputBuilders list + Windmill.WorkItemCommitRequest.Builder newOutputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(newWork.getWorkItem().getKey()) + .setShardingKey(newWork.getWorkItem().getShardingKey()) + .setWorkToken(newWork.getWorkItem().getWorkToken()) + .setCacheToken(newWork.getWorkItem().getCacheToken()); + + this.outputBuilder = newOutputBuilder; + this.outputBuilders.add(newOutputBuilder); + newWork.setOnFailureListener(() -> this.failedWork = newWork); + this.executedWorks.add(newWork); + + // 2. Detect and log hot keys for dynamically polled keys + if (newWork.getWorkItem().hasHotKeyInfo() && hotKeyLogger != null && stepName != null) { + Windmill.HotKeyInfo hotKeyInfo = newWork.getWorkItem().getHotKeyInfo(); + Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000); + if (newKey != null && hotKeyLoggingEnabled) { + hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge, newKey); + } else { + hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge); + } + } + + // Note: We do NOT clear sideInputCache here, allowing Key B to reuse warm side inputs! + + // 3. Re-initialize state cache and state/timer internals across all step contexts + Instant processingTime = + computeProcessingTime(newWork.getWorkItem().getTimers().getTimersList()); + Collection stepContexts = getAllStepContexts(); + if (!stepContexts.isEmpty()) { + WindmillStateCache.ForKey cacheForKey = + stateCache.forKey( + getComputationKey(), newWork.getWorkItem().getCacheToken(), getWorkToken()); + WindmillStateReader newReader = newWork.createWindmillStateReader(); + this.activeStateReader = newReader; + for (StepContext stepContext : stepContexts) { + stepContext.start(newReader, processingTime, cacheForKey, newWork.watermarks()); + } + } else { + this.activeStateReader = null; + } + } + + public List getExecutedWorks() { + return executedWorks; + } + + public long getStateBytesRead() { + return stateBytesRead; + } + + public List getOutputBuilders() { + return outputBuilders; + } + + public Map> getAccumulatedCallbacks() { + return accumulatedCallbacks; + } + + public @Nullable Object getKey() { + return key; + } + + public Work getWork() { + return work; } String getStateFamily(NameContext nameContext) { 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 075a1a8a4250..e4aeec3b7182 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 @@ -35,7 +35,7 @@ public abstract class WindmillReaderIteratorBase extends NativeReader.NativeReaderIterator> { private final StreamingModeExecutionContext context; - private final Windmill.WorkItem work; + private Windmill.WorkItem work; private int bundleIndex = 0; private int messageIndex = -1; private @Nullable WindowedValue current = null; @@ -57,15 +57,27 @@ public boolean start() throws IOException { @Override public boolean advance() throws IOException { if (context.workIsFailed()) { - throw new WorkItemCancelledException(context.getWorkItem().getShardingKey()); + throw new WorkItemCancelledException( + checkNotNull(context.getFailedWork()).getWorkItem().getShardingKey()); } while (true) { if (bundleIndex >= work.getMessageBundlesCount()) { + // If elements are exhausted, try advancing the execution context to the next key in the + // group + if (context.advance()) { + // Transition succeeded! Update iterator references to the new work item + this.work = context.getWork().getWorkItem(); + this.bundleIndex = 0; + this.messageIndex = -1; + continue; + } + + // All work items are exhausted. Iterator returns false. current = null; - context.finishKey(); return false; } + Windmill.InputMessageBundle bundle = work.getMessageBundles(bundleIndex); ++messageIndex; if (messageIndex >= bundle.getMessagesCount()) { @@ -73,6 +85,7 @@ public boolean advance() throws IOException { ++bundleIndex; continue; } + try { current = checkNotNull(decodeMessage(bundle.getMessages(messageIndex))); return true; 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 488684769bd9..51d00ce507df 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 @@ -151,51 +151,64 @@ public NativeReaderIterator>> iterator() throw && Iterables.isEmpty(keyedWorkItem.elementsIterable())); final WindowedValue> value = new ValueInEmptyWindows<>(keyedWorkItem); - // Return a noop iterator when current workitem is an empty workitem. - if (isEmptyWorkItem) { - return new NativeReaderIterator>>() { - @Override - public boolean start() throws IOException { - context.finishKey(); - return false; + return new NativeReaderIterator>>() { + private @Nullable WindowedValue> current = null; + private boolean started = false; + + @Override + public boolean start() throws IOException { + if (context.workIsFailed()) { + throw new WorkItemCancelledException( + checkStateNotNull(context.getFailedWork()).getWorkItem().getShardingKey()); } - - @Override - public boolean advance() throws IOException { + if (started) { return false; } - - @Override - public WindowedValue> getCurrent() { - throw new NoSuchElementException(); + started = true; + if (isEmptyWorkItem) { + return advance(); // Try to transition immediately if the first key is empty! } - }; - } else { - return new NativeReaderIterator>>() { - private @Nullable WindowedValue> current = null; - - @Override - public boolean start() throws IOException { - current = value; - return true; + current = value; + return true; + } + + @Override + public boolean advance() throws IOException { + if (context.workIsFailed()) { + throw new WorkItemCancelledException( + checkStateNotNull(context.getFailedWork()).getWorkItem().getShardingKey()); } - @Override - public boolean advance() throws IOException { - current = null; - context.finishKey(); - return false; + if (context.advance()) { + @SuppressWarnings("unchecked") + K newKey = (K) context.getKey(); + KeyedWorkItem newKeyedWorkItem = + new WindmillKeyedWorkItem<>( + newKey, + context.getWork().getWorkItem(), + windowCoder, + windowsCoder, + valueCoder, + context.getWindmillTagEncoding(), + context.getDrainMode(), + skipUndecodableElements.isAccessible() + && Boolean.TRUE.equals(skipUndecodableElements.get())); + current = new ValueInEmptyWindows<>(newKeyedWorkItem); + return true; } - @Override - public WindowedValue> getCurrent() { - if (current == null) { - throw new NoSuchElementException(); - } - return value; + current = null; + return false; + } + + @Override + public WindowedValue> getCurrent() { + if (current == null) { + throw new NoSuchElementException(); } - }; - } + return current; + } + }; } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemCancelledException.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemCancelledException.java index a12a5075c5ee..85e63cc0064c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemCancelledException.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemCancelledException.java @@ -17,21 +17,33 @@ */ package org.apache.beam.runners.dataflow.worker; +import java.util.Optional; +import javax.annotation.Nullable; + /** Indicates that the work item was cancelled and should not be retried. */ @SuppressWarnings({ "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) public class WorkItemCancelledException extends RuntimeException { + private final @Nullable Long shardingKey; + public WorkItemCancelledException(long sharding_key) { super("Work item cancelled for key " + sharding_key); + this.shardingKey = sharding_key; } public WorkItemCancelledException(String message, Throwable cause) { super(message, cause); + this.shardingKey = null; } public WorkItemCancelledException(Throwable cause) { super(cause); + this.shardingKey = null; + } + + public Optional getShardingKey() { + return Optional.ofNullable(shardingKey); } /** Returns whether an exception was caused by a {@link WorkItemCancelledException}. */ 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 29d5fb3561a1..e1fa860236a9 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 @@ -457,7 +457,7 @@ public NativeReaderIterator>> iterator() thro context.setActiveReader(reader); - return new UnboundedReaderIterator<>(reader, context, started, options); + return new UnboundedReaderIterator(reader, started); } @Override @@ -487,6 +487,179 @@ private UnboundedSource parseSource(int index } return (UnboundedSource) rawSource; } + + private class UnboundedReaderIterator + extends NativeReader.NativeReaderIterator>> { + // Do not close reader. The reader is cached in StreamingModeExecutionContext.readerCache, and + // will be reused until the cache is evicted, expired or invalidated. + // See UnboundedReader#iterator(). + private UnboundedSource.UnboundedReader reader; + private boolean started; + private final Instant endTime; + private final int maxElems; + private final FluentBackoff backoffFactory; + private int elemsRead = 0; + + private UnboundedReaderIterator(UnboundedSource.UnboundedReader reader, boolean started) { + this.reader = reader; + this.started = started; + DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class); + long maxReadTimeMs = debugOptions.getUnboundedReaderMaxReadTimeMs(); + this.endTime = Instant.now().plus(Duration.millis(maxReadTimeMs)); + this.maxElems = debugOptions.getUnboundedReaderMaxElements(); + this.backoffFactory = + FluentBackoff.DEFAULT + .withInitialBackoff(Duration.millis(10)) + .withMaxCumulativeBackoff( + Duration.millis(debugOptions.getUnboundedReaderMaxWaitForElementsMs())); + } + + private void initNextKey() throws IOException { + @SuppressWarnings("unchecked") + UnboundedSource.UnboundedReader nextReader = + (UnboundedSource.UnboundedReader) context.getCachedReader(); + this.started = nextReader != null; + if (nextReader == null) { + String key = context.getSerializedKey().toStringUtf8(); + // Key is expected to be a zero-padded integer representing the split index. + int splitIndex = Integer.parseInt(key.substring(0, 16), 16) - 1; + + UnboundedSource splitSource = parseSource(splitIndex); + + UnboundedSource.@Nullable CheckpointMark checkpoint = null; + if (splitSource.getCheckpointMarkCoder() != null) { + checkpoint = context.getReaderCheckpoint(splitSource.getCheckpointMarkCoder()); + } + + nextReader = splitSource.createReader(options, checkpoint); + } + + context.setActiveReader(nextReader); + this.reader = nextReader; + } + + @Override + public boolean start() throws IOException { + if (context.workIsFailed()) { + throw new WorkItemCancelledException( + org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions + .checkNotNull(context.getFailedWork()) + .getWorkItem() + .getShardingKey()); + } + if (started) { + // This is a reader that has been restored from the unbounded reader cache. + // It has already been started, so this call to start() should delegate + // to advance() instead. + return advance(); + } + try { + if (!reader.start()) { + if (context.advance()) { + initNextKey(); + elemsRead = 0; + return advance(); + } + context.finishKey(); + return false; + } + } catch (Exception e) { + throw new IOException( + "Failed to start reading from source: " + reader.getCurrentSource(), e); + } + elemsRead++; + return true; + } + + @Override + public boolean advance() throws IOException { + // Limits are placed on how much data we allow to return, how long we process the input + // before checkpointing and how long we block for input to be available. This ensures + // 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( + org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions + .checkNotNull(context.getFailedWork()) + .getWorkItem() + .getShardingKey()); + } + if (elemsRead >= maxElems + || Instant.now().isAfter(endTime) + || context.isSinkFullHintSet()) { + if (context.advance()) { + initNextKey(); + elemsRead = 0; + backoff = backoffFactory.backoff(); + continue; + } + context.finishKey(); + return false; + } + + boolean hasElement; + try { + if (started) { + hasElement = reader.advance(); + } else { + hasElement = reader.start(); + started = true; + } + } catch (Exception e) { + throw new IOException("Failed to read from source: " + reader.getCurrentSource(), e); + } + + if (hasElement) { + elemsRead++; + return true; + } + + long nextBackoff = backoff.nextBackOffMillis(); + if (nextBackoff == BackOff.STOP) { + if (context.advance()) { + initNextKey(); + elemsRead = 0; + backoff = backoffFactory.backoff(); + continue; + } + context.finishKey(); + return false; + } + Uninterruptibles.sleepUninterruptibly(nextBackoff, TimeUnit.MILLISECONDS); + } + } + + @Override + public WindowedValue> getCurrent() throws NoSuchElementException { + WindowedValue result = + WindowedValues.timestampedValueInGlobalWindow( + reader.getCurrent(), reader.getCurrentTimestamp()); + return result.withValue( + new ValueWithRecordId<>(result.getValue(), reader.getCurrentRecordId())); + } + + @Override + public void close() { + // Don't close reader. + } + + @Override + public NativeReader.Progress getProgress() { + return null; + } + + @Override + public NativeReader.DynamicSplitResult requestDynamicSplit( + NativeReader.DynamicSplitRequest request) { + return null; + } + + @Override + public double getRemainingParallelism() { + return Double.NaN; + } + } } /** @@ -781,121 +954,4 @@ public double getRemainingParallelism() { return Double.NaN; } } - - private static class UnboundedReaderIterator - extends NativeReader.NativeReaderIterator>> { - // Do not close reader. The reader is cached in StreamingModeExecutionContext.readerCache, and - // will be reused until the cache is evicted, expired or invalidated. - // See UnboundedReader#iterator(). - private final UnboundedSource.UnboundedReader reader; - private final StreamingModeExecutionContext context; - private final boolean started; - private final Instant endTime; - private final int maxElems; - private final FluentBackoff backoffFactory; - private int elemsRead = 0; - - private UnboundedReaderIterator( - UnboundedSource.UnboundedReader reader, - StreamingModeExecutionContext context, - boolean started, - PipelineOptions options) { - this.reader = reader; - this.context = context; - this.started = started; - DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class); - long maxReadTimeMs = debugOptions.getUnboundedReaderMaxReadTimeMs(); - this.endTime = Instant.now().plus(Duration.millis(maxReadTimeMs)); - this.maxElems = debugOptions.getUnboundedReaderMaxElements(); - this.backoffFactory = - FluentBackoff.DEFAULT - .withInitialBackoff(Duration.millis(10)) - .withMaxCumulativeBackoff( - Duration.millis(debugOptions.getUnboundedReaderMaxWaitForElementsMs())); - } - - @Override - public boolean start() throws IOException { - if (started) { - // This is a reader that has been restored from the unbounded reader cache. - // It has already been started, so this call to start() should delegate - // to advance() instead. - return advance(); - } - try { - if (!reader.start()) { - context.finishKey(); - return false; - } - } catch (Exception e) { - throw new IOException( - "Failed to start reading from source: " + reader.getCurrentSource(), e); - } - elemsRead++; - return true; - } - - @Override - public boolean advance() throws IOException { - // Limits are placed on how much data we allow to return, how long we process the input - // before checkpointing and how long we block for input to be available. This ensures - // 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.finishKey(); - return false; - } - try { - if (reader.advance()) { - elemsRead++; - return true; - } - } catch (Exception e) { - throw new IOException("Failed to advance source: " + reader.getCurrentSource(), e); - } - long nextBackoff = backoff.nextBackOffMillis(); - if (nextBackoff == BackOff.STOP) { - context.finishKey(); - return false; - } - Uninterruptibles.sleepUninterruptibly(nextBackoff, TimeUnit.MILLISECONDS); - } - } - - @Override - public WindowedValue> getCurrent() throws NoSuchElementException { - WindowedValue result = - WindowedValues.timestampedValueInGlobalWindow( - reader.getCurrent(), reader.getCurrentTimestamp()); - return result.withValue( - new ValueWithRecordId<>(result.getValue(), reader.getCurrentRecordId())); - } - - @Override - public void close() { - // Don't close reader. - } - - @Override - public NativeReader.Progress getProgress() { - return null; - } - - @Override - public NativeReader.DynamicSplitResult requestDynamicSplit( - NativeReader.DynamicSplitRequest request) { - return null; - } - - @Override - public double getRemainingParallelism() { - return Double.NaN; - } - } } 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 b4f3a22a7f52..b83bc0d2cce7 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 @@ -24,8 +24,10 @@ import org.apache.beam.runners.core.metrics.ExecutionStateTracker; import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutor; import org.apache.beam.runners.dataflow.worker.DataflowWorkExecutor; +import org.apache.beam.runners.dataflow.worker.HotKeyLogger; import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; +import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; @@ -72,9 +74,37 @@ public final void executeWork( Work work, WindmillStateReader stateReader, SideInputStateFetcher sideInputStateFetcher, - Windmill.WorkItemCommitRequest.Builder outputBuilder) + Windmill.WorkItemCommitRequest.Builder outputBuilder, + BoundedQueueExecutor workQueueExecutor, + BoundedQueueExecutorWorkHandle budgetHandle, + HotKeyLogger hotKeyLogger, + boolean hotKeyLoggingEnabled, + String stepName, + int maxKeyGroupBatchSize, + long maxKeyGroupBatchTimeNanos, + long maxKeyGroupBatchBytes, + @Nullable String sourceBytesProcessCounterName, + StreamingModeExecutionContext.KeySwitchListener keySwitchListener) throws Exception { - context().start(key, work, stateReader, sideInputStateFetcher, outputBuilder, workExecutor()); + context() + .start( + key, + work, + stateReader, + sideInputStateFetcher, + outputBuilder, + workExecutor(), + workQueueExecutor, + budgetHandle, + hotKeyLogger, + hotKeyLoggingEnabled, + stepName, + keyCoder().orElse(null), + maxKeyGroupBatchSize, + maxKeyGroupBatchTimeNanos, + maxKeyGroupBatchBytes, + keySwitchListener, + sourceBytesProcessCounterName); workExecutor().execute(); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java index b0fda0e0516b..6029c34b76e0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java @@ -82,6 +82,7 @@ public final class Work implements RefreshableWork { private volatile TimedState currentState; private volatile boolean isFailed; private volatile String processingThreadName = ""; + private volatile @Nullable Runnable onFailureListener = null; private final boolean drainMode; private Work( @@ -247,8 +248,19 @@ public void setProcessingThreadName(String processingThreadName) { } @Override - public void setFailed() { + public synchronized void setFailed() { this.isFailed = true; + Runnable listener = onFailureListener; + if (listener != null) { + listener.run(); + } + } + + public synchronized void setOnFailureListener(@Nullable Runnable listener) { + this.onFailureListener = listener; + if (isFailed && listener != null) { + listener.run(); + } } public boolean isCommitPending() { @@ -273,6 +285,10 @@ public void queueCommit(WorkItemCommitRequest commitRequest, ComputationState co processingContext.workCommitter().accept(Commit.create(commitRequest, computationState, this)); } + public Consumer workCommitter() { + return processingContext.workCommitter(); + } + public WindmillStateReader createWindmillStateReader() { return WindmillStateReader.forWork(this); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java index 526b67890783..d62d41ae0037 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java @@ -108,6 +108,13 @@ boolean commitWorkItem( Windmill.WorkItemCommitRequest request, Consumer onDone); + default boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + throw new UnsupportedOperationException(); + } + /** Flushes any pending work items to the wire. */ void flush(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java index b840d22a3434..fc951cb6633f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java @@ -18,11 +18,14 @@ package org.apache.beam.runners.dataflow.worker.windmill.client.commits; import com.google.auto.value.AutoValue; +import java.util.Optional; import org.apache.beam.runners.dataflow.worker.streaming.ComputationState; import org.apache.beam.runners.dataflow.worker.streaming.Work; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** Value class for a queued commit. */ @Internal @@ -32,7 +35,20 @@ public abstract class Commit { public static Commit create( WorkItemCommitRequest request, ComputationState computationState, Work work) { Preconditions.checkArgument(request.getSerializedSize() > 0); - return new AutoValue_Commit(request, computationState, work); + return new AutoValue_Commit( + request, computationState, work, Optional.empty(), ImmutableList.of(work)); + } + + public static Commit createMultiKey( + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest, + ComputationState computationState, + ImmutableList workBatch) { + return new AutoValue_Commit( + WorkItemCommitRequest.getDefaultInstance(), + computationState, + workBatch.get(0), + Optional.of(multiKeyRequest), + workBatch); } public final String computationId() { @@ -45,7 +61,23 @@ public final String computationId() { public abstract Work work(); + public abstract Optional multiKeyRequest(); + + public abstract ImmutableList workBatch(); + + public final boolean isFailed() { + for (Work w : workBatch()) { + if (w.isFailed()) { + return true; + } + } + return false; + } + public final int getSize() { + if (multiKeyRequest().isPresent()) { + return multiKeyRequest().get().getSerializedSize(); + } return request().getSerializedSize(); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java index b68f53121b86..55207d6bcba9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java @@ -100,7 +100,7 @@ public void start() { @Override public void commit(Commit commit) { - if (commit.work().isFailed()) { + if (commit.isFailed()) { failCommit(commit); } else { commitQueue.put(commit); @@ -147,8 +147,15 @@ private void drainCommitQueue() { } private void failCommit(Commit commit) { - commit.work().setFailed(); - onCommitComplete.accept(CompleteCommit.forFailedWork(commit)); + for (Work w : commit.workBatch()) { + w.setFailed(); + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), + w.getShardedKey(), + w.id(), + org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus.ABORTED)); + } } @Override @@ -173,8 +180,8 @@ private void streamingCommitLoop() { // take() blocks until a value is available in the commitQueue. Preconditions.checkNotNull(initialCommit); - if (initialCommit.work().isFailed()) { - onCommitComplete.accept(CompleteCommit.forFailedWork(initialCommit)); + if (initialCommit.isFailed()) { + failCommit(initialCommit); initialCommit = null; continue; } @@ -202,20 +209,40 @@ private void streamingCommitLoop() { /** Adds the commit to the batch if it fits, returning true if it is consumed. */ private boolean tryAddToCommitBatch(Commit commit, CommitWorkStream.RequestBatcher batcher) { Preconditions.checkNotNull(commit); - commit.work().setState(Work.State.COMMITTING); + for (Work w : commit.workBatch()) { + w.setState(Work.State.COMMITTING); + } activeCommitBytes.addAndGet(commit.getSize()); - boolean isCommitAccepted = - batcher.commitWorkItem( - commit.computationId(), - commit.request(), - commitStatus -> { - onCommitComplete.accept(CompleteCommit.create(commit, commitStatus)); - activeCommitBytes.addAndGet(-commit.getSize()); - }); + boolean isCommitAccepted; + if (commit.multiKeyRequest().isPresent()) { + isCommitAccepted = + batcher.commitMultiKeyWorkItem( + commit.computationId(), + commit.multiKeyRequest().get(), + commitStatus -> { + for (Work w : commit.workBatch()) { + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), w.getShardedKey(), w.id(), commitStatus)); + } + activeCommitBytes.addAndGet(-commit.getSize()); + }); + } else { + isCommitAccepted = + batcher.commitWorkItem( + commit.computationId(), + commit.request(), + commitStatus -> { + onCommitComplete.accept(CompleteCommit.create(commit, commitStatus)); + activeCommitBytes.addAndGet(-commit.getSize()); + }); + } // Since the commit was not accepted, revert the changes made above. if (!isCommitAccepted) { - commit.work().setState(Work.State.COMMIT_QUEUED); + for (Work w : commit.workBatch()) { + w.setState(Work.State.COMMIT_QUEUED); + } activeCommitBytes.addAndGet(-commit.getSize()); } @@ -246,8 +273,8 @@ private boolean tryAddToCommitBatch(Commit commit, CommitWorkStream.RequestBatch } // Drop commits for failed work. Such commits will be dropped by Windmill anyway. - if (commit.work().isFailed()) { - onCommitComplete.accept(CompleteCommit.forFailedWork(commit)); + if (commit.isFailed()) { + failCommit(commit); continue; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java index d24676652fd8..0882de996cac 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java @@ -35,6 +35,7 @@ import java.util.function.Function; import javax.annotation.Nullable; import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitRequestChunk; @@ -270,7 +271,7 @@ private void flushInternal(Map requests) if (requests.size() == 1) { Map.Entry elem = requests.entrySet().iterator().next(); - if (elem.getValue().request().getSerializedSize() + if (elem.getValue().serializedCommit().size() > AbstractWindmillStream.RPC_STREAM_CHUNK_SIZE) { issueMultiChunkRequest(elem.getKey(), elem.getValue()); } else { @@ -289,6 +290,7 @@ private void issueSingleRequest(long id, PendingRequest pendingRequest) .setComputationId(pendingRequest.computationId()) .setRequestId(id) .setShardingKey(pendingRequest.shardingKey()) + .setCommitType(pendingRequest.commitType()) .setSerializedWorkItemCommit(pendingRequest.serializedCommit()); StreamingCommitWorkRequest chunk = requestBuilder.build(); synchronized (this) { @@ -318,7 +320,8 @@ private void issueBatchedRequest(Map requests) chunkBuilder .setRequestId(entry.getKey()) .setShardingKey(request.shardingKey()) - .setSerializedWorkItemCommit(request.serializedCommit()); + .setSerializedWorkItemCommit(request.serializedCommit()) + .setCommitType(request.commitType()); } StreamingCommitWorkRequest request = requestBuilder.build(); synchronized (this) { @@ -360,7 +363,8 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) .setRequestId(id) .setSerializedWorkItemCommit(chunk) .setComputationId(pendingRequest.computationId()) - .setShardingKey(pendingRequest.shardingKey()); + .setShardingKey(pendingRequest.shardingKey()) + .setCommitType(pendingRequest.commitType()); int remaining = serializedCommit.size() - end; if (remaining > 0) { chunkBuilder.setRemainingBytesForWorkItem(remaining); @@ -378,34 +382,34 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) @AutoValue abstract static class PendingRequest { - - private static PendingRequest create( - String computationId, WorkItemCommitRequest request, Consumer onDone) { - return new AutoValue_GrpcCommitWorkStream_PendingRequest(computationId, request, onDone); + static PendingRequest create( + String computationId, + long shardingKey, + ByteString serializedCommit, + StreamingCommitRequestChunk.CommitType commitType, + Consumer onDone) { + return new AutoValue_GrpcCommitWorkStream_PendingRequest( + computationId, shardingKey, serializedCommit, commitType, onDone); } abstract String computationId(); - abstract WorkItemCommitRequest request(); + abstract long shardingKey(); + + abstract ByteString serializedCommit(); + + abstract StreamingCommitRequestChunk.CommitType commitType(); abstract Consumer onDone(); private long getBytes() { - return (long) request().getSerializedSize() + computationId().length(); - } - - private ByteString serializedCommit() { - return request().toByteString(); + return (long) serializedCommit().size() + computationId().length(); } private void completeWithStatus(CommitStatus commitStatus) { onDone().accept(commitStatus); } - private long shardingKey() { - return request().getShardingKey(); - } - private void abort() { completeWithStatus(CommitStatus.ABORTED); } @@ -462,7 +466,33 @@ public boolean commitWorkItem( return false; } - PendingRequest request = PendingRequest.create(computation, commitRequest, onDone); + PendingRequest request = + PendingRequest.create( + computation, + commitRequest.getShardingKey(), + commitRequest.toByteString(), + StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_SINGLE_KEY, + onDone); + add(idGenerator.incrementAndGet(), request); + return true; + } + + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest commitRequest, + Consumer onDone) { + if (!canAccept(commitRequest.getSerializedSize() + computation.length())) { + return false; + } + + PendingRequest request = + PendingRequest.create( + computation, + commitRequest.getKeyGroup().getLow(), + commitRequest.toByteString(), + StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_MULTI_KEY, + onDone); add(idGenerator.incrementAndGet(), request); return true; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateReader.java index c609bed4eae0..e739e1a7e19c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateReader.java @@ -39,6 +39,7 @@ import org.apache.beam.runners.dataflow.worker.KeyTokenInvalidException; import org.apache.beam.runners.dataflow.worker.WindmillTimeUtils; import org.apache.beam.runners.dataflow.worker.WorkItemCancelledException; +import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataRequest; @@ -588,7 +589,7 @@ private KeyedGetDataRequest createRequest(Iterable> toFetch) { private void consumeResponse(KeyedGetDataResponse response, Set> toFetch) { bytesRead += response.getSerializedSize(); if (response.getFailed()) { - throw new KeyTokenInvalidException(key.toStringUtf8()); + throw new KeyTokenInvalidException(ShardedKey.create(key, shardingKey), key.toStringUtf8()); } if (!key.equals(response.getKey())) { 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 dbcb0f09c0a9..a9f0e6d4555a 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 @@ -21,6 +21,10 @@ import com.google.api.services.dataflow.model.MapTask; import com.google.auto.value.AutoValue; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ScheduledExecutorService; @@ -28,11 +32,13 @@ import java.util.function.Function; import java.util.function.Supplier; import javax.annotation.concurrent.ThreadSafe; +import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; import org.apache.beam.runners.dataflow.worker.DataflowExecutionStateSampler; import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutorFactory; import org.apache.beam.runners.dataflow.worker.HotKeyLogger; import org.apache.beam.runners.dataflow.worker.ReaderCache; +import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext; import org.apache.beam.runners.dataflow.worker.WorkItemCancelledException; import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC; import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle; @@ -60,6 +66,7 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.fn.IdGenerator; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; @@ -78,6 +85,19 @@ public class StreamingWorkScheduler { private static final Logger LOG = LoggerFactory.getLogger(StreamingWorkScheduler.class); + private static class BatchExecutionException extends RuntimeException { + private final List executedWorks; + + BatchExecutionException(Throwable cause, List executedWorks) { + super(cause); + this.executedWorks = executedWorks; + } + + List getExecutedWorks() { + return executedWorks; + } + } + private final DataflowWorkerHarnessOptions options; private final Supplier clock; private final ComputationWorkExecutorFactory computationWorkExecutorFactory; @@ -90,10 +110,12 @@ public class StreamingWorkScheduler { private final ConcurrentMap stageInfoMap; private final DataflowExecutionStateSampler sampler; private final StreamingGlobalConfigHandle globalConfigHandle; + private final BoundedQueueExecutor workExecutor; public StreamingWorkScheduler( DataflowWorkerHarnessOptions options, Supplier clock, + BoundedQueueExecutor workExecutor, ComputationWorkExecutorFactory computationWorkExecutorFactory, SideInputStateFetcherFactory sideInputStateFetcherFactory, FailureTracker failureTracker, @@ -106,6 +128,7 @@ public StreamingWorkScheduler( StreamingGlobalConfigHandle globalConfigHandle) { this.options = options; this.clock = clock; + this.workExecutor = workExecutor; this.computationWorkExecutorFactory = computationWorkExecutorFactory; this.sideInputStateFetcherFactory = sideInputStateFetcherFactory; this.failureTracker = failureTracker; @@ -148,6 +171,7 @@ public static StreamingWorkScheduler create( return new StreamingWorkScheduler( options, clock, + workExecutor, computationWorkExecutorFactory, SideInputStateFetcherFactory.fromOptions(options), failureTracker, @@ -246,7 +270,7 @@ private void processWork( } private void processWork( - ComputationState computationState, Work work, BoundedQueueExecutorWorkHandle unusedHandle) { + ComputationState computationState, Work work, BoundedQueueExecutorWorkHandle handle) { Windmill.WorkItem workItem = work.getWorkItem(); String computationId = computationState.getComputationId(); ByteString key = workItem.getKey(); @@ -272,55 +296,132 @@ private void processWork( stageInfoMap.computeIfAbsent( mapTask.getStageName(), s -> StageInfo.create(s, mapTask.getSystemName())); + List worksToCleanup = null; try { if (work.isFailed()) { throw new WorkItemCancelledException(workItem.getShardingKey()); } - // Execute the user code for the Work. - ExecuteWorkResult executeWorkResult = executeWork(work, stageInfo, computationState); - Windmill.WorkItemCommitRequest.Builder commitRequest = executeWorkResult.commitWorkRequest(); + // Execute the user code for the Work batch. + ExecuteWorkResult executeWorkResult = executeWork(work, stageInfo, computationState, handle); + List workBatch = executeWorkResult.workBatch(); + worksToCleanup = workBatch; + List outputBuilders = + executeWorkResult.outputBuilders(); + Map> accumulatedCallbacks = + executeWorkResult.accumulatedCallbacks(); + + // Cache accumulated flat-map of all callbacks + commitFinalizer.cacheCommitFinalizers(accumulatedCallbacks); + + if (workBatch.size() > 1) { + // Multi-key commit packing + Windmill.MultiKeyWorkItemCommitRequest.Builder multiKeyBuilder = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder(); + + // Group by the KeyGroup (from primary work item) + Work primaryWork = workBatch.get(0); + if (primaryWork.getKeyGroup().isPresent()) { + Work.KeyGroup keyGroup = primaryWork.getKeyGroup().get(); + multiKeyBuilder.setKeyGroup( + Windmill.Uint128Proto.newBuilder() + .setHigh(keyGroup.high()) + .setLow(keyGroup.low()) + .build()); + } - // Validate the commit request, possibly requesting truncation if the commitSize is too large. - Windmill.WorkItemCommitRequest validatedCommitRequest = - validateCommitRequestSize(commitRequest.build(), computationId, workItem); + for (int i = 0; i < workBatch.size(); i++) { + Windmill.WorkItemCommitRequest.Builder builder = outputBuilders.get(i); + Work w = workBatch.get(i); + builder.addAllPerWorkItemLatencyAttributions(w.getLatencyAttributions(sampler)); + + // Aggregate ONLY finalize IDs to the top level + multiKeyBuilder.addAllFinalizeIds(builder.getFinalizeIdsList()); + // Clear only finalize IDs from individual request + builder.clearFinalizeIds(); + // Keep output_messages and pubsub_messages scoped inside builder + multiKeyBuilder.addRequests(builder.build()); + } - // Queue the commit. - work.queueCommit(validatedCommitRequest, computationState); - recordProcessingStats(commitRequest, workItem, executeWorkResult); - LOG.debug("Processing done for work token: {}", workItem.getWorkToken()); + // Transition states of all completed works in the batch to COMMIT_QUEUED and submit + for (Work w : workBatch) { + w.setState(Work.State.COMMIT_QUEUED); + } + + // Package and submit the commit batch transactionally + primaryWork + .workCommitter() + .accept( + Commit.createMultiKey( + multiKeyBuilder.build(), computationState, ImmutableList.copyOf(workBatch))); + } else { + // Standard single-key commit path (fully backward-compatible with Appliance path) + Windmill.WorkItemCommitRequest.Builder commitRequest = outputBuilders.get(0); + Windmill.WorkItemCommitRequest validatedCommitRequest = + validateCommitRequestSize( + commitRequest.build(), computationId, workBatch.get(0).getWorkItem()); + workBatch.get(0).setState(Work.State.COMMIT_QUEUED); + validatedCommitRequest = + validatedCommitRequest + .toBuilder() + .addAllPerWorkItemLatencyAttributions( + workBatch.get(0).getLatencyAttributions(sampler)) + .build(); + workBatch.get(0).queueCommit(validatedCommitRequest, computationState); + } + + recordProcessingStats(workBatch, outputBuilders, executeWorkResult.stateBytesRead()); + LOG.debug("Processing done for work batch size: {}", workBatch.size()); } catch (Throwable t) { - // OutOfMemoryError that are caught will be rethrown and trigger jvm termination. + // Handle batch failure rollback and rescheduling try { - workFailureProcessor.logAndProcessFailure( + List failedBatch = ImmutableList.of(work); + Throwable errorToProcess = t; + if (t instanceof BatchExecutionException) { + failedBatch = ((BatchExecutionException) t).getExecutedWorks(); + Throwable cause = t.getCause(); + if (cause != null) { + errorToProcess = cause; + } + } + worksToCleanup = failedBatch; + + List executableWorks = new ArrayList<>(); + for (Work w : failedBatch) { + executableWorks.add( + ExecutableWork.create(w, (retry, h) -> processWork(computationState, retry, h))); + } + + workFailureProcessor.logAndProcessFailureBatch( computationId, - ExecutableWork.create(work, (retry, h) -> processWork(computationState, retry, h)), - t, + executableWorks, + errorToProcess, invalidWork -> computationState.completeWorkAndScheduleNextWorkForKey( invalidWork.getShardedKey(), invalidWork.id())); } catch (OutOfMemoryError oom) { throw oom; } catch (Throwable t2) { - LOG.warn("Failed to process work failure safely for work {}", work.id(), t2); + LOG.warn("Failed to process work failure safely for work batch", t2); throw ExceptionUtils.propagate(t2); } } finally { - // Update total processing time counters. Updating in finally clause ensures that - // work items causing exceptions are also accounted in time spent. long processingTimeMsecs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - processingStartTimeNanos); stageInfo.totalProcessingMsecs().addValue(processingTimeMsecs); - // Attribute all the processing to timers if the work item contains any timers. - // Tests show that work items rarely contain both timers and message bundles. It should - // be a fairly close approximation. - // Another option: Derive time split between messages and timers based on recent totals. - // either here or in DFE. if (work.getWorkItem().hasTimers()) { stageInfo.timerProcessingMsecs().addValue(processingTimeMsecs); } + if (worksToCleanup != null) { + for (Work w : worksToCleanup) { + w.setOnFailureListener(null); + } + } else { + work.setOnFailureListener(null); + } + resetWorkLoggingContext(work.getLatencyTrackingId()); work.setProcessingThreadName(""); } @@ -354,24 +455,34 @@ private Windmill.WorkItemCommitRequest validateCommitRequestSize( } private void recordProcessingStats( - Windmill.WorkItemCommitRequest.Builder outputBuilder, - Windmill.WorkItem workItem, - ExecuteWorkResult executeWorkResult) { - // Compute shuffle and state byte statistics these will be flushed asynchronously. - long stateBytesWritten = - outputBuilder - .clearOutputMessages() - .clearPerWorkItemLatencyAttributions() - .build() - .getSerializedSize(); - - streamingCounters.windmillShuffleBytesRead().addValue(computeShuffleBytesRead(workItem)); - streamingCounters.windmillStateBytesRead().addValue(executeWorkResult.stateBytesRead()); - streamingCounters.windmillStateBytesWritten().addValue(stateBytesWritten); + List workBatch, + List outputBuilders, + long totalStateBytesRead) { + long totalStateBytesWritten = 0; + long totalShuffleBytesRead = 0; + for (int i = 0; i < workBatch.size(); i++) { + Windmill.WorkItem workItem = workBatch.get(i).getWorkItem(); + Windmill.WorkItemCommitRequest.Builder outputBuilder = outputBuilders.get(i); + long stateBytesWritten = + outputBuilder + .clearOutputMessages() + .clearPerWorkItemLatencyAttributions() + .build() + .getSerializedSize(); + totalStateBytesWritten += stateBytesWritten; + totalShuffleBytesRead += computeShuffleBytesRead(workItem); + } + streamingCounters.windmillShuffleBytesRead().addValue(totalShuffleBytesRead); + streamingCounters.windmillStateBytesRead().addValue(totalStateBytesRead); + streamingCounters.windmillStateBytesWritten().addValue(totalStateBytesWritten); } private ExecuteWorkResult executeWork( - Work work, StageInfo stageInfo, ComputationState computationState) throws Exception { + Work work, + StageInfo stageInfo, + ComputationState computationState, + BoundedQueueExecutorWorkHandle handle) + throws Exception { Windmill.WorkItem workItem = work.getWorkItem(); ByteString key = workItem.getKey(); Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem); @@ -388,18 +499,40 @@ private ExecuteWorkResult executeWork( SideInputStateFetcher localSideInputStateFetcher = sideInputStateFetcherFactory.createSideInputStateFetcher(work::fetchSideInput); - // If the read output KVs, then we can decode Windmill's byte key into userland - // key object and provide it to the execution context for use with per-key state. - // Otherwise, we pass null. - // - // The coder type that will be present is: - // WindowedValueCoder(TimerOrElementCoder(KvCoder)) Optional> keyCoder = computationWorkExecutor.keyCoder(); @SuppressWarnings("deprecation") @Nullable final Object executionKey = !keyCoder.isPresent() ? null : keyCoder.get().decode(key.newInput(), Coder.Context.OUTER); + // Parse limits from experiments + String batchSizeStr = + org.apache.beam.sdk.options.ExperimentalOptions.getExperimentValue( + options, "max_key_group_batch_size"); + int maxKeyGroupBatchSize = batchSizeStr != null ? Integer.parseInt(batchSizeStr) : 100; + + String batchTimeStr = + org.apache.beam.sdk.options.ExperimentalOptions.getExperimentValue( + options, "max_key_group_batch_time_ms"); + long maxKeyGroupBatchTimeNanos = + TimeUnit.MILLISECONDS.toNanos(batchTimeStr != null ? Long.parseLong(batchTimeStr) : 100); + + String batchBytesStr = + org.apache.beam.sdk.options.ExperimentalOptions.getExperimentValue( + options, "max_key_group_batch_bytes"); + long maxKeyGroupBatchBytes = + batchBytesStr != null ? Long.parseLong(batchBytesStr) : 10L * 1024 * 1024; + + // MDC and samplers aligner callback + StreamingModeExecutionContext.KeySwitchListener keySwitchListener = + (oldWork, newWork) -> { + resetWorkLoggingContext(oldWork.getLatencyTrackingId()); + setUpWorkLoggingContext( + newWork.getLatencyTrackingId(), computationState.getComputationId()); + newWork.setProcessingThreadName(Thread.currentThread().getName()); + oldWork.setProcessingThreadName(""); + }; + if (workItem.hasHotKeyInfo()) { Windmill.HotKeyInfo hotKeyInfo = workItem.getHotKeyInfo(); Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000); @@ -417,57 +550,74 @@ private ExecuteWorkResult executeWork( // Blocks while executing work. computationWorkExecutor.executeWork( - executionKey, work, stateReader, localSideInputStateFetcher, outputBuilder); - - if (work.isFailed()) { - throw new WorkItemCancelledException(workItem.getShardingKey()); + executionKey, + work, + stateReader, + localSideInputStateFetcher, + outputBuilder, + workExecutor, + handle, + hotKeyLogger, + options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging"), + getShuffleTaskStepName(computationState.getMapTask()), + maxKeyGroupBatchSize, + maxKeyGroupBatchTimeNanos, + maxKeyGroupBatchBytes, + computationState.sourceBytesProcessCounterName(), + keySwitchListener); + + StreamingModeExecutionContext context = computationWorkExecutor.context(); + if (context.workIsFailed()) { + throw new WorkItemCancelledException( + Preconditions.checkNotNull(context.getFailedWork()).getWorkItem().getShardingKey()); } - // Reports source bytes processed to WorkItemCommitRequest if available. - try { - long sourceBytesProcessed = - computationWorkExecutor.computeSourceBytesProcessed( - computationState.sourceBytesProcessCounterName()); - outputBuilder.setSourceBytesProcessed(sourceBytesProcessed); - } catch (Exception e) { - LOG.error("{}", e.toString()); - } - - commitFinalizer.cacheCommitFinalizers(computationWorkExecutor.context().flushState()); + // Retrieve executed works, output builders, and accumulated callbacks from execution context + ImmutableList workBatch = ImmutableList.copyOf(context.getExecutedWorks()); + ImmutableList outputBuilders = + ImmutableList.copyOf(context.getOutputBuilders()); + Map> accumulatedCallbacks = + new HashMap<>(context.getAccumulatedCallbacks()); // Release the execution state for another thread to use. computationState.releaseComputationWorkExecutor(computationWorkExecutor); computationWorkExecutor = null; - work.setState(Work.State.COMMIT_QUEUED); - outputBuilder.addAllPerWorkItemLatencyAttributions(work.getLatencyAttributions(sampler)); - return ExecuteWorkResult.create( - outputBuilder, stateReader.getBytesRead() + localSideInputStateFetcher.getBytesRead()); + workBatch, + outputBuilders, + accumulatedCallbacks, + context.getStateBytesRead() + localSideInputStateFetcher.getBytesRead()); } catch (Throwable t) { + List executedWorks = ImmutableList.of(work); if (computationWorkExecutor != null) { - // If processing failed due to a thrown exception, close the executionState. Do not - // return/release the executionState back to computationState as that will lead to this - // executionState instance being reused. - LOG.debug("Invalidating executor after work item {} failed", workItem.getWorkToken(), t); + executedWorks = ImmutableList.copyOf(computationWorkExecutor.context().getExecutedWorks()); + if (executedWorks.isEmpty()) { + executedWorks = ImmutableList.of(work); + } + LOG.debug("Invalidating executor after work failure", t); computationWorkExecutor.invalidate(); } - - // Re-throw the exception, it will be caught and handled by workFailureProcessor downstream. - throw t; + throw new BatchExecutionException(t, executedWorks); } } @AutoValue abstract static class ExecuteWorkResult { - - private static ExecuteWorkResult create( - Windmill.WorkItemCommitRequest.Builder commitWorkRequest, long stateBytesRead) { + static ExecuteWorkResult create( + ImmutableList workBatch, + ImmutableList outputBuilders, + Map> accumulatedCallbacks, + long stateBytesRead) { return new AutoValue_StreamingWorkScheduler_ExecuteWorkResult( - commitWorkRequest, stateBytesRead); + workBatch, outputBuilders, accumulatedCallbacks, stateBytesRead); } - abstract Windmill.WorkItemCommitRequest.Builder commitWorkRequest(); + abstract ImmutableList workBatch(); + + abstract ImmutableList outputBuilders(); + + abstract Map> accumulatedCallbacks(); abstract long stateBytesRead(); } 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 18c8e9b8d83c..163c3f3cbd21 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 @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures; +import java.util.List; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import java.util.function.Supplier; @@ -124,6 +125,58 @@ public void logAndProcessFailure( } } + public void logAndProcessFailureBatch( + String computationId, + List executableWorks, + Throwable t, + Consumer onInvalidWork) + throws Throwable { + List worksToRetryLocally = new java.util.ArrayList<>(); + + for (ExecutableWork executableWork : executableWorks) { + switch (evaluateRetry(computationId, executableWork.work(), t)) { + case DO_NOT_RETRY: + onInvalidWork.accept(executableWork.work()); + break; + case RETRY_LOCALLY: + worksToRetryLocally.add(executableWork); + break; + case RETHROW_THROWABLE: + throw t; + } + } + + if (!worksToRetryLocally.isEmpty()) { + // Sleep ONCE for the entire batch delay to avoid sequential thread blocks + Uninterruptibles.sleepUninterruptibly(retryLocallyDelayMs, TimeUnit.MILLISECONDS); + for (ExecutableWork ew : worksToRetryLocally) { + workUnitExecutor.forceExecute(ew, ew.work().getSerializedWorkItemSize()); + } + } + } + + private static @Nullable KeyTokenInvalidException getKeyTokenInvalidException( + @Nullable Throwable t) { + while (t != null) { + if (t instanceof KeyTokenInvalidException) { + return (KeyTokenInvalidException) t; + } + t = t.getCause(); + } + return null; + } + + private static @Nullable WorkItemCancelledException getWorkItemCancelledException( + @Nullable Throwable t) { + while (t != null) { + if (t instanceof WorkItemCancelledException) { + return (WorkItemCancelledException) t; + } + t = t.getCause(); + } + return null; + } + private String tryToDumpHeap() { return heapDumper .dumpAndGetHeap() @@ -147,19 +200,46 @@ private RetryEvaluation evaluateRetry(String computationId, Work work, Throwable @Nullable final Throwable cause = t.getCause(); Throwable parsedException = (t instanceof UserCodeException && cause != null) ? cause : t; if (KeyTokenInvalidException.isKeyTokenInvalidException(parsedException)) { - LOG.debug( - "Execution of work for computation '{}' on sharding key '{}' failed due to token expiration. " - + "Work will not be retried locally.", - computationId, - work.getWorkItem().getShardingKey()); + KeyTokenInvalidException invalidException = getKeyTokenInvalidException(parsedException); + if (invalidException != null && invalidException.getShardedKey().isPresent()) { + if (work.getShardedKey().equals(invalidException.getShardedKey().get())) { + LOG.debug( + "Execution of work for computation '{}' on sharding key '{}' failed due to token expiration. " + + "Work will not be retried locally.", + computationId, + work.getWorkItem().getShardingKey()); + return RetryEvaluation.DO_NOT_RETRY; + } else { + LOG.debug( + "Execution of work for computation '{}' on sharding key '{}' aborted due to token mismatch on another key. " + + "Work will be retried locally.", + computationId, + work.getWorkItem().getShardingKey()); + return RetryEvaluation.RETRY_LOCALLY; + } + } return RetryEvaluation.DO_NOT_RETRY; } if (WorkItemCancelledException.isWorkItemCancelledException(parsedException)) { - LOG.debug( - "Execution of work for computation '{}' on sharding key '{}' failed. " - + "Work will not be retried locally.", - computationId, - work.getWorkItem().getShardingKey()); + WorkItemCancelledException cancelledException = + getWorkItemCancelledException(parsedException); + if (cancelledException != null && cancelledException.getShardingKey().isPresent()) { + if (work.getWorkItem().getShardingKey() == cancelledException.getShardingKey().get()) { + LOG.debug( + "Execution of work for computation '{}' on sharding key '{}' failed. " + + "Work will not be retried locally.", + computationId, + work.getWorkItem().getShardingKey()); + return RetryEvaluation.DO_NOT_RETRY; + } else { + LOG.debug( + "Execution of work for computation '{}' on sharding key '{}' aborted due to cancellation of another key. " + + "Work will be retried locally.", + computationId, + work.getWorkItem().getShardingKey()); + return RetryEvaluation.RETRY_LOCALLY; + } + } return RetryEvaluation.DO_NOT_RETRY; } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java index 5be8ec0a6c72..250968c0bb28 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java @@ -400,6 +400,7 @@ public void shutdown() {} public RequestBatcher batcher() { return new RequestBatcher() { final List requests = new ArrayList<>(); + final List multiKeyRequests = new ArrayList<>(); @Override public boolean commitWorkItem( @@ -423,6 +424,18 @@ public boolean commitWorkItem( return true; } + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + LOG.debug("commitWorkStream::commitMultiKeyWorkItem: {}", request); + if (multiKeyRequests.size() > 5) return false; + multiKeyRequests.add(new MultiKeyRequestAndDone(request, onDone)); + flush(); + return true; + } + @Override public void flush() { for (RequestAndDone elem : requests) { @@ -445,6 +458,26 @@ public void flush() { .orElse(Windmill.CommitStatus.OK)); } requests.clear(); + + for (MultiKeyRequestAndDone elem : multiKeyRequests) { + Windmill.CommitStatus status = Windmill.CommitStatus.OK; + for (WorkItemCommitRequest req : elem.request.getRequestsList()) { + commitsReceived.put(req.getWorkToken(), req); + Windmill.CommitStatus itemStatus = + Optional.ofNullable( + streamingCommitsToOffer.remove( + WorkId.builder() + .setWorkToken(req.getWorkToken()) + .setCacheToken(req.getCacheToken()) + .build())) + .orElse(Windmill.CommitStatus.OK); + if (itemStatus != Windmill.CommitStatus.OK) { + status = itemStatus; + } + } + elem.onDone.accept(status); + } + multiKeyRequests.clear(); } class RequestAndDone { @@ -456,6 +489,18 @@ class RequestAndDone { this.onDone = onDone; } } + + class MultiKeyRequestAndDone { + final Consumer onDone; + final Windmill.MultiKeyWorkItemCommitRequest request; + + MultiKeyRequestAndDone( + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + this.request = request; + this.onDone = onDone; + } + } }; } 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 d58f20076994..95e822b5b491 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 @@ -1327,7 +1327,7 @@ public void testKeyCommitTooLargeException() throws Exception { makeExpectedTruncationRequestOutput( 1, "large_key", DEFAULT_SHARDING_KEY, largeCommit.getEstimatedWorkItemCommitBytes()) .build(), - largeCommit); + removeDynamicFields(largeCommit)); // Check this explicitly since the estimated commit bytes weren't actually // checked against an expected value in the previous step @@ -3507,8 +3507,8 @@ public void testExceptionInvalidatesCache() throws Exception { } // Ensure that the invalidated dofn had tearDown called on them. - assertEquals(1, TestExceptionInvalidatesCacheFn.tearDownCallCount.get()); - assertEquals(2, TestExceptionInvalidatesCacheFn.setupCallCount.get()); + assertEquals(2, TestExceptionInvalidatesCacheFn.tearDownCallCount.get()); + assertEquals(3, TestExceptionInvalidatesCacheFn.setupCallCount.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 216ca5386675..3249b13ef832 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 @@ -51,17 +51,21 @@ 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.KeySwitchListener; 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; 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.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.Watermarks; 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.BoundedQueueExecutor; 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; @@ -464,4 +468,502 @@ public void testSetBacklogBytes() { assertEquals(1234, outputBuilder.getSourceBacklogBytes()); } + + @Test + public void testAdvanceKeySwitching() throws Exception { + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setWorkToken(100L) + .setCacheToken(1000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workA = + createMockWork( + workItemA, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setWorkToken(200L) + .setCacheToken(2000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workB = + createMockWork( + workItemB, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + BoundedQueueExecutor mockExecutor = mock(BoundedQueueExecutor.class); + BoundedQueueExecutorWorkHandle mockBudget = mock(BoundedQueueExecutorWorkHandle.class); + ExecutableWork executableWorkB = ExecutableWork.create(workB, (w, h) -> {}); + + org.mockito.Mockito.when( + mockExecutor.pollWork( + org.mockito.Mockito.eq(COMPUTATION_ID), + org.mockito.Mockito.eq( + org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup.create(1L, 2L)), + org.mockito.Mockito.eq(mockBudget))) + .thenReturn(java.util.Optional.of(executableWorkB)); + + Windmill.WorkItemCommitRequest.Builder outputBuilderA = + Windmill.WorkItemCommitRequest.newBuilder(); + + java.util.concurrent.atomic.AtomicReference oldWorkRef = + new java.util.concurrent.atomic.AtomicReference<>(); + java.util.concurrent.atomic.AtomicReference newWorkRef = + new java.util.concurrent.atomic.AtomicReference<>(); + KeySwitchListener listener = + (oldW, newW) -> { + oldWorkRef.set(oldW); + newWorkRef.set(newW); + }; + + executionContext.start( + "keyA", + workA, + stateReader, + sideInputStateFetcher, + outputBuilderA, + workExecutor, + mockExecutor, + mockBudget, + new HotKeyLogger(), + false, + "step1", + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + 5, // maxKeyGroupBatchSize + 1000000000L, // maxKeyGroupBatchTimeNanos (1s) + 10000000L, // maxKeyGroupBatchBytes (10MB) + listener); + + // 1. Verify primary key info + assertEquals("keyA", executionContext.getKey()); + assertEquals(workA, executionContext.getWork()); + assertEquals(1, executionContext.getExecutedWorks().size()); + assertEquals(1, executionContext.getOutputBuilders().size()); + + // 2. Advance context (trigger key switch) + boolean advanced = executionContext.advance(); + assertTrue(advanced); + + // 3. Verify that the context has transitioned to keyB + assertEquals("keyB", executionContext.getKey()); + assertEquals(workB, executionContext.getWork()); + assertEquals(2, executionContext.getExecutedWorks().size()); + assertEquals(2, executionContext.getOutputBuilders().size()); + + // Verify listener was called + assertEquals(workA, oldWorkRef.get()); + assertEquals(workB, newWorkRef.get()); + + // 4. Advance again (should return false since pollWork returns empty) + org.mockito.Mockito.when( + mockExecutor.pollWork( + org.mockito.Mockito.eq(COMPUTATION_ID), + org.mockito.Mockito.eq( + org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup.create(1L, 2L)), + org.mockito.Mockito.eq(mockBudget))) + .thenReturn(java.util.Optional.empty()); + + boolean advancedAgain = executionContext.advance(); + assertFalse(advancedAgain); + } + + @Test + public void testAdvanceLimitThresholds() throws Exception { + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setWorkToken(100L) + .setCacheToken(1000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workA = + createMockWork( + workItemA, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setWorkToken(200L) + .setCacheToken(2000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workB = + createMockWork( + workItemB, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + BoundedQueueExecutor mockExecutor = mock(BoundedQueueExecutor.class); + BoundedQueueExecutorWorkHandle mockBudget = mock(BoundedQueueExecutorWorkHandle.class); + ExecutableWork executableWorkB = ExecutableWork.create(workB, (w, h) -> {}); + + org.mockito.Mockito.when( + mockExecutor.pollWork( + org.mockito.Mockito.eq(COMPUTATION_ID), + org.mockito.Mockito.eq( + org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup.create(1L, 2L)), + org.mockito.Mockito.eq(mockBudget))) + .thenReturn(java.util.Optional.of(executableWorkB)); + + Windmill.WorkItemCommitRequest.Builder outputBuilderA = + Windmill.WorkItemCommitRequest.newBuilder(); + + // Case 1: maxKeyGroupBatchSize is 0 + executionContext.start( + "keyA", + workA, + stateReader, + sideInputStateFetcher, + outputBuilderA, + workExecutor, + mockExecutor, + mockBudget, + new HotKeyLogger(), + false, + "step1", + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + 0, // maxKeyGroupBatchSize = 0 (batch size threshold hit immediately!) + 1000000000L, + 10000000L, + (k, c) -> {}); + + assertFalse(executionContext.advance()); + + // Case 2: maxKeyGroupBatchBytes limit is exceeded + Windmill.WorkItemCommitRequest.Builder outputBuilderAForBytes = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("some_non_empty_key_to_increase_size")); + + executionContext.start( + "keyA", + workA, + stateReader, + sideInputStateFetcher, + outputBuilderAForBytes, + workExecutor, + mockExecutor, + mockBudget, + new HotKeyLogger(), + false, + "step1", + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + 5, + 1000000000L, + 1L, // maxKeyGroupBatchBytes = 1 byte (exceeded!) + (k, c) -> {}); + + assertFalse(executionContext.advance()); + } + + @Test + public void testFinishKeyReentrantSafety() { + Windmill.WorkItemCommitRequest.Builder outputBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + executionContext.start( + "key", + createMockWork( + Windmill.WorkItem.newBuilder().setKey(ByteString.EMPTY).setWorkToken(17L).build(), + Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), + stateReader, + sideInputStateFetcher, + outputBuilder, + workExecutor); + + // First call + executionContext.finishKey(); + // Second call - should not throw any Exception + executionContext.finishKey(); + } + + @Test + public void testWorkIsFailed_heartbeatFailureOnPolledKey() throws Exception { + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setWorkToken(100L) + .setCacheToken(1000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workA = + createMockWork( + workItemA, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setWorkToken(200L) + .setCacheToken(2000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workB = + createMockWork( + workItemB, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + BoundedQueueExecutor mockExecutor = mock(BoundedQueueExecutor.class); + BoundedQueueExecutorWorkHandle mockBudget = mock(BoundedQueueExecutorWorkHandle.class); + ExecutableWork executableWorkB = ExecutableWork.create(workB, (w, h) -> {}); + + org.mockito.Mockito.when( + mockExecutor.pollWork( + org.mockito.Mockito.eq(COMPUTATION_ID), + org.mockito.Mockito.eq( + org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup.create(1L, 2L)), + org.mockito.Mockito.eq(mockBudget))) + .thenReturn(java.util.Optional.of(executableWorkB)); + + Windmill.WorkItemCommitRequest.Builder outputBuilderA = + Windmill.WorkItemCommitRequest.newBuilder(); + + executionContext.start( + "keyA", + workA, + stateReader, + sideInputStateFetcher, + outputBuilderA, + workExecutor, + mockExecutor, + mockBudget, + new HotKeyLogger(), + false, + "step1", + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + 5, + 1000000000L, + 10000000L, + (k, c) -> {}); + + // Heartbeat fails on primary work A + workA.setFailed(); + + assertTrue(executionContext.workIsFailed()); + assertEquals(workA, executionContext.getFailedWork()); + } + + @Test + public void testAdvance_abortsOnFailedExecutedKey() throws Exception { + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setWorkToken(100L) + .setCacheToken(1000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workA = + createMockWork( + workItemA, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setWorkToken(200L) + .setCacheToken(2000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workB = + createMockWork( + workItemB, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + BoundedQueueExecutor mockExecutor = mock(BoundedQueueExecutor.class); + BoundedQueueExecutorWorkHandle mockBudget = mock(BoundedQueueExecutorWorkHandle.class); + ExecutableWork executableWorkB = ExecutableWork.create(workB, (w, h) -> {}); + + org.mockito.Mockito.when( + mockExecutor.pollWork( + org.mockito.Mockito.eq(COMPUTATION_ID), + org.mockito.Mockito.eq( + org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup.create(1L, 2L)), + org.mockito.Mockito.eq(mockBudget))) + .thenReturn(java.util.Optional.of(executableWorkB)); + + Windmill.WorkItemCommitRequest.Builder outputBuilderA = + Windmill.WorkItemCommitRequest.newBuilder(); + + executionContext.start( + "keyA", + workA, + stateReader, + sideInputStateFetcher, + outputBuilderA, + workExecutor, + mockExecutor, + mockBudget, + new HotKeyLogger(), + false, + "step1", + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + 5, + 1000000000L, + 10000000L, + (k, c) -> {}); + + // Simulate heartbeat failure on key A before switching to key B + workA.setFailed(); + + // Advance should fail immediately and throw WorkItemCancelledException + boolean thrown = false; + try { + executionContext.advance(); + org.junit.Assert.fail("Expected WorkItemCancelledException"); + } catch (WorkItemCancelledException e) { + thrown = true; + } + assertTrue(thrown); + } + + @Test + @SuppressWarnings({"rawtypes", "unchecked"}) + public void testInvalidateCache_clearsAllExecutedKeys() throws Exception { + CounterSet counterSet = new CounterSet(); + ConcurrentHashMap stateNameMap = new ConcurrentHashMap<>(); + stateNameMap.put(NameContextsForTests.nameContextForTest().userName(), "testStateFamily"); + + // Mock stateCache + WindmillStateCache.ForComputation mockStateCache = + mock(WindmillStateCache.ForComputation.class); + + // Use a real ReaderCache so we can verify interactions directly + ReaderCache testReaderCache = + new ReaderCache(Duration.standardMinutes(1), Executors.newSingleThreadExecutor()); + + StreamingModeExecutionContext testContext = + new StreamingModeExecutionContext( + counterSet, + COMPUTATION_ID, + testReaderCache, + stateNameMap, + mockStateCache, + StreamingStepMetricsContainer.createRegistry(), + new DataflowExecutionStateTracker( + ExecutionStateSampler.newForTest(), + executionStateRegistry.getState( + NameContext.forStage("stage"), "other", null, NoopProfileScope.NOOP), + counterSet, + PipelineOptionsFactory.create(), + "test-work-item-id"), + executionStateRegistry, + globalConfigHandle, + Long.MAX_VALUE, + /*throwExceptionOnLargeOutput=*/ false); + + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setWorkToken(100L) + .setCacheToken(1000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workA = + createMockWork( + workItemA, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setWorkToken(200L) + .setCacheToken(2000L) + .setKeyGroup(Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build()) + .build(); + Work workB = + createMockWork( + workItemB, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + BoundedQueueExecutor mockExecutor = mock(BoundedQueueExecutor.class); + BoundedQueueExecutorWorkHandle mockBudget = mock(BoundedQueueExecutorWorkHandle.class); + ExecutableWork executableWorkB = ExecutableWork.create(workB, (w, h) -> {}); + + org.mockito.Mockito.when( + mockExecutor.pollWork( + org.mockito.Mockito.eq(COMPUTATION_ID), + org.mockito.Mockito.eq( + org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup.create(1L, 2L)), + org.mockito.Mockito.eq(mockBudget))) + .thenReturn(java.util.Optional.of(executableWorkB)); + + Windmill.WorkItemCommitRequest.Builder outputBuilderA = + Windmill.WorkItemCommitRequest.newBuilder(); + + testContext.start( + "keyA", + workA, + stateReader, + sideInputStateFetcher, + outputBuilderA, + workExecutor, + mockExecutor, + mockBudget, + new HotKeyLogger(), + false, + "step1", + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + 5, + 1000000000L, + 10000000L, + (k, c) -> {}); + + // Cache reader for keyA + org.apache.beam.sdk.io.UnboundedSource.UnboundedReader readerA = + mock(org.apache.beam.sdk.io.UnboundedSource.UnboundedReader.class); + org.apache.beam.sdk.io.UnboundedSource mockSourceA = + mock(org.apache.beam.sdk.io.UnboundedSource.class); + org.mockito.Mockito.when(readerA.getCurrentSource()).thenReturn(mockSourceA); + org.mockito.Mockito.when(readerA.getWatermark()).thenReturn(new Instant(1000)); + testContext.setActiveReader(readerA); + + // Advance to keyB (which calls finishKey() for keyA, caching readerA) + boolean advanced = testContext.advance(); + assertTrue(advanced); + + // Cache reader for keyB + org.apache.beam.sdk.io.UnboundedSource.UnboundedReader readerB = + mock(org.apache.beam.sdk.io.UnboundedSource.UnboundedReader.class); + org.apache.beam.sdk.io.UnboundedSource mockSourceB = + mock(org.apache.beam.sdk.io.UnboundedSource.class); + org.mockito.Mockito.when(readerB.getCurrentSource()).thenReturn(mockSourceB); + org.mockito.Mockito.when(readerB.getWatermark()).thenReturn(new Instant(1000)); + testContext.setActiveReader(readerB); + + // Call finishKey() for keyB, caching readerB + testContext.finishKey(); + + // Verify both readers are now cached in testReaderCache + org.apache.beam.runners.dataflow.worker.WindmillComputationKey compKeyA = + org.apache.beam.runners.dataflow.worker.WindmillComputationKey.create( + COMPUTATION_ID, workA.getShardedKey()); + org.apache.beam.runners.dataflow.worker.WindmillComputationKey compKeyB = + org.apache.beam.runners.dataflow.worker.WindmillComputationKey.create( + COMPUTATION_ID, workB.getShardedKey()); + + // We acquire readerA and readerB to ensure they were cached + org.apache.beam.sdk.io.UnboundedSource.UnboundedReader cachedA = + testReaderCache.acquireReader(compKeyA, 1000L, 100L + 1); + org.apache.beam.sdk.io.UnboundedSource.UnboundedReader cachedB = + testReaderCache.acquireReader(compKeyB, 2000L, 200L + 1); + assertEquals(readerA, cachedA); + assertEquals(readerB, cachedB); + + // Put them back into cache + testReaderCache.cacheReader(compKeyA, 1000L, 100L, readerA); + testReaderCache.cacheReader(compKeyB, 2000L, 200L, readerB); + + // Call invalidateCache() + testContext.invalidateCache(); + + // Verify both readers are invalidated from the cache (acquireReader returns null) + org.apache.beam.sdk.io.UnboundedSource.UnboundedReader clearedA = + testReaderCache.acquireReader(compKeyA, 1000L, 100L + 1); + org.apache.beam.sdk.io.UnboundedSource.UnboundedReader clearedB = + testReaderCache.acquireReader(compKeyB, 2000L, 200L + 1); + org.junit.Assert.assertNull(clearedA); + org.junit.Assert.assertNull(clearedB); + + // Verify that stateCache.invalidate was called for both keys + org.mockito.Mockito.verify(mockStateCache) + .invalidate( + org.mockito.Mockito.eq(workItemA.getKey()), + org.mockito.Mockito.eq(workItemA.getShardingKey())); + org.mockito.Mockito.verify(mockStateCache) + .invalidate( + org.mockito.Mockito.eq(workItemB.getKey()), + org.mockito.Mockito.eq(workItemB.getShardingKey())); + } } 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 539c38eeb1da..7780a49728a2 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 @@ -30,6 +30,7 @@ import java.util.Arrays; import java.util.List; import java.util.concurrent.ThreadLocalRandom; +import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.options.ValueProvider; @@ -94,6 +95,9 @@ public void testWorkItemCancelledException() throws IOException { Windmill.WorkItem.newBuilder().setKey(ByteString.EMPTY).setWorkToken(0L).build(); when(mockContext.getWorkItem()).thenReturn(workItem); + Work mockFailedWork = createMockWork(workItem); + when(mockContext.getFailedWork()).thenReturn(mockFailedWork); + try (TestWindmillReaderIterator iter = new TestWindmillReaderIterator(mockContext, ValueProvider.StaticValueProvider.of(false))) { iter.start(); @@ -122,6 +126,12 @@ public void testFinishKeyCalled() throws Exception { .build()) .build(); when(mockContext.getWorkItem()).thenReturn(workItem); + when(mockContext.advance()) + .thenAnswer( + inv -> { + mockContext.finishKey(); + return false; + }); try (TestWindmillReaderIterator iter = new TestWindmillReaderIterator(mockContext, ValueProvider.StaticValueProvider.of(false))) { @@ -131,6 +141,78 @@ public void testFinishKeyCalled() throws Exception { } } + @Test + public void testAdvanceKeyChaining() throws Exception { + StreamingModeExecutionContext mockContext = mock(StreamingModeExecutionContext.class); + when(mockContext.workIsFailed()).thenReturn(false); + + // Work item A (1 message) + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setWorkToken(100L) + .addMessageBundles( + Windmill.InputMessageBundle.newBuilder() + .setSourceComputationId("foo") + .addMessages( + Windmill.Message.newBuilder() + .setTimestamp(1000) + .setData(ByteString.EMPTY) + .build()) + .build()) + .build(); + when(mockContext.getWorkItem()).thenReturn(workItemA); + + // Work item B (1 message) + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setWorkToken(200L) + .addMessageBundles( + Windmill.InputMessageBundle.newBuilder() + .setSourceComputationId("foo") + .addMessages( + Windmill.Message.newBuilder() + .setTimestamp(2000) + .setData(ByteString.EMPTY) + .build()) + .build()) + .build(); + + Work mockWorkB = createMockWork(workItemB); + + // Set up context.advance() to mock transition + when(mockContext.advance()) + .thenAnswer( + new org.mockito.stubbing.Answer() { + private int count = 0; + + @Override + public Boolean answer(org.mockito.invocation.InvocationOnMock invocation) { + if (count == 0) { + count++; + when(mockContext.getWork()).thenReturn(mockWorkB); + return true; + } + return false; + } + }); + + try (TestWindmillReaderIterator iter = + new TestWindmillReaderIterator(mockContext, ValueProvider.StaticValueProvider.of(false))) { + assertTrue(iter.start()); + assertEquals(1000L, iter.getCurrent().getValue().longValue()); + + // Advance should trigger context.advance(), transition to workItemB, and decode message from + // workItemB (timestamp 2000) + assertTrue(iter.advance()); + assertEquals(2000L, iter.getCurrent().getValue().longValue()); + + // Next advance should exhaust it and return false + assertFalse(iter.advance()); + } + } + private void testForMessageBundleCounts(int... messageBundleCounts) throws IOException { testForMessageBundleCounts(false, messageBundleCounts); } @@ -179,4 +261,53 @@ private void testForMessageBundleCounts(boolean skipErrors, int... messageBundle assertEquals(Arrays.toString(messageBundleCounts) + skipErrors, expected, actual); } } + + @Test + public void testAdvance_abortsOnHeartbeatFailure() throws Exception { + StreamingModeExecutionContext mockContext = mock(StreamingModeExecutionContext.class); + when(mockContext.workIsFailed()).thenReturn(true); + + Windmill.WorkItem failedWorkItem = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setWorkToken(100L) + .setShardingKey(999L) + .build(); + Work mockFailedWork = createMockWork(failedWorkItem); + when(mockContext.getFailedWork()).thenReturn(mockFailedWork); + + try (TestWindmillReaderIterator iter = + new TestWindmillReaderIterator(mockContext, ValueProvider.StaticValueProvider.of(false))) { + boolean thrown = false; + try { + iter.advance(); + org.junit.Assert.fail("Expected WorkItemCancelledException"); + } catch (WorkItemCancelledException e) { + thrown = true; + assertTrue(e.getShardingKey().isPresent()); + assertEquals(999L, e.getShardingKey().get().longValue()); + } + assertTrue(thrown); + } + } + + private static Work createMockWork(Windmill.WorkItem workItem) { + return Work.create( + workItem, + workItem.getSerializedSize(), + org.apache.beam.runners.dataflow.worker.streaming.Watermarks.builder() + .setInputDataWatermark(new org.joda.time.Instant(1000)) + .build(), + Work.createProcessingContext( + "computationId", + mock( + org.apache.beam.runners.dataflow.worker.windmill.client.getdata.GetDataClient + .class), + ignored -> {}, + mock( + org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender + .class)), + false, + org.joda.time.Instant::now); + } } 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 d5cf2948d928..126e8d44603e 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 @@ -83,11 +83,14 @@ import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.util.PropertyNames; import org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowExecutionStateTracker; +import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.KeySwitchListener; import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionStateRegistry; import org.apache.beam.runners.dataflow.worker.WorkerCustomSources.SplittableOnlyBoundedSource; 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.ExecutableWork; 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; @@ -95,6 +98,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; 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; @@ -1057,4 +1061,123 @@ public void testFailedWorkItemsAbort() throws Exception { } assertThat(numReads, equalTo(2)); } + + @Test + public void testUnboundedReaderIterator_multiKeySwitching() throws Exception { + CounterSet counterSet = new CounterSet(); + StreamingModeExecutionStateRegistry executionStateRegistry = + new StreamingModeExecutionStateRegistry(); + ReaderCache readerCache = new ReaderCache(Duration.standardMinutes(1), Runnable::run); + StreamingGlobalConfigHandle globalConfigHandle = + new FixedGlobalConfigHandle(StreamingGlobalConfig.builder().build()); + StreamingModeExecutionContext context = + new StreamingModeExecutionContext( + counterSet, + COMPUTATION_ID, + readerCache, + /*stateNameMap=*/ ImmutableMap.of(), + /*stateCache=*/ null, + StreamingStepMetricsContainer.createRegistry(), + new DataflowExecutionStateTracker( + ExecutionStateSampler.newForTest(), + executionStateRegistry.getState( + NameContext.forStage("stageName"), "other", null, NoopProfileScope.NOOP), + counterSet, + PipelineOptionsFactory.create(), + "test-work-item-id"), + executionStateRegistry, + globalConfigHandle, + Long.MAX_VALUE, + /*throwExceptionOnLargeOutput=*/ false); + + options.setNumWorkers(5); + int maxElements = 2; + DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class); + debugOptions.setUnboundedReaderMaxElements(maxElements); + + Windmill.Uint128Proto keyGroup = + Windmill.Uint128Proto.newBuilder().setHigh(1L).setLow(2L).build(); + + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("0000000000000001")) + .setWorkToken(1) + .setCacheToken(1) + .setKeyGroup(keyGroup) + .setSourceState(Windmill.SourceState.newBuilder().setState(ByteString.EMPTY).build()) + .build(); + + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("0000000000000002")) + .setWorkToken(2) + .setCacheToken(1) + .setKeyGroup(keyGroup) + .setSourceState(Windmill.SourceState.newBuilder().setState(ByteString.EMPTY).build()) + .build(); + + Work workA = + createMockWork( + workItemA, Watermarks.builder().setInputDataWatermark(new Instant(0)).build()); + Work workB = + createMockWork( + workItemB, Watermarks.builder().setInputDataWatermark(new Instant(0)).build()); + + BoundedQueueExecutor mockExecutor = mock(BoundedQueueExecutor.class); + BoundedQueueExecutorWorkHandle mockBudget = mock(BoundedQueueExecutorWorkHandle.class); + ExecutableWork executableWorkB = ExecutableWork.create(workB, (w, h) -> {}); + + org.mockito.Mockito.when( + mockExecutor.pollWork( + org.mockito.Mockito.eq(COMPUTATION_ID), + org.mockito.Mockito.eq( + org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup.create(1L, 2L)), + org.mockito.Mockito.eq(mockBudget))) + .thenReturn(java.util.Optional.of(executableWorkB)); + + context.start( + "keyA", + workA, + mock(WindmillStateReader.class), + mock(SideInputStateFetcher.class), + Windmill.WorkItemCommitRequest.newBuilder(), + mock(WorkExecutor.class), + mockExecutor, + mockBudget, + new HotKeyLogger(), + false, + "step1", + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + 5, // maxKeyGroupBatchSize + 1000000000L, // maxKeyGroupBatchTimeNanos + 10000000L, // maxKeyGroupBatchBytes + mock(KeySwitchListener.class)); + + @SuppressWarnings({"unchecked", "rawtypes"}) + NativeReader>>> reader = + (NativeReader) + WorkerCustomSources.create( + (CloudObject) + serializeToCloudSource(new TestCountingSource(Integer.MAX_VALUE), options) + .getSpec(), + options, + context); + + NativeReaderIterator>>> readerIterator = + reader.iterator(); + + assertTrue(readerIterator.start()); + WindowedValue>> val1 = readerIterator.getCurrent(); + assertEquals(Integer.valueOf(0), val1.getValue().getValue().getKey()); + + assertTrue(readerIterator.advance()); + WindowedValue>> val2 = readerIterator.getCurrent(); + assertEquals(Integer.valueOf(0), val2.getValue().getValue().getKey()); + + assertTrue(readerIterator.advance()); + WindowedValue>> val3 = readerIterator.getCurrent(); + assertEquals(Integer.valueOf(1), val3.getValue().getValue().getKey()); + + assertEquals("0000000000000002", context.getKey().toString()); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java index 01197622c24d..72289b3c266a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java @@ -474,4 +474,67 @@ public void testStop_drainsCommitQueue_concurrentCommit() waitForExpectedSetSize(completeCommits, sentCommits.intValue()); } + + @Test + public void testCommit_multiKeyCommitFailedWork() { + Set completeCommits = Collections.newSetFromMap(new ConcurrentHashMap<>()); + workCommitter = createWorkCommitter(completeCommits::add); + + Work workA = createMockWork(101L); + Work workB = createMockWork(102L); + Work workC = createMockWork(103L); + + // Mark non-primary key B as failed + workB.setFailed(); + + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workA.getWorkItem().getKey()) + .setShardingKey(workA.getWorkItem().getShardingKey()) + .setWorkToken(workA.getWorkItem().getWorkToken()) + .setCacheToken(workA.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workB.getWorkItem().getKey()) + .setShardingKey(workB.getWorkItem().getShardingKey()) + .setWorkToken(workB.getWorkItem().getWorkToken()) + .setCacheToken(workB.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workC.getWorkItem().getKey()) + .setShardingKey(workC.getWorkItem().getShardingKey()) + .setWorkToken(workC.getWorkItem().getWorkToken()) + .setCacheToken(workC.getWorkItem().getCacheToken()) + .build()) + .build(); + + Commit commit = + Commit.createMultiKey( + multiKeyRequest, + createComputationState("computationId"), + org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.of( + workA, workB, workC)); + + workCommitter.start(); + workCommitter.commit(commit); + + // The entire batch must be aborted immediately without making network calls + waitForExpectedSetSize(completeCommits, 3); + + // Verify all three works are aborted individually + assertThat(completeCommits) + .containsExactly( + CompleteCommit.create( + "computationId", workA.getShardedKey(), workA.id(), Windmill.CommitStatus.ABORTED), + CompleteCommit.create( + "computationId", workB.getShardedKey(), workB.id(), Windmill.CommitStatus.ABORTED), + CompleteCommit.create( + "computationId", workC.getShardedKey(), workC.id(), Windmill.CommitStatus.ABORTED)); + + workCommitter.stop(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java index 51bd4816b031..35265e3c147d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java @@ -213,4 +213,155 @@ public void logAndProcessFailure_retriesOnUncaughtUnhandledException_streamingAp runWork.await(); assertThat(invalidWork).isEmpty(); } + + @Test + public void logAndProcessFailureBatch_isolatesFailureForKeyTokenInvalidException() + throws Throwable { + Set executedWork = new HashSet<>(); + Set invalidWork = new HashSet<>(); + + // Setup Work A (the invalid key) + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setShardingKey(1000L) + .setWorkToken(100L) + .build(); + ExecutableWork workA = + ExecutableWork.create( + Work.create( + workItemA, + workItemA.getSerializedSize(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), + Work.createProcessingContext( + "computationId", + new FakeGetDataClient(), + ignored -> {}, + mock(HeartbeatSender.class)), + false, + Instant::now), + (w, h) -> executedWork.add(w)); + + // Setup Work B (the healthy key in the batch) + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setShardingKey(2000L) + .setWorkToken(200L) + .build(); + CountDownLatch runWork = new CountDownLatch(1); + ExecutableWork workB = + ExecutableWork.create( + Work.create( + workItemB, + workItemB.getSerializedSize(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), + Work.createProcessingContext( + "computationId", + new FakeGetDataClient(), + ignored -> {}, + mock(HeartbeatSender.class)), + false, + Instant::now), + (w, h) -> { + executedWork.add(w); + runWork.countDown(); + }); + + WorkFailureProcessor workFailureProcessor = + createWorkFailureProcessor(streamingEngineFailureReporter()); + + java.util.List batch = java.util.Arrays.asList(workA, workB); + + // Trigger batch failure where KeyA has token invalid exception + org.apache.beam.runners.dataflow.worker.streaming.ShardedKey shardedKeyA = + org.apache.beam.runners.dataflow.worker.streaming.ShardedKey.create( + ByteString.copyFromUtf8("keyA"), 1000L); + KeyTokenInvalidException exception = new KeyTokenInvalidException(shardedKeyA, "keyA"); + + workFailureProcessor.logAndProcessFailureBatch( + DEFAULT_COMPUTATION_ID, batch, exception, invalidWork::add); + + // Wait for asynchronous execution to complete + runWork.await(5, TimeUnit.SECONDS); + + // Work A should be marked as invalid (DO_NOT_RETRY) and not rescheduled + assertThat(invalidWork).containsExactly(workA.work()); + // Work B (healthy key) should be rescheduled and executed (RETRY_LOCALLY) + assertThat(executedWork).containsExactly(workB.work()); + } + + @Test + public void logAndProcessFailureBatch_isolatesFailureForWorkItemCancelledException() + throws Throwable { + Set executedWork = new HashSet<>(); + Set invalidWork = new HashSet<>(); + + // Setup Work A (the cancelled key) + Windmill.WorkItem workItemA = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyA")) + .setShardingKey(1000L) + .setWorkToken(100L) + .build(); + ExecutableWork workA = + ExecutableWork.create( + Work.create( + workItemA, + workItemA.getSerializedSize(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), + Work.createProcessingContext( + "computationId", + new FakeGetDataClient(), + ignored -> {}, + mock(HeartbeatSender.class)), + false, + Instant::now), + (w, h) -> executedWork.add(w)); + + // Setup Work B (the healthy key in the batch) + Windmill.WorkItem workItemB = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("keyB")) + .setShardingKey(2000L) + .setWorkToken(200L) + .build(); + CountDownLatch runWork = new CountDownLatch(1); + ExecutableWork workB = + ExecutableWork.create( + Work.create( + workItemB, + workItemB.getSerializedSize(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), + Work.createProcessingContext( + "computationId", + new FakeGetDataClient(), + ignored -> {}, + mock(HeartbeatSender.class)), + false, + Instant::now), + (w, h) -> { + executedWork.add(w); + runWork.countDown(); + }); + + WorkFailureProcessor workFailureProcessor = + createWorkFailureProcessor(streamingEngineFailureReporter()); + + java.util.List batch = java.util.Arrays.asList(workA, workB); + + // Trigger batch failure where KeyA is cancelled (takes shardingKey) + WorkItemCancelledException exception = new WorkItemCancelledException(1000L); + + workFailureProcessor.logAndProcessFailureBatch( + DEFAULT_COMPUTATION_ID, batch, exception, invalidWork::add); + + // Wait for asynchronous execution to complete + runWork.await(5, TimeUnit.SECONDS); + + // Work A should be marked as invalid (DO_NOT_RETRY) and not rescheduled + assertThat(invalidWork).containsExactly(workA.work()); + // Work B (healthy key) should be rescheduled and executed (RETRY_LOCALLY) + assertThat(executedWork).containsExactly(workB.work()); + } } diff --git a/sdks/java/core/build.gradle b/sdks/java/core/build.gradle index f532e9d14166..6ccd2a0a39b7 100644 --- a/sdks/java/core/build.gradle +++ b/sdks/java/core/build.gradle @@ -152,3 +152,12 @@ project.tasks.compileTestJava { test { useJUnit() } + +jar { + zip64 = false +} + +shadowJar { + zip64 = false +} + From df77e7493365e92541b6c84f7a44f527389311a8 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 1 Jun 2026 07:14:43 +0000 Subject: [PATCH 10/11] chore: Save progress of Phase 1-6 before spawning coders for Phase 7 --- .../worker/StreamingModeExecutionContext.java | 96 ++++++++++++++++++- .../streaming/ComputationWorkExecutor.java | 11 +-- .../ComputationWorkExecutorFactory.java | 22 ++++- .../processing/StreamingWorkScheduler.java | 29 +----- 4 files changed, 113 insertions(+), 45 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 5bca60b5661d..a303b05cdcc8 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 @@ -217,6 +217,40 @@ public StreamingModeExecutionContext( StreamingGlobalConfigHandle globalConfigHandle, long sinkByteLimit, boolean throwExceptionOnLargeOutput) { + this( + counterFactory, + computationId, + readerCache, + stateNameMap, + stateCache, + metricsContainerRegistry, + executionStateTracker, + executionStateRegistry, + globalConfigHandle, + sinkByteLimit, + throwExceptionOnLargeOutput, + new HotKeyLogger(), + /* hotKeyLoggingEnabled= */ false, + /* stepName= */ null, + /* sourceBytesProcessCounterName= */ null); + } + + public StreamingModeExecutionContext( + CounterFactory counterFactory, + String computationId, + ReaderCache readerCache, + Map stateNameMap, + ForComputation stateCache, + MetricsContainerRegistry metricsContainerRegistry, + DataflowExecutionStateTracker executionStateTracker, + StreamingModeExecutionStateRegistry executionStateRegistry, + StreamingGlobalConfigHandle globalConfigHandle, + long sinkByteLimit, + boolean throwExceptionOnLargeOutput, + @Nullable HotKeyLogger hotKeyLogger, + boolean hotKeyLoggingEnabled, + @Nullable String stepName, + @Nullable String sourceBytesProcessCounterName) { super( counterFactory, metricsContainerRegistry, @@ -231,6 +265,10 @@ public StreamingModeExecutionContext( this.stateCache = stateCache; this.backlogBytes = UnboundedReader.BACKLOG_UNKNOWN; this.throwExceptionOnLargeOutput = throwExceptionOnLargeOutput; + this.hotKeyLogger = hotKeyLogger; + this.hotKeyLoggingEnabled = hotKeyLoggingEnabled; + this.stepName = stepName; + this.sourceBytesProcessCounterName = sourceBytesProcessCounterName; } @VisibleForTesting @@ -351,6 +389,7 @@ public void start( /* sourceBytesProcessCounterName= */ null); } + @Deprecated public void start( @Nullable Object key, Work work, @@ -369,14 +408,45 @@ public void start( long maxKeyGroupBatchBytes, KeySwitchListener keySwitchListener, @Nullable String sourceBytesProcessCounterName) { + this.hotKeyLogger = hotKeyLogger; + this.hotKeyLoggingEnabled = hotKeyLoggingEnabled; + this.stepName = stepName; + this.sourceBytesProcessCounterName = sourceBytesProcessCounterName; + start( + key, + work, + stateReader, + sideInputStateFetcher, + outputBuilder, + workExecutor, + workQueueExecutor, + budgetHandle, + keyCoder, + maxKeyGroupBatchSize, + maxKeyGroupBatchTimeNanos, + maxKeyGroupBatchBytes, + keySwitchListener); + } + + public void start( + @Nullable Object key, + Work work, + WindmillStateReader stateReader, + SideInputStateFetcher sideInputStateFetcher, + Windmill.WorkItemCommitRequest.Builder outputBuilder, + WorkExecutor workExecutor, + BoundedQueueExecutor workQueueExecutor, + BoundedQueueExecutorWorkHandle budgetHandle, + @Nullable Coder keyCoder, + int maxKeyGroupBatchSize, + long maxKeyGroupBatchTimeNanos, + long maxKeyGroupBatchBytes, + KeySwitchListener keySwitchListener) { this.key = key; this.work = work; this.workExecutor = workExecutor; this.workQueueExecutor = workQueueExecutor; this.budgetHandle = budgetHandle; - this.hotKeyLogger = hotKeyLogger; - this.hotKeyLoggingEnabled = hotKeyLoggingEnabled; - this.stepName = stepName; this.keyCoder = keyCoder; this.maxKeyGroupBatchSize = maxKeyGroupBatchSize; @@ -415,7 +485,25 @@ public void start( this.activeStateReader = stateReader; this.stateBytesRead = 0; - this.sourceBytesProcessCounterName = sourceBytesProcessCounterName; + + LOG.warn( + "JETSKI_DEBUG: hasHotKeyInfo={}, hotKeyLogger={}, stepName={}, hotKeyLoggingEnabled={}, key={}", + work.getWorkItem().hasHotKeyInfo(), + this.hotKeyLogger != null, + this.stepName, + this.hotKeyLoggingEnabled, + key); + + // Move primary key hotkey logging from StreamingWorkScheduler into start(...) + if (work.getWorkItem().hasHotKeyInfo() && this.hotKeyLogger != null && this.stepName != null) { + Windmill.HotKeyInfo hotKeyInfo = work.getWorkItem().getHotKeyInfo(); + Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000); + if (key != null && this.hotKeyLoggingEnabled) { + this.hotKeyLogger.logHotKeyDetection(this.stepName, hotKeyAge, key); + } else { + this.hotKeyLogger.logHotKeyDetection(this.stepName, hotKeyAge); + } + } Instant processingTime = computeProcessingTime(work.getWorkItem().getTimers().getTimersList()); 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 b83bc0d2cce7..9ad0c69bbe9d 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 @@ -24,7 +24,6 @@ import org.apache.beam.runners.core.metrics.ExecutionStateTracker; import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutor; import org.apache.beam.runners.dataflow.worker.DataflowWorkExecutor; -import org.apache.beam.runners.dataflow.worker.HotKeyLogger; import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; @@ -77,13 +76,9 @@ public final void executeWork( Windmill.WorkItemCommitRequest.Builder outputBuilder, BoundedQueueExecutor workQueueExecutor, BoundedQueueExecutorWorkHandle budgetHandle, - HotKeyLogger hotKeyLogger, - boolean hotKeyLoggingEnabled, - String stepName, int maxKeyGroupBatchSize, long maxKeyGroupBatchTimeNanos, long maxKeyGroupBatchBytes, - @Nullable String sourceBytesProcessCounterName, StreamingModeExecutionContext.KeySwitchListener keySwitchListener) throws Exception { context() @@ -96,15 +91,11 @@ public final void executeWork( workExecutor(), workQueueExecutor, budgetHandle, - hotKeyLogger, - hotKeyLoggingEnabled, - stepName, keyCoder().orElse(null), maxKeyGroupBatchSize, maxKeyGroupBatchTimeNanos, maxKeyGroupBatchBytes, - keySwitchListener, - sourceBytesProcessCounterName); + keySwitchListener); workExecutor().execute(); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index 269799903300..9e2f002b909a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -29,6 +29,7 @@ import org.apache.beam.runners.dataflow.worker.DataflowExecutionStateSampler; import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutor; import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutorFactory; +import org.apache.beam.runners.dataflow.worker.HotKeyLogger; import org.apache.beam.runners.dataflow.worker.IntrinsicMapTaskExecutorFactory; import org.apache.beam.runners.dataflow.worker.ReaderCache; import org.apache.beam.runners.dataflow.worker.ReaderRegistry; @@ -97,6 +98,7 @@ final class ComputationWorkExecutorFactory { private final IdGenerator idGenerator; private final StreamingGlobalConfigHandle globalConfigHandle; private final boolean throwExceptionOnLargeOutput; + private final HotKeyLogger hotKeyLogger; ComputationWorkExecutorFactory( DataflowWorkerHarnessOptions options, @@ -106,7 +108,8 @@ final class ComputationWorkExecutorFactory { DataflowExecutionStateSampler sampler, CounterSet pendingDeltaCounters, IdGenerator idGenerator, - StreamingGlobalConfigHandle globalConfigHandle) { + StreamingGlobalConfigHandle globalConfigHandle, + HotKeyLogger hotKeyLogger) { this.options = options; this.mapTaskExecutorFactory = mapTaskExecutorFactory; this.readerCache = readerCache; @@ -124,6 +127,7 @@ final class ComputationWorkExecutorFactory { : StreamingDataflowWorker.MAX_SINK_BYTES; this.throwExceptionOnLargeOutput = hasExperiment(options, THROW_EXCEPTIONS_ON_LARGE_OUTPUT_EXPERIMENT); + this.hotKeyLogger = hotKeyLogger; } private static Nodes.ParallelInstructionNode extractReadNode( @@ -191,8 +195,12 @@ ComputationWorkExecutor createComputationWorkExecutor( DataflowExecutionContext.DataflowExecutionStateTracker executionStateTracker = createExecutionStateTracker(stageInfo, mapTask, workLatencyTrackingId); + boolean hotKeyLoggingEnabled = + options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging"); + String stepName = computationState.getMapTask().getInstructions().get(0).getName(); StreamingModeExecutionContext context = - createExecutionContext(computationState, stageInfo, executionStateTracker); + createExecutionContext( + computationState, stageInfo, executionStateTracker, hotKeyLoggingEnabled, stepName); DataflowMapTaskExecutor mapTaskExecutor = createMapTaskExecutor(context, mapTask, mapTaskNetwork); ReadOperation readOperation = getValidatedReadOperation(mapTaskExecutor); @@ -255,7 +263,9 @@ ComputationWorkExecutor createComputationWorkExecutor( private StreamingModeExecutionContext createExecutionContext( ComputationState computationState, StageInfo stageInfo, - DataflowExecutionContext.DataflowExecutionStateTracker executionStateTracker) { + DataflowExecutionContext.DataflowExecutionStateTracker executionStateTracker, + boolean hotKeyLoggingEnabled, + String stepName) { String computationId = computationState.getComputationId(); return new StreamingModeExecutionContext( pendingDeltaCounters, @@ -268,7 +278,11 @@ private StreamingModeExecutionContext createExecutionContext( stageInfo.executionStateRegistry(), globalConfigHandle, maxSinkBytes, - throwExceptionOnLargeOutput); + throwExceptionOnLargeOutput, + hotKeyLogger, + hotKeyLoggingEnabled, + stepName, + computationState.sourceBytesProcessCounterName()); } private DataflowMapTaskExecutor createMapTaskExecutor( 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 a9f0e6d4555a..845edfe04561 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 @@ -17,8 +17,6 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.work.processing; -import static org.apache.beam.sdk.options.ExperimentalOptions.hasExperiment; - import com.google.api.services.dataflow.model.MapTask; import com.google.auto.value.AutoValue; import java.util.ArrayList; @@ -69,7 +67,6 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.checkerframework.checker.nullness.qual.Nullable; -import org.joda.time.Duration; import org.joda.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -106,7 +103,6 @@ List getExecutedWorks() { private final WorkFailureProcessor workFailureProcessor; private final StreamingCommitFinalizer commitFinalizer; private final StreamingCounters streamingCounters; - private final HotKeyLogger hotKeyLogger; private final ConcurrentMap stageInfoMap; private final DataflowExecutionStateSampler sampler; private final StreamingGlobalConfigHandle globalConfigHandle; @@ -122,7 +118,6 @@ public StreamingWorkScheduler( WorkFailureProcessor workFailureProcessor, StreamingCommitFinalizer commitFinalizer, StreamingCounters streamingCounters, - HotKeyLogger hotKeyLogger, ConcurrentMap stageInfoMap, DataflowExecutionStateSampler sampler, StreamingGlobalConfigHandle globalConfigHandle) { @@ -135,7 +130,6 @@ public StreamingWorkScheduler( this.workFailureProcessor = workFailureProcessor; this.commitFinalizer = commitFinalizer; this.streamingCounters = streamingCounters; - this.hotKeyLogger = hotKeyLogger; this.stageInfoMap = stageInfoMap; this.sampler = sampler; this.globalConfigHandle = globalConfigHandle; @@ -166,7 +160,8 @@ public static StreamingWorkScheduler create( sampler, streamingCounters.pendingDeltaCounters(), idGenerator, - globalConfigHandle); + globalConfigHandle, + hotKeyLogger); return new StreamingWorkScheduler( options, @@ -178,7 +173,6 @@ public static StreamingWorkScheduler create( workFailureProcessor, StreamingCommitFinalizer.create(workExecutor, commitFinalizerCleanupExecutor), streamingCounters, - hotKeyLogger, stageInfoMap, sampler, globalConfigHandle); @@ -533,21 +527,6 @@ private ExecuteWorkResult executeWork( oldWork.setProcessingThreadName(""); }; - if (workItem.hasHotKeyInfo()) { - Windmill.HotKeyInfo hotKeyInfo = workItem.getHotKeyInfo(); - Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000); - - String stepName = getShuffleTaskStepName(computationState.getMapTask()); - if (executionKey != null - && (options.isHotKeyLoggingEnabled() - || hasExperiment(options, "enable_hot_key_logging")) - && keyCoder.isPresent()) { - hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge, executionKey); - } else { - hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge); - } - } - // Blocks while executing work. computationWorkExecutor.executeWork( executionKey, @@ -557,13 +536,9 @@ private ExecuteWorkResult executeWork( outputBuilder, workExecutor, handle, - hotKeyLogger, - options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging"), - getShuffleTaskStepName(computationState.getMapTask()), maxKeyGroupBatchSize, maxKeyGroupBatchTimeNanos, maxKeyGroupBatchBytes, - computationState.sourceBytesProcessCounterName(), keySwitchListener); StreamingModeExecutionContext context = computationWorkExecutor.context(); From 6c78923e378a0fe715fec6b7001f203986694273 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 1 Jun 2026 07:36:33 +0000 Subject: [PATCH 11/11] context update --- .../worker/StreamingModeExecutionContext.java | 44 ++-- .../streaming/ComputationWorkExecutor.java | 7 +- .../ComputationWorkExecutorFactory.java | 13 +- .../processing/StreamingWorkScheduler.java | 26 +- .../worker/StreamingDataflowWorkerTest.java | 226 +++++++++++++++++- .../StreamingModeExecutionContextTest.java | 72 +++--- .../worker/WorkerCustomSourcesTest.java | 3 - 7 files changed, 294 insertions(+), 97 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 a303b05cdcc8..d62e6c04d0c3 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 @@ -330,14 +330,12 @@ public void start( Work work, WindmillStateReader stateReader, SideInputStateFetcher sideInputStateFetcher, - Windmill.WorkItemCommitRequest.Builder outputBuilder, WorkExecutor workExecutor) { start( key, work, stateReader, sideInputStateFetcher, - outputBuilder, workExecutor, /* workQueueExecutor= */ null, /* budgetHandle= */ null, @@ -357,7 +355,6 @@ public void start( Work work, WindmillStateReader stateReader, SideInputStateFetcher sideInputStateFetcher, - Windmill.WorkItemCommitRequest.Builder outputBuilder, WorkExecutor workExecutor, BoundedQueueExecutor workQueueExecutor, BoundedQueueExecutorWorkHandle budgetHandle, @@ -374,7 +371,6 @@ public void start( work, stateReader, sideInputStateFetcher, - outputBuilder, workExecutor, workQueueExecutor, budgetHandle, @@ -395,7 +391,6 @@ public void start( Work work, WindmillStateReader stateReader, SideInputStateFetcher sideInputStateFetcher, - Windmill.WorkItemCommitRequest.Builder outputBuilder, WorkExecutor workExecutor, BoundedQueueExecutor workQueueExecutor, BoundedQueueExecutorWorkHandle budgetHandle, @@ -417,7 +412,6 @@ public void start( work, stateReader, sideInputStateFetcher, - outputBuilder, workExecutor, workQueueExecutor, budgetHandle, @@ -433,7 +427,6 @@ public void start( Work work, WindmillStateReader stateReader, SideInputStateFetcher sideInputStateFetcher, - Windmill.WorkItemCommitRequest.Builder outputBuilder, WorkExecutor workExecutor, BoundedQueueExecutor workQueueExecutor, BoundedQueueExecutorWorkHandle budgetHandle, @@ -442,7 +435,17 @@ public void start( long maxKeyGroupBatchTimeNanos, long maxKeyGroupBatchBytes, KeySwitchListener keySwitchListener) { - this.key = key; + if (key != null) { + this.key = key; + } else if (keyCoder != null) { + try { + this.key = keyCoder.decode(work.getWorkItem().getKey().newInput(), Coder.Context.OUTER); + } catch (IOException e) { + throw new RuntimeException("Failed to decode primary key during start", e); + } + } else { + this.key = null; + } this.work = work; this.workExecutor = workExecutor; this.workQueueExecutor = workQueueExecutor; @@ -466,8 +469,15 @@ public void start( work.setOnFailureListener(() -> this.failedWork = work); this.executedWorks.add(work); - this.outputBuilders.add(outputBuilder); - this.outputBuilder = outputBuilder; + + Windmill.WorkItemCommitRequest.Builder primaryOutputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(work.getWorkItem().getKey()) + .setShardingKey(work.getWorkItem().getShardingKey()) + .setWorkToken(work.getWorkItem().getWorkToken()) + .setCacheToken(work.getWorkItem().getCacheToken()); + this.outputBuilders.add(primaryOutputBuilder); + this.outputBuilder = primaryOutputBuilder; this.finishKeyCalled = false; this.computationKey = WindmillComputationKey.create(computationId, work.getShardedKey()); @@ -486,20 +496,12 @@ public void start( this.activeStateReader = stateReader; this.stateBytesRead = 0; - LOG.warn( - "JETSKI_DEBUG: hasHotKeyInfo={}, hotKeyLogger={}, stepName={}, hotKeyLoggingEnabled={}, key={}", - work.getWorkItem().hasHotKeyInfo(), - this.hotKeyLogger != null, - this.stepName, - this.hotKeyLoggingEnabled, - key); - // Move primary key hotkey logging from StreamingWorkScheduler into start(...) - if (work.getWorkItem().hasHotKeyInfo() && this.hotKeyLogger != null && this.stepName != null) { + if (work.getWorkItem().hasHotKeyInfo() && this.hotKeyLogger != null) { Windmill.HotKeyInfo hotKeyInfo = work.getWorkItem().getHotKeyInfo(); Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000); - if (key != null && this.hotKeyLoggingEnabled) { - this.hotKeyLogger.logHotKeyDetection(this.stepName, hotKeyAge, key); + if (this.key != null && this.hotKeyLoggingEnabled) { + this.hotKeyLogger.logHotKeyDetection(this.stepName, hotKeyAge, this.key); } else { this.hotKeyLogger.logHotKeyDetection(this.stepName, hotKeyAge); } 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 9ad0c69bbe9d..d0cb12953ea5 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 @@ -29,11 +29,9 @@ import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; -import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateReader; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.coders.Coder; -import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -69,11 +67,9 @@ public static ComputationWorkExecutor.Builder builder() { * Executes DoFns for the Work. Blocks the calling thread until DoFn(s) have completed execution. */ public final void executeWork( - @Nullable Object key, Work work, WindmillStateReader stateReader, SideInputStateFetcher sideInputStateFetcher, - Windmill.WorkItemCommitRequest.Builder outputBuilder, BoundedQueueExecutor workQueueExecutor, BoundedQueueExecutorWorkHandle budgetHandle, int maxKeyGroupBatchSize, @@ -83,11 +79,10 @@ public final void executeWork( throws Exception { context() .start( - key, + null, work, stateReader, sideInputStateFetcher, - outputBuilder, workExecutor(), workQueueExecutor, budgetHandle, diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index 9e2f002b909a..d1f1ec1fcdf7 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -20,6 +20,7 @@ import static org.apache.beam.runners.dataflow.DataflowRunner.hasExperiment; import com.google.api.services.dataflow.model.MapTask; +import com.google.api.services.dataflow.model.ParallelInstruction; import java.util.function.Function; import org.apache.beam.runners.dataflow.internal.CustomSources; import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; @@ -197,7 +198,17 @@ ComputationWorkExecutor createComputationWorkExecutor( createExecutionStateTracker(stageInfo, mapTask, workLatencyTrackingId); boolean hotKeyLoggingEnabled = options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging"); - String stepName = computationState.getMapTask().getInstructions().get(0).getName(); + ParallelInstruction instruction = computationState.getMapTask().getInstructions().get(0); + String stepName = instruction.getName(); + if (stepName == null) { + stepName = instruction.getOriginalName(); + } + if (stepName == null) { + stepName = instruction.getSystemName(); + } + if (stepName == null) { + stepName = "unknown-step"; + } StreamingModeExecutionContext context = createExecutionContext( computationState, stageInfo, executionStateTracker, hotKeyLoggingEnabled, stepName); 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 845edfe04561..c536624f4670 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 @@ -23,7 +23,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -61,12 +60,10 @@ import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker; import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.WorkFailureProcessor; import org.apache.beam.sdk.annotations.Internal; -import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.fn.IdGenerator; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; -import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -216,8 +213,7 @@ private static String getShuffleTaskStepName(MapTask mapTask) { } /** Resets logging context of the Thread executing the {@link Work} for logging. */ - private void resetWorkLoggingContext(String workLatencyTrackingId) { - sampler.resetForWorkId(workLatencyTrackingId); + private void resetWorkLoggingContext() { DataflowWorkerLoggingMDC.setWorkId(null); DataflowWorkerLoggingMDC.setStageName(null); } @@ -327,7 +323,9 @@ private void processWork( for (int i = 0; i < workBatch.size(); i++) { Windmill.WorkItemCommitRequest.Builder builder = outputBuilders.get(i); Work w = workBatch.get(i); - builder.addAllPerWorkItemLatencyAttributions(w.getLatencyAttributions(sampler)); + if (i == 0) { + builder.addAllPerWorkItemLatencyAttributions(w.getLatencyAttributions(sampler)); + } // Aggregate ONLY finalize IDs to the top level multiKeyBuilder.addAllFinalizeIds(builder.getFinalizeIdsList()); @@ -416,7 +414,8 @@ private void processWork( work.setOnFailureListener(null); } - resetWorkLoggingContext(work.getLatencyTrackingId()); + resetWorkLoggingContext(); + sampler.resetForWorkId(work.getLatencyTrackingId()); work.setProcessingThreadName(""); } } @@ -477,9 +476,6 @@ private ExecuteWorkResult executeWork( ComputationState computationState, BoundedQueueExecutorWorkHandle handle) throws Exception { - Windmill.WorkItem workItem = work.getWorkItem(); - ByteString key = workItem.getKey(); - Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem); ComputationWorkExecutor computationWorkExecutor = computationState .acquireComputationWorkExecutor() @@ -493,12 +489,6 @@ private ExecuteWorkResult executeWork( SideInputStateFetcher localSideInputStateFetcher = sideInputStateFetcherFactory.createSideInputStateFetcher(work::fetchSideInput); - Optional> keyCoder = computationWorkExecutor.keyCoder(); - @SuppressWarnings("deprecation") - @Nullable - final Object executionKey = - !keyCoder.isPresent() ? null : keyCoder.get().decode(key.newInput(), Coder.Context.OUTER); - // Parse limits from experiments String batchSizeStr = org.apache.beam.sdk.options.ExperimentalOptions.getExperimentValue( @@ -520,7 +510,7 @@ private ExecuteWorkResult executeWork( // MDC and samplers aligner callback StreamingModeExecutionContext.KeySwitchListener keySwitchListener = (oldWork, newWork) -> { - resetWorkLoggingContext(oldWork.getLatencyTrackingId()); + resetWorkLoggingContext(); setUpWorkLoggingContext( newWork.getLatencyTrackingId(), computationState.getComputationId()); newWork.setProcessingThreadName(Thread.currentThread().getName()); @@ -529,11 +519,9 @@ private ExecuteWorkResult executeWork( // Blocks while executing work. computationWorkExecutor.executeWork( - executionKey, work, stateReader, localSideInputStateFetcher, - outputBuilder, workExecutor, handle, maxKeyGroupBatchSize, 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 95e822b5b491..b759b306bb76 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 @@ -29,6 +29,7 @@ import static org.hamcrest.Matchers.lessThan; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; @@ -206,6 +207,7 @@ import org.joda.time.Instant; import org.junit.After; import org.junit.Assert; +import org.junit.Assume; import org.junit.Before; import org.junit.Ignore; import org.junit.Rule; @@ -571,11 +573,16 @@ private Windmill.GetWorkResponse buildInput(String input, byte[] metadata) throw Windmill.GetWorkResponse.Builder builder = Windmill.GetWorkResponse.newBuilder(); TextFormat.merge(input, builder); if (metadata != null) { - Windmill.InputMessageBundle.Builder messageBundleBuilder = - builder.getWorkBuilder(0).getWorkBuilder(0).getMessageBundlesBuilder(0); - for (Windmill.Message.Builder messageBuilder : - messageBundleBuilder.getMessagesBuilderList()) { - messageBuilder.setMetadata(addPaneTag(PaneInfo.NO_FIRING, metadata)); + for (Windmill.ComputationWorkItems.Builder compBuilder : builder.getWorkBuilderList()) { + for (Windmill.WorkItem.Builder workBuilder : compBuilder.getWorkBuilderList()) { + for (Windmill.InputMessageBundle.Builder messageBundleBuilder : + workBuilder.getMessageBundlesBuilderList()) { + for (Windmill.Message.Builder messageBuilder : + messageBundleBuilder.getMessagesBuilderList()) { + messageBuilder.setMetadata(addPaneTag(PaneInfo.NO_FIRING, metadata)); + } + } + } } } @@ -5019,6 +5026,215 @@ public void processElement(ProcessContext c) { } } + private Windmill.GetWorkResponse makeInputWithKeyGroup( + int index, long timestamp, String key, long shardingKey, long keyGroupHigh, long keyGroupLow) + throws Exception { + return buildInput( + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"" + + key + + "\"" + + " sharding_key: " + + shardingKey + + " work_token: " + + index + + " cache_token: " + + (index + 1) + + " key_group {" + + " high: " + + keyGroupHigh + + " low: " + + keyGroupLow + + " }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: " + + timestamp + + " data: \"data" + + index + + "\"" + + " }" + + " }" + + " }" + + "}", + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + } + + private Windmill.GetWorkResponse makeInputWithMultiKeys( + int index1, + long timestamp1, + String key1, + long shardingKey1, + int index2, + long timestamp2, + String key2, + long shardingKey2, + long keyGroupHigh, + long keyGroupLow) + throws Exception { + return buildInput( + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"" + + key1 + + "\"" + + " sharding_key: " + + shardingKey1 + + " work_token: " + + index1 + + " cache_token: " + + (index1 + 1) + + " key_group {" + + " high: " + + keyGroupHigh + + " low: " + + keyGroupLow + + " }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: " + + timestamp1 + + " data: \"data" + + index1 + + "\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"" + + key2 + + "\"" + + " sharding_key: " + + shardingKey2 + + " work_token: " + + index2 + + " cache_token: " + + (index2 + 1) + + " key_group {" + + " high: " + + keyGroupHigh + + " low: " + + keyGroupLow + + " }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: " + + timestamp2 + + " data: \"data" + + index2 + + "\"" + + " }" + + " }" + + " }" + + "}", + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + } + + @Test + public void testPerWorkItemLatencyAttributions_MultiKeyBundleOnlyPrimaryReports() + throws Exception { + Assume.assumeTrue(streamingEngine); + + final int workToken1 = 9001; + final int workToken2 = 9002; + final long shardingKey1 = 101L; + final long shardingKey2 = 102L; + final long keyGroupHigh = 1L; + final long keyGroupLow = 2L; + + FakeClock clock = new FakeClock(); + // Inject processing latency on the fake clock in the worker via FakeSlowDoFn. + List instructions = + Arrays.asList( + makeSourceInstruction(StringUtf8Coder.of()), + makeDoFnInstruction( + new FakeSlowDoFn(clock, Duration.millis(1000)), 0, StringUtf8Coder.of()), + makeSinkInstruction(StringUtf8Coder.of(), 0)); + + // We use 1 worker thread so they are processed sequentially, but they will be queued + // in the BoundedQueueExecutor and batched because they are in the same key group. + StreamingDataflowWorker worker = + makeWorker( + defaultWorkerParams( + "--numberOfWorkerHarnessThreads=1", + "--experiments=max_key_group_batch_size=5,unstable_enable_multi_key_bundle") + .setInstructions(instructions) + .setClock(clock) + .setExecutorSupplier(clock::newFakeScheduledExecutor) + .build()); + worker.start(); + + ActiveWorkRefreshSink awrSink = + new ActiveWorkRefreshSink(StreamingDataflowWorkerTest::emptyDataResponder); + server.whenGetDataCalled().answerByDefault(awrSink::getData).delayEachResponseBy(Duration.ZERO); + + // Queue both work items in the FakeWindmillServer before starting. + // They both have the same key group (high=1, low=2) but different keys ("key1", "key2") and + // sharding keys (101, 102). + server + .whenGetWorkCalled() + .thenReturn( + makeInputWithMultiKeys( + workToken1, + 0L /* timestamp */, + "key1", + shardingKey1, + workToken2, + 1000L /* timestamp */, + "key2", + shardingKey2, + keyGroupHigh, + keyGroupLow)); + + // Wait for both commits to complete. They should be committed together in a multi-key bundle. + Map commits = server.waitForAndGetCommits(2); + + worker.stop(); + + WorkItemCommitRequest commit1 = commits.get((long) workToken1); + WorkItemCommitRequest commit2 = commits.get((long) workToken2); + + assertNotNull(commit1); + assertNotNull(commit2); + + // The first work item (primary) should have latency attributions populated. + assertTrue(commit1.getPerWorkItemLatencyAttributionsCount() > 0); + // Verify that we have ACTIVE state latency (which was slow due to FakeSlowDoFn). + boolean hasActiveLA = false; + for (LatencyAttribution la : commit1.getPerWorkItemLatencyAttributionsList()) { + if (la.getState() == State.ACTIVE) { + hasActiveLA = true; + assertEquals(2000L, la.getTotalDurationMillis()); + } + } + assertTrue("Primary commit should have ACTIVE latency attribution", hasActiveLA); + + // The second work item (secondary) should NOT have any latency attributions reported. + assertEquals(0, commit2.getPerWorkItemLatencyAttributionsCount()); + } + @AutoValue abstract static class StreamingDataflowWorkerTestParams { 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 3249b13ef832..9c332243bbad 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 @@ -159,8 +159,6 @@ COMPUTATION_ID, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.cla @Test public void testTimerInternalsSetTimer() throws Exception { - Windmill.WorkItemCommitRequest.Builder outputBuilder = - Windmill.WorkItemCommitRequest.newBuilder(); NameContext nameContext = NameContextsForTests.nameContextForTest(); DataflowOperationContext operationContext = executionContext.createOperationContext(nameContext); @@ -174,7 +172,6 @@ public void testTimerInternalsSetTimer() throws Exception { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder, workExecutor); TimerInternals timerInternals = stepContext.timerInternals(); @@ -189,6 +186,7 @@ public void testTimerInternalsSetTimer() throws Exception { executionContext.finishKey(); executionContext.flushState(); + Windmill.WorkItemCommitRequest.Builder outputBuilder = executionContext.getOutputBuilder(); Windmill.Timer timer = outputBuilder.buildPartial().getOutputTimers(0); assertThat(timer.getTag().toStringUtf8(), equalTo("/skey+0:5000")); assertThat(timer.getTimestamp(), equalTo(TimeUnit.MILLISECONDS.toMicros(5000))); @@ -197,9 +195,6 @@ public void testTimerInternalsSetTimer() throws Exception { @Test public void testTimerInternalsProcessingTimeSkew() { - Windmill.WorkItemCommitRequest.Builder outputBuilder = - Windmill.WorkItemCommitRequest.newBuilder(); - NameContext nameContext = NameContextsForTests.nameContextForTest(); DataflowOperationContext operationContext = executionContext.createOperationContext(nameContext); @@ -226,7 +221,6 @@ public void testTimerInternalsProcessingTimeSkew() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder, workExecutor); TimerInternals timerInternals = stepContext.timerInternals(); assertTrue(timerTimestamp.isBefore(timerInternals.currentProcessingTime())); @@ -425,8 +419,6 @@ public void testStateTagEncodingBasedOnConfig() { for (Boolean isV2Encoding : Lists.newArrayList(Boolean.TRUE, Boolean.FALSE)) { Class expectedEncoding = isV2Encoding ? WindmillTagEncodingV2.class : WindmillTagEncodingV1.class; - Windmill.WorkItemCommitRequest.Builder outputBuilder = - Windmill.WorkItemCommitRequest.newBuilder(); globalConfigHandle.setConfig( StreamingGlobalConfig.builder().setEnableStateTagEncodingV2(isV2Encoding).build()); executionContext.start( @@ -436,7 +428,6 @@ public void testStateTagEncodingBasedOnConfig() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder, workExecutor); assertEquals(expectedEncoding, executionContext.getWindmillTagEncoding().getClass()); } @@ -444,8 +435,6 @@ public void testStateTagEncodingBasedOnConfig() { @Test public void testSetBacklogBytes() { - Windmill.WorkItemCommitRequest.Builder outputBuilder = - Windmill.WorkItemCommitRequest.newBuilder(); NameContext nameContext = NameContextsForTests.nameContextForTest(); DataflowOperationContext operationContext = executionContext.createOperationContext(nameContext); @@ -459,14 +448,13 @@ public void testSetBacklogBytes() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder, workExecutor); stepContext.setBacklogBytes(1234.0); executionContext.finishKey(); executionContext.flushState(); - assertEquals(1234, outputBuilder.getSourceBacklogBytes()); + assertEquals(1234, executionContext.getOutputBuilder().getSourceBacklogBytes()); } @Test @@ -505,9 +493,6 @@ public void testAdvanceKeySwitching() throws Exception { org.mockito.Mockito.eq(mockBudget))) .thenReturn(java.util.Optional.of(executableWorkB)); - Windmill.WorkItemCommitRequest.Builder outputBuilderA = - Windmill.WorkItemCommitRequest.newBuilder(); - java.util.concurrent.atomic.AtomicReference oldWorkRef = new java.util.concurrent.atomic.AtomicReference<>(); java.util.concurrent.atomic.AtomicReference newWorkRef = @@ -523,7 +508,6 @@ public void testAdvanceKeySwitching() throws Exception { workA, stateReader, sideInputStateFetcher, - outputBuilderA, workExecutor, mockExecutor, mockBudget, @@ -605,16 +589,12 @@ public void testAdvanceLimitThresholds() throws Exception { org.mockito.Mockito.eq(mockBudget))) .thenReturn(java.util.Optional.of(executableWorkB)); - Windmill.WorkItemCommitRequest.Builder outputBuilderA = - Windmill.WorkItemCommitRequest.newBuilder(); - // Case 1: maxKeyGroupBatchSize is 0 executionContext.start( "keyA", workA, stateReader, sideInputStateFetcher, - outputBuilderA, workExecutor, mockExecutor, mockBudget, @@ -630,16 +610,11 @@ public void testAdvanceLimitThresholds() throws Exception { assertFalse(executionContext.advance()); // Case 2: maxKeyGroupBatchBytes limit is exceeded - Windmill.WorkItemCommitRequest.Builder outputBuilderAForBytes = - Windmill.WorkItemCommitRequest.newBuilder() - .setKey(ByteString.copyFromUtf8("some_non_empty_key_to_increase_size")); - executionContext.start( "keyA", workA, stateReader, sideInputStateFetcher, - outputBuilderAForBytes, workExecutor, mockExecutor, mockBudget, @@ -657,8 +632,6 @@ public void testAdvanceLimitThresholds() throws Exception { @Test public void testFinishKeyReentrantSafety() { - Windmill.WorkItemCommitRequest.Builder outputBuilder = - Windmill.WorkItemCommitRequest.newBuilder(); executionContext.start( "key", createMockWork( @@ -666,7 +639,6 @@ public void testFinishKeyReentrantSafety() { Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()), stateReader, sideInputStateFetcher, - outputBuilder, workExecutor); // First call @@ -711,15 +683,11 @@ public void testWorkIsFailed_heartbeatFailureOnPolledKey() throws Exception { org.mockito.Mockito.eq(mockBudget))) .thenReturn(java.util.Optional.of(executableWorkB)); - Windmill.WorkItemCommitRequest.Builder outputBuilderA = - Windmill.WorkItemCommitRequest.newBuilder(); - executionContext.start( "keyA", workA, stateReader, sideInputStateFetcher, - outputBuilderA, workExecutor, mockExecutor, mockBudget, @@ -775,15 +743,11 @@ public void testAdvance_abortsOnFailedExecutedKey() throws Exception { org.mockito.Mockito.eq(mockBudget))) .thenReturn(java.util.Optional.of(executableWorkB)); - Windmill.WorkItemCommitRequest.Builder outputBuilderA = - Windmill.WorkItemCommitRequest.newBuilder(); - executionContext.start( "keyA", workA, stateReader, sideInputStateFetcher, - outputBuilderA, workExecutor, mockExecutor, mockBudget, @@ -879,15 +843,11 @@ public void testInvalidateCache_clearsAllExecutedKeys() throws Exception { org.mockito.Mockito.eq(mockBudget))) .thenReturn(java.util.Optional.of(executableWorkB)); - Windmill.WorkItemCommitRequest.Builder outputBuilderA = - Windmill.WorkItemCommitRequest.newBuilder(); - testContext.start( "keyA", workA, stateReader, sideInputStateFetcher, - outputBuilderA, workExecutor, mockExecutor, mockBudget, @@ -966,4 +926,32 @@ public void testInvalidateCache_clearsAllExecutedKeys() throws Exception { org.mockito.Mockito.eq(workItemB.getKey()), org.mockito.Mockito.eq(workItemB.getShardingKey())); } + + @Test + public void testStart_internalKeyDecoding() throws Exception { + Windmill.WorkItem workItem = + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("decodedKey")) + .setWorkToken(17L) + .build(); + Work work = + createMockWork( + workItem, Watermarks.builder().setInputDataWatermark(new Instant(1000)).build()); + + executionContext.start( + /* key= */ null, + work, + stateReader, + sideInputStateFetcher, + workExecutor, + /* workQueueExecutor= */ null, + /* budgetHandle= */ null, + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + /* maxKeyGroupBatchSize= */ 1, + /* maxKeyGroupBatchTimeNanos= */ 0L, + /* maxKeyGroupBatchBytes= */ 0L, + /* keySwitchListener= */ (k, c) -> {}); + + assertEquals("decodedKey", executionContext.getKey()); + } } 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 126e8d44603e..f1296c71d9df 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 @@ -648,7 +648,6 @@ public void testReadUnboundedReader() throws Exception { Watermarks.builder().setInputDataWatermark(new Instant(0)).build()), mock(WindmillStateReader.class), mock(SideInputStateFetcher.class), - Windmill.WorkItemCommitRequest.newBuilder(), mock(WorkExecutor.class)); @SuppressWarnings({"unchecked", "rawtypes"}) @@ -1029,7 +1028,6 @@ public void testFailedWorkItemsAbort() throws Exception { dummyWork, mock(WindmillStateReader.class), mock(SideInputStateFetcher.class), - Windmill.WorkItemCommitRequest.newBuilder(), mock(WorkExecutor.class)); @SuppressWarnings({"unchecked", "rawtypes"}) @@ -1140,7 +1138,6 @@ public void testUnboundedReaderIterator_multiKeySwitching() throws Exception { workA, mock(WindmillStateReader.class), mock(SideInputStateFetcher.class), - Windmill.WorkItemCommitRequest.newBuilder(), mock(WorkExecutor.class), mockExecutor, mockBudget,