arunpandianp commented on code in PR #40593:
URL: https://github.com/apache/beam/pull/40593#discussion_r4296958813
##########
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:
Keeping the null checks at one place inside getWork, getHandle and updated
visibility of the methods.
--
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]
