From ad68a3c26ecb344026a7abd49ce71d7e88d0e327 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Tue, 2 Jun 2026 09:01:33 +0000 Subject: [PATCH 1/9] [Dataflow Streaming] Introduce KeyGroupWorkQueue and integrate with BoundedQueueExecutor - Add Uint128Proto and key_group to WorkItem in windmill.proto. - Update Work model to support key groups. - Implement KeyGroupWorkQueue for grouping tasks by key group. KepGroupWorkQueue is a FIFO queue that allows polling elements based on global order and also by order within a key group. - Integrate BoundedQueueExecutor with KepGroupWorkQueue and support targeted polling. --- .../worker/StreamingDataflowWorker.java | 8 +- .../worker/streaming/ExecutableWork.java | 13 +- .../dataflow/worker/streaming/Work.java | 60 +++ .../worker/util/BoundedQueueExecutor.java | 26 +- .../worker/util/KeyGroupWorkQueue.java | 463 ++++++++++++++++++ .../worker/StreamingDataflowWorkerTest.java | 12 +- .../worker/util/BoundedQueueExecutorTest.java | 151 +++++- .../worker/util/KeyGroupWorkQueueTest.java | 461 +++++++++++++++++ .../StreamingCommitFinalizerTest.java | 3 +- .../failures/WorkFailureProcessorTest.java | 3 +- .../work/refresh/ActiveWorkRefresherTest.java | 3 +- .../windmill/src/main/proto/windmill.proto | 7 + 12 files changed, 1192 insertions(+), 18 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java create mode 100644 runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java 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..4d04c97a1eca 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,9 @@ 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"; + // Don't use. Experiment guarding multi key bundles. The feature is work in progress and + // incomplete. + private static final String UNSTABLE_ENABLE_MULTI_KEY_BUNDLE = "unstable_enable_multi_key_bundle"; private final WindmillStateCache stateCache; private AtomicReference statusPages = new AtomicReference<>(); @@ -1017,6 +1020,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, UNSTABLE_ENABLE_MULTI_KEY_BUNDLE); return new BoundedQueueExecutor( chooseMaxThreads(options), THREAD_EXPIRATION_TIME_SEC, @@ -1024,7 +1029,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 ecaa673f5570..432227c9253f 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; @@ -62,11 +63,11 @@ public void run(BoundedQueueExecutorWorkHandle handle) { } } - public final WorkId id() { + public WorkId id() { return work().id(); } - public final Windmill.WorkItem getWorkItem() { + public Windmill.WorkItem getWorkItem() { return work().getWorkItem(); } @@ -74,4 +75,12 @@ public final Windmill.WorkItem getWorkItem() { public String toString() { return "ExecutableWork{" + id() + "}"; } + + public String getComputationId() { + return work().getComputationId(); + } + + public 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 cb01e1e508ce..9dbaf1bb519c 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()) + '-' @@ -383,6 +391,14 @@ private boolean isCommitPending() { abstract Instant startTime(); } + public String getComputationId() { + return processingContext.computationId(); + } + + public Optional getKeyGroup() { + return keyGroup; + } + @AutoValue public abstract static class ProcessingContext { @@ -416,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 c6fd96e0a4cb..57046147e204 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,6 +20,7 @@ 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; @@ -29,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,7 +87,8 @@ public BoundedQueueExecutor( int maximumElementsOutstanding, long maximumBytesOutstanding, ThreadFactory threadFactory, - boolean useFairMonitor) { + boolean useFairMonitor, + boolean useKeyGroupWorkQueue) { this.maximumPoolSize = initialMaximumPoolSize; monitor = new Monitor(useFairMonitor); executor = @@ -94,7 +97,7 @@ public BoundedQueueExecutor( initialMaximumPoolSize, keepAliveTime, unit, - new LinkedBlockingQueue<>(), + useKeyGroupWorkQueue ? new KeyGroupWorkQueue() : new LinkedBlockingQueue<>(), threadFactory) { @Override protected void beforeExecute(Thread t, Runnable r) { @@ -313,7 +316,7 @@ public synchronized void close() { } } - private static final class QueuedWork implements Runnable { + static final class QueuedWork implements Runnable { private final ExecutableWork work; private final BoundedQueueExecutorWorkHandleImpl handle; @@ -378,6 +381,23 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) return new BoundedQueueExecutorWorkHandleImpl(elements, bytes); } + /** 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; + if (!(executor.getQueue() instanceof KeyGroupWorkQueue)) { + return Optional.empty(); + } + QueuedWork queuedWork = + ((KeyGroupWorkQueue) executor.getQueue()).pollWork(computationId, keyGroup); + if (queuedWork == null) { + return Optional.empty(); + } + internalHandle.merge(queuedWork.getHandle()); + return 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/KeyGroupWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java new file mode 100644 index 000000000000..50494edf70bf --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java @@ -0,0 +1,463 @@ +/* + * 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.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.util.AbstractQueue; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +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. In addition to global FIFO ordering, the queue + * supports polling work by computation + key group in FIFO order + */ +class KeyGroupWorkQueue extends AbstractQueue implements BlockingQueue { + + static class Node { + final @Nullable Runnable task; + final @Nullable String computationId; + final Work.@Nullable KeyGroup keyGroup; + + // prevNode, nextNode are used for the global order across all queued Runnables + @Nullable Node prevNode; + @Nullable Node nextNode; + + // prevKeyGroupNode and nextKeyGroupNode are used for the keyGroup level lists linking + // QueuedWork with same keyGroup + @Nullable Node prevKeyGroupNode; + @Nullable Node nextKeyGroupNode; + + Node(@Nullable Runnable task) { + 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; + } + } + } + + /** Double linked list implementing key group level queue */ + private static class KeyGroupWorkList { + final Node head; + final Node tail; + + KeyGroupWorkList() { + head = new Node(null); + tail = new Node(null); + head.nextKeyGroupNode = tail; + tail.prevKeyGroupNode = head; + } + + boolean isEmpty() { + return head.nextKeyGroupNode == tail; + } + + void append(Node node) { + @Nullable Node last = tail.prevKeyGroupNode; + if (last == null) { + throw new NullPointerException("tail.prevComp is null"); + } + node.prevKeyGroupNode = last; + node.nextKeyGroupNode = tail; + last.nextKeyGroupNode = node; + tail.prevKeyGroupNode = node; + } + + void remove(Node node) { + @Nullable Node prev = node.prevKeyGroupNode; + @Nullable Node next = node.nextKeyGroupNode; + if (prev != null && next != null) { + prev.nextKeyGroupNode = next; + next.prevKeyGroupNode = prev; + node.prevKeyGroupNode = null; + node.nextKeyGroupNode = 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); + + private final Map keyGroupQueueMap = new HashMap<>(); + + private int size = 0; + + public KeyGroupWorkQueue() { + globalHead.nextNode = globalTail; + globalTail.prevNode = globalHead; + } + + private void unlinkNode(Node node) { + // 1. Unlink from global list + Node prevG = node.prevNode; + Node nextG = node.nextNode; + if (prevG != null && nextG != null) { + prevG.nextNode = nextG; + nextG.prevNode = prevG; + } + node.prevNode = null; + node.nextNode = null; + + // 2. Unlink from key group list + if (node.computationId != null) { + QueueKey key = QueueKey.create(node.computationId, node.keyGroup); + KeyGroupWorkList keyGroupQueue = keyGroupQueueMap.get(key); + if (keyGroupQueue != null) { + keyGroupQueue.remove(node); + if (keyGroupQueue.isEmpty()) { + keyGroupQueueMap.remove(key); + } + } + } + --size; + } + + private @Nullable Node removeFirstGlobal() { + @Nullable Node first = globalHead.nextNode; + if (first == null || first == globalTail) { + return null; + } + unlinkNode(first); + return first; + } + + /** + * Remove and Return QueuedWork for the computationId, keyGroup in the FIFO order Returns null, if + * there are no matches. + */ + public @Nullable QueuedWork pollWork(String computationId, Work.KeyGroup keyGroup) { + if (computationId == null || keyGroup == null) { + return null; + } + lock.lock(); + try { + QueueKey key = QueueKey.create(computationId, keyGroup); + KeyGroupWorkList keyGroupWorkList = keyGroupQueueMap.get(key); + if (keyGroupWorkList == null || keyGroupWorkList.isEmpty()) { + return null; + } + + // Retrieve the first pending task for this computation and keyGroup in O(1) + @Nullable Node firstNode = keyGroupWorkList.head.nextKeyGroupNode; + if (firstNode == null || firstNode == keyGroupWorkList.tail) { + return null; + } + unlinkNode(firstNode); + + Runnable task = firstNode.task; + if (task == null) { + return null; + } + return (QueuedWork) task; + } finally { + lock.unlock(); + } + } + + @Override + public boolean offer(Runnable runnable) { + if (runnable == null) throw new NullPointerException(); + lock.lock(); + try { + Node node = new Node(runnable); + + // Append to global list tail + @Nullable Node lastG = globalTail.prevNode; + if (lastG == null) { + throw new NullPointerException("globalTail.prevNode is null"); + } + node.prevNode = lastG; + node.nextNode = globalTail; + lastG.nextNode = node; + globalTail.prevNode = node; + + // Append to key group list if applicable + if (node.computationId != null) { + QueueKey key = QueueKey.create(node.computationId, node.keyGroup); + KeyGroupWorkList keyGroupWorkList = + keyGroupQueueMap.computeIfAbsent(key, k -> new KeyGroupWorkList()); + keyGroupWorkList.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 @Nullable Runnable poll() { + lock.lock(); + try { + @Nullable 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(); + } + @Nullable Node node = removeFirstGlobal(); + checkStateNotNull(node, "Queue is empty but size was " + size); + Runnable task = node.task; + checkStateNotNull(task, "Encountered null task in queue"); + return task; + } finally { + lock.unlock(); + } + } + + @Override + public @Nullable 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); + } + @Nullable Node node = removeFirstGlobal(); + return (node != null) ? node.task : null; + } finally { + lock.unlock(); + } + } + + @Override + public @Nullable Runnable peek() { + lock.lock(); + try { + @Nullable Node first = globalHead.nextNode; + if (first == null || 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 + @Nullable Node curr = globalHead.nextNode; + while (curr != null && curr != globalTail) { + if (o.equals(curr.task)) { + unlinkNode(curr); + return true; + } + curr = curr.nextNode; + } + return false; + } finally { + lock.unlock(); + } + } + + @Override + public boolean contains(Object o) { + if (o == null) return false; + lock.lock(); + try { + @Nullable Node curr = globalHead.nextNode; + while (curr != null && curr != globalTail) { + if (o.equals(curr.task)) { + return true; + } + curr = curr.nextNode; + } + 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; + @Nullable Node curr = globalHead.nextNode; + while (curr != null && curr != globalTail && added < maxElements) { + @Nullable Node next = curr.nextNode; + unlinkNode(curr); + Runnable task = curr.task; + if (task != null) { + c.add(task); + ++added; + } + curr = next; + } + return added; + } finally { + lock.unlock(); + } + } + + @Override + public void clear() { + lock.lock(); + try { + @Nullable Node curr = globalHead.nextNode; + while (curr != null && curr != globalTail) { + @Nullable Node next = curr.nextNode; + unlinkNode(curr); + curr = next; + } + } finally { + lock.unlock(); + } + } + + @Override + public Iterator iterator() { + lock.lock(); + try { + List snapshot = new ArrayList<>(size); + @Nullable Node curr = globalHead.nextNode; + while (curr != null && curr != globalTail) { + if (curr.task != null) { + snapshot.add(curr.task); + } + curr = curr.nextNode; + } + return Collections.unmodifiableList(snapshot).iterator(); + } finally { + 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); + } + + @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/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index d58f20076994..5bcdffcc2564 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 @@ -3036,7 +3036,8 @@ public void testMaxThreadMetric() throws Exception { .setNameFormat("DataflowWorkUnits-%d") .setDaemon(true) .build(), - /*useFairMonitor=*/ false); + /*useFairMonitor=*/ false, + /*useKeyGroupWorkQueue=*/ false); ComputationState computationState = new ComputationState( @@ -3097,7 +3098,8 @@ public void testActiveThreadMetric() throws Exception { .setNameFormat("DataflowWorkUnits-%d") .setDaemon(true) .build(), - /*useFairMonitor=*/ false); + /*useFairMonitor=*/ false, + /*useKeyGroupWorkQueue=*/ false); ComputationState computationState = new ComputationState( @@ -3167,7 +3169,8 @@ public void testOutstandingBytesMetric() throws Exception { .setNameFormat("DataflowWorkUnits-%d") .setDaemon(true) .build(), - /*useFairMonitor=*/ false); + /*useFairMonitor=*/ false, + /*useKeyGroupWorkQueue=*/ false); ComputationState computationState = new ComputationState( @@ -3241,7 +3244,8 @@ public void testOutstandingBundlesMetric() throws Exception { .setNameFormat("DataflowWorkUnits-%d") .setDaemon(true) .build(), - /*useFairMonitor=*/ false); + /*useFairMonitor=*/ false, + /*useKeyGroupWorkQueue=*/ false); ComputationState computationState = new 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..aa9ce32b3c5c 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,13 +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( @@ -80,10 +98,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) -> { @@ -116,7 +131,8 @@ public void setUp() { .setNameFormat("DataflowWorkUnits-%d") .setDaemon(true) .build(), - useFairMonitor); + useFairMonitor, + /*useKeyGroupWorkQueue=*/ false); } @Test @@ -413,4 +429,129 @@ public void testRenderSummaryHtml() { + "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, + 60, + TimeUnit.SECONDS, + 100, + 10000000, + new ThreadFactoryBuilder().setNameFormat("testStealing-%d").setDaemon(true).build(), + useFairMonitor, + /*useKeyGroupWorkQueue=*/ true); + + // 1. Create blocker task to occupy the worker thread + 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); + } + }); + + 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 work1 = createWorkWithCompIdAndKeyGroup("compA", keyGroup1, ignored -> {}); + ExecutableWork work2 = + createWorkWithCompIdAndKeyGroup( + "compA", + keyGroup2, + ignored -> { + targetStart.countDown(); + }); + + // Enqueue tasks (they will wait in the queue because the thread is blocked) + testExecutor.execute(work1, 100); + testExecutor.execute(work2, 150); + + // Total outstanding elements must be 3 (blocker + work1 + work2) + assertEquals(3, testExecutor.elementsOutstanding()); + + // Steal work2 using pollWork with compA and keyGroup2 + try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { + java.util.Optional stolen = + testExecutor.pollWork("compA", keyGroup2, stealHandle); + assertTrue(stolen.isPresent()); + assertEquals(work2, stolen.get()); + + // Run the stolen task + stolen.get().run(stealHandle); + targetStart.await(); + } + + // 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()); + } + + // Unblock the blocker and shut down + blockerStop.countDown(); + testExecutor.shutdown(); + } + + @Test + public void testPollWorkWithLinkedBlockingQueue() throws Exception { + BoundedQueueExecutor testExecutor = + new BoundedQueueExecutor( + 1, + 60, + TimeUnit.SECONDS, + 100, + 10000000, + new ThreadFactoryBuilder().setNameFormat("testLinkedQueue-%d").setDaemon(true).build(), + useFairMonitor, + /* useKeyGroupWorkQueue= */ false); + + 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); + } + }); + + testExecutor.execute(blockerWork, 0); + blockerStart.await(); + + Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1); + ExecutableWork work = createWorkWithCompIdAndKeyGroup("compA", keyGroup, ignored -> {}); + testExecutor.execute(work, 100); + + try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { + java.util.Optional stolen = + testExecutor.pollWork("compA", keyGroup, stealHandle); + assertFalse(stolen.isPresent()); + } + + blockerStop.countDown(); + testExecutor.shutdown(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java new file mode 100644 index 000000000000..45ebab4bca28 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java @@ -0,0 +1,461 @@ +/* + * 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.junit.Assert.fail; +import static org.mockito.Mockito.mock; + +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +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) +public class KeyGroupWorkQueueTest { + + 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, + /*useKeyGroupWorkQueue=*/ true); + } + + 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( + 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() {} + + @Override + public String toString() { + return "MockRunnable(" + id + ")"; + } + } + + @Test + public void testBasicOfferAndPoll() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + 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() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + 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() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + 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() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + MockRunnable task1 = new MockRunnable("1"); + MockRunnable task2 = new MockRunnable("2"); + queue.offer(task1); + queue.offer(task2); + + 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(); + fail("Iterator must be immutable"); + } catch (UnsupportedOperationException e) { + // Expected + } + } + + @Test + public void testPollWorkTargeted() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + + 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", DEFAULT_KEY_GROUP); + 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() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + QueuedWork workA1 = createQueuedWork("compA", 100); + queue.offer(workA1); + + // Steal A1 + QueuedWork polled = queue.pollWork("compA", DEFAULT_KEY_GROUP); + 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", DEFAULT_KEY_GROUP); + assertNotNull(polledB); + assertTrue(queue.isEmpty()); + } + + @Test + public void testConcurrentStress() throws InterruptedException, ExecutionException { + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + 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); + List> futures = new ArrayList<>(); + + // Start producers + for (int i = 0; i < producerThreads; i++) { + futures.add( + 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) { + throw new RuntimeException(e); + } finally { + doneLatch.countDown(); + } + })); + } + + // Start consumers (mix of poll and pollWork) + for (int i = 0; i < consumerThreads; i++) { + final int consumerId = i; + futures.add( + 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, DEFAULT_KEY_GROUP); + } else { + // Global poll + task = queue.poll(); + } + if (task != null) { + consumedCount.incrementAndGet(); + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + doneLatch.countDown(); + } + })); + } + + startLatch.countDown(); + assertTrue(doneLatch.await(10, TimeUnit.SECONDS)); + + // Check for exceptions in threads + for (Future future : futures) { + future.get(); + } + + executorService.shutdown(); + assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS)); + + assertEquals(0, queue.size()); + assertTrue(queue.isEmpty()); + } + + @Test + public void testTakeBlocksAndWakesUp() throws InterruptedException { + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + final MockRunnable task = new MockRunnable("take-task"); + final AtomicReference result = new 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 KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + final MockRunnable task = new MockRunnable("poll-task"); + final AtomicReference result = new 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(500, TimeUnit.MILLISECONDS)); + } catch (InterruptedException e) { + // Ignore + } finally { + finished.countDown(); + } + }); + t1.setDaemon(true); + t1.start(); + + assertTrue(started.await(2, TimeUnit.SECONDS)); + Thread.sleep(100); + 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 AtomicReference result2 = new 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(100); + assertEquals(Thread.State.TIMED_WAITING, t2.getState()); + + queue.offer(task); + + assertTrue(finished2.await(2, TimeUnit.SECONDS)); + assertEquals(task, result2.get()); + } + + @Test + public void testPollWorkWithKeyGroup() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + + 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()); + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizerTest.java index 07b4b14fd115..ef0d8e434858 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizerTest.java @@ -62,7 +62,8 @@ public void setUp() { .setNameFormat("FinalizationCallback-%d") .setDaemon(true) .build(), - /*useFairMonitor=*/ false); + /*useFairMonitor=*/ false, + /*useKeyGroupWorkQueue=*/ false); cleanupExecutor = Executors.newScheduledThreadPool( 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..0610ed44c27f 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 @@ -64,7 +64,8 @@ private static WorkFailureProcessor createWorkFailureProcessor( .setNameFormat("DataflowWorkUnits-%d") .setDaemon(true) .build(), - /*useFairMonitor=*/ false); + /*useFairMonitor=*/ false, + /*useKeyGroupWorkQueue=*/ false); return WorkFailureProcessor.forTesting(workExecutor, failureTracker, Optional::empty, clock, 0); } 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 88a82c6f76b6..f32282056e4f 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 @@ -80,7 +80,8 @@ private static BoundedQueueExecutor workExecutor() { 1, 10000000, new ThreadFactoryBuilder().setNameFormat("DataflowWorkUnits-%d").setDaemon(true).build(), - /*useFairMonitor=*/ false); + /*useFairMonitor=*/ false, + /*useKeyGroupWorkQueue=*/ false); } private static ComputationState createComputationState(int computationIdSuffix) { 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..aaa09c105fc3 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; } From 5fc07cc29479eecf2cd46e9a57ca0c23ea4c20e9 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 5 Jun 2026 08:15:53 +0000 Subject: [PATCH 2/9] address comments --- .../worker/streaming/ExecutableWork.java | 3 +- .../dataflow/worker/streaming/Work.java | 20 +++- .../worker/util/BoundedQueueExecutor.java | 21 ++-- .../worker/util/KeyGroupWorkQueue.java | 112 +++++++++--------- .../worker/util/BoundedQueueExecutorTest.java | 23 ++-- .../worker/util/KeyGroupWorkQueueTest.java | 76 ++++++++---- 6 files changed, 144 insertions(+), 111 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 432227c9253f..7748a554f0fc 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,6 @@ 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; @@ -80,7 +79,7 @@ public String getComputationId() { return work().getComputationId(); } - public Optional getKeyGroup() { + public Work.KeyGroup 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 9dbaf1bb519c..086203b7d7fc 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 @@ -76,7 +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 KeyGroup keyGroup; private final String latencyTrackingId; private final long serializedWorkItemSize; private volatile TimedState currentState; @@ -106,9 +106,8 @@ private Work( this.id = WorkId.of(workItem); this.keyGroup = workItem.hasKeyGroup() - ? Optional.of( - KeyGroup.create(workItem.getKeyGroup().getHigh(), workItem.getKeyGroup().getLow())) - : Optional.empty(); + ? KeyGroup.create(workItem.getKeyGroup().getHigh(), workItem.getKeyGroup().getLow()) + : KeyGroup.DEFAULT; this.latencyTrackingId = Long.toHexString(workItem.getShardingKey()) + '-' @@ -395,7 +394,7 @@ public String getComputationId() { return processingContext.computationId(); } - public Optional getKeyGroup() { + public KeyGroup getKeyGroup() { return keyGroup; } @@ -433,7 +432,16 @@ private Optional fetchKeyedState(KeyedGetDataRequest reque } } + /** + * WorkItems with same key group and computation are eligible to be executed together in a + * multi-key bundle. + */ public static final class KeyGroup { + + // The default 0 key group. Work items with 0 keyGroup will always be executed + // separately and not in a multi-key bundle + public static final KeyGroup DEFAULT = new KeyGroup(0, 0); + private final long high; private final long low; @@ -473,7 +481,7 @@ public int hashCode() { @Override public String toString() { - return "KeyGroup{" + "high=" + high + ", low=" + low + '}'; + return String.format("%016x%016x", high, 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 57046147e204..44fced15be77 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,6 @@ 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; @@ -35,6 +34,7 @@ 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; +import org.checkerframework.checker.nullness.qual.Nullable; /** An executor for executing work on windmill items. */ @SuppressWarnings({ @@ -80,6 +80,8 @@ private static class Budget { @GuardedBy("this") private long totalTimeMaxActiveThreadsUsed; + private final @Nullable KeyGroupWorkQueue keyGroupWorkQueue; + public BoundedQueueExecutor( int initialMaximumPoolSize, long keepAliveTime, @@ -89,6 +91,7 @@ public BoundedQueueExecutor( ThreadFactory threadFactory, boolean useFairMonitor, boolean useKeyGroupWorkQueue) { + this.keyGroupWorkQueue = useKeyGroupWorkQueue ? new KeyGroupWorkQueue(useFairMonitor) : null; this.maximumPoolSize = initialMaximumPoolSize; monitor = new Monitor(useFairMonitor); executor = @@ -97,7 +100,7 @@ public BoundedQueueExecutor( initialMaximumPoolSize, keepAliveTime, unit, - useKeyGroupWorkQueue ? new KeyGroupWorkQueue() : new LinkedBlockingQueue<>(), + keyGroupWorkQueue != null ? keyGroupWorkQueue : new LinkedBlockingQueue<>(), threadFactory) { @Override protected void beforeExecute(Thread t, Runnable r) { @@ -381,21 +384,19 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) return new BoundedQueueExecutorWorkHandleImpl(elements, bytes); } - /** Poll work for a specific computationId and keyGroup. */ - public Optional pollWork( + public @Nullable ExecutableWork pollWork( String computationId, Work.KeyGroup keyGroup, BoundedQueueExecutorWorkHandle handle) { checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl); BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; - if (!(executor.getQueue() instanceof KeyGroupWorkQueue)) { - return Optional.empty(); + if (keyGroupWorkQueue == null) { + return null; } - QueuedWork queuedWork = - ((KeyGroupWorkQueue) executor.getQueue()).pollWork(computationId, keyGroup); + QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup); if (queuedWork == null) { - return Optional.empty(); + return null; } internalHandle.merge(queuedWork.getHandle()); - return Optional.of(queuedWork.getWork()); + return queuedWork.getWork(); } private void decrementCounters(int elements, long bytes) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java index 50494edf70bf..a90eed07be23 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java @@ -17,24 +17,26 @@ */ package org.apache.beam.runners.dataflow.worker.util; +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import java.util.AbstractQueue; -import java.util.ArrayList; import java.util.Collection; -import java.util.Collections; import java.util.HashMap; import java.util.Iterator; -import java.util.List; 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 javax.annotation.concurrent.GuardedBy; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor.QueuedWork; +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.jspecify.annotations.NonNull; /** * A custom, thread-safe doubly-linked BlockingQueue. In addition to global FIFO ordering, the queue @@ -46,6 +48,8 @@ static class Node { final @Nullable Runnable task; final @Nullable String computationId; final Work.@Nullable KeyGroup keyGroup; + // cached keyGroupList if the Node is part of one. + @Nullable KeyGroupWorkList keyGroupList; // prevNode, nextNode are used for the global order across all queued Runnables @Nullable Node prevNode; @@ -60,7 +64,7 @@ static class Node { this.task = task; if (task instanceof QueuedWork) { this.computationId = ((QueuedWork) task).getWork().getComputationId(); - this.keyGroup = ((QueuedWork) task).getWork().getKeyGroup().orElse(null); + this.keyGroup = ((QueuedWork) task).getWork().getKeyGroup(); } else { this.computationId = null; this.keyGroup = null; @@ -70,12 +74,10 @@ static class Node { /** Double linked list implementing key group level queue */ private static class KeyGroupWorkList { - final Node head; - final Node tail; + final Node head = new Node(null); + final Node tail = new Node(null); KeyGroupWorkList() { - head = new Node(null); - tail = new Node(null); head.nextKeyGroupNode = tail; tail.prevKeyGroupNode = head; } @@ -107,47 +109,56 @@ void remove(Node node) { } } - private final ReentrantLock lock = new ReentrantLock(); - private final Condition notEmpty = lock.newCondition(); + private final ReentrantLock lock; + private final Condition notEmpty; // Sentinels for the global list + @GuardedBy("lock") private final Node globalHead = new Node(null); + + @GuardedBy("lock") private final Node globalTail = new Node(null); + @GuardedBy("lock") private final Map keyGroupQueueMap = new HashMap<>(); + @GuardedBy("lock") private int size = 0; - public KeyGroupWorkQueue() { + public KeyGroupWorkQueue(boolean fair) { + this.lock = new ReentrantLock(fair); + this.notEmpty = lock.newCondition(); globalHead.nextNode = globalTail; globalTail.prevNode = globalHead; } + @GuardedBy("lock") private void unlinkNode(Node node) { + // An existing node should always have previous and next since we have sentinels // 1. Unlink from global list - Node prevG = node.prevNode; - Node nextG = node.nextNode; - if (prevG != null && nextG != null) { - prevG.nextNode = nextG; - nextG.prevNode = prevG; - } + Node prevG = checkArgumentNotNull(node.prevNode); + Node nextG = checkArgumentNotNull(node.nextNode); + prevG.nextNode = nextG; + nextG.prevNode = prevG; node.prevNode = null; node.nextNode = null; // 2. Unlink from key group list - if (node.computationId != null) { - QueueKey key = QueueKey.create(node.computationId, node.keyGroup); - KeyGroupWorkList keyGroupQueue = keyGroupQueueMap.get(key); - if (keyGroupQueue != null) { - keyGroupQueue.remove(node); - if (keyGroupQueue.isEmpty()) { - keyGroupQueueMap.remove(key); - } + KeyGroupWorkList list = node.keyGroupList; + if (list != null) { + list.remove(node); + if (list.isEmpty()) { + String compId = checkStateNotNull(node.computationId); + Work.KeyGroup keyGroup = checkStateNotNull(node.keyGroup); + QueueKey key = new QueueKey(compId, keyGroup); + keyGroupQueueMap.remove(key); } + node.keyGroupList = null; } --size; } + @GuardedBy("lock") private @Nullable Node removeFirstGlobal() { @Nullable Node first = globalHead.nextNode; if (first == null || first == globalTail) { @@ -162,12 +173,13 @@ private void unlinkNode(Node node) { * there are no matches. */ public @Nullable QueuedWork pollWork(String computationId, Work.KeyGroup keyGroup) { - if (computationId == null || keyGroup == null) { + Preconditions.checkArgument(computationId != null && keyGroup != null); + if (keyGroup.equals(Work.KeyGroup.DEFAULT)) { return null; } + QueueKey key = new QueueKey(computationId, keyGroup); lock.lock(); try { - QueueKey key = QueueKey.create(computationId, keyGroup); KeyGroupWorkList keyGroupWorkList = keyGroupQueueMap.get(key); if (keyGroupWorkList == null || keyGroupWorkList.isEmpty()) { return null; @@ -180,39 +192,33 @@ private void unlinkNode(Node node) { } unlinkNode(firstNode); - Runnable task = firstNode.task; - if (task == null) { - return null; - } - return (QueuedWork) task; + return (QueuedWork) firstNode.task; } finally { lock.unlock(); } } @Override - public boolean offer(Runnable runnable) { - if (runnable == null) throw new NullPointerException(); + public boolean offer(@NonNull Runnable runnable) { + Node node = new Node(checkStateNotNull(runnable)); lock.lock(); try { - Node node = new Node(runnable); - // Append to global list tail - @Nullable Node lastG = globalTail.prevNode; - if (lastG == null) { - throw new NullPointerException("globalTail.prevNode is null"); - } + Node lastG = checkStateNotNull(globalTail.prevNode); node.prevNode = lastG; node.nextNode = globalTail; lastG.nextNode = node; globalTail.prevNode = node; // Append to key group list if applicable - if (node.computationId != null) { - QueueKey key = QueueKey.create(node.computationId, node.keyGroup); + String compId = node.computationId; + Work.KeyGroup keyGroup = node.keyGroup; + if (compId != null && keyGroup != null && !keyGroup.equals(Work.KeyGroup.DEFAULT)) { + QueueKey key = new QueueKey(compId, keyGroup); KeyGroupWorkList keyGroupWorkList = keyGroupQueueMap.computeIfAbsent(key, k -> new KeyGroupWorkList()); keyGroupWorkList.append(node); + node.keyGroupList = keyGroupWorkList; } ++size; @@ -253,9 +259,7 @@ public Runnable take() throws InterruptedException { } @Nullable Node node = removeFirstGlobal(); checkStateNotNull(node, "Queue is empty but size was " + size); - Runnable task = node.task; - checkStateNotNull(task, "Encountered null task in queue"); - return task; + return checkStateNotNull(node.task, "Encountered null task in queue"); } finally { lock.unlock(); } @@ -405,15 +409,15 @@ public void clear() { public Iterator iterator() { lock.lock(); try { - List snapshot = new ArrayList<>(size); + ImmutableList.Builder builder = ImmutableList.builderWithExpectedSize(size); @Nullable Node curr = globalHead.nextNode; while (curr != null && curr != globalTail) { if (curr.task != null) { - snapshot.add(curr.task); + builder.add(curr.task); } curr = curr.nextNode; } - return Collections.unmodifiableList(snapshot).iterator(); + return builder.build().iterator(); } finally { lock.unlock(); } @@ -421,17 +425,13 @@ public Iterator iterator() { static final class QueueKey { private final String computationId; - private final Work.@Nullable KeyGroup keyGroup; + private final Work.KeyGroup keyGroup; - private QueueKey(String computationId, Work.@Nullable KeyGroup keyGroup) { - this.computationId = Objects.requireNonNull(computationId); + QueueKey(String computationId, Work.KeyGroup keyGroup) { + this.computationId = computationId; this.keyGroup = keyGroup; } - public static QueueKey create(String computationId, Work.@Nullable KeyGroup keyGroup) { - return new QueueKey(computationId, keyGroup); - } - @Override public boolean equals(@Nullable Object o) { if (this == o) { @@ -441,7 +441,7 @@ public boolean equals(@Nullable Object o) { return false; } QueueKey other = (QueueKey) o; - return computationId.equals(other.computationId) && Objects.equals(keyGroup, other.keyGroup); + return computationId.equals(other.computationId) && keyGroup.equals(other.keyGroup); } @Override 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 aa9ce32b3c5c..a98102751fb2 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 @@ -20,6 +20,8 @@ import static org.hamcrest.Matchers.greaterThan; 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.assertThat; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; @@ -487,22 +489,20 @@ public void testPollWork() throws Exception { // Steal work2 using pollWork with compA and keyGroup2 try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { - java.util.Optional stolen = - testExecutor.pollWork("compA", keyGroup2, stealHandle); - assertTrue(stolen.isPresent()); - assertEquals(work2, stolen.get()); + ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup2, stealHandle); + assertNotNull(stolen); + assertEquals(work2, stolen); // Run the stolen task - stolen.get().run(stealHandle); + stolen.run(stealHandle); targetStart.await(); } // 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()); + ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup1, stealHandle); + assertNotNull(stolen); + assertEquals(work1, stolen); } // Unblock the blocker and shut down @@ -546,9 +546,8 @@ public void testPollWorkWithLinkedBlockingQueue() throws Exception { testExecutor.execute(work, 100); try (BoundedQueueExecutorWorkHandleImpl stealHandle = testExecutor.createBudgetHandle(0, 0L)) { - java.util.Optional stolen = - testExecutor.pollWork("compA", keyGroup, stealHandle); - assertFalse(stolen.isPresent()); + ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup, stealHandle); + assertNull(stolen); } blockerStop.countDown(); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java index 45ebab4bca28..98023901570e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java @@ -45,6 +45,7 @@ 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.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.junit.Before; import org.junit.Test; @@ -70,26 +71,28 @@ public void setUp() { /*useKeyGroupWorkQueue=*/ true); } - private static final Work.KeyGroup DEFAULT_KEY_GROUP = Work.KeyGroup.create(1, 2); + private static final Work.KeyGroup TEST_KEY_GROUP = Work.KeyGroup.create(1, 2); private QueuedWork createQueuedWork(String computationId, long workBytes) { - return createQueuedWork(computationId, DEFAULT_KEY_GROUP, workBytes); + return createQueuedWork(computationId, TEST_KEY_GROUP, workBytes); } private QueuedWork createQueuedWork( - String computationId, Work.KeyGroup keyGroup, long workBytes) { - WorkItem workItem = + String computationId, Work.@Nullable KeyGroup keyGroup, long workBytes) { + WorkItem.Builder workItemBuilder = 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(); + .setCacheToken(1); + if (keyGroup != null) { + workItemBuilder.setKeyGroup( + org.apache.beam.runners.dataflow.worker.windmill.Windmill.Uint128Proto.newBuilder() + .setHigh(keyGroup.high()) + .setLow(keyGroup.low()) + .build()); + } + WorkItem workItem = workItemBuilder.build(); ExecutableWork work = ExecutableWork.create( Work.create( @@ -125,7 +128,7 @@ public String toString() { @Test public void testBasicOfferAndPoll() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); assertTrue(queue.isEmpty()); assertEquals(0, queue.size()); @@ -144,7 +147,7 @@ public void testBasicOfferAndPoll() { @Test public void testRemove() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); MockRunnable task1 = new MockRunnable("1"); MockRunnable task2 = new MockRunnable("2"); @@ -159,7 +162,7 @@ public void testRemove() { @Test public void testDrainTo() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); MockRunnable task1 = new MockRunnable("1"); MockRunnable task2 = new MockRunnable("2"); queue.offer(task1); @@ -175,7 +178,7 @@ public void testDrainTo() { @Test public void testIteratorSafeTraversalAndImmutable() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); MockRunnable task1 = new MockRunnable("1"); MockRunnable task2 = new MockRunnable("2"); queue.offer(task1); @@ -202,7 +205,7 @@ public void testIteratorSafeTraversalAndImmutable() { @Test public void testPollWorkTargeted() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); QueuedWork workA1 = createQueuedWork("compA", 100); QueuedWork workB1 = createQueuedWork("compB", 200); @@ -215,7 +218,7 @@ public void testPollWorkTargeted() { assertEquals(3, queue.size()); // Targeted poll A - QueuedWork polledA1 = queue.pollWork("compA", DEFAULT_KEY_GROUP); + QueuedWork polledA1 = queue.pollWork("compA", TEST_KEY_GROUP); assertNotNull(polledA1); assertEquals("compA", polledA1.getWork().getComputationId()); assertEquals(100, polledA1.getHandle().bytes()); @@ -234,12 +237,12 @@ public void testPollWorkTargeted() { @Test public void testMemoryPruningLeavesZeroLeaks() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); QueuedWork workA1 = createQueuedWork("compA", 100); queue.offer(workA1); // Steal A1 - QueuedWork polled = queue.pollWork("compA", DEFAULT_KEY_GROUP); + QueuedWork polled = queue.pollWork("compA", TEST_KEY_GROUP); assertNotNull(polled); assertTrue(queue.isEmpty()); @@ -249,14 +252,14 @@ public void testMemoryPruningLeavesZeroLeaks() { assertEquals(1, queue.size()); // Steal B1 - QueuedWork polledB = queue.pollWork("compB", DEFAULT_KEY_GROUP); + QueuedWork polledB = queue.pollWork("compB", TEST_KEY_GROUP); assertNotNull(polledB); assertTrue(queue.isEmpty()); } @Test public void testConcurrentStress() throws InterruptedException, ExecutionException { - final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); final int producerThreads = 4; final int consumerThreads = 4; final int tasksPerProducer = 1000; @@ -301,7 +304,7 @@ public void testConcurrentStress() throws InterruptedException, ExecutionExcepti if (consumerId % 2 == 0) { // Targeted poll String compId = "comp-" + (consumedCount.get() % 5); - task = queue.pollWork(compId, DEFAULT_KEY_GROUP); + task = queue.pollWork(compId, TEST_KEY_GROUP); } else { // Global poll task = queue.poll(); @@ -335,7 +338,7 @@ public void testConcurrentStress() throws InterruptedException, ExecutionExcepti @Test public void testTakeBlocksAndWakesUp() throws InterruptedException { - final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); final MockRunnable task = new MockRunnable("take-task"); final AtomicReference result = new AtomicReference<>(); final CountDownLatch started = new CountDownLatch(1); @@ -369,7 +372,7 @@ public void testTakeBlocksAndWakesUp() throws InterruptedException { @Test public void testPollWithTimeout() throws InterruptedException { - final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); final MockRunnable task = new MockRunnable("poll-task"); final AtomicReference result = new AtomicReference<>(); final CountDownLatch started = new CountDownLatch(1); @@ -430,7 +433,7 @@ public void testPollWithTimeout() throws InterruptedException { @Test public void testPollWorkWithKeyGroup() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); Work.KeyGroup keyGroup1 = Work.KeyGroup.create(1, 1); Work.KeyGroup keyGroup2 = Work.KeyGroup.create(1, 2); @@ -458,4 +461,27 @@ public void testPollWorkWithKeyGroup() { assertEquals(workA1, polledA1); assertTrue(queue.isEmpty()); } + + @Test + public void testDefaultKeyGroupNotIndexedInKeyGroupQueue() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + + // Work item with explicit KeyGroup.DEFAULT + QueuedWork workExplicitDefault = createQueuedWork("compA", Work.KeyGroup.DEFAULT, 100); + // Work item with implicit KeyGroup.DEFAULT (null passed to helper) + QueuedWork workImplicitDefault = createQueuedWork("compA", null, 150); + + queue.offer(workExplicitDefault); + queue.offer(workImplicitDefault); + + assertEquals(2, queue.size()); + + // Verify they cannot be retrieved via pollWork with KeyGroup.DEFAULT + assertNull(queue.pollWork("compA", Work.KeyGroup.DEFAULT)); + + // Verify they can still be retrieved via global poll in FIFO order + assertEquals(workExplicitDefault, queue.poll()); + assertEquals(workImplicitDefault, queue.poll()); + assertTrue(queue.isEmpty()); + } } From b0f8235eb91daa1c6603ff1f7cea2a9e1e0c6e0e Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 5 Jun 2026 10:25:38 +0000 Subject: [PATCH 3/9] address comments --- .../dataflow/worker/streaming/Work.java | 2 +- .../worker/util/BoundedQueueExecutor.java | 3 + .../worker/util/KeyGroupWorkQueue.java | 91 +++++++++---------- .../worker/util/KeyGroupWorkQueueTest.java | 23 ----- 4 files changed, 49 insertions(+), 70 deletions(-) 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 086203b7d7fc..0a09a5a13309 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 @@ -438,7 +438,7 @@ private Optional fetchKeyedState(KeyedGetDataRequest reque */ public static final class KeyGroup { - // The default 0 key group. Work items with 0 keyGroup will always be executed + // Work items equaling to the default keyGroup will always be executed // separately and not in a multi-key bundle public static final KeyGroup DEFAULT = new KeyGroup(0, 0); 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 44fced15be77..71200a33c21e 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 @@ -30,6 +30,7 @@ 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.runners.dataflow.worker.streaming.Work.KeyGroup; 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; @@ -80,6 +81,7 @@ private static class Budget { @GuardedBy("this") private long totalTimeMaxActiveThreadsUsed; + // If set the keyGroupWorkQueue is used by the underlying executor. private final @Nullable KeyGroupWorkQueue keyGroupWorkQueue; public BoundedQueueExecutor( @@ -387,6 +389,7 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) public @Nullable ExecutableWork pollWork( String computationId, Work.KeyGroup keyGroup, BoundedQueueExecutorWorkHandle handle) { checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl); + checkArgument(computationId != null && !KeyGroup.DEFAULT.equals(keyGroup)); BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; if (keyGroupWorkQueue == null) { return null; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java index a90eed07be23..f26c6ca1ee14 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java @@ -19,6 +19,7 @@ import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import java.util.AbstractQueue; import java.util.Collection; @@ -32,8 +33,8 @@ import java.util.concurrent.locks.ReentrantLock; import javax.annotation.concurrent.GuardedBy; import org.apache.beam.runners.dataflow.worker.streaming.Work; +import org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor.QueuedWork; -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.jspecify.annotations.NonNull; @@ -44,8 +45,14 @@ */ class KeyGroupWorkQueue extends AbstractQueue implements BlockingQueue { + public static final Runnable SENTINEL_RUNNABLE = + () -> { + throw new IllegalStateException("sentinel runnable called"); + }; + static class Node { - final @Nullable Runnable task; + // If keyGroup is non-null, task is an instance of QueuedWork + final Runnable task; final @Nullable String computationId; final Work.@Nullable KeyGroup keyGroup; // cached keyGroupList if the Node is part of one. @@ -60,7 +67,7 @@ static class Node { @Nullable Node prevKeyGroupNode; @Nullable Node nextKeyGroupNode; - Node(@Nullable Runnable task) { + Node(Runnable task) { this.task = task; if (task instanceof QueuedWork) { this.computationId = ((QueuedWork) task).getWork().getComputationId(); @@ -74,8 +81,8 @@ static class Node { /** Double linked list implementing key group level queue */ private static class KeyGroupWorkList { - final Node head = new Node(null); - final Node tail = new Node(null); + final Node head = new Node(SENTINEL_RUNNABLE); + final Node tail = new Node(SENTINEL_RUNNABLE); KeyGroupWorkList() { head.nextKeyGroupNode = tail; @@ -87,10 +94,7 @@ boolean isEmpty() { } void append(Node node) { - @Nullable Node last = tail.prevKeyGroupNode; - if (last == null) { - throw new NullPointerException("tail.prevComp is null"); - } + Node last = checkStateNotNull(tail.prevKeyGroupNode); node.prevKeyGroupNode = last; node.nextKeyGroupNode = tail; last.nextKeyGroupNode = node; @@ -114,10 +118,10 @@ void remove(Node node) { // Sentinels for the global list @GuardedBy("lock") - private final Node globalHead = new Node(null); + private final Node globalHead = new Node(SENTINEL_RUNNABLE); @GuardedBy("lock") - private final Node globalTail = new Node(null); + private final Node globalTail = new Node(SENTINEL_RUNNABLE); @GuardedBy("lock") private final Map keyGroupQueueMap = new HashMap<>(); @@ -160,8 +164,8 @@ private void unlinkNode(Node node) { @GuardedBy("lock") private @Nullable Node removeFirstGlobal() { - @Nullable Node first = globalHead.nextNode; - if (first == null || first == globalTail) { + Node first = checkStateNotNull(globalHead.nextNode); + if (first == globalTail) { return null; } unlinkNode(first); @@ -169,14 +173,13 @@ private void unlinkNode(Node node) { } /** - * Remove and Return QueuedWork for the computationId, keyGroup in the FIFO order Returns null, if - * there are no matches. + * Remove and Return QueuedWork for the computationId, keyGroup in the FIFO order. Returns null, + * if there are no matches. + * + * @param keyGroup should not be equal to KeyGroup.DEFAULT */ public @Nullable QueuedWork pollWork(String computationId, Work.KeyGroup keyGroup) { - Preconditions.checkArgument(computationId != null && keyGroup != null); - if (keyGroup.equals(Work.KeyGroup.DEFAULT)) { - return null; - } + checkArgument(computationId != null && !KeyGroup.DEFAULT.equals(keyGroup)); QueueKey key = new QueueKey(computationId, keyGroup); lock.lock(); try { @@ -186,8 +189,8 @@ private void unlinkNode(Node node) { } // Retrieve the first pending task for this computation and keyGroup in O(1) - @Nullable Node firstNode = keyGroupWorkList.head.nextKeyGroupNode; - if (firstNode == null || firstNode == keyGroupWorkList.tail) { + Node firstNode = checkStateNotNull(keyGroupWorkList.head.nextKeyGroupNode); + if (firstNode == keyGroupWorkList.tail) { return null; } unlinkNode(firstNode); @@ -259,7 +262,7 @@ public Runnable take() throws InterruptedException { } @Nullable Node node = removeFirstGlobal(); checkStateNotNull(node, "Queue is empty but size was " + size); - return checkStateNotNull(node.task, "Encountered null task in queue"); + return node.task; } finally { lock.unlock(); } @@ -287,8 +290,8 @@ public Runnable take() throws InterruptedException { public @Nullable Runnable peek() { lock.lock(); try { - @Nullable Node first = globalHead.nextNode; - if (first == null || first == globalTail) { + Node first = checkStateNotNull(globalHead.nextNode); + if (first == globalTail) { return null; } return first.task; @@ -323,13 +326,13 @@ public boolean remove(Object o) { lock.lock(); try { // Walk the global queue in O(N) to find and unlink the node - @Nullable Node curr = globalHead.nextNode; - while (curr != null && curr != globalTail) { + Node curr = checkStateNotNull(globalHead.nextNode); + while (curr != globalTail) { if (o.equals(curr.task)) { unlinkNode(curr); return true; } - curr = curr.nextNode; + curr = checkStateNotNull(curr.nextNode); } return false; } finally { @@ -342,12 +345,12 @@ public boolean contains(Object o) { if (o == null) return false; lock.lock(); try { - @Nullable Node curr = globalHead.nextNode; - while (curr != null && curr != globalTail) { + Node curr = checkStateNotNull(globalHead.nextNode); + while (curr != globalTail) { if (o.equals(curr.task)) { return true; } - curr = curr.nextNode; + curr = checkStateNotNull(curr.nextNode); } return false; } finally { @@ -373,15 +376,13 @@ public int drainTo(Collection c, int maxElements) { lock.lock(); try { int added = 0; - @Nullable Node curr = globalHead.nextNode; - while (curr != null && curr != globalTail && added < maxElements) { - @Nullable Node next = curr.nextNode; + Node curr = checkStateNotNull(globalHead.nextNode); + while (curr != globalTail && added < maxElements) { + Node next = checkStateNotNull(curr.nextNode); unlinkNode(curr); Runnable task = curr.task; - if (task != null) { - c.add(task); - ++added; - } + c.add(task); + ++added; curr = next; } return added; @@ -394,9 +395,9 @@ public int drainTo(Collection c, int maxElements) { public void clear() { lock.lock(); try { - @Nullable Node curr = globalHead.nextNode; - while (curr != null && curr != globalTail) { - @Nullable Node next = curr.nextNode; + Node curr = checkStateNotNull(globalHead.nextNode); + while (curr != globalTail) { + Node next = checkStateNotNull(curr.nextNode); unlinkNode(curr); curr = next; } @@ -410,12 +411,10 @@ public Iterator iterator() { lock.lock(); try { ImmutableList.Builder builder = ImmutableList.builderWithExpectedSize(size); - @Nullable Node curr = globalHead.nextNode; - while (curr != null && curr != globalTail) { - if (curr.task != null) { - builder.add(curr.task); - } - curr = curr.nextNode; + Node curr = checkStateNotNull(globalHead.nextNode); + while (curr != globalTail) { + builder.add(curr.task); + curr = checkStateNotNull(curr.nextNode); } return builder.build().iterator(); } finally { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java index 98023901570e..d034c4a571b5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java @@ -461,27 +461,4 @@ public void testPollWorkWithKeyGroup() { assertEquals(workA1, polledA1); assertTrue(queue.isEmpty()); } - - @Test - public void testDefaultKeyGroupNotIndexedInKeyGroupQueue() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); - - // Work item with explicit KeyGroup.DEFAULT - QueuedWork workExplicitDefault = createQueuedWork("compA", Work.KeyGroup.DEFAULT, 100); - // Work item with implicit KeyGroup.DEFAULT (null passed to helper) - QueuedWork workImplicitDefault = createQueuedWork("compA", null, 150); - - queue.offer(workExplicitDefault); - queue.offer(workImplicitDefault); - - assertEquals(2, queue.size()); - - // Verify they cannot be retrieved via pollWork with KeyGroup.DEFAULT - assertNull(queue.pollWork("compA", Work.KeyGroup.DEFAULT)); - - // Verify they can still be retrieved via global poll in FIFO order - assertEquals(workExplicitDefault, queue.poll()); - assertEquals(workImplicitDefault, queue.poll()); - assertTrue(queue.isEmpty()); - } } From 7f7d972981db34eb153f4c53d708b333795c5647 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 5 Jun 2026 10:32:28 +0000 Subject: [PATCH 4/9] Add fairQueue test for KeyGroupWorkQueue --- .../worker/util/KeyGroupWorkQueueTest.java | 35 ++++++++++++------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java index d034c4a571b5..42079361391c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java @@ -26,6 +26,7 @@ import static org.mockito.Mockito.mock; import java.util.ArrayList; +import java.util.Arrays; import java.util.Iterator; import java.util.List; import java.util.concurrent.CountDownLatch; @@ -50,11 +51,19 @@ import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameters; -@RunWith(JUnit4.class) +@RunWith(Parameterized.class) public class KeyGroupWorkQueueTest { + @Parameters(name = "fairQueue={0}") + public static Iterable data() { + return Arrays.asList(new Object[][] {{true}, {false}}); + } + + @Parameterized.Parameter public boolean fairQueue; + private BoundedQueueExecutor executor; @Before @@ -67,7 +76,7 @@ public void setUp() { 100, 10000000, new ThreadFactoryBuilder().setNameFormat("Test-%d").setDaemon(true).build(), - false, + fairQueue, /*useKeyGroupWorkQueue=*/ true); } @@ -128,7 +137,7 @@ public String toString() { @Test public void testBasicOfferAndPoll() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); assertTrue(queue.isEmpty()); assertEquals(0, queue.size()); @@ -147,7 +156,7 @@ public void testBasicOfferAndPoll() { @Test public void testRemove() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); MockRunnable task1 = new MockRunnable("1"); MockRunnable task2 = new MockRunnable("2"); @@ -162,7 +171,7 @@ public void testRemove() { @Test public void testDrainTo() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); MockRunnable task1 = new MockRunnable("1"); MockRunnable task2 = new MockRunnable("2"); queue.offer(task1); @@ -178,7 +187,7 @@ public void testDrainTo() { @Test public void testIteratorSafeTraversalAndImmutable() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); MockRunnable task1 = new MockRunnable("1"); MockRunnable task2 = new MockRunnable("2"); queue.offer(task1); @@ -205,7 +214,7 @@ public void testIteratorSafeTraversalAndImmutable() { @Test public void testPollWorkTargeted() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); QueuedWork workA1 = createQueuedWork("compA", 100); QueuedWork workB1 = createQueuedWork("compB", 200); @@ -237,7 +246,7 @@ public void testPollWorkTargeted() { @Test public void testMemoryPruningLeavesZeroLeaks() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); QueuedWork workA1 = createQueuedWork("compA", 100); queue.offer(workA1); @@ -259,7 +268,7 @@ public void testMemoryPruningLeavesZeroLeaks() { @Test public void testConcurrentStress() throws InterruptedException, ExecutionException { - final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); final int producerThreads = 4; final int consumerThreads = 4; final int tasksPerProducer = 1000; @@ -338,7 +347,7 @@ public void testConcurrentStress() throws InterruptedException, ExecutionExcepti @Test public void testTakeBlocksAndWakesUp() throws InterruptedException { - final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); final MockRunnable task = new MockRunnable("take-task"); final AtomicReference result = new AtomicReference<>(); final CountDownLatch started = new CountDownLatch(1); @@ -372,7 +381,7 @@ public void testTakeBlocksAndWakesUp() throws InterruptedException { @Test public void testPollWithTimeout() throws InterruptedException { - final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); final MockRunnable task = new MockRunnable("poll-task"); final AtomicReference result = new AtomicReference<>(); final CountDownLatch started = new CountDownLatch(1); @@ -433,7 +442,7 @@ public void testPollWithTimeout() throws InterruptedException { @Test public void testPollWorkWithKeyGroup() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(false); + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); Work.KeyGroup keyGroup1 = Work.KeyGroup.create(1, 1); Work.KeyGroup keyGroup2 = Work.KeyGroup.create(1, 2); From 4749b08add7e4e6ba2faaeca676a9df0c0716272 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Fri, 5 Jun 2026 11:55:16 +0000 Subject: [PATCH 5/9] address comments --- .../apache/beam/runners/dataflow/worker/streaming/Work.java | 3 +++ .../runners/dataflow/worker/util/BoundedQueueExecutor.java | 2 +- .../beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java | 4 ++-- 3 files changed, 6 insertions(+), 3 deletions(-) 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 0a09a5a13309..53ed30fdedbb 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 @@ -451,6 +451,9 @@ private KeyGroup(long high, long low) { } public static KeyGroup create(long high, long low) { + if (high == 0 && low == 0) { + return DEFAULT; + } return new KeyGroup(high, 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 71200a33c21e..1445492968fc 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 @@ -389,7 +389,7 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) public @Nullable ExecutableWork pollWork( String computationId, Work.KeyGroup keyGroup, BoundedQueueExecutorWorkHandle handle) { checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl); - checkArgument(computationId != null && !KeyGroup.DEFAULT.equals(keyGroup)); + checkArgument(computationId != null && keyGroup != null && !keyGroup.equals(KeyGroup.DEFAULT)); BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle; if (keyGroupWorkQueue == null) { return null; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java index f26c6ca1ee14..823178d96584 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java @@ -179,7 +179,7 @@ private void unlinkNode(Node node) { * @param keyGroup should not be equal to KeyGroup.DEFAULT */ public @Nullable QueuedWork pollWork(String computationId, Work.KeyGroup keyGroup) { - checkArgument(computationId != null && !KeyGroup.DEFAULT.equals(keyGroup)); + checkArgument(computationId != null && keyGroup != null && !keyGroup.equals(KeyGroup.DEFAULT)); QueueKey key = new QueueKey(computationId, keyGroup); lock.lock(); try { @@ -216,7 +216,7 @@ public boolean offer(@NonNull Runnable runnable) { // Append to key group list if applicable String compId = node.computationId; Work.KeyGroup keyGroup = node.keyGroup; - if (compId != null && keyGroup != null && !keyGroup.equals(Work.KeyGroup.DEFAULT)) { + if (compId != null && keyGroup != null && !keyGroup.equals(KeyGroup.DEFAULT)) { QueueKey key = new QueueKey(compId, keyGroup); KeyGroupWorkList keyGroupWorkList = keyGroupQueueMap.computeIfAbsent(key, k -> new KeyGroupWorkList()); From 6430094da0bc8f378109968da78f33b2b1faea2b Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Sun, 7 Jun 2026 09:22:04 +0000 Subject: [PATCH 6/9] Address comments --- .../worker/util/BoundedQueueExecutor.java | 5 +- .../worker/util/KeyGroupWorkQueueTest.java | 96 +++++++++---------- 2 files changed, 49 insertions(+), 52 deletions(-) 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 1445492968fc..8964246c1160 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 @@ -38,9 +38,6 @@ import org.checkerframework.checker.nullness.qual.Nullable; /** An executor for executing work on windmill items. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class BoundedQueueExecutor { private final ThreadPoolExecutor executor; @@ -394,7 +391,7 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long bytes) if (keyGroupWorkQueue == null) { return null; } - QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup); + @Nullable QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup); if (queuedWork == null) { return null; } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java index 42079361391c..91b3e38900b3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java @@ -25,6 +25,7 @@ import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; +import java.lang.Thread.State; import java.util.ArrayList; import java.util.Arrays; import java.util.Iterator; @@ -119,10 +120,10 @@ private QueuedWork createQueuedWork( return new QueuedWork(work, executor.createBudgetHandle(1, workBytes)); } - private static class MockRunnable implements Runnable { + private static class NoOpRunnable implements Runnable { final String id; - MockRunnable(String id) { + NoOpRunnable(String id) { this.id = id; } @@ -131,7 +132,7 @@ public void run() {} @Override public String toString() { - return "MockRunnable(" + id + ")"; + return "NoOpRunnable(" + id + ")"; } } @@ -141,8 +142,8 @@ public void testBasicOfferAndPoll() { assertTrue(queue.isEmpty()); assertEquals(0, queue.size()); - MockRunnable task1 = new MockRunnable("1"); - MockRunnable task2 = new MockRunnable("2"); + NoOpRunnable task1 = new NoOpRunnable("1"); + NoOpRunnable task2 = new NoOpRunnable("2"); assertTrue(queue.offer(task1)); assertTrue(queue.offer(task2)); @@ -157,8 +158,8 @@ public void testBasicOfferAndPoll() { @Test public void testRemove() { KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); - MockRunnable task1 = new MockRunnable("1"); - MockRunnable task2 = new MockRunnable("2"); + NoOpRunnable task1 = new NoOpRunnable("1"); + NoOpRunnable task2 = new NoOpRunnable("2"); queue.offer(task1); queue.offer(task2); @@ -172,8 +173,8 @@ public void testRemove() { @Test public void testDrainTo() { KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); - MockRunnable task1 = new MockRunnable("1"); - MockRunnable task2 = new MockRunnable("2"); + NoOpRunnable task1 = new NoOpRunnable("1"); + NoOpRunnable task2 = new NoOpRunnable("2"); queue.offer(task1); queue.offer(task2); @@ -188,8 +189,8 @@ public void testDrainTo() { @Test public void testIteratorSafeTraversalAndImmutable() { KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); - MockRunnable task1 = new MockRunnable("1"); - MockRunnable task2 = new MockRunnable("2"); + NoOpRunnable task1 = new NoOpRunnable("1"); + NoOpRunnable task2 = new NoOpRunnable("2"); queue.offer(task1); queue.offer(task2); @@ -225,45 +226,27 @@ public void testPollWorkTargeted() { queue.offer(workA2); assertEquals(3, queue.size()); + assertFalse(queue.isEmpty()); // Targeted poll A QueuedWork polledA1 = queue.pollWork("compA", TEST_KEY_GROUP); assertNotNull(polledA1); assertEquals("compA", polledA1.getWork().getComputationId()); assertEquals(100, polledA1.getHandle().bytes()); - - // Verify size decremented assertEquals(2, queue.size()); + assertFalse(queue.isEmpty()); // Poll next should be B1 (since A1 was stolen, B1 is now first global) assertEquals(workB1, queue.poll()); assertEquals(1, queue.size()); + assertFalse(queue.isEmpty()); // Last should be A2 assertEquals(workA2, queue.poll()); + assertEquals(0, queue.size()); assertTrue(queue.isEmpty()); - } - - @Test - public void testMemoryPruningLeavesZeroLeaks() { - KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); - QueuedWork workA1 = createQueuedWork("compA", 100); - queue.offer(workA1); - - // Steal A1 - QueuedWork polled = queue.pollWork("compA", TEST_KEY_GROUP); - 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", TEST_KEY_GROUP); - assertNotNull(polledB); - assertTrue(queue.isEmpty()); + assertEquals(null, queue.poll()); } @Test @@ -317,6 +300,9 @@ public void testConcurrentStress() throws InterruptedException, ExecutionExcepti } else { // Global poll task = queue.poll(); + if (task == null) { + task = queue.poll(10, TimeUnit.MICROSECONDS); + } } if (task != null) { consumedCount.incrementAndGet(); @@ -331,7 +317,7 @@ public void testConcurrentStress() throws InterruptedException, ExecutionExcepti } startLatch.countDown(); - assertTrue(doneLatch.await(10, TimeUnit.SECONDS)); + assertTrue(doneLatch.await(30, TimeUnit.SECONDS)); // Check for exceptions in threads for (Future future : futures) { @@ -339,7 +325,7 @@ public void testConcurrentStress() throws InterruptedException, ExecutionExcepti } executorService.shutdown(); - assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(executorService.awaitTermination(30, TimeUnit.SECONDS)); assertEquals(0, queue.size()); assertTrue(queue.isEmpty()); @@ -348,7 +334,7 @@ public void testConcurrentStress() throws InterruptedException, ExecutionExcepti @Test public void testTakeBlocksAndWakesUp() throws InterruptedException { final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); - final MockRunnable task = new MockRunnable("take-task"); + final NoOpRunnable task = new NoOpRunnable("take-task"); final AtomicReference result = new AtomicReference<>(); final CountDownLatch started = new CountDownLatch(1); final CountDownLatch finished = new CountDownLatch(1); @@ -368,21 +354,22 @@ public void testTakeBlocksAndWakesUp() throws InterruptedException { t.setDaemon(true); t.start(); - assertTrue(started.await(2, TimeUnit.SECONDS)); - // Give thread a moment to enter await() - Thread.sleep(100); + assertTrue(started.await(30, TimeUnit.SECONDS)); + while (t.getState() != State.WAITING) { + Thread.sleep(1); + } assertEquals(Thread.State.WAITING, t.getState()); queue.offer(task); - assertTrue(finished.await(2, TimeUnit.SECONDS)); + assertTrue(finished.await(30, TimeUnit.SECONDS)); assertEquals(task, result.get()); } @Test public void testPollWithTimeout() throws InterruptedException { final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); - final MockRunnable task = new MockRunnable("poll-task"); + final NoOpRunnable task = new NoOpRunnable("poll-task"); final AtomicReference result = new AtomicReference<>(); final CountDownLatch started = new CountDownLatch(1); final CountDownLatch finished = new CountDownLatch(1); @@ -403,11 +390,13 @@ public void testPollWithTimeout() throws InterruptedException { t1.setDaemon(true); t1.start(); - assertTrue(started.await(2, TimeUnit.SECONDS)); - Thread.sleep(100); + assertTrue(started.await(30, TimeUnit.SECONDS)); + while (t1.getState() != State.TIMED_WAITING) { + Thread.sleep(1); + } assertEquals(Thread.State.TIMED_WAITING, t1.getState()); - assertTrue(finished.await(2, TimeUnit.SECONDS)); + assertTrue(finished.await(30, TimeUnit.SECONDS)); assertNull(result.get()); // 2. Verify timed poll receives task offered concurrently @@ -430,13 +419,15 @@ public void testPollWithTimeout() throws InterruptedException { t2.setDaemon(true); t2.start(); - assertTrue(started2.await(2, TimeUnit.SECONDS)); - Thread.sleep(100); + assertTrue(started2.await(30, TimeUnit.SECONDS)); + while (t2.getState() != State.TIMED_WAITING) { + Thread.sleep(1); + } assertEquals(Thread.State.TIMED_WAITING, t2.getState()); queue.offer(task); - assertTrue(finished2.await(2, TimeUnit.SECONDS)); + assertTrue(finished2.await(30, TimeUnit.SECONDS)); assertEquals(task, result2.get()); } @@ -446,6 +437,7 @@ public void testPollWorkWithKeyGroup() { Work.KeyGroup keyGroup1 = Work.KeyGroup.create(1, 1); Work.KeyGroup keyGroup2 = Work.KeyGroup.create(1, 2); + Work.KeyGroup keyGroupNotExist = Work.KeyGroup.create(3, 4); QueuedWork workA1 = createQueuedWork("compA", keyGroup1, 100); QueuedWork workA2 = createQueuedWork("compA", keyGroup2, 150); @@ -455,6 +447,10 @@ public void testPollWorkWithKeyGroup() { assertEquals(2, queue.size()); + QueuedWork polledNotExist = queue.pollWork("compA", keyGroupNotExist); + assertNull(polledNotExist); + assertEquals(2, queue.size()); + // Poll with keyGroup2 first - should return workA2 QueuedWork polledA2 = queue.pollWork("compA", keyGroup2); assertNotNull(polledA2); @@ -469,5 +465,9 @@ public void testPollWorkWithKeyGroup() { assertNotNull(polledA1); assertEquals(workA1, polledA1); assertTrue(queue.isEmpty()); + + polledNotExist = queue.pollWork("compA", keyGroupNotExist); + assertNull(polledNotExist); + assertTrue(queue.isEmpty()); } } From efbac1ca6ee90573ea103fd797e01b3f20b55dcc Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Sun, 7 Jun 2026 09:57:22 +0000 Subject: [PATCH 7/9] Address comments --- .../worker/util/KeyGroupWorkQueueTest.java | 127 +++++++++++------- 1 file changed, 79 insertions(+), 48 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java index 91b3e38900b3..994aa2030f3f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java @@ -36,7 +36,6 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; import org.apache.beam.runners.dataflow.worker.streaming.Watermarks; @@ -251,73 +250,103 @@ public void testPollWorkTargeted() { @Test public void testConcurrentStress() throws InterruptedException, ExecutionException { - final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); - final int producerThreads = 4; - final int consumerThreads = 4; - final int tasksPerProducer = 1000; - final int totalTasks = producerThreads * tasksPerProducer; + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); + int producerThreads = 4; + int consumerThreads = 4; + int tasksPerProducer = 1000; + int totalTasks = producerThreads * tasksPerProducer; + Runnable poisonPill = + new Runnable() { + @Override + public void run() {} + + @Override + public String toString() { + return "POISON_PILL"; + } + }; 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); + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch producersDoneLatch = new CountDownLatch(producerThreads); + CountDownLatch consumersDoneLatch = new CountDownLatch(consumerThreads); + CountDownLatch consumedLatch = new CountDownLatch(totalTasks); List> futures = new ArrayList<>(); - // Start producers - for (int i = 0; i < producerThreads; i++) { + // Start consumers + for (int i = 0; i < consumerThreads; i++) { + int consumerId = i; futures.add( executorService.submit( () -> { try { startLatch.await(); - for (int j = 0; j < tasksPerProducer; j++) { - String compId = "comp-" + (j % 5); - queue.offer(createQueuedWork(compId, 10)); + int iteration = consumerId % 4; + while (true) { + int strategy = iteration; + iteration = (iteration + 1) % 4; + Runnable task = null; + if (strategy == 0) { + String compId = "comp-" + (consumedLatch.getCount() % 5); + task = queue.pollWork(compId, TEST_KEY_GROUP); + } else if (strategy == 1) { + task = queue.poll(); + } else if (strategy == 2) { + task = queue.poll(10, TimeUnit.MICROSECONDS); + } else if (strategy == 3) { + task = queue.take(); + } + + if (task == poisonPill) { + break; + } + if (task != null) { + consumedLatch.countDown(); + } } } catch (Exception e) { throw new RuntimeException(e); } finally { - doneLatch.countDown(); + consumersDoneLatch.countDown(); } })); } - // Start consumers (mix of poll and pollWork) - for (int i = 0; i < consumerThreads; i++) { - final int consumerId = i; + // Start producers + for (int i = 0; i < producerThreads; i++) { futures.add( 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, TEST_KEY_GROUP); - } else { - // Global poll - task = queue.poll(); - if (task == null) { - task = queue.poll(10, TimeUnit.MICROSECONDS); - } - } - if (task != null) { - consumedCount.incrementAndGet(); - } + for (int j = 0; j < tasksPerProducer; j++) { + String compId = "comp-" + (j % 5); + queue.offer(createQueuedWork(compId, 10)); } } catch (Exception e) { throw new RuntimeException(e); } finally { - doneLatch.countDown(); + producersDoneLatch.countDown(); } })); } + // Release the start latch to start the test startLatch.countDown(); - assertTrue(doneLatch.await(30, TimeUnit.SECONDS)); + + // Wait for all tasks to be consumed + assertTrue(consumedLatch.await(30, TimeUnit.SECONDS)); + + // Send poison pills to stop all consumers + for (int i = 0; i < consumerThreads; i++) { + queue.offer(poisonPill); + } + + // Wait for consumers to finish + assertTrue(consumersDoneLatch.await(30, TimeUnit.SECONDS)); + // Wait for producers to finish + assertTrue(producersDoneLatch.await(30, TimeUnit.SECONDS)); // Check for exceptions in threads for (Future future : futures) { @@ -355,10 +384,7 @@ public void testTakeBlocksAndWakesUp() throws InterruptedException { t.start(); assertTrue(started.await(30, TimeUnit.SECONDS)); - while (t.getState() != State.WAITING) { - Thread.sleep(1); - } - assertEquals(Thread.State.WAITING, t.getState()); + waitForThreadState(t, State.WAITING); queue.offer(task); @@ -391,10 +417,7 @@ public void testPollWithTimeout() throws InterruptedException { t1.start(); assertTrue(started.await(30, TimeUnit.SECONDS)); - while (t1.getState() != State.TIMED_WAITING) { - Thread.sleep(1); - } - assertEquals(Thread.State.TIMED_WAITING, t1.getState()); + waitForThreadState(t1, State.TIMED_WAITING); assertTrue(finished.await(30, TimeUnit.SECONDS)); assertNull(result.get()); @@ -420,10 +443,7 @@ public void testPollWithTimeout() throws InterruptedException { t2.start(); assertTrue(started2.await(30, TimeUnit.SECONDS)); - while (t2.getState() != State.TIMED_WAITING) { - Thread.sleep(1); - } - assertEquals(Thread.State.TIMED_WAITING, t2.getState()); + waitForThreadState(t2, State.TIMED_WAITING); queue.offer(task); @@ -470,4 +490,15 @@ public void testPollWorkWithKeyGroup() { assertNull(polledNotExist); assertTrue(queue.isEmpty()); } + + private void waitForThreadState(Thread t, State state) throws InterruptedException { + long timeoutMs = 30000; + long start = System.currentTimeMillis(); + while (t.getState() != state) { + if (System.currentTimeMillis() - start > timeoutMs) { + fail("Thread did not reach " + state + " state within " + timeoutMs + "ms"); + } + Thread.sleep(1); + } + } } From 15b7c6bdd2947f59679382c746a40d09cd0d4176 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 8 Jun 2026 01:40:47 +0000 Subject: [PATCH 8/9] fix import --- runners/google-cloud-dataflow-java/worker/build.gradle | 1 + .../beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/build.gradle b/runners/google-cloud-dataflow-java/worker/build.gradle index 21879861e9d6..29569ea9c72b 100644 --- a/runners/google-cloud-dataflow-java/worker/build.gradle +++ b/runners/google-cloud-dataflow-java/worker/build.gradle @@ -212,6 +212,7 @@ dependencies { implementation library.java.jackson_core implementation library.java.jackson_databind implementation library.java.joda_time + implementation library.java.jspecify implementation library.java.opentelemetry_context implementation library.java.slf4j_api implementation library.java.vendored_grpc_1_69_0 diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java index 823178d96584..d151157ec68f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java @@ -36,8 +36,8 @@ import org.apache.beam.runners.dataflow.worker.streaming.Work.KeyGroup; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor.QueuedWork; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; -import org.jspecify.annotations.NonNull; /** * A custom, thread-safe doubly-linked BlockingQueue. In addition to global FIFO ordering, the queue From 31f432ecf9e099da1172ef550b9d198ed15a24af Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 8 Jun 2026 01:48:51 +0000 Subject: [PATCH 9/9] fix import --- runners/google-cloud-dataflow-java/worker/build.gradle | 1 - 1 file changed, 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/build.gradle b/runners/google-cloud-dataflow-java/worker/build.gradle index 29569ea9c72b..21879861e9d6 100644 --- a/runners/google-cloud-dataflow-java/worker/build.gradle +++ b/runners/google-cloud-dataflow-java/worker/build.gradle @@ -212,7 +212,6 @@ dependencies { implementation library.java.jackson_core implementation library.java.jackson_databind implementation library.java.joda_time - implementation library.java.jspecify implementation library.java.opentelemetry_context implementation library.java.slf4j_api implementation library.java.vendored_grpc_1_69_0