This is an automated email from the ASF dual-hosted git repository.
derrickaw pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 18db4896cf1 [Multikey] [Dataflow Streaming] Replace KeyGroupWorkQueue
with a sharded KeyGroupIndex (#40593)
18db4896cf1 is described below
commit 18db4896cf142b11f4f6e91d57ad4183d76c6115
Author: Arun Pandian <[email protected]>
AuthorDate: Mon Oct 19 13:48:38 2026 -0700
[Multikey] [Dataflow Streaming] Replace KeyGroupWorkQueue with a sharded
KeyGroupIndex (#40593)
* [Dataflow Streaming] Replace KeyGroupWorkQueue with a sharded
KeyGroupWorkIndex
`KeyGroupWorkQueue` guarded a doubly-linked global queue and per-key-group
queues with a single `ReentrantLock`, creating lock contention across worker
and `GetWork` threads when multi-key bundles are enabled.
Use `LinkedBlockingQueue` as the `BoundedQueueExecutor` work queue
unconditionally, and index stealable work in `KeyGroupWorkIndex` across 64
fixed `HashMap<QueueKey, ArrayDeque<QueuedWork>>` shards synchronized per
shard. An executor thread and `pollWork` coordinate on a `QueuedWork` via a
single atomic state transition, clearing work and budget references immediately
upon claim. Claimed entries at the head of a key-group queue are drained in
amortized O(1) time during index mutation [...]
* Address comments + refactorings + change purge logic
* Address comments
* Clarify comments regarding maximumElementsOutstanding
---
.../dataflow/worker/StreamingDataflowWorker.java | 7 +-
.../runners/dataflow/worker/streaming/Work.java | 3 +-
.../dataflow/worker/util/BoundedQueueExecutor.java | 224 +++++++--
.../dataflow/worker/util/KeyGroupWorkIndex.java | 190 ++++++++
.../dataflow/worker/util/KeyGroupWorkQueue.java | 472 -------------------
.../worker/StreamingDataflowWorkerTest.java | 12 +-
.../worker/util/BoundedQueueExecutorTest.java | 428 +++++++++++------
.../worker/util/KeyGroupWorkIndexTest.java | 319 +++++++++++++
.../worker/util/KeyGroupWorkQueueTest.java | 523 ---------------------
.../processing/StreamingCommitFinalizerTest.java | 2 +-
.../failures/WorkFailureProcessorTest.java | 2 +-
.../work/refresh/ActiveWorkRefresherTest.java | 2 +-
12 files changed, 991 insertions(+), 1193 deletions(-)
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
index fd00cea775f..dd96142c41a 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
@@ -1055,7 +1055,6 @@ public final class StreamingDataflowWorker {
DataflowWorkerHarnessOptions options, MultiKeyBundleOptions
multiKeyBundleOptions) {
boolean useFairMonitor =
DataflowRunner.hasExperiment(options,
BOUNDED_QUEUE_EXECUTOR_USE_FAIR_MONITOR_EXPERIMENT);
- boolean useKeyGroupWorkQueue =
multiKeyBundleOptions.multiKeyBundleEnabled();
return new BoundedQueueExecutor(
chooseMaxThreads(options),
THREAD_EXPIRATION_TIME_SEC,
@@ -1064,7 +1063,7 @@ public final class StreamingDataflowWorker {
chooseMaxBytesOutstanding(options),
new
ThreadFactoryBuilder().setNameFormat("DataflowWorkUnits-%d").setDaemon(true).build(),
useFairMonitor,
- useKeyGroupWorkQueue);
+ multiKeyBundleOptions.multiKeyBundleEnabled());
}
public static void main(String[] args) throws Exception {
@@ -1134,8 +1133,8 @@ public final class StreamingDataflowWorker {
}
@VisibleForTesting
- public boolean workExecutorIsEmpty() {
- return workUnitExecutor.executorQueueIsEmpty();
+ public boolean workExecutorIsEmptyForTesting() {
+ return workUnitExecutor.executorQueueIsEmptyForTesting();
}
@VisibleForTesting
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 2acee9410fa..724668b689f 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,7 +25,6 @@ import java.util.EnumMap;
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.concurrent.atomic.AtomicBoolean;
@@ -554,7 +553,7 @@ public final class Work implements RefreshableWork {
@Override
public int hashCode() {
- return Objects.hash(high, low);
+ return 31 * Long.hashCode(high) + Long.hashCode(low);
}
@Override
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 046d8cae9f9..1d10d074b11 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,8 @@ import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import javax.annotation.concurrent.GuardedBy;
import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
@@ -41,13 +43,20 @@ import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurren
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. */
+/**
+ * An executor for executing work on windmill items.
+ *
+ * <p>When multi-key bundles are enabled, batchable work is also indexed in
{@link
+ * KeyGroupWorkIndex} so running bundles can steal it via {@link #pollWork}.
Stolen work leaves a
+ * no-op stale shell in the executor queue, which is discarded when dequeued
or purged in bulk if
+ * the queue grows too large.
+ */
public class BoundedQueueExecutor {
private final ThreadPoolExecutor executor;
private final long maximumBytesOutstanding;
- // Used to guard elementsOutstanding and bytesOutstanding.
+ // Used to guard writes to elementsOutstanding and bytesOutstanding.
private final Monitor monitor;
private static class Budget {
@@ -64,11 +73,10 @@ public class BoundedQueueExecutor {
private final ConcurrentLinkedQueue<Budget> decrementQueue = new
ConcurrentLinkedQueue<>();
private final Object decrementQueueDrainLock = new Object();
private final AtomicBoolean isDecrementBatchPending = new
AtomicBoolean(false);
- private int elementsOutstanding = 0;
+ private final AtomicInteger elementsOutstanding = new AtomicInteger();
private long bytesOutstanding = 0;
- @GuardedBy("this")
- private int maximumElementsOutstanding;
+ private volatile int maximumElementsOutstanding;
@GuardedBy("this")
private int activeCount;
@@ -82,8 +90,10 @@ public class BoundedQueueExecutor {
@GuardedBy("this")
private long totalTimeMaxActiveThreadsUsed;
- // If set the keyGroupWorkQueue is used by the underlying executor.
- private final @Nullable KeyGroupWorkQueue keyGroupWorkQueue;
+ private final LinkedBlockingQueue<Runnable> workQueue = new
LinkedBlockingQueue<>();
+ private final @Nullable KeyGroupWorkIndex keyGroupWorkIndex;
+
+ private final AtomicBoolean purgingStaleShells = new AtomicBoolean(false);
public BoundedQueueExecutor(
int initialMaximumPoolSize,
@@ -93,8 +103,8 @@ public class BoundedQueueExecutor {
long maximumBytesOutstanding,
ThreadFactory threadFactory,
boolean useFairMonitor,
- boolean useKeyGroupWorkQueue) {
- this.keyGroupWorkQueue = useKeyGroupWorkQueue ? new
KeyGroupWorkQueue(useFairMonitor) : null;
+ boolean multiKeyBundleEnabled) {
+ this.keyGroupWorkIndex = multiKeyBundleEnabled ? new KeyGroupWorkIndex() :
null;
this.maximumPoolSize = initialMaximumPoolSize;
monitor = new Monitor(useFairMonitor);
executor =
@@ -103,7 +113,7 @@ public class BoundedQueueExecutor {
initialMaximumPoolSize,
keepAliveTime,
unit,
- keyGroupWorkQueue != null ? keyGroupWorkQueue : new
LinkedBlockingQueue<>(),
+ workQueue,
threadFactory) {
@Override
protected void beforeExecute(Thread t, Runnable r) {
@@ -139,9 +149,9 @@ public class BoundedQueueExecutor {
new Guard(monitor) {
@Override
public boolean isSatisfied() {
- return elementsOutstanding == 0
- || (bytesAvailable() >= workBytes
- && elementsOutstanding < maximumElementsOutstanding());
+ int outstanding = elementsOutstanding.get();
+ return outstanding == 0
+ || (bytesAvailable() >= workBytes && outstanding <
maximumElementsOutstanding());
}
});
executeMonitorHeld(work, workBytes);
@@ -181,8 +191,21 @@ public class BoundedQueueExecutor {
}
}
- public boolean executorQueueIsEmpty() {
- return executor.getQueue().isEmpty();
+ @VisibleForTesting
+ public boolean executorQueueIsEmptyForTesting() {
+ return workQueue.stream().allMatch(BoundedQueueExecutor::isStaleShell);
+ }
+
+ /** Returns the raw number of entries in the executor queue, including stale
shells. */
+ @VisibleForTesting
+ int rawQueueSizeForTesting() {
+ return workQueue.size();
+ }
+
+ @VisibleForTesting
+ boolean keyGroupWorkIndexIsEmpty() {
+ @Nullable KeyGroupWorkIndex index = keyGroupWorkIndex;
+ return index == null || index.isEmpty();
}
public synchronized long allThreadsActiveTime() {
@@ -203,19 +226,14 @@ public class BoundedQueueExecutor {
}
public int elementsOutstanding() {
- monitor.enter();
- try {
- return elementsOutstanding;
- } finally {
- monitor.leave();
- }
+ return elementsOutstanding.get();
}
public long maximumBytesOutstanding() {
return maximumBytesOutstanding;
}
- public synchronized int maximumElementsOutstanding() {
+ public int maximumElementsOutstanding() {
return maximumElementsOutstanding;
}
@@ -238,7 +256,7 @@ public class BoundedQueueExecutor {
builder.append("<br>/n");
builder.append("Work Queue Size: ");
- builder.append(elementsOutstanding);
+ builder.append(elementsOutstanding.get());
builder.append("/");
builder.append(maximumElementsOutstanding());
builder.append("<br>/n");
@@ -328,49 +346,132 @@ public class BoundedQueueExecutor {
}
}
+ /**
+ * A {@link Runnable} for an {@link ExecutableWork} that can be claimed
either by an executor
+ * thread in {@link #run} or by a multi-key bundle in {@link #pollWork}.
+ *
+ * <p>Thread safety: {@link #state} is {@code volatile} and transitions at
most once from {@link
+ * #QUEUED} to {@link #STARTED} or {@link #STOLEN} via {@link #STATE}{@code
.compareAndSet}.
+ * Non-final fields {@link #work} and {@link #handle} are written before
{@code state} is
+ * initialized in the constructor, and are subsequently read and cleared
only by the single thread
+ * that wins that {@code compareAndSet}. All other accesses read only {@code
final} fields or
+ * {@code volatile state}.
+ */
static final class QueuedWork implements Runnable {
- private final ExecutableWork work;
- private final BoundedQueueExecutorWorkHandleImpl handle;
-
- public QueuedWork(ExecutableWork work, BoundedQueueExecutorWorkHandleImpl
handle) {
+ private static final int QUEUED = 0;
+ private static final int STARTED = 1; // Claimed by an executor thread in
run().
+ private static final int STOLEN = 2; // Claimed by a multi-key bundle in
pollWork().
+
+ private static final AtomicIntegerFieldUpdater<QueuedWork> STATE =
+ AtomicIntegerFieldUpdater.newUpdater(QueuedWork.class, "state");
+
+ // Non-null iff the work is eligible to be stolen by a multi-key bundle.
+ private final @Nullable KeyGroupWorkIndex keyGroupWorkIndex;
+ private final KeyGroupWorkIndex.@Nullable QueueKey queueKey;
+ // Cleared upon claim so stale shells and claimed index entries do not
retain work references.
+ private @Nullable ExecutableWork work;
+ private @Nullable BoundedQueueExecutorWorkHandleImpl handle;
+ private volatile int state = QUEUED;
+
+ QueuedWork(
+ ExecutableWork work,
+ BoundedQueueExecutorWorkHandleImpl handle,
+ @Nullable KeyGroupWorkIndex keyGroupWorkIndex) {
this.work = work;
this.handle = handle;
+ Work w = work.work();
+ if (keyGroupWorkIndex != null
+ && !w.getKeyGroup().equals(Work.KeyGroup.DEFAULT)
+ && !w.isMultiKeyBatchingDisabled()) {
+ this.keyGroupWorkIndex = keyGroupWorkIndex;
+ this.queueKey = new KeyGroupWorkIndex.QueueKey(w.getComputationId(),
w.getKeyGroup());
+ } else {
+ this.keyGroupWorkIndex = null;
+ this.queueKey = null;
+ }
+ }
+
+ @VisibleForTesting
+ ExecutableWork getWork() {
+ return checkStateNotNull(work, "QueuedWork was already claimed.");
+ }
+
+ private BoundedQueueExecutorWorkHandleImpl getHandle() {
+ return checkStateNotNull(handle, "QueuedWork was already claimed.");
+ }
+
+ @VisibleForTesting
+ boolean isStealable() {
+ return keyGroupWorkIndex != null;
+ }
+
+ KeyGroupWorkIndex.QueueKey getQueueKey() {
+ return checkStateNotNull(queueKey);
}
- public ExecutableWork getWork() {
- return work;
+ boolean isClaimed() {
+ return state != QUEUED;
}
- public BoundedQueueExecutorWorkHandleImpl getHandle() {
- return handle;
+ private boolean isStolen() {
+ return state == STOLEN;
+ }
+
+ @VisibleForTesting
+ boolean tryClaimForExecution() {
+ return STATE.compareAndSet(this, QUEUED, STARTED);
+ }
+
+ boolean tryClaimForSteal() {
+ return STATE.compareAndSet(this, QUEUED, STOLEN);
+ }
+
+ private void clearClaimedWork() {
+ work = null;
+ handle = null;
}
@Override
public void run() {
- checkArgument(!handle.isClosed());
- try (handle) {
- work.run(handle);
+ if (!tryClaimForExecution()) {
+ return; // Already stolen by a multi-key bundle.
+ }
+ ExecutableWork workToRun = getWork();
+ BoundedQueueExecutorWorkHandleImpl workHandle = getHandle();
+ clearClaimedWork();
+ @Nullable KeyGroupWorkIndex index = keyGroupWorkIndex;
+ if (index != null) {
+ index.drainClaimed(this);
+ }
+ checkArgument(!workHandle.isClosed());
+ try (workHandle) {
+ workToRun.run(workHandle);
}
}
}
private void executeMonitorHeld(ExecutableWork work, long workBytes) {
- ++elementsOutstanding;
+ elementsOutstanding.incrementAndGet();
bytesOutstanding += workBytes;
monitor.leave();
BoundedQueueExecutorWorkHandleImpl handle =
new BoundedQueueExecutorWorkHandleImpl(work.work(), workBytes);
+ QueuedWork queuedWork = new QueuedWork(work, handle, keyGroupWorkIndex);
try {
- executor.execute(new QueuedWork(work, handle));
+ executor.execute(queuedWork);
} catch (Throwable t) {
handle.close();
throw ExceptionUtils.safeWrapThrowableAsException(t);
}
+ if (queuedWork.isStealable()) {
+ // Index after queuing so work rejected after shutdown cannot be stolen.
+ checkStateNotNull(keyGroupWorkIndex).add(queuedWork);
+ }
}
private void executeMonitorHeld(Runnable work) {
- ++elementsOutstanding;
+ elementsOutstanding.incrementAndGet();
monitor.leave();
try {
@@ -403,23 +504,56 @@ public class BoundedQueueExecutor {
checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl);
checkStateNotNull(onFailedWorkHandler);
BoundedQueueExecutorWorkHandleImpl internalHandle =
(BoundedQueueExecutorWorkHandleImpl) handle;
- if (keyGroupWorkQueue == null) {
+ @Nullable KeyGroupWorkIndex index = keyGroupWorkIndex;
+ if (index == null) {
return null;
}
while (true) {
- @Nullable QueuedWork queuedWork =
keyGroupWorkQueue.pollWork(computationId, keyGroup);
- if (queuedWork == null) {
+ @Nullable QueuedWork stolen = index.pollAndClaim(computationId,
keyGroup);
+ if (stolen == null) {
return null;
}
- Work work = queuedWork.getWork().work();
+ ExecutableWork stolenWork = stolen.getWork();
+ BoundedQueueExecutorWorkHandleImpl stolenHandle = stolen.getHandle();
+ stolen.clearClaimedWork();
+ onWorkStolen();
+ Work work = stolenWork.work();
if (work.isFailed()) {
- queuedWork.getHandle().close();
+ stolenHandle.close();
onFailedWorkHandler.onFailedWork(work);
continue;
}
- internalHandle.merge(queuedWork.getHandle());
- return queuedWork.getWork();
+ internalHandle.merge(stolenHandle);
+ return stolenWork;
+ }
+ }
+
+ /**
+ * Purges stale shells in bulk if they accumulate while all executor threads
are busy running
+ * multi-key bundles.
+ *
+ * <p>{@code forceExecute} can queue more than {@code
maximumElementsOutstanding} unstarted items
+ * (e.g. retries or finalizations). {@code elementsOutstanding} counts
non-shell work items, so
+ * comparing {@code workQueue.size()} against {@code elementsOutstanding +
+ * maximumElementsOutstanding} means there is at least {@code
maximumElementsOutstanding} closed
+ * stale shells before running an O(n) {@code removeIf}.
+ */
+ private void onWorkStolen() {
+ if (workQueue.size() <= (long) elementsOutstanding.get() +
maximumElementsOutstanding
+ || !purgingStaleShells.compareAndSet(false, true)) {
+ return;
}
+ try {
+ if (!executor.isShutdown()) {
+ workQueue.removeIf(BoundedQueueExecutor::isStaleShell);
+ }
+ } finally {
+ purgingStaleShells.set(false);
+ }
+ }
+
+ private static boolean isStaleShell(Runnable runnable) {
+ return runnable instanceof QueuedWork && ((QueuedWork)
runnable).isStolen();
}
private void decrementCounters(int elements, long bytes) {
@@ -457,7 +591,7 @@ public class BoundedQueueExecutor {
}
monitor.enter();
- elementsOutstanding -= elementsToDecrement;
+ elementsOutstanding.addAndGet(-elementsToDecrement);
bytesOutstanding -= bytesToDecrement;
monitor.leave();
}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkIndex.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkIndex.java
new file mode 100644
index 00000000000..10cdab083b5
--- /dev/null
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkIndex.java
@@ -0,0 +1,190 @@
+/*
+ * 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.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
+import java.util.ArrayDeque;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+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.annotations.VisibleForTesting;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Sharded FIFO index of {@link QueuedWork} by computation and {@link
Work.KeyGroup} for work
+ * stealing by multi-key bundles.
+ */
+@ThreadSafe
+final class KeyGroupWorkIndex {
+
+ private static final int NUM_SHARDS = 64;
+ private static final int INITIAL_QUEUE_CAPACITY = 4;
+
+ private final ImmutableList<Shard> shards;
+
+ KeyGroupWorkIndex() {
+ ImmutableList.Builder<Shard> builder =
ImmutableList.builderWithExpectedSize(NUM_SHARDS);
+ for (int i = 0; i < NUM_SHARDS; i++) {
+ builder.add(new Shard());
+ }
+ this.shards = builder.build();
+ }
+
+ private Shard shardFor(QueueKey queueKey) {
+ return shards.get(Math.floorMod(queueKey.hashCode(), NUM_SHARDS));
+ }
+
+ /** Adds {@code work} to its key-group queue if it has not been claimed yet.
*/
+ void add(QueuedWork work) {
+ if (work.isClaimed()) {
+ return;
+ }
+ QueueKey queueKey = work.getQueueKey();
+ shardFor(queueKey).add(queueKey, work);
+ }
+
+ /**
+ * Drains claimed work from the head of {@code work}'s key-group queue. If
an earlier item is
+ * still unclaimed, {@code work} remains in the queue with its references
cleared and is drained
+ * in amortized O(1) time once the earlier item is claimed.
+ */
+ void drainClaimed(QueuedWork completedWork) {
+ checkArgument(completedWork.isClaimed());
+ QueueKey queueKey = completedWork.getQueueKey();
+ shardFor(queueKey).drainClaimed(queueKey);
+ }
+
+ /** Claims and returns the oldest unclaimed work for the computation and key
group, or null. */
+ @Nullable QueuedWork pollAndClaim(String computationId, Work.KeyGroup
keyGroup) {
+ QueueKey queueKey = new QueueKey(computationId, keyGroup);
+ return shardFor(queueKey).pollAndClaim(queueKey);
+ }
+
+ @VisibleForTesting
+ boolean isEmpty() {
+ for (Shard shard : shards) {
+ if (!shard.isEmpty()) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ /** A shard of key-group queues; the map and its queues are guarded by the
shard's monitor. */
+ private static final class Shard {
+
+ @GuardedBy("this")
+ private final Map<QueueKey, ArrayDeque<QueuedWork>> queuesByKeyGroup = new
HashMap<>();
+
+ private synchronized void add(QueueKey queueKey, QueuedWork work) {
+ if (work.isClaimed()) {
+ return;
+ }
+ ArrayDeque<QueuedWork> queue =
+ queuesByKeyGroup.computeIfAbsent(
+ queueKey, unused -> new ArrayDeque<>(INITIAL_QUEUE_CAPACITY));
+ drainClaimedFromHead(queue);
+ queue.addLast(work);
+ }
+
+ private synchronized void drainClaimed(QueueKey queueKey) {
+ @Nullable ArrayDeque<QueuedWork> queue = queuesByKeyGroup.get(queueKey);
+ if (queue == null) {
+ return;
+ }
+ drainClaimedFromHead(queue);
+ if (queue.isEmpty()) {
+ queuesByKeyGroup.remove(queueKey);
+ }
+ }
+
+ private synchronized @Nullable QueuedWork pollAndClaim(QueueKey queueKey) {
+ @Nullable ArrayDeque<QueuedWork> queue = queuesByKeyGroup.get(queueKey);
+ if (queue == null) {
+ return null;
+ }
+ @Nullable QueuedWork claimed = null;
+ while (claimed == null && !queue.isEmpty()) {
+ QueuedWork candidate = queue.removeFirst();
+ if (candidate.tryClaimForSteal()) {
+ claimed = candidate;
+ }
+ }
+ drainClaimedFromHead(queue);
+ if (queue.isEmpty()) {
+ queuesByKeyGroup.remove(queueKey);
+ }
+ return claimed;
+ }
+
+ private synchronized boolean isEmpty() {
+ return queuesByKeyGroup.isEmpty();
+ }
+
+ private static void drainClaimedFromHead(ArrayDeque<QueuedWork> queue) {
+ while (!queue.isEmpty() && queue.peekFirst().isClaimed()) {
+ queue.removeFirst();
+ }
+ }
+ }
+
+ static final class QueueKey {
+ private final String computationId;
+ private final Work.KeyGroup keyGroup;
+ private final int hashCode;
+
+ QueueKey(String computationId, Work.KeyGroup keyGroup) {
+ this.computationId = computationId;
+ this.keyGroup = keyGroup;
+ this.hashCode = 31 * computationId.hashCode() + keyGroup.hashCode();
+ }
+
+ @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) &&
keyGroup.equals(other.keyGroup);
+ }
+
+ @Override
+ public int hashCode() {
+ return hashCode;
+ }
+
+ @Override
+ public String toString() {
+ return "QueueKey{"
+ + "computationId='"
+ + computationId
+ + '\''
+ + ", keyGroup="
+ + keyGroup
+ + '}';
+ }
+ }
+}
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
deleted file mode 100644
index dd409616ab9..00000000000
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java
+++ /dev/null
@@ -1,472 +0,0 @@
-/*
- * 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.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 static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
-
-import java.util.AbstractQueue;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.Map;
-import java.util.Objects;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.locks.Condition;
-import java.util.concurrent.locks.ReentrantLock;
-import 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.collect.ImmutableList;
-import org.checkerframework.checker.nullness.qual.NonNull;
-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<Runnable> implements
BlockingQueue<Runnable> {
-
- public static final Runnable SENTINEL_RUNNABLE =
- () -> {
- throw new IllegalStateException("sentinel runnable called");
- };
-
- static class Node {
- // 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.
- @Nullable KeyGroupWorkList keyGroupList;
-
- // 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;
-
- private static boolean isMultiKeyBatchingDisabled(Runnable task) {
- return !(task instanceof QueuedWork)
- || ((QueuedWork) task).getWork().isMultiKeyBatchingDisabled();
- }
-
- Node(Runnable task) {
- this.task = task;
- if (!isMultiKeyBatchingDisabled(task)) {
- this.computationId = ((QueuedWork) task).getWork().getComputationId();
- this.keyGroup = ((QueuedWork) task).getWork().getKeyGroup();
- } else {
- this.computationId = null;
- this.keyGroup = null;
- }
- }
- }
-
- /** Double linked list implementing key group level queue */
- private static class KeyGroupWorkList {
- final Node head = new Node(SENTINEL_RUNNABLE);
- final Node tail = new Node(SENTINEL_RUNNABLE);
-
- KeyGroupWorkList() {
- head.nextKeyGroupNode = tail;
- tail.prevKeyGroupNode = head;
- }
-
- boolean isEmpty() {
- return head.nextKeyGroupNode == tail;
- }
-
- void append(Node node) {
- Node last = checkStateNotNull(tail.prevKeyGroupNode);
- 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;
- private final Condition notEmpty;
-
- // Sentinels for the global list
- @GuardedBy("lock")
- private final Node globalHead = new Node(SENTINEL_RUNNABLE);
-
- @GuardedBy("lock")
- private final Node globalTail = new Node(SENTINEL_RUNNABLE);
-
- @GuardedBy("lock")
- private final Map<QueueKey, KeyGroupWorkList> keyGroupQueueMap = new
HashMap<>();
-
- @GuardedBy("lock")
- private int size = 0;
-
- 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 = 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
- 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() {
- Node first = checkStateNotNull(globalHead.nextNode);
- if (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.
- *
- * @param keyGroup should not be equal to KeyGroup.DEFAULT
- */
- public @Nullable QueuedWork pollWork(String computationId, Work.KeyGroup
keyGroup) {
- checkArgument(computationId != null && keyGroup != null &&
!keyGroup.equals(KeyGroup.DEFAULT));
- QueueKey key = new QueueKey(computationId, keyGroup);
- lock.lock();
- try {
- 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)
- Node firstNode =
checkStateNotNull(keyGroupWorkList.head.nextKeyGroupNode);
- if (firstNode == keyGroupWorkList.tail) {
- return null;
- }
-
- // MultiKeyBatchingDisabled items should not be in keyGroupWorkList
- checkState(!Node.isMultiKeyBatchingDisabled(firstNode.task));
-
- unlinkNode(firstNode);
-
- return (QueuedWork) firstNode.task;
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public boolean offer(@NonNull Runnable runnable) {
- Node node = new Node(checkStateNotNull(runnable));
- lock.lock();
- try {
- // Append to global list tail
- Node lastG = checkStateNotNull(globalTail.prevNode);
- node.prevNode = lastG;
- node.nextNode = globalTail;
- lastG.nextNode = node;
- globalTail.prevNode = node;
-
- // Append to key group list if applicable
- String compId = node.computationId;
- Work.KeyGroup keyGroup = node.keyGroup;
- if (compId != null && keyGroup != null &&
!keyGroup.equals(KeyGroup.DEFAULT)) {
- QueueKey key = new QueueKey(compId, keyGroup);
- KeyGroupWorkList keyGroupWorkList =
- keyGroupQueueMap.computeIfAbsent(key, k -> new KeyGroupWorkList());
- keyGroupWorkList.append(node);
- node.keyGroupList = keyGroupWorkList;
- }
-
- ++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);
- return node.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 {
- Node first = checkStateNotNull(globalHead.nextNode);
- if (first == globalTail) {
- return null;
- }
- return first.task;
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public int size() {
- lock.lock();
- try {
- return size;
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public boolean isEmpty() {
- lock.lock();
- try {
- return size == 0;
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public boolean remove(Object o) {
- if (o == null) return false;
- lock.lock();
- try {
- // Walk the global queue in O(N) to find and unlink the node
- Node curr = checkStateNotNull(globalHead.nextNode);
- while (curr != globalTail) {
- if (o.equals(curr.task)) {
- unlinkNode(curr);
- return true;
- }
- curr = checkStateNotNull(curr.nextNode);
- }
- return false;
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public boolean contains(Object o) {
- if (o == null) return false;
- lock.lock();
- try {
- Node curr = checkStateNotNull(globalHead.nextNode);
- while (curr != globalTail) {
- if (o.equals(curr.task)) {
- return true;
- }
- curr = checkStateNotNull(curr.nextNode);
- }
- return false;
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public int remainingCapacity() {
- return Integer.MAX_VALUE;
- }
-
- @Override
- public int drainTo(Collection<? super Runnable> c) {
- return drainTo(c, Integer.MAX_VALUE);
- }
-
- @Override
- public int drainTo(Collection<? super Runnable> c, int maxElements) {
- if (c == null) throw new NullPointerException();
- if (c == this) throw new IllegalArgumentException();
- if (maxElements <= 0) return 0;
- lock.lock();
- try {
- int added = 0;
- Node curr = checkStateNotNull(globalHead.nextNode);
- while (curr != globalTail && added < maxElements) {
- Node next = checkStateNotNull(curr.nextNode);
- unlinkNode(curr);
- Runnable task = curr.task;
- c.add(task);
- ++added;
- curr = next;
- }
- return added;
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public void clear() {
- lock.lock();
- try {
- Node curr = checkStateNotNull(globalHead.nextNode);
- while (curr != globalTail) {
- Node next = checkStateNotNull(curr.nextNode);
- unlinkNode(curr);
- curr = next;
- }
- } finally {
- lock.unlock();
- }
- }
-
- @Override
- public Iterator<Runnable> iterator() {
- lock.lock();
- try {
- ImmutableList.Builder<Runnable> builder =
ImmutableList.builderWithExpectedSize(size);
- Node curr = checkStateNotNull(globalHead.nextNode);
- while (curr != globalTail) {
- builder.add(curr.task);
- curr = checkStateNotNull(curr.nextNode);
- }
- return builder.build().iterator();
- } finally {
- lock.unlock();
- }
- }
-
- static final class QueueKey {
- private final String computationId;
- private final Work.KeyGroup keyGroup;
-
- QueueKey(String computationId, Work.KeyGroup keyGroup) {
- this.computationId = computationId;
- this.keyGroup = 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) &&
keyGroup.equals(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 b679ca40a95..f3371b54673 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
@@ -2405,10 +2405,10 @@ public class StreamingDataflowWorkerTest {
// Wait until the worker has given up.
int maxTries = 100;
- while (maxTries-- > 0 && !worker.workExecutorIsEmpty()) {
+ while (maxTries-- > 0 && !worker.workExecutorIsEmptyForTesting()) {
Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
}
- assertTrue(worker.workExecutorIsEmpty());
+ assertTrue(worker.workExecutorIsEmptyForTesting());
// Spam worker updates a few times.
maxTries = 10;
@@ -4005,7 +4005,7 @@ public class StreamingDataflowWorkerTest {
.setDaemon(true)
.build(),
/* useFairMonitor= */ false,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ false);
ComputationState computationState =
new ComputationState(
@@ -4067,7 +4067,7 @@ public class StreamingDataflowWorkerTest {
.setDaemon(true)
.build(),
/* useFairMonitor= */ false,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ false);
ComputationState computationState =
new ComputationState(
@@ -4138,7 +4138,7 @@ public class StreamingDataflowWorkerTest {
.setDaemon(true)
.build(),
/* useFairMonitor= */ false,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ false);
ComputationState computationState =
new ComputationState(
@@ -4213,7 +4213,7 @@ public class StreamingDataflowWorkerTest {
.setDaemon(true)
.build(),
/* useFairMonitor= */ false,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ 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 c1408db5b4a..4446d1ea2cb 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
@@ -23,6 +23,7 @@ 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.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
@@ -31,6 +32,9 @@ import java.util.Arrays;
import java.util.Collection;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicIntegerArray;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
@@ -99,28 +103,10 @@ public class BoundedQueueExecutorTest {
private static ExecutableWork createWorkWithCompIdAndKeyGroupAndWorkToken(
String computationId, Work.KeyGroup keyGroup, long workToken,
Consumer<Work> executeWorkFn) {
- WorkItem workItem =
- WorkItem.newBuilder()
- .setKey(ByteString.EMPTY)
- .setShardingKey(1)
- .setWorkToken(workToken)
- .setCacheToken(1)
- .setKeyGroup(
- Windmill.Uint128Proto.newBuilder()
- .setHigh(keyGroup.high())
- .setLow(keyGroup.low())
- .build())
- .build();
- return ExecutableWork.create(
- Work.create(
- workItem,
- workItem.getSerializedSize(),
- Watermarks.builder().setInputDataWatermark(Instant.now()).build(),
- Work.createProcessingContext(
- computationId, new FakeGetDataClient(), ignored -> {},
mock(HeartbeatSender.class)),
- false,
- Instant::now,
- ImmutableList.of()),
+ return createWork(
+ createProcessingContext(computationId),
+ keyGroup,
+ workToken,
(work, handle) -> executeWorkFn.accept(work));
}
@@ -128,11 +114,25 @@ public class BoundedQueueExecutorTest {
String computationId,
Work.KeyGroup keyGroup,
BiConsumer<Work, BoundedQueueExecutorWorkHandle> executeWorkFn) {
+ return createWork(
+ createProcessingContext(computationId), keyGroup, /* workToken= */ 33,
executeWorkFn);
+ }
+
+ private static Work.ProcessingContext createProcessingContext(String
computationId) {
+ return Work.createProcessingContext(
+ computationId, new FakeGetDataClient(), ignored -> {},
mock(HeartbeatSender.class));
+ }
+
+ private static ExecutableWork createWork(
+ Work.ProcessingContext processingContext,
+ Work.KeyGroup keyGroup,
+ long workToken,
+ BiConsumer<Work, BoundedQueueExecutorWorkHandle> executeWorkFn) {
WorkItem workItem =
WorkItem.newBuilder()
.setKey(ByteString.EMPTY)
.setShardingKey(1)
- .setWorkToken(33)
+ .setWorkToken(workToken)
.setCacheToken(1)
.setKeyGroup(
Windmill.Uint128Proto.newBuilder()
@@ -145,8 +145,7 @@ public class BoundedQueueExecutorTest {
workItem,
workItem.getSerializedSize(),
Watermarks.builder().setInputDataWatermark(Instant.now()).build(),
- Work.createProcessingContext(
- computationId, new FakeGetDataClient(), ignored -> {},
mock(HeartbeatSender.class)),
+ processingContext,
false,
Instant::now,
ImmutableList.of()),
@@ -179,7 +178,7 @@ public class BoundedQueueExecutorTest {
.setDaemon(true)
.build(),
useFairMonitor,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ false);
}
@Test
@@ -201,13 +200,13 @@ public class BoundedQueueExecutorTest {
// m1 and m2 have started and all threads are occupied so m3 will be
queued and not executed.
executor.execute(m3, 1);
assertFalse(processStart3.await(1000, TimeUnit.MILLISECONDS));
- assertFalse(executor.executorQueueIsEmpty());
+ assertFalse(executor.executorQueueIsEmptyForTesting());
// Stop m1 so there is an available thread for m3 to run.
processStop1.countDown();
processStart3.await();
// m3 started.
- assertTrue(executor.executorQueueIsEmpty());
+ assertTrue(executor.executorQueueIsEmptyForTesting());
processStop2.countDown();
processStop3.countDown();
executor.shutdown();
@@ -236,7 +235,7 @@ public class BoundedQueueExecutorTest {
assertFalse(processStart2.await(1000, TimeUnit.MILLISECONDS));
// m2 will wait for monitor instead of being queued.
assertEquals(Thread.State.WAITING, m2Runner.getState());
- assertTrue(executor.executorQueueIsEmpty());
+ assertTrue(executor.executorQueueIsEmptyForTesting());
// Stop m1 so there are available bytes for m2 to run.
processStop1.countDown();
@@ -487,55 +486,23 @@ public class BoundedQueueExecutorTest {
@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);
+ BoundedQueueExecutor testExecutor = createMultiKeyBundleExecutor(1, 100);
// 1. Create blocker task to occupy the worker thread
- CountDownLatch blockerStart = new CountDownLatch(1);
CountDownLatch blockerStop = new CountDownLatch(1);
- AtomicReference<BoundedQueueExecutorWorkHandle> blockerHandleRef = new
AtomicReference<>();
- ExecutableWork blockerWork =
- createWorkWithHandle(
- "blockerComp",
- DEFAULT_KEY_GROUP,
- (work, handle) -> {
- blockerHandleRef.set(handle);
- blockerStart.countDown();
- try {
- blockerStop.await();
- } catch (InterruptedException e) {
- throw new RuntimeException(e);
- }
- });
-
- testExecutor.execute(blockerWork, 0);
- blockerStart.await();
- BoundedQueueExecutorWorkHandleImpl stealHandle =
- (BoundedQueueExecutorWorkHandleImpl) blockerHandleRef.get();
- assertNotNull(stealHandle);
+ BoundedQueueExecutorWorkHandleImpl stealHandle =
occupyThread(testExecutor, blockerStop);
// 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 -> {});
+ AtomicInteger work1Runs = new AtomicInteger();
+ AtomicInteger work2Runs = new AtomicInteger();
+ ExecutableWork work1 =
+ createWorkWithCompIdAndKeyGroup("compA", keyGroup1, ignored ->
work1Runs.incrementAndGet());
ExecutableWork work2 =
- createWorkWithCompIdAndKeyGroup(
- "compA",
- keyGroup2,
- ignored -> {
- targetStart.countDown();
- });
+ createWorkWithCompIdAndKeyGroup("compA", keyGroup2, ignored ->
work2Runs.incrementAndGet());
// Enqueue tasks (they will wait in the queue because the thread is
blocked)
testExecutor.execute(work1, 100);
@@ -543,70 +510,66 @@ public class BoundedQueueExecutorTest {
// Total outstanding elements must be 3 (blocker + work1 + work2)
assertEquals(3, testExecutor.elementsOutstanding());
+ assertFalse(testExecutor.executorQueueIsEmptyForTesting());
// Steal work2 using pollWork with compA and keyGroup2
- ExecutableWork stolen = testExecutor.pollWork("compA", keyGroup2,
stealHandle, ignored -> {});
- assertNotNull(stolen);
+ ExecutableWork stolen =
+ testExecutor.pollWork("compA", keyGroup2, stealHandle,
FAILING_FAILED_WORK_HANDLER);
assertEquals(work2, stolen);
// Run the stolen task
stolen.run(stealHandle);
- targetStart.await();
+ assertEquals(1, work2Runs.get());
// Steal work1 using pollWork with compA and keyGroup1
ExecutableWork stolen1 =
testExecutor.pollWork("compA", keyGroup1, stealHandle,
FAILING_FAILED_WORK_HANDLER);
- assertNotNull(stolen1);
assertEquals(work1, stolen1);
+ stolen1.run(stealHandle);
+ assertEquals(1, work1Runs.get());
- // Unblock the blocker and shut down
+ // All queued work was stolen; only stale shells remain in the executor
queue
+ assertTrue(testExecutor.keyGroupWorkIndexIsEmpty());
+ assertTrue(testExecutor.executorQueueIsEmptyForTesting());
+ assertEquals(2, testExecutor.rawQueueSizeForTesting());
+ assertEquals(3, testExecutor.elementsOutstanding());
+ assertEquals(250, testExecutor.bytesOutstanding());
+
+ // Closing the merged handle releases the budget for the blocker and both
stolen tasks
+ stealHandle.close();
+ assertEquals(0, testExecutor.elementsOutstanding());
+ assertEquals(0, testExecutor.bytesOutstanding());
+
+ // Unblock the blocker and shut down; dequeued stale shells are no-ops
blockerStop.countDown();
testExecutor.shutdown();
+ assertEquals(1, work1Runs.get());
+ assertEquals(1, work2Runs.get());
+ assertEquals(0, testExecutor.rawQueueSizeForTesting());
+ assertEquals(0, testExecutor.elementsOutstanding());
+ assertEquals(0, testExecutor.bytesOutstanding());
}
@Test
- public void testPollWorkWithLinkedBlockingQueue() throws Exception {
+ public void testPollWork_multiKeyBundleDisabled() throws Exception {
BoundedQueueExecutor testExecutor =
new BoundedQueueExecutor(
1,
- 60,
+ DEFAULT_THREAD_EXPIRATION_SEC,
TimeUnit.SECONDS,
100,
- 10000000,
- new
ThreadFactoryBuilder().setNameFormat("testLinkedQueue-%d").setDaemon(true).build(),
+ MAXIMUM_BYTES_OUTSTANDING,
+ new
ThreadFactoryBuilder().setNameFormat("testDisabled-%d").setDaemon(true).build(),
useFairMonitor,
- /* useKeyGroupWorkQueue= */ false);
-
- CountDownLatch blockerStart = new CountDownLatch(1);
+ /* multiKeyBundleEnabled= */ false);
CountDownLatch blockerStop = new CountDownLatch(1);
- AtomicReference<BoundedQueueExecutorWorkHandle> blockerHandleRef = new
AtomicReference<>();
- ExecutableWork blockerWork =
- createWorkWithHandle(
- "blockerComp",
- DEFAULT_KEY_GROUP,
- (work, handle) -> {
- blockerHandleRef.set(handle);
- blockerStart.countDown();
- try {
- blockerStop.await();
- } catch (InterruptedException e) {
- throw new RuntimeException(e);
- }
- });
-
- testExecutor.execute(blockerWork, 0);
- blockerStart.await();
- BoundedQueueExecutorWorkHandleImpl stealHandle =
- (BoundedQueueExecutorWorkHandleImpl) blockerHandleRef.get();
- assertNotNull(stealHandle);
+ BoundedQueueExecutorWorkHandleImpl stealHandle =
occupyThread(testExecutor, blockerStop);
Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1);
ExecutableWork work = createWorkWithCompIdAndKeyGroup("compA", keyGroup,
ignored -> {});
testExecutor.execute(work, 100);
- ExecutableWork stolen =
- testExecutor.pollWork("compA", keyGroup, stealHandle,
FAILING_FAILED_WORK_HANDLER);
- assertNull(stolen);
+ assertNull(testExecutor.pollWork("compA", keyGroup, stealHandle,
FAILING_FAILED_WORK_HANDLER));
blockerStop.countDown();
testExecutor.shutdown();
@@ -614,39 +577,10 @@ public class BoundedQueueExecutorTest {
@Test
public void testPollWork_skipsFailedWorkAndCallsOnFailedWorkHandler() throws
Exception {
- BoundedQueueExecutor testExecutor =
- new BoundedQueueExecutor(
- 1,
- 60,
- TimeUnit.SECONDS,
- 100,
- 10000000,
- new
ThreadFactoryBuilder().setNameFormat("testPollWork-%d").setDaemon(true).build(),
- useFairMonitor,
- /* useKeyGroupWorkQueue= */ true);
-
- CountDownLatch blockerStart = new CountDownLatch(1);
+ BoundedQueueExecutor testExecutor = createMultiKeyBundleExecutor(1, 100);
CountDownLatch blockerStop = new CountDownLatch(1);
- AtomicReference<BoundedQueueExecutorWorkHandle> blockerHandleRef = new
AtomicReference<>();
- ExecutableWork blockerWork =
- createWorkWithHandle(
- "compA",
- DEFAULT_KEY_GROUP,
- (work, handle) -> {
- blockerHandleRef.set(handle);
- blockerStart.countDown();
- try {
- blockerStop.await();
- } catch (InterruptedException e) {
- throw new RuntimeException(e);
- }
- });
-
- testExecutor.execute(blockerWork, 10);
- blockerStart.await();
BoundedQueueExecutorWorkHandleImpl stealHandle =
- (BoundedQueueExecutorWorkHandleImpl) blockerHandleRef.get();
- assertNotNull(stealHandle);
+ occupyThread(testExecutor, blockerStop, /* bytes= */ 10);
Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1);
FailedWorkHandler onFailedWorkHandler = mock(FailedWorkHandler.class);
@@ -658,7 +592,7 @@ public class BoundedQueueExecutorTest {
ExecutableWork work3 =
createWorkWithCompIdAndKeyGroupAndWorkToken("compA", keyGroup, 103,
ignored -> {});
- // Enqueue both tasks (they will wait in the queue because the thread is
blocked).
+ // Enqueue tasks (they will wait in the queue because the thread is
blocked).
testExecutor.execute(work1, 100);
testExecutor.execute(work2, 150);
testExecutor.execute(work3, 200);
@@ -674,9 +608,7 @@ public class BoundedQueueExecutorTest {
// onFailedWorkHandler callback, and return work2.
ExecutableWork stolen =
testExecutor.pollWork("compA", keyGroup, stealHandle,
onFailedWorkHandler);
- assertNotNull(stolen);
assertEquals(work2, stolen);
-
verify(onFailedWorkHandler).onFailedWork(work1.work());
// Verify stealHandle merged (blockerWork: 10 bytes, work2: 150 bytes).
@@ -693,9 +625,229 @@ public class BoundedQueueExecutorTest {
assertNull(testExecutor.pollWork("compA", keyGroup, stealHandle,
onFailedWorkHandler));
blockerStop.countDown();
- stealHandle.close();
+ testExecutor.shutdown();
+ assertEquals(0, testExecutor.elementsOutstanding());
+ assertEquals(0, testExecutor.bytesOutstanding());
+ }
+
+ @Test
+ public void testPollWork_doesNotStealWorkStartedByExecutor() throws
Exception {
+ BoundedQueueExecutor testExecutor = createMultiKeyBundleExecutor(1, 1000);
+ Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1);
+ CountDownLatch release = new CountDownLatch(1);
+ BoundedQueueExecutorWorkHandleImpl handle =
+ occupyThread(testExecutor, "compA", keyGroup, release, /* bytes= */
100);
+
+ assertTrue(testExecutor.keyGroupWorkIndexIsEmpty());
+ assertNull(testExecutor.pollWork("compA", keyGroup, handle,
FAILING_FAILED_WORK_HANDLER));
+
+ release.countDown();
+ testExecutor.shutdown();
+ assertEquals(0, testExecutor.elementsOutstanding());
+ assertEquals(0, testExecutor.bytesOutstanding());
+ }
+
+ @Test
+ public void testPollWork_doesNotStealWorkWithMultiKeyBatchingDisabled()
throws Exception {
+ BoundedQueueExecutor testExecutor = createMultiKeyBundleExecutor(1, 1000);
+ CountDownLatch releaseBlocker = new CountDownLatch(1);
+ BoundedQueueExecutorWorkHandleImpl stealHandle =
occupyThread(testExecutor, releaseBlocker);
+
+ Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1);
+ AtomicBoolean disabledWorkRan = new AtomicBoolean();
+ ExecutableWork disabledWork =
+ createWorkWithCompIdAndKeyGroup("compA", keyGroup, ignored ->
disabledWorkRan.set(true));
+ disabledWork.work().setMultiKeyBatchingDisabled(true);
+ ExecutableWork enabledWork = createWorkWithCompIdAndKeyGroup("compA",
keyGroup, ignored -> {});
+ testExecutor.execute(disabledWork, 100);
+ testExecutor.execute(enabledWork, 100);
+
+ assertEquals(
+ enabledWork,
+ testExecutor.pollWork("compA", keyGroup, stealHandle,
FAILING_FAILED_WORK_HANDLER));
+ assertNull(testExecutor.pollWork("compA", keyGroup, stealHandle,
FAILING_FAILED_WORK_HANDLER));
+
+ releaseBlocker.countDown();
+ testExecutor.shutdown();
+ assertTrue(disabledWorkRan.get());
+ assertEquals(0, testExecutor.elementsOutstanding());
+ }
+
+ @Test
+ public void testPollWork_purgesStaleShellsToBoundExecutorQueue() throws
Exception {
+ int maximumElementsOutstanding = 4;
+ BoundedQueueExecutor testExecutor = createMultiKeyBundleExecutor(2,
maximumElementsOutstanding);
+ CountDownLatch releaseBundles = new CountDownLatch(1);
+ BoundedQueueExecutorWorkHandleImpl bundle1 = occupyThread(testExecutor,
releaseBundles);
+ BoundedQueueExecutorWorkHandleImpl bundle2 = occupyThread(testExecutor,
releaseBundles);
+
+ Work.KeyGroup keyGroup = Work.KeyGroup.create(1, 1);
+ int numWorks = 20;
+ for (int i = 0; i < numWorks; i++) {
+ testExecutor.forceExecute(
+ createWorkWithCompIdAndKeyGroup("compA", keyGroup, ignored -> {}),
10);
+ }
+
+ // Each steal leaves a stale shell in the executor queue. While bundle1 is
open, its stolen
+ // work is still counted in elementsOutstanding, so there is nothing to
purge.
+ int numStolenByBundle1 = maximumElementsOutstanding + 2;
+ for (int i = 0; i < numStolenByBundle1; i++) {
+ assertNotNull(testExecutor.pollWork("compA", keyGroup, bundle1,
FAILING_FAILED_WORK_HANDLER));
+ }
+ // 2 = bundle1 + bundle2.
+ assertEquals(2 + numWorks, testExecutor.elementsOutstanding());
+ assertEquals(numWorks, testExecutor.rawQueueSizeForTesting());
+
+ // Closing bundle1 leaves its 6 stale shells uncounted. The executor queue
now exceeds
+ // elementsOutstanding (bundle2 + unstarted work) by more than
maximumElementsOutstanding, so
+ // the next steal purges every stale shell.
+ bundle1.close();
+ int numStolenByBundle2 = 1;
+ assertNotNull(testExecutor.pollWork("compA", keyGroup, bundle2,
FAILING_FAILED_WORK_HANDLER));
+ // 1 = bundle2.
+ assertEquals(1 + numWorks - numStolenByBundle1,
testExecutor.elementsOutstanding());
+ assertEquals(
+ numWorks - numStolenByBundle1 - numStolenByBundle2,
testExecutor.rawQueueSizeForTesting());
+
+ releaseBundles.countDown();
+ testExecutor.shutdown();
+ assertEquals(0, testExecutor.rawQueueSizeForTesting());
+ assertEquals(0, testExecutor.elementsOutstanding());
+ assertEquals(0, testExecutor.bytesOutstanding());
+ }
+
+ @Test
+ public void testConcurrentExecuteAndPollWork_processesEachWorkExactlyOnce()
throws Exception {
+ int numThreads = 4;
+ int numWorks = 4000;
+ int maxStolenPerBundle = 10;
+ BoundedQueueExecutor testExecutor =
+ createMultiKeyBundleExecutor(numThreads, /*
maximumElementsOutstanding= */ 8);
+ Work.ProcessingContext processingContext =
createProcessingContext("compA");
+ ImmutableList<Work.KeyGroup> keyGroups =
+ ImmutableList.of(
+ Work.KeyGroup.create(1, 1),
+ Work.KeyGroup.create(1, 2),
+ Work.KeyGroup.create(2, 1),
+ Work.KeyGroup.create(2, 2));
+
+ AtomicIntegerArray processedCounts = new AtomicIntegerArray(numWorks);
+ AtomicInteger numStolen = new AtomicInteger();
+ CountDownLatch allProcessed = new CountDownLatch(numWorks);
+ BiConsumer<Work, BoundedQueueExecutorWorkHandle> bundleFn =
+ (work, handle) -> {
+ processedCounts.incrementAndGet((int)
work.getWorkItem().getWorkToken());
+ allProcessed.countDown();
+ for (int i = 0; i < maxStolenPerBundle; i++) {
+ ExecutableWork stolen =
+ testExecutor.pollWork(
+ work.getComputationId(),
+ work.getKeyGroup(),
+ handle,
+ FAILING_FAILED_WORK_HANDLER);
+ if (stolen == null) {
+ break;
+ }
+ numStolen.incrementAndGet();
+ processedCounts.incrementAndGet((int)
stolen.work().getWorkItem().getWorkToken());
+ allProcessed.countDown();
+ }
+ };
+
+ // Hold all threads until half of the work is queued so bundles steal from
a backlog.
+ CountDownLatch releaseBlockers = new CountDownLatch(1);
+ for (int i = 0; i < numThreads; i++) {
+ occupyThread(testExecutor, releaseBlockers);
+ }
+ for (int workToken = 0; workToken < numWorks; workToken++) {
+ if (workToken == numWorks / 2) {
+ releaseBlockers.countDown();
+ }
+ testExecutor.forceExecute(
+ createWork(
+ processingContext, keyGroups.get(workToken % keyGroups.size()),
workToken, bundleFn),
+ 1);
+ }
+ allProcessed.await();
+ testExecutor.shutdown();
+
+ for (int workToken = 0; workToken < numWorks; workToken++) {
+ assertEquals("work token " + workToken, 1,
processedCounts.get(workToken));
+ }
+ assertThat(numStolen.get(), greaterThan(0));
+ assertEquals(0, testExecutor.rawQueueSizeForTesting());
+ assertTrue(testExecutor.keyGroupWorkIndexIsEmpty());
assertEquals(0, testExecutor.elementsOutstanding());
assertEquals(0, testExecutor.bytesOutstanding());
+ }
+
+ @Test
+ public void testExecuteAfterShutdown_rejectedWorkIsNotStealable() throws
Exception {
+ BoundedQueueExecutor testExecutor = createMultiKeyBundleExecutor(1, 1000);
testExecutor.shutdown();
+
+ ExecutableWork work =
+ createWorkWithCompIdAndKeyGroup(
+ "compA", Work.KeyGroup.create(1, 1), ignored -> Assert.fail());
+ assertThrows(RuntimeException.class, () -> testExecutor.execute(work,
100));
+
+ assertTrue(testExecutor.keyGroupWorkIndexIsEmpty());
+ assertEquals(0, testExecutor.elementsOutstanding());
+ assertEquals(0, testExecutor.bytesOutstanding());
+ }
+
+ private BoundedQueueExecutor createMultiKeyBundleExecutor(
+ int maxThreads, int maximumElementsOutstanding) {
+ return new BoundedQueueExecutor(
+ maxThreads,
+ DEFAULT_THREAD_EXPIRATION_SEC,
+ TimeUnit.SECONDS,
+ maximumElementsOutstanding,
+ MAXIMUM_BYTES_OUTSTANDING,
+ new
ThreadFactoryBuilder().setNameFormat("testMultiKeyBundle-%d").setDaemon(true).build(),
+ useFairMonitor,
+ /* multiKeyBundleEnabled= */ true);
+ }
+
+ private static BoundedQueueExecutorWorkHandleImpl occupyThread(
+ BoundedQueueExecutor testExecutor, CountDownLatch release) throws
InterruptedException {
+ return occupyThread(testExecutor, release, /* bytes= */ 0);
+ }
+
+ private static BoundedQueueExecutorWorkHandleImpl occupyThread(
+ BoundedQueueExecutor testExecutor, CountDownLatch release, long bytes)
+ throws InterruptedException {
+ return occupyThread(testExecutor, "blockerComp", Work.KeyGroup.DEFAULT,
release, bytes);
+ }
+
+ private static BoundedQueueExecutorWorkHandleImpl occupyThread(
+ BoundedQueueExecutor testExecutor,
+ String computationId,
+ Work.KeyGroup keyGroup,
+ CountDownLatch release,
+ long bytes)
+ throws InterruptedException {
+ CountDownLatch started = new CountDownLatch(1);
+ AtomicReference<BoundedQueueExecutorWorkHandle> handle = new
AtomicReference<>();
+ testExecutor.execute(
+ createWorkWithHandle(
+ computationId,
+ keyGroup,
+ (work, workHandle) -> {
+ handle.set(workHandle);
+ started.countDown();
+ awaitUninterruptibly(release);
+ }),
+ bytes);
+ started.await();
+ return (BoundedQueueExecutorWorkHandleImpl) handle.get();
+ }
+
+ private static void awaitUninterruptibly(CountDownLatch latch) {
+ try {
+ latch.await();
+ } catch (InterruptedException e) {
+ throw new RuntimeException(e);
+ }
}
}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkIndexTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkIndexTest.java
new file mode 100644
index 00000000000..f41b0da4eec
--- /dev/null
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkIndexTest.java
@@ -0,0 +1,319 @@
+/*
+ * 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.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicIntegerArray;
+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;
+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.collect.ImmutableList;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
+import org.joda.time.Instant;
+import org.junit.After;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class KeyGroupWorkIndexTest {
+
+ private static final String COMP_A = "compA";
+ private static final String COMP_B = "compB";
+ private static final Work.KeyGroup KEY_GROUP_1 = Work.KeyGroup.create(1, 1);
+ private static final Work.KeyGroup KEY_GROUP_2 = Work.KeyGroup.create(1, 2);
+
+ @Rule public transient Timeout globalTimeout = Timeout.seconds(300);
+
+ // Creates the budget handles of the QueuedWork created by the tests. Its
executor is never used.
+ private final BoundedQueueExecutor owner =
+ new BoundedQueueExecutor(
+ 1,
+ 60,
+ TimeUnit.SECONDS,
+ 100,
+ 10000000,
+ new ThreadFactoryBuilder()
+ .setNameFormat("KeyGroupWorkIndexTest-%d")
+ .setDaemon(true)
+ .build(),
+ /* useFairMonitor= */ false,
+ /* multiKeyBundleEnabled= */ true);
+ private final KeyGroupWorkIndex index = new KeyGroupWorkIndex();
+ private long nextWorkToken = 0;
+
+ @After
+ public void tearDown() throws InterruptedException {
+ owner.shutdown();
+ }
+
+ private QueuedWork createQueuedWork(String computationId, Work.KeyGroup
keyGroup) {
+ long workToken = nextWorkToken++;
+ WorkItem workItem =
+ WorkItem.newBuilder()
+ .setKey(ByteString.EMPTY)
+ .setShardingKey(workToken)
+ .setWorkToken(workToken)
+ .setCacheToken(1)
+ .setKeyGroup(
+ Windmill.Uint128Proto.newBuilder()
+ .setHigh(keyGroup.high())
+ .setLow(keyGroup.low())
+ .build())
+ .build();
+ Work.ProcessingContext processingContext =
+ Work.createProcessingContext(
+ computationId, new FakeGetDataClient(), ignored -> {},
mock(HeartbeatSender.class));
+ ExecutableWork work =
+ ExecutableWork.create(
+ Work.create(
+ workItem,
+ workItem.getSerializedSize(),
+
Watermarks.builder().setInputDataWatermark(Instant.now()).build(),
+ processingContext,
+ false,
+ Instant::now,
+ ImmutableList.of()),
+ (w, handle) -> {});
+ QueuedWork queuedWork =
+ new QueuedWork(work, owner.createBudgetHandle(work.work(), /* bytes=
*/ 0), index);
+ assertTrue(queuedWork.isStealable());
+ return queuedWork;
+ }
+
+ private static int workToken(QueuedWork queuedWork) {
+ return (int) queuedWork.getWork().work().getWorkItem().getWorkToken();
+ }
+
+ @Test
+ public void testPollAndClaim_returnsWorkInFifoOrder() {
+ QueuedWork work1 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ QueuedWork work2 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ QueuedWork work3 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ index.add(work1);
+ index.add(work2);
+ index.add(work3);
+
+ assertSame(work1, index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertSame(work2, index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertSame(work3, index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertNull(index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertTrue(work1.isClaimed());
+ assertTrue(work2.isClaimed());
+ assertTrue(work3.isClaimed());
+ assertTrue(index.isEmpty());
+ }
+
+ @Test
+ public void testPollAndClaim_isolatesComputationsAndKeyGroups() {
+ QueuedWork workA1 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ QueuedWork workA2 = createQueuedWork(COMP_A, KEY_GROUP_2);
+ QueuedWork workB1 = createQueuedWork(COMP_B, KEY_GROUP_1);
+ index.add(workA1);
+ index.add(workA2);
+ index.add(workB1);
+
+ assertNull(index.pollAndClaim(COMP_A, Work.KeyGroup.create(3, 4)));
+ assertNull(index.pollAndClaim(COMP_B, KEY_GROUP_2));
+ assertSame(workA2, index.pollAndClaim(COMP_A, KEY_GROUP_2));
+ assertNull(index.pollAndClaim(COMP_A, KEY_GROUP_2));
+ assertSame(workB1, index.pollAndClaim(COMP_B, KEY_GROUP_1));
+ assertNull(index.pollAndClaim(COMP_B, KEY_GROUP_1));
+ assertSame(workA1, index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertNull(index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertTrue(index.isEmpty());
+ }
+
+ @Test
+ public void testPollAndClaim_skipsWorkClaimedByExecutor() {
+ QueuedWork work1 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ QueuedWork work2 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ index.add(work1);
+ index.add(work2);
+
+ assertTrue(work1.tryClaimForExecution());
+
+ assertSame(work2, index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertNull(index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertTrue(index.isEmpty());
+ index.drainClaimed(work1);
+ assertTrue(index.isEmpty());
+
+ QueuedWork work3 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ QueuedWork work4 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ index.add(work3);
+ index.add(work4);
+
+ assertTrue(work4.tryClaimForExecution());
+ index.drainClaimed(work4);
+ assertFalse(index.isEmpty());
+
+ assertSame(work3, index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertTrue(index.isEmpty());
+ }
+
+ @Test
+ public void testAdd_skipsClaimedWork() {
+ QueuedWork work = createQueuedWork(COMP_A, KEY_GROUP_1);
+ assertTrue(work.tryClaimForExecution());
+
+ index.add(work);
+
+ assertTrue(index.isEmpty());
+ assertNull(index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ }
+
+ @Test
+ public void testDrainClaimed_drainsFromHead() {
+ QueuedWork work1 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ QueuedWork work2 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ index.add(work1);
+ index.add(work2);
+
+ assertTrue(work2.tryClaimForExecution());
+ index.drainClaimed(work2);
+ assertFalse(index.isEmpty());
+ assertTrue(work1.tryClaimForExecution());
+ index.drainClaimed(work1);
+ assertTrue(index.isEmpty());
+
+ QueuedWork work3 = createQueuedWork(COMP_A, KEY_GROUP_1);
+ index.add(work3);
+ assertSame(work3, index.pollAndClaim(COMP_A, KEY_GROUP_1));
+ assertTrue(index.isEmpty());
+ }
+
+ @Test
+ public void
testConcurrentAddDrainAndPollAndClaim_claimsEachWorkExactlyOnce() throws
Exception {
+ int numProducers = 2;
+ int numExecutors = 2;
+ int numStealers = 2;
+ int numWorksPerProducer = 10000;
+ ImmutableList<String> computations = ImmutableList.of(COMP_A, COMP_B);
+ ImmutableList<Work.KeyGroup> keyGroups =
+ ImmutableList.of(KEY_GROUP_1, KEY_GROUP_2, Work.KeyGroup.create(2, 1));
+
+ List<List<QueuedWork>> worksByProducer = new ArrayList<>();
+ for (int producer = 0; producer < numProducers; producer++) {
+ List<QueuedWork> works = new ArrayList<>();
+ for (int i = 0; i < numWorksPerProducer; i++) {
+ works.add(
+ createQueuedWork(
+ computations.get(i % computations.size()), keyGroups.get(i %
keyGroups.size())));
+ }
+ worksByProducer.add(works);
+ }
+ int numWorks = numProducers * numWorksPerProducer;
+ // Simulates the executor queue: work is queued before it is indexed.
+ ConcurrentLinkedQueue<QueuedWork> executorQueue = new
ConcurrentLinkedQueue<>();
+ AtomicIntegerArray executedCounts = new AtomicIntegerArray(numWorks);
+ AtomicIntegerArray stolenCounts = new AtomicIntegerArray(numWorks);
+ AtomicBoolean producersDone = new AtomicBoolean(false);
+ AtomicBoolean executorsDone = new AtomicBoolean(false);
+
+ ExecutorService threads =
+ Executors.newFixedThreadPool(numProducers + numExecutors +
numStealers);
+ List<Future<?>> producerFutures = new ArrayList<>();
+ for (List<QueuedWork> works : worksByProducer) {
+ producerFutures.add(
+ threads.submit(
+ () -> {
+ for (QueuedWork work : works) {
+ executorQueue.add(work);
+ index.add(work);
+ }
+ }));
+ }
+ List<Future<?>> executorFutures = new ArrayList<>();
+ for (int i = 0; i < numExecutors; i++) {
+ executorFutures.add(
+ threads.submit(
+ () -> {
+ while (!producersDone.get() || !executorQueue.isEmpty()) {
+ QueuedWork work = executorQueue.poll();
+ if (work != null && work.tryClaimForExecution()) {
+ index.drainClaimed(work);
+ executedCounts.incrementAndGet(workToken(work));
+ }
+ }
+ }));
+ }
+ List<Future<?>> stealerFutures = new ArrayList<>();
+ for (int i = 0; i < numStealers; i++) {
+ stealerFutures.add(
+ threads.submit(
+ () -> {
+ while (!executorsDone.get()) {
+ ThreadLocalRandom random = ThreadLocalRandom.current();
+ QueuedWork work =
+ index.pollAndClaim(
+
computations.get(random.nextInt(computations.size())),
+ keyGroups.get(random.nextInt(keyGroups.size())));
+ if (work != null) {
+ stolenCounts.incrementAndGet(workToken(work));
+ }
+ }
+ }));
+ }
+
+ for (Future<?> future : producerFutures) {
+ future.get();
+ }
+ producersDone.set(true);
+ for (Future<?> future : executorFutures) {
+ future.get();
+ }
+ executorsDone.set(true);
+ for (Future<?> future : stealerFutures) {
+ future.get();
+ }
+ threads.shutdown();
+ assertTrue(threads.awaitTermination(30, TimeUnit.SECONDS));
+
+ for (int workToken = 0; workToken < numWorks; workToken++) {
+ assertEquals(
+ "work token " + workToken,
+ 1,
+ executedCounts.get(workToken) + stolenCounts.get(workToken));
+ }
+ assertTrue(index.isEmpty());
+ }
+}
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
deleted file mode 100644
index a58851cf370..00000000000
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java
+++ /dev/null
@@ -1,523 +0,0 @@
-/*
- * 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.lang.Thread.State;
-import java.util.ArrayList;
-import java.util.Arrays;
-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.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.collect.ImmutableList;
-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;
-import org.junit.runner.RunWith;
-import org.junit.runners.Parameterized;
-import org.junit.runners.Parameterized.Parameters;
-
-@RunWith(Parameterized.class)
-public class KeyGroupWorkQueueTest {
-
- @Parameters(name = "fairQueue={0}")
- public static Iterable<Object[]> data() {
- return Arrays.asList(new Object[][] {{true}, {false}});
- }
-
- @Parameterized.Parameter public boolean fairQueue;
- private BoundedQueueExecutor executor;
-
- @Before
- public void setUp() {
- executor =
- new BoundedQueueExecutor(
- 2,
- 60,
- TimeUnit.SECONDS,
- 100,
- 10000000,
- new
ThreadFactoryBuilder().setNameFormat("Test-%d").setDaemon(true).build(),
- fairQueue,
- /* useKeyGroupWorkQueue= */ true);
- }
-
- private static final Work.KeyGroup TEST_KEY_GROUP = Work.KeyGroup.create(1,
2);
-
- private QueuedWork createQueuedWork(String computationId, long workBytes) {
- return createQueuedWork(computationId, TEST_KEY_GROUP, workBytes);
- }
-
- private QueuedWork createQueuedWork(
- String computationId, Work.@Nullable KeyGroup keyGroup, long workBytes) {
- WorkItem.Builder workItemBuilder =
- WorkItem.newBuilder()
- .setKey(ByteString.EMPTY)
- .setShardingKey(1)
- .setWorkToken(33)
- .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(
- workItem,
- workItem.getSerializedSize(),
-
Watermarks.builder().setInputDataWatermark(Instant.now()).build(),
- Work.createProcessingContext(
- computationId,
- new FakeGetDataClient(),
- ignored -> {},
- mock(HeartbeatSender.class)),
- false,
- Instant::now,
- ImmutableList.of()),
- (w, h) -> {});
- return new QueuedWork(work, executor.createBudgetHandle(work.work(),
workBytes));
- }
-
- private static class NoOpRunnable implements Runnable {
- final String id;
-
- NoOpRunnable(String id) {
- this.id = id;
- }
-
- @Override
- public void run() {}
-
- @Override
- public String toString() {
- return "NoOpRunnable(" + id + ")";
- }
- }
-
- @Test
- public void testBasicOfferAndPoll() {
- KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue);
- assertTrue(queue.isEmpty());
- assertEquals(0, queue.size());
-
- NoOpRunnable task1 = new NoOpRunnable("1");
- NoOpRunnable task2 = new NoOpRunnable("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(fairQueue);
- NoOpRunnable task1 = new NoOpRunnable("1");
- NoOpRunnable task2 = new NoOpRunnable("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(fairQueue);
- NoOpRunnable task1 = new NoOpRunnable("1");
- NoOpRunnable task2 = new NoOpRunnable("2");
- queue.offer(task1);
- queue.offer(task2);
-
- List<Runnable> 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(fairQueue);
- NoOpRunnable task1 = new NoOpRunnable("1");
- NoOpRunnable task2 = new NoOpRunnable("2");
- queue.offer(task1);
- queue.offer(task2);
-
- Iterator<Runnable> 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(fairQueue);
-
- 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());
- 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());
- 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());
-
- assertEquals(null, queue.poll());
- }
-
- @Test
- public void testConcurrentStress() throws InterruptedException,
ExecutionException {
- 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);
- CountDownLatch startLatch = new CountDownLatch(1);
- CountDownLatch producersDoneLatch = new CountDownLatch(producerThreads);
- CountDownLatch consumersDoneLatch = new CountDownLatch(consumerThreads);
- CountDownLatch consumedLatch = new CountDownLatch(totalTasks);
- List<Future<?>> futures = new ArrayList<>();
-
- // Start consumers
- for (int i = 0; i < consumerThreads; i++) {
- int consumerId = i;
- futures.add(
- executorService.submit(
- () -> {
- try {
- startLatch.await();
- 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 {
- consumersDoneLatch.countDown();
- }
- }));
- }
- // 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 {
- producersDoneLatch.countDown();
- }
- }));
- }
-
- // Release the start latch to start the test
- startLatch.countDown();
-
- // 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) {
- future.get();
- }
-
- executorService.shutdown();
- assertTrue(executorService.awaitTermination(30, TimeUnit.SECONDS));
-
- assertEquals(0, queue.size());
- assertTrue(queue.isEmpty());
- }
-
- @Test
- public void testTakeBlocksAndWakesUp() throws InterruptedException {
- final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue);
- final NoOpRunnable task = new NoOpRunnable("take-task");
- final AtomicReference<Runnable> 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(30, TimeUnit.SECONDS));
- waitForThreadState(t, State.WAITING);
-
- queue.offer(task);
-
- assertTrue(finished.await(30, TimeUnit.SECONDS));
- assertEquals(task, result.get());
- }
-
- @Test
- public void testPollWithTimeout() throws InterruptedException {
- final KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue);
- final NoOpRunnable task = new NoOpRunnable("poll-task");
- final AtomicReference<Runnable> 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(30, TimeUnit.SECONDS));
- waitForThreadState(t1, State.TIMED_WAITING);
-
- assertTrue(finished.await(30, 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<Runnable> 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(30, TimeUnit.SECONDS));
- waitForThreadState(t2, State.TIMED_WAITING);
-
- queue.offer(task);
-
- assertTrue(finished2.await(30, TimeUnit.SECONDS));
- assertEquals(task, result2.get());
- }
-
- @Test
- public void testPollWorkWithKeyGroup() {
- KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue);
-
- 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);
-
- queue.offer(workA1);
- queue.offer(workA2);
-
- 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);
- 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());
- polledNotExist = queue.pollWork("compA", keyGroupNotExist);
- assertNull(polledNotExist);
- assertTrue(queue.isEmpty());
- }
-
- @Test
- public void testOffer_multiKeyBatchingDisabled_notInsertedInKeyGroupQueue() {
- KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue);
- QueuedWork workDisabled = createQueuedWork("compA", 100);
- workDisabled.getWork().work().setMultiKeyBatchingDisabled(true);
- QueuedWork workEnabled = createQueuedWork("compA", 200);
-
- queue.offer(workDisabled);
- queue.offer(workEnabled);
- assertEquals(2, queue.size());
-
- QueuedWork polledWork = queue.pollWork("compA", TEST_KEY_GROUP);
- assertNotNull(polledWork);
- assertEquals(workEnabled, polledWork);
- assertEquals(1, queue.size());
-
- assertNull(queue.pollWork("compA", TEST_KEY_GROUP));
- assertEquals(workDisabled, queue.poll());
- 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);
- }
- }
-}
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 a11fb1f3ad3..688bad933e0 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
@@ -63,7 +63,7 @@ public class StreamingCommitFinalizerTest {
.setDaemon(true)
.build(),
/* useFairMonitor= */ false,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ 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 6f440c9cf30..f11add07f6f 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
@@ -68,7 +68,7 @@ public class WorkFailureProcessorTest {
.setDaemon(true)
.build(),
/* useFairMonitor= */ false,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ 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 017a22ca684..8ad5896dfab 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
@@ -75,7 +75,7 @@ public class ActiveWorkRefresherTest {
10000000,
new
ThreadFactoryBuilder().setNameFormat("DataflowWorkUnits-%d").setDaemon(true).build(),
/* useFairMonitor= */ false,
- /* useKeyGroupWorkQueue= */ false);
+ /* multiKeyBundleEnabled= */ false);
}
private static ComputationState createComputationState(int
computationIdSuffix) {
