This is an automated email from the ASF dual-hosted git repository.
scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 98ca01f415c [Dataflow Streaming] Prepare BoundedQueueExecutor for
MultiKey bundles (#38592)
98ca01f415c is described below
commit 98ca01f415cd090274718f7f2dbe2d73552d1116
Author: Arun Pandian <[email protected]>
AuthorDate: Tue Jun 2 00:43:26 2026 -0700
[Dataflow Streaming] Prepare BoundedQueueExecutor for MultiKey bundles
(#38592)
---
...rk.java => BoundedQueueExecutorWorkHandle.java} | 34 +----
.../dataflow/worker/streaming/ExecutableWork.java | 51 +++++--
.../dataflow/worker/util/BoundedQueueExecutor.java | 166 +++++++++++++++++++--
.../dataflow/worker/util/ExceptionUtils.java | 43 ++++++
.../work/processing/StreamingCommitFinalizer.java | 2 +-
.../work/processing/StreamingWorkScheduler.java | 18 ++-
.../worker/StreamingDataflowWorkerTest.java | 4 +-
.../worker/streaming/ActiveWorkStateTest.java | 4 +-
.../streaming/ComputationStateCacheTest.java | 2 +-
.../worker/util/BoundedQueueExecutorTest.java | 107 +++++++++----
.../failures/WorkFailureProcessorTest.java | 4 +-
.../work/refresh/ActiveWorkRefresherTest.java | 4 +-
12 files changed, 344 insertions(+), 95 deletions(-)
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java
similarity index 53%
copy from
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java
copy to
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java
index db279f06663..1ca53496694 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java
@@ -17,32 +17,8 @@
*/
package org.apache.beam.runners.dataflow.worker.streaming;
-import com.google.auto.value.AutoValue;
-import java.util.function.Consumer;
-import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
-
-/** {@link Work} instance and a processing function used to process the work.
*/
-@AutoValue
-public abstract class ExecutableWork implements Runnable {
-
- public static ExecutableWork create(Work work, Consumer<Work> executeWorkFn)
{
- return new AutoValue_ExecutableWork(work, executeWorkFn);
- }
-
- public abstract Work work();
-
- public abstract Consumer<Work> executeWorkFn();
-
- @Override
- public void run() {
- executeWorkFn().accept(work());
- }
-
- public final WorkId id() {
- return work().id();
- }
-
- public final Windmill.WorkItem getWorkItem() {
- return work().getWorkItem();
- }
-}
+/**
+ * A handle to use when requesting pulling more work from @BoundedQueueExecutor
+ * via @BoundedQueueExecutor.pollWork
+ */
+public interface BoundedQueueExecutorWorkHandle {}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java
index db279f06663..ecaa673f557 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java
@@ -17,25 +17,49 @@
*/
package org.apache.beam.runners.dataflow.worker.streaming;
-import com.google.auto.value.AutoValue;
-import java.util.function.Consumer;
+import java.util.Objects;
+import java.util.function.BiConsumer;
+import org.apache.beam.runners.dataflow.worker.util.ExceptionUtils;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
/** {@link Work} instance and a processing function used to process the work.
*/
-@AutoValue
-public abstract class ExecutableWork implements Runnable {
+public final class ExecutableWork {
- public static ExecutableWork create(Work work, Consumer<Work> executeWorkFn)
{
- return new AutoValue_ExecutableWork(work, executeWorkFn);
+ private final Work work;
+ private final BiConsumer<Work, BoundedQueueExecutorWorkHandle> executeWorkFn;
+
+ private ExecutableWork(
+ Work work, BiConsumer<Work, BoundedQueueExecutorWorkHandle>
executeWorkFn) {
+ this.work = Objects.requireNonNull(work);
+ this.executeWorkFn = Objects.requireNonNull(executeWorkFn);
}
- public abstract Work work();
+ /**
+ * Creates an {@link ExecutableWork} instance.
+ *
+ * @param executeWorkFn The function executing the work. It'll be called
along with a
+ * BoundedQueueExecutorWorkHandle. The handle needs to be passed to
BoundedQueueExecutor when
+ * requesting more work to process inline.
+ */
+ public static ExecutableWork create(
+ Work work, BiConsumer<Work, BoundedQueueExecutorWorkHandle>
executeWorkFn) {
+ return new ExecutableWork(work, executeWorkFn);
+ }
- public abstract Consumer<Work> executeWorkFn();
+ public Work work() {
+ return work;
+ }
- @Override
- public void run() {
- executeWorkFn().accept(work());
+ public BiConsumer<Work, BoundedQueueExecutorWorkHandle> executeWorkFn() {
+ return executeWorkFn;
+ }
+
+ public void run(BoundedQueueExecutorWorkHandle handle) {
+ try {
+ executeWorkFn().accept(work(), handle);
+ } catch (Throwable t) {
+ throw ExceptionUtils.safeWrapThrowableAsException(t);
+ }
}
public final WorkId id() {
@@ -45,4 +69,9 @@ public abstract class ExecutableWork implements Runnable {
public final Windmill.WorkItem getWorkItem() {
return work().getWorkItem();
}
+
+ @Override
+ public String toString() {
+ return "ExecutableWork{" + id() + "}";
+ }
}
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 9079c3cc69b..c6fd96e0a4c 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
@@ -17,6 +17,9 @@
*/
package org.apache.beam.runners.dataflow.worker.util;
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadFactory;
@@ -24,6 +27,10 @@ import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.concurrent.GuardedBy;
+import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
+import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor.Guard;
@@ -38,7 +45,19 @@ public class BoundedQueueExecutor {
// Used to guard elementsOutstanding and bytesOutstanding.
private final Monitor monitor;
- private final ConcurrentLinkedQueue<Long> decrementQueue = new
ConcurrentLinkedQueue<>();
+
+ private static class Budget {
+
+ final int elements;
+ final long bytes;
+
+ Budget(int elements, long bytes) {
+ this.elements = elements;
+ this.bytes = bytes;
+ }
+ }
+
+ private final ConcurrentLinkedQueue<Budget> decrementQueue = new
ConcurrentLinkedQueue<>();
private final Object decrementQueueDrainLock = new Object();
private final AtomicBoolean isDecrementBatchPending = new
AtomicBoolean(false);
private int elementsOutstanding = 0;
@@ -106,7 +125,7 @@ public class BoundedQueueExecutor {
// Before adding a Work to the queue, check that there are enough bytes of
space or no other
// outstanding elements of work.
- public void execute(Runnable work, long workBytes) {
+ public void execute(ExecutableWork work, long workBytes) {
monitor.enterWhenUninterruptibly(
new Guard(monitor) {
@Override
@@ -119,12 +138,18 @@ public class BoundedQueueExecutor {
executeMonitorHeld(work, workBytes);
}
- // Forcibly add something to the queue, ignoring the length limit.
- public void forceExecute(Runnable work, long workBytes) {
+ // Forcibly add ExecutableWork to the queue, ignoring the limits.
+ public void forceExecute(ExecutableWork work, long workBytes) {
monitor.enter();
executeMonitorHeld(work, workBytes);
}
+ /** Forcibly execute a Runnable callback with 0 bytes of size. */
+ public void forceExecute(Runnable runnable) {
+ monitor.enter();
+ executeMonitorHeld(runnable);
+ }
+
// Set the maximum/core pool size of the executor.
public synchronized void setMaximumPoolSize(int maximumPoolSize, int
maximumElementsOutstanding) {
// For ThreadPoolExecutor, the maximum pool size should always greater
than or equal to core
@@ -221,8 +246,115 @@ public class BoundedQueueExecutor {
}
}
- private void executeMonitorHeld(Runnable work, long workBytes) {
+ /**
+ * A handle to use when requesting pulling more work from
@BoundedQueueExecutor
+ * via @BoundedQueueExecutor.pollWork. A single handle aggregates all
budgets from work pulled for
+ * inline execution and releases the budget after the multi work bundle is
complete.
+ */
+ final class BoundedQueueExecutorWorkHandleImpl
+ implements BoundedQueueExecutorWorkHandle, AutoCloseable {
+
+ @GuardedBy("this")
+ private int elements;
+
+ @GuardedBy("this")
+ private long bytes;
+
+ @GuardedBy("this")
+ private boolean closed = false;
+
+ private BoundedQueueExecutorWorkHandleImpl(int elements, long bytes) {
+ checkArgument(elements >= 0 && bytes >= 0);
+ this.elements = elements;
+ this.bytes = bytes;
+ }
+
+ /**
+ * Merges the budget from another handle into this handle.
+ *
+ * <p>This transfers the budget (elements and bytes) from the {@code
other} handle to this
+ * handle, and marks the {@code other} handle as closed to prevent it from
releasing the budget
+ * again if it is closed.
+ */
+ public void merge(BoundedQueueExecutorWorkHandleImpl other) {
+ checkArgumentNotNull(other);
+ synchronized (this) {
+ Preconditions.checkState(!closed, "Cannot merge into a closed handle");
+ synchronized (other) {
+ Preconditions.checkState(!other.closed, "Cannot merge a closed
handle");
+ this.elements += other.elements;
+ this.bytes += other.bytes;
+ other.closed = true;
+ other.elements = 0;
+ other.bytes = 0;
+ }
+ }
+ }
+
+ public synchronized boolean isClosed() {
+ return closed;
+ }
+
+ @VisibleForTesting
+ synchronized int elements() {
+ return elements;
+ }
+
+ @VisibleForTesting
+ synchronized long bytes() {
+ return bytes;
+ }
+
+ @Override
+ public synchronized void close() {
+ if (closed) return;
+ closed = true;
+ decrementCounters(this.elements, this.bytes);
+ }
+ }
+
+ private static final class QueuedWork implements Runnable {
+
+ private final ExecutableWork work;
+ private final BoundedQueueExecutorWorkHandleImpl handle;
+
+ public QueuedWork(ExecutableWork work, BoundedQueueExecutorWorkHandleImpl
handle) {
+ this.work = work;
+ this.handle = handle;
+ }
+
+ public ExecutableWork getWork() {
+ return work;
+ }
+
+ public BoundedQueueExecutorWorkHandleImpl getHandle() {
+ return handle;
+ }
+
+ @Override
+ public void run() {
+ checkArgument(!handle.isClosed());
+ try (handle) {
+ work.run(handle);
+ }
+ }
+ }
+
+ private void executeMonitorHeld(ExecutableWork work, long workBytes) {
+ ++elementsOutstanding;
bytesOutstanding += workBytes;
+ monitor.leave();
+ BoundedQueueExecutorWorkHandleImpl handle =
+ new BoundedQueueExecutorWorkHandleImpl(1, workBytes);
+ try {
+ executor.execute(new QueuedWork(work, handle));
+ } catch (Throwable t) {
+ handle.close();
+ throw ExceptionUtils.safeWrapThrowableAsException(t);
+ }
+ }
+
+ private void executeMonitorHeld(Runnable work) {
++elementsOutstanding;
monitor.leave();
@@ -232,21 +364,25 @@ public class BoundedQueueExecutor {
try {
work.run();
} finally {
- decrementCounters(workBytes);
+ decrementCounters(1, 0L);
}
});
- } catch (RuntimeException e) {
- // If the execute() call threw an exception, decrement counters here.
- decrementCounters(workBytes);
- throw e;
+ } catch (Throwable t) {
+ decrementCounters(1, 0L);
+ throw ExceptionUtils.safeWrapThrowableAsException(t);
}
}
- private void decrementCounters(long workBytes) {
+ @VisibleForTesting
+ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(int elements, long
bytes) {
+ return new BoundedQueueExecutorWorkHandleImpl(elements, bytes);
+ }
+
+ private void decrementCounters(int elements, long bytes) {
// All threads queue decrements and one thread grabs the monitor and
updates
// counters. We do this to reduce contention on monitor which is locked by
// GetWork thread
- decrementQueue.add(workBytes);
+ decrementQueue.add(new Budget(elements, bytes));
boolean submittedToExistingBatch = isDecrementBatchPending.getAndSet(true);
if (submittedToExistingBatch) {
// There is already a thread about to drain the decrement queue
@@ -265,12 +401,12 @@ public class BoundedQueueExecutor {
long bytesToDecrement = 0;
int elementsToDecrement = 0;
while (true) {
- Long pollResult = decrementQueue.poll();
+ Budget pollResult = decrementQueue.poll();
if (pollResult == null) {
break;
}
- bytesToDecrement += pollResult;
- ++elementsToDecrement;
+ bytesToDecrement += pollResult.bytes;
+ elementsToDecrement += pollResult.elements;
}
if (elementsToDecrement == 0) {
return;
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ExceptionUtils.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ExceptionUtils.java
new file mode 100644
index 00000000000..d04a385d6e6
--- /dev/null
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ExceptionUtils.java
@@ -0,0 +1,43 @@
+/*
+ * 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 javax.annotation.CheckReturnValue;
+import org.apache.beam.sdk.annotations.Internal;
+
+/** Utility methods for simplifying work with exceptions and throwables. */
+@Internal
+public final class ExceptionUtils {
+
+ private ExceptionUtils() {}
+
+ /**
+ * Returns the {@code throwable} as-is if it is an instance of {@link
RuntimeException} throws if
+ * it is an {@link Error}, or returns the {@code throwable} wrapped in a
{@code RuntimeException}.
+ */
+ @CheckReturnValue
+ public static RuntimeException safeWrapThrowableAsException(Throwable
throwable) {
+ if (throwable instanceof RuntimeException) {
+ return (RuntimeException) throwable;
+ } else if (throwable instanceof Error) {
+ throw (Error) throwable;
+ } else {
+ return new RuntimeException(throwable);
+ }
+ }
+}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java
index 5a66545ab33..22573bf1ced 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingCommitFinalizer.java
@@ -156,7 +156,7 @@ final class StreamingCommitFinalizer {
}
for (Runnable callback : callbacksToExecute) {
try {
- finalizationExecutor.forceExecute(callback, 0);
+ finalizationExecutor.forceExecute(callback);
} catch (OutOfMemoryError oom) {
throw oom;
} catch (Throwable t) {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
index 1428037d9ca..364608be82c 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
@@ -35,6 +35,7 @@ import org.apache.beam.runners.dataflow.worker.HotKeyLogger;
import org.apache.beam.runners.dataflow.worker.ReaderCache;
import org.apache.beam.runners.dataflow.worker.WorkItemCancelledException;
import
org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC;
+import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
import
org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
@@ -47,6 +48,7 @@ import
org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounte
import
org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher;
import
org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory;
import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
+import org.apache.beam.runners.dataflow.worker.util.ExceptionUtils;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import
org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution;
import org.apache.beam.runners.dataflow.worker.windmill.client.commits.Commit;
@@ -218,7 +220,7 @@ public class StreamingWorkScheduler {
ExecutableWork.create(
Work.create(
workItem, serializedWorkItemSize, watermarks,
processingContext, drainMode, clock),
- work -> processWork(computationState, work,
getWorkStreamLatencies)));
+ (work, handle) -> processWork(computationState, work,
getWorkStreamLatencies, handle)));
}
/** Adds any applied finalize ids to the commit finalizer to have their
callbacks executed. */
@@ -232,16 +234,19 @@ public class StreamingWorkScheduler {
* internally if processing fails due to uncaught {@link Exception}(s).
*
* @implNote This will block the calling thread during execution of user
DoFns.
+ * @param handle handled to pass to BoundedQueueExecutor.pollWork, currently
unused
*/
private void processWork(
ComputationState computationState,
Work work,
- ImmutableList<LatencyAttribution> getWorkStreamLatencies) {
+ ImmutableList<LatencyAttribution> getWorkStreamLatencies,
+ BoundedQueueExecutorWorkHandle handle) {
work.recordGetWorkStreamLatencies(getWorkStreamLatencies);
- processWork(computationState, work);
+ processWork(computationState, work, handle);
}
- private void processWork(ComputationState computationState, Work work) {
+ private void processWork(
+ ComputationState computationState, Work work,
BoundedQueueExecutorWorkHandle unusedHandle) {
Windmill.WorkItem workItem = work.getWorkItem();
String computationId = computationState.getComputationId();
ByteString key = workItem.getKey();
@@ -289,7 +294,7 @@ public class StreamingWorkScheduler {
try {
workFailureProcessor.logAndProcessFailure(
computationId,
- ExecutableWork.create(work, retry -> processWork(computationState,
retry)),
+ ExecutableWork.create(work, (retry, h) ->
processWork(computationState, retry, h)),
t,
invalidWork ->
computationState.completeWorkAndScheduleNextWorkForKey(
@@ -297,7 +302,8 @@ public class StreamingWorkScheduler {
} catch (OutOfMemoryError oom) {
throw oom;
} catch (Throwable t2) {
- throw new RuntimeException(t2);
+ LOG.warn("Failed to process work failure safely for work {}",
work.id(), t2);
+ throw ExceptionUtils.safeWrapThrowableAsException(t2);
}
} finally {
// Update total processing time counters. Updating in finally clause
ensures that
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 ff82a1ab5c4..d58f2007699 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
@@ -373,7 +373,9 @@ public class StreamingDataflowWorkerTest {
computationId, new FakeGetDataClient(), ignored -> {},
mock(HeartbeatSender.class)),
false,
Instant::now),
- processWorkFn);
+ (work, handle) -> {
+ processWorkFn.accept(work);
+ });
}
private byte[] intervalWindowBytes(IntervalWindow window) throws Exception {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
index 865ae261280..0f14efdd0c0 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
@@ -73,7 +73,7 @@ public class ActiveWorkStateTest {
createWorkProcessingContext(),
false,
Instant::now),
- ignored -> {});
+ (work, handle) -> {});
}
private static ExecutableWork expiredWork(Windmill.WorkItem workItem) {
@@ -85,7 +85,7 @@ public class ActiveWorkStateTest {
createWorkProcessingContext(),
false,
() -> Instant.EPOCH),
- ignored -> {});
+ (work, handle) -> {});
}
private static Work.ProcessingContext createWorkProcessingContext() {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java
index 1c8b8fca131..30ad97140e1 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java
@@ -77,7 +77,7 @@ public class ComputationStateCacheTest {
mock(HeartbeatSender.class)),
false,
Instant::now),
- ignored -> {});
+ (work, handle) -> {});
}
@Before
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 d7ea039bb80..55fe82c7163 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
@@ -32,6 +32,7 @@ import java.util.function.Consumer;
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.BoundedQueueExecutorWorkHandleImpl;
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;
@@ -85,18 +86,21 @@ public class BoundedQueueExecutorTest {
mock(HeartbeatSender.class)),
false,
Instant::now),
- executeWorkFn);
+ (work, handle) -> {
+ executeWorkFn.accept(work);
+ });
}
- private Runnable createSleepProcessWorkFn(CountDownLatch start,
CountDownLatch stop) {
- return () -> {
- start.countDown();
- try {
- stop.await();
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- };
+ private ExecutableWork createSleepProcessWork(CountDownLatch start,
CountDownLatch stop) {
+ return createWork(
+ ignored -> {
+ start.countDown();
+ try {
+ stop.await();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
}
@Before
@@ -123,9 +127,9 @@ public class BoundedQueueExecutorTest {
CountDownLatch processStop2 = new CountDownLatch(1);
CountDownLatch processStart3 = new CountDownLatch(1);
CountDownLatch processStop3 = new CountDownLatch(1);
- Runnable m1 = createSleepProcessWorkFn(processStart1, processStop1);
- Runnable m2 = createSleepProcessWorkFn(processStart2, processStop2);
- Runnable m3 = createSleepProcessWorkFn(processStart3, processStop3);
+ ExecutableWork m1 = createSleepProcessWork(processStart1, processStop1);
+ ExecutableWork m2 = createSleepProcessWork(processStart2, processStop2);
+ ExecutableWork m3 = createSleepProcessWork(processStart3, processStop3);
executor.execute(m1, 1);
processStart1.await();
@@ -152,8 +156,8 @@ public class BoundedQueueExecutorTest {
CountDownLatch processStop1 = new CountDownLatch(1);
CountDownLatch processStart2 = new CountDownLatch(1);
CountDownLatch processStop2 = new CountDownLatch(1);
- Runnable m1 = createSleepProcessWorkFn(processStart1, processStop1);
- Runnable m2 = createSleepProcessWorkFn(processStart2, processStop2);
+ ExecutableWork m1 = createSleepProcessWork(processStart1, processStop1);
+ ExecutableWork m2 = createSleepProcessWork(processStart2, processStop2);
executor.execute(m1, 10000000);
processStart1.await();
@@ -187,9 +191,9 @@ public class BoundedQueueExecutorTest {
CountDownLatch processStart2 = new CountDownLatch(1);
CountDownLatch processStart3 = new CountDownLatch(1);
CountDownLatch stop = new CountDownLatch(1);
- Runnable m1 = createSleepProcessWorkFn(processStart1, stop);
- Runnable m2 = createSleepProcessWorkFn(processStart2, stop);
- Runnable m3 = createSleepProcessWorkFn(processStart3, stop);
+ ExecutableWork m1 = createSleepProcessWork(processStart1, stop);
+ ExecutableWork m2 = createSleepProcessWork(processStart2, stop);
+ ExecutableWork m3 = createSleepProcessWork(processStart3, stop);
// Initial state.
assertEquals(0, executor.activeCount());
@@ -225,9 +229,9 @@ public class BoundedQueueExecutorTest {
CountDownLatch processStart2 = new CountDownLatch(1);
CountDownLatch processStart3 = new CountDownLatch(1);
CountDownLatch stop = new CountDownLatch(1);
- Runnable m1 = createSleepProcessWorkFn(processStart1, stop);
- Runnable m2 = createSleepProcessWorkFn(processStart2, stop);
- Runnable m3 = createSleepProcessWorkFn(processStart3, stop);
+ ExecutableWork m1 = createSleepProcessWork(processStart1, stop);
+ ExecutableWork m2 = createSleepProcessWork(processStart2, stop);
+ ExecutableWork m3 = createSleepProcessWork(processStart3, stop);
// Initial state.
assertEquals(0, executor.activeCount());
@@ -264,9 +268,9 @@ public class BoundedQueueExecutorTest {
CountDownLatch processStart2 = new CountDownLatch(1);
CountDownLatch processStart3 = new CountDownLatch(1);
CountDownLatch stop = new CountDownLatch(1);
- Runnable m1 = createSleepProcessWorkFn(processStart1, stop);
- Runnable m2 = createSleepProcessWorkFn(processStart2, stop);
- Runnable m3 = createSleepProcessWorkFn(processStart3, stop);
+ ExecutableWork m1 = createSleepProcessWork(processStart1, stop);
+ ExecutableWork m2 = createSleepProcessWork(processStart2, stop);
+ ExecutableWork m3 = createSleepProcessWork(processStart3, stop);
// Initial state.
assertEquals(0, executor.activeCount());
@@ -308,9 +312,9 @@ public class BoundedQueueExecutorTest {
CountDownLatch processStop2 = new CountDownLatch(1);
CountDownLatch processStart3 = new CountDownLatch(1);
CountDownLatch processStop3 = new CountDownLatch(1);
- Runnable m1 = createSleepProcessWorkFn(processStart1, processStop1);
- Runnable m2 = createSleepProcessWorkFn(processStart2, processStop2);
- Runnable m3 = createSleepProcessWorkFn(processStart3, processStop3);
+ ExecutableWork m1 = createSleepProcessWork(processStart1, processStop1);
+ ExecutableWork m2 = createSleepProcessWork(processStart2, processStop2);
+ ExecutableWork m3 = createSleepProcessWork(processStart3, processStop3);
// Initial state.
assertEquals(0, executor.activeCount());
@@ -351,6 +355,55 @@ public class BoundedQueueExecutorTest {
executor.shutdown();
}
+ @Test
+ public void testRunnableExceptionPropagationDecrementsCounters() throws
Exception {
+ CountDownLatch processStart = new CountDownLatch(1);
+ CountDownLatch processStop = new CountDownLatch(1);
+
+ Runnable work =
+ () -> {
+ processStart.countDown();
+ try {
+ processStop.await();
+ } catch (InterruptedException e) {
+ throw new RuntimeException(e);
+ }
+ throw new RuntimeException("Simulated finalizer processing
exception");
+ };
+
+ executor.forceExecute(work);
+ processStart.await();
+
+ assertEquals(1, executor.elementsOutstanding());
+
+ processStop.countDown();
+
+ // Wait until outstanding elements are released
+ while (executor.elementsOutstanding() != 0) {
+ Thread.sleep(10);
+ }
+
+ assertEquals(0, executor.elementsOutstanding());
+ }
+
+ @Test
+ public void testHandleMerge() throws Exception {
+ BoundedQueueExecutorWorkHandleImpl handle1 =
executor.createBudgetHandle(1, 100L);
+ BoundedQueueExecutorWorkHandleImpl handle2 =
executor.createBudgetHandle(2, 200L);
+
+ handle1.merge(handle2);
+
+ // Verify that handle2 has 0 budget and is closed.
+ assertEquals(0, handle2.elements());
+ assertEquals(0, handle2.bytes());
+ assertTrue(handle2.isClosed());
+
+ // Verify that handle1 has the combined budget and is not closed.
+ assertEquals(3, handle1.elements());
+ assertEquals(300L, handle1.bytes());
+ assertFalse(handle1.isClosed());
+ }
+
@Test
public void testRenderSummaryHtml() {
String expectedSummaryHtml =
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 68a11895fa1..51bd4816b03 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
@@ -98,7 +98,9 @@ public class WorkFailureProcessorTest {
mock(HeartbeatSender.class)),
false,
clock),
- processWorkFn);
+ (work, handle) -> {
+ processWorkFn.accept(work);
+ });
}
private static ExecutableWork createWork(Consumer<Work> processWorkFn) {
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 054db878c86..88a82c6f76b 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
@@ -137,7 +137,9 @@ public class ActiveWorkRefresherTest {
"computationId", new FakeGetDataClient(), ignored -> {},
heartbeatSender),
false,
ActiveWorkRefresherTest::aLongTimeAgo),
- processWork);
+ (work, handle) -> {
+ processWork.accept(work);
+ });
}
@Test
