scwhittle commented on code in PR #40593:
URL: https://github.com/apache/beam/pull/40593#discussion_r4292549671
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java:
##########
@@ -328,49 +346,129 @@ 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}.
+ *
+ * <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;
+ }
}
public ExecutableWork getWork() {
Review Comment:
I think that this and getHandle are maybe only used by the other methods and
could be removed?
See if other methods could be made private or removed as well.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java:
##########
@@ -403,25 +501,52 @@ 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.
+ */
+ private void onWorkStolen() {
+ if (workQueue.size() <= (long) elementsOutstanding.get() +
maximumElementsOutstanding
Review Comment:
could we remove checking elementsOutstanding and revert changing it to an
atomic?
seems like we could just use
workQueue.size() < 2 * maximumElementsOutstanding
That seems like it woudl trigger only if maximumElementsOutstanding stale
entries are there
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java:
##########
@@ -742,25 +688,26 @@ public void
testPollWork_purgeBacksOffWhileQueueHoldsMostlyUnstartedWork() throw
createWorkWithCompIdAndKeyGroup("compA", keyGroup, ignored -> {}),
10);
}
- // The first steal purges its own stale shell; the remaining steals back
off until the queue
- // grows by more than maximumElementsOutstanding.
- for (int i = 0; i < numWorks; i++) {
- assertNotNull(
- testExecutor.pollWork("compA", keyGroup, stealHandle,
FAILING_FAILED_WORK_HANDLER));
- }
- assertEquals(numWorks - 1, testExecutor.executorQueueSize());
-
- for (int i = 0; i <= maximumElementsOutstanding; 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));
}
- assertNotNull(
- testExecutor.pollWork("compA", keyGroup, stealHandle,
FAILING_FAILED_WORK_HANDLER));
- assertEquals(maximumElementsOutstanding, testExecutor.executorQueueSize());
+ assertEquals(numWorks, testExecutor.rawQueueSizeForTesting());
Review Comment:
can we verify elementsOutstanding here and below not just the rawQeueuSize?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java:
##########
Review Comment:
how about getQueueKey? at call-sites it is confusing with performing queuing
action
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
