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 fd00cea775f2..dd96142c41a9 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 @@ private static BoundedQueueExecutor createWorkUnitExecutor( 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 @@ private static BoundedQueueExecutor createWorkUnitExecutor( 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 @@ void reportPeriodicWorkerUpdatesForTest() { } @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 2acee9410fa3..724668b689f1 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.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 boolean equals(@Nullable Object o) { @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 046d8cae9f9d..1d10d074b115 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.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.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. + * + *

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 @@ private static class Budget { private final ConcurrentLinkedQueue 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 @@ private static class Budget { @GuardedBy("this") private long totalTimeMaxActiveThreadsUsed; - // If set the keyGroupWorkQueue is used by the underlying executor. - private final @Nullable KeyGroupWorkQueue keyGroupWorkQueue; + private final LinkedBlockingQueue workQueue = new LinkedBlockingQueue<>(); + private final @Nullable KeyGroupWorkIndex keyGroupWorkIndex; + + private final AtomicBoolean purgingStaleShells = new AtomicBoolean(false); public BoundedQueueExecutor( int initialMaximumPoolSize, @@ -93,8 +103,8 @@ public 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 BoundedQueueExecutor( initialMaximumPoolSize, keepAliveTime, unit, - keyGroupWorkQueue != null ? keyGroupWorkQueue : new LinkedBlockingQueue<>(), + workQueue, threadFactory) { @Override protected void beforeExecute(Thread t, Runnable r) { @@ -139,9 +149,9 @@ public void execute(ExecutableWork work, long workBytes) { 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 void shutdown() throws InterruptedException { } } - 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 long bytesOutstanding() { } 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 String summaryHtml() { builder.append("
/n"); builder.append("Work Queue Size: "); - builder.append(elementsOutstanding); + builder.append(elementsOutstanding.get()); builder.append("/"); builder.append(maximumElementsOutstanding()); builder.append("
/n"); @@ -328,49 +346,132 @@ public synchronized void close() { } } + /** + * 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}. + * + *

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 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 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(Work work, long bytes) { 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. + * + *

{@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 @@ private void decrementCounters(int elements, long bytes) { } 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 000000000000..10cdab083b5f --- /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 shards; + + KeyGroupWorkIndex() { + ImmutableList.Builder 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> queuesByKeyGroup = new HashMap<>(); + + private synchronized void add(QueueKey queueKey, QueuedWork work) { + if (work.isClaimed()) { + return; + } + ArrayDeque queue = + queuesByKeyGroup.computeIfAbsent( + queueKey, unused -> new ArrayDeque<>(INITIAL_QUEUE_CAPACITY)); + drainClaimedFromHead(queue); + queue.addLast(work); + } + + private synchronized void drainClaimed(QueueKey queueKey) { + @Nullable ArrayDeque 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 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 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 dd409616ab9d..000000000000 --- 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 implements BlockingQueue { - - 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 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 c) { - return drainTo(c, Integer.MAX_VALUE); - } - - @Override - public int drainTo(Collection c, int maxElements) { - if (c == null) throw new NullPointerException(); - if (c == this) throw new IllegalArgumentException(); - if (maxElements <= 0) return 0; - lock.lock(); - try { - int added = 0; - Node curr = 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 iterator() { - lock.lock(); - try { - ImmutableList.Builder 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 b679ca40a956..f3371b54673f 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 void testExceptions() throws Exception { // 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 void testMaxThreadMetric() throws Exception { .setDaemon(true) .build(), /* useFairMonitor= */ false, - /* useKeyGroupWorkQueue= */ false); + /* multiKeyBundleEnabled= */ false); ComputationState computationState = new ComputationState( @@ -4067,7 +4067,7 @@ public void testActiveThreadMetric() throws Exception { .setDaemon(true) .build(), /* useFairMonitor= */ false, - /* useKeyGroupWorkQueue= */ false); + /* multiKeyBundleEnabled= */ false); ComputationState computationState = new ComputationState( @@ -4138,7 +4138,7 @@ public void testOutstandingBytesMetric() throws Exception { .setDaemon(true) .build(), /* useFairMonitor= */ false, - /* useKeyGroupWorkQueue= */ false); + /* multiKeyBundleEnabled= */ false); ComputationState computationState = new ComputationState( @@ -4213,7 +4213,7 @@ public void testOutstandingBundlesMetric() throws Exception { .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 c1408db5b4a6..4446d1ea2cbb 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.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.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 @@ private static ExecutableWork createWorkWithCompIdAndKeyGroup( private static ExecutableWork createWorkWithCompIdAndKeyGroupAndWorkToken( String computationId, Work.KeyGroup keyGroup, long workToken, Consumer 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 @@ private static ExecutableWork createWorkWithHandle( String computationId, Work.KeyGroup keyGroup, BiConsumer 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 executeWorkFn) { WorkItem workItem = WorkItem.newBuilder() .setKey(ByteString.EMPTY) .setShardingKey(1) - .setWorkToken(33) + .setWorkToken(workToken) .setCacheToken(1) .setKeyGroup( Windmill.Uint128Proto.newBuilder() @@ -145,8 +145,7 @@ private static ExecutableWork createWorkWithHandle( 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 void setUp() { .setDaemon(true) .build(), useFairMonitor, - /* useKeyGroupWorkQueue= */ false); + /* multiKeyBundleEnabled= */ false); } @Test @@ -201,13 +200,13 @@ public void testScheduleWorkWhenExceedMaximumPoolSize() throws Exception { // 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 void testScheduleWorkWhenExceedMaximumBytesOutstanding() throws Exception 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 void testRenderSummaryHtml() { @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 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 void testPollWork() throws Exception { // 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 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 void testPollWorkWithLinkedBlockingQueue() throws Exception { @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 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 void testPollWork_skipsFailedWorkAndCallsOnFailedWorkHandler() throws Exc 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 void testPollWork_skipsFailedWorkAndCallsOnFailedWorkHandler() throws Exc // 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 void testPollWork_skipsFailedWorkAndCallsOnFailedWorkHandler() throws Exc 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 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 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 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 000000000000..f41b0da4eece --- /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 computations = ImmutableList.of(COMP_A, COMP_B); + ImmutableList keyGroups = + ImmutableList.of(KEY_GROUP_1, KEY_GROUP_2, Work.KeyGroup.create(2, 1)); + + List> worksByProducer = new ArrayList<>(); + for (int producer = 0; producer < numProducers; producer++) { + List 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 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> producerFutures = new ArrayList<>(); + for (List works : worksByProducer) { + producerFutures.add( + threads.submit( + () -> { + for (QueuedWork work : works) { + executorQueue.add(work); + index.add(work); + } + })); + } + List> 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> 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 a58851cf3702..000000000000 --- 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 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 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 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> 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 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 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 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 a11fb1f3ad33..688bad933e0f 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 void setUp() { .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 6f440c9cf30b..f11add07f6f1 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 @@ private static WorkFailureProcessor createWorkFailureProcessor( .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 017a22ca6842..8ad5896dfabb 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 @@ private static BoundedQueueExecutor workExecutor() { 10000000, new ThreadFactoryBuilder().setNameFormat("DataflowWorkUnits-%d").setDaemon(true).build(), /* useFairMonitor= */ false, - /* useKeyGroupWorkQueue= */ false); + /* multiKeyBundleEnabled= */ false); } private static ComputationState createComputationState(int computationIdSuffix) {