Skip to content

[Dataflow Streaming] [Multi Key] Drop failed work in BoundedQueueExecutor::pollWork #38920

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
arunpandianp wants to merge 11 commits into
base: master
Choose a base branch
from
Open

[Dataflow Streaming] [Multi Key] Drop failed work in BoundedQueueExecutor::pollWork #38920

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
2c12786
Drop failed work in BoundedQueueExecutor::pollWork
arunpandianp Jun 11, 2026
3670a03
address comment
arunpandianp Jun 11, 2026
546739c
Merge remote-tracking branch 'beam/master' into multikey_queue_filter…
arunpandianp Aug 5, 2026
2ddc10e
Revert "address comment"
arunpandianp Aug 5, 2026
fc15302
Revert "Drop failed work in BoundedQueueExecutor::pollWork"
arunpandianp Aug 5, 2026
7938f7f
[Dataflow Streaming] Remove finalizeCommits from processWork
arunpandianp Aug 5, 2026
ebfe88a
Merge branch 'fixfinalizer' into multikey_queue_filter_failed
arunpandianp Aug 5, 2026
fce54b0
Plumb ComputationState to ProcessingContext
arunpandianp Aug 6, 2026
840ecda
Drop failed workitems during pollwork
arunpandianp Aug 6, 2026
bbda033
Improve tests
arunpandianp Aug 6, 2026
bc9b4d0
Merge remote-tracking branch 'beam/master' into multikey_queue_filter…
arunpandianp Aug 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 16 additions & 18 deletions ...worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,28 +405,25 @@ private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHar
.setBytes(MAX_GET_WORK_FETCH_BYTES)
.build(),
windmillStreamFactory,
(workItem,
(computationState,
workItem,
serializedWorkItemSize,
watermarks,
processingContext,
drainMode,
appliedFinalizeIds,
getWorkStreamLatencies) ->
checkNotNull(computationStateCache)
.get(processingContext.computationId())
.ifPresent(
computationState -> {
memoryMonitor.waitForResources("GetWork");
streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
streamingWorkScheduler.scheduleWork(
computationState,
workItem,
serializedWorkItemSize,
watermarks,
processingContext,
drainMode,
getWorkStreamLatencies);
}),
getWorkStreamLatencies) -> {
memoryMonitor.waitForResources("GetWork");
streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
streamingWorkScheduler.scheduleWork(
computationState,
workItem,
serializedWorkItemSize,
watermarks,
processingContext,
drainMode,
getWorkStreamLatencies);
},
ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache),
GetWorkBudgetDistributors.distributeEvenly(),
checkNotNull(dispatcherClient),
Expand All @@ -441,7 +438,8 @@ private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHar
.setCommitWorkStreamFactory(
() -> CloseableStream.create(commitWorkStream, () -> {}))
.build(),
getDataMetricTracker);
getDataMetricTracker,
checkNotNull(this.computationStateCache)::get);
ChannelzServlet channelzServlet =
createChannelzServlet(
options, fanOutStreamingEngineWorkerHarness::currentWindmillEndpoints);
Expand Down
5 changes: 2 additions & 3 deletions ...rker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,18 +28,17 @@
import java.util.Optional;
import java.util.Queue;
import java.util.function.BiConsumer;
import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem;
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache.ForComputation;
import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
import org.apache.beam.sdk.annotations.Internal;
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.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
Expand Down Expand Up @@ -78,7 +77,7 @@ public final class ActiveWorkState {

private ActiveWorkState(
Map<Long, LinkedHashMap<WorkId, ExecutableWork>> activeWork,
ForComputation computationStateCache) {
WindmillStateCache.ForComputation computationStateCache) {
this.activeWork = activeWork;
this.computationStateCache = computationStateCache;
this.activeGetWorkBudget = GetWorkBudget.noBudget();
Expand Down
2 changes: 1 addition & 1 deletion ...ker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,14 @@
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentLinkedQueue;
import javax.annotation.Nullable;
import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
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.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;

/**
Expand Down
26 changes: 19 additions & 7 deletions ...low-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,22 +144,26 @@ public static Work create(
}

public static ProcessingContext createProcessingContext(
String computationId,
ComputationState computationState,
GetDataClient getDataClient,
Consumer<Commit> workCommitter,
HeartbeatSender heartbeatSender) {
return ProcessingContext.create(
computationId, getDataClient, workCommitter, heartbeatSender, /* backendWorkerToken= */ "");
computationState,
getDataClient,
workCommitter,
heartbeatSender,
/* backendWorkerToken= */ "");
}

public static ProcessingContext createProcessingContext(
String computationId,
ComputationState computationState,
GetDataClient getDataClient,
Consumer<Commit> workCommitter,
HeartbeatSender heartbeatSender,
String backendWorkerToken) {
return ProcessingContext.create(
computationId, getDataClient, workCommitter, heartbeatSender, backendWorkerToken);
computationState, getDataClient, workCommitter, heartbeatSender, backendWorkerToken);
}

private static LatencyAttribution.Builder createLatencyAttributionWithActiveLatencyBreakdown(
Expand Down Expand Up @@ -207,6 +211,10 @@ public long getSerializedWorkItemSize() {
return serializedWorkItemSize;
}

public ComputationState getComputationState() {
return processingContext.computationState();
}

public String getComputationId() {
return processingContext.computationId();
}
Expand Down Expand Up @@ -457,17 +465,21 @@ public KeyGroup getKeyGroup() {
public abstract static class ProcessingContext {

private static ProcessingContext create(
String computationId,
ComputationState computationState,
GetDataClient getDataClient,
Consumer<Commit> workCommitter,
HeartbeatSender heartbeatSender,
String backendWorkerToken) {
return new AutoValue_Work_ProcessingContext(
computationId, getDataClient, heartbeatSender, workCommitter, backendWorkerToken);
computationState, getDataClient, heartbeatSender, workCommitter, backendWorkerToken);
}

/** Computation that the {@link Work} belongs to. */
public abstract String computationId();
public abstract ComputationState computationState();

public String computationId() {
return computationState().getComputationId();
}

/** Handles GetData requests to streaming backend. */
public abstract GetDataClient getDataClient();
Expand Down
23 changes: 15 additions & 8 deletions ...he/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair;
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetWorkRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
Expand Down Expand Up @@ -96,6 +97,7 @@ public final class FanOutStreamingEngineWorkerHarness implements StreamingWorker
private final GetWorkBudget totalGetWorkBudget;
private final Function<WindmillStream.CommitWorkStream, WorkCommitter> workCommitterFactory;
private final ThrottlingGetDataMetricTracker getDataMetricTracker;
private final Function<String, Optional<ComputationState>> computationStateFetcher;
private final ExecutorService windmillStreamManager;
private final ExecutorService workerMetadataConsumer;
private final Object metadataLock = new Object();
Expand Down Expand Up @@ -131,7 +133,8 @@ private FanOutStreamingEngineWorkerHarness(
GrpcDispatcherClient dispatcherClient,
Function<WindmillStream.CommitWorkStream, WorkCommitter> workCommitterFactory,
ThrottlingGetDataMetricTracker getDataMetricTracker,
ExecutorService workerMetadataConsumer) {
ExecutorService workerMetadataConsumer,
Function<String, Optional<ComputationState>> computationStateFetcher) {
this.jobHeader = jobHeader;
this.getDataMetricTracker = getDataMetricTracker;
this.started = false;
Expand All @@ -150,6 +153,7 @@ private FanOutStreamingEngineWorkerHarness(
this.activeMetadataType = WindmillEndpoints.Type.UNKNOWN;
this.pendingMetadataType = WindmillEndpoints.Type.UNKNOWN;
this.workCommitterFactory = workCommitterFactory;
this.computationStateFetcher = computationStateFetcher;
}

/**
Expand All @@ -166,7 +170,8 @@ public static FanOutStreamingEngineWorkerHarness create(
GetWorkBudgetDistributor getWorkBudgetDistributor,
GrpcDispatcherClient dispatcherClient,
Function<WindmillStream.CommitWorkStream, WorkCommitter> workCommitterFactory,
ThrottlingGetDataMetricTracker getDataMetricTracker) {
ThrottlingGetDataMetricTracker getDataMetricTracker,
Function<String, Optional<ComputationState>> computationStateFetcher) {
return new FanOutStreamingEngineWorkerHarness(
jobHeader,
totalGetWorkBudget,
Expand All @@ -178,9 +183,8 @@ public static FanOutStreamingEngineWorkerHarness create(
workCommitterFactory,
getDataMetricTracker,
Executors.newSingleThreadExecutor(
new ThreadFactoryBuilder()
.setNameFormat(WORKER_METADATA_CONSUMER_THREAD_NAME)
.build()));
new ThreadFactoryBuilder().setNameFormat(WORKER_METADATA_CONSUMER_THREAD_NAME).build()),
computationStateFetcher);
}

@VisibleForTesting
Expand All @@ -193,7 +197,8 @@ static FanOutStreamingEngineWorkerHarness forTesting(
GetWorkBudgetDistributor getWorkBudgetDistributor,
GrpcDispatcherClient dispatcherClient,
Function<WindmillStream.CommitWorkStream, WorkCommitter> workCommitterFactory,
ThrottlingGetDataMetricTracker getDataMetricTracker) {
ThrottlingGetDataMetricTracker getDataMetricTracker,
Function<String, Optional<ComputationState>> computationStateFetcher) {
FanOutStreamingEngineWorkerHarness fanOutStreamingEngineWorkProvider =
new FanOutStreamingEngineWorkerHarness(
jobHeader,
Expand All @@ -210,7 +215,8 @@ static FanOutStreamingEngineWorkerHarness forTesting(
// blocked by the consumeWorkerMetadata() task. Test suites run in different
// environments and non-determinism has lead to past flakiness. See
// https://github.com/apache/beam/issues/28957.
MoreExecutors.newDirectExecutorService());
MoreExecutors.newDirectExecutorService(),
computationStateFetcher);
fanOutStreamingEngineWorkProvider.start();
return fanOutStreamingEngineWorkProvider;
}
Expand Down Expand Up @@ -448,7 +454,8 @@ private WindmillStreamSender createAndStartWindmillStreamSender(Endpoint endpoin
getDataStream ->
StreamGetDataClient.create(
getDataStream, this::getGlobalDataStream, getDataMetricTracker),
workCommitterFactory);
workCommitterFactory,
computationStateFetcher);
windmillStreamSender.start();
return windmillStreamSender;
}
Expand Down
4 changes: 2 additions & 2 deletions .../org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ private void streamingEngineDispatchLoop(
.setOutputDataWatermark(workItem.getOutputDataWatermark())
.build(),
Work.createProcessingContext(
computationId,
computationState,
getDataClient,
workCommitter::commit,
heartbeatSender),
Expand Down Expand Up @@ -250,7 +250,7 @@ private void applianceDispatchLoop(Supplier<Windmill.GetWorkResponse> getWorkFn)
workItem.getSerializedSize(),
watermarks.setOutputDataWatermark(workItem.getOutputDataWatermark()).build(),
Work.createProcessingContext(
computationId, getDataClient, workCommitter::commit, heartbeatSender),
computationState, getDataClient, workCommitter::commit, heartbeatSender),
computationWork.getDrainMode(),
/* getWorkStreamLatencies= */ ImmutableList.of());
}
Expand Down
14 changes: 10 additions & 4 deletions .../java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSender.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;

import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
Expand All @@ -27,6 +28,7 @@
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetWorkRequest;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillConnection;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream;
Expand Down Expand Up @@ -75,7 +77,8 @@ private WindmillStreamSender(
GrpcWindmillStreamFactory streamingEngineStreamFactory,
WorkItemScheduler workItemScheduler,
Function<GetDataStream, GetDataClient> getDataClientFactory,
Function<CommitWorkStream, WorkCommitter> workCommitterFactory) {
Function<CommitWorkStream, WorkCommitter> workCommitterFactory,
Function<String, Optional<ComputationState>> computationStateFetcher) {
this.started = new AtomicBoolean(false);
this.getWorkBudget = getWorkBudget;

Expand All @@ -91,7 +94,8 @@ private WindmillStreamSender(
FixedStreamHeartbeatSender.create(getDataStream),
getDataClientFactory.apply(getDataStream),
workCommitter,
workItemScheduler);
workItemScheduler,
computationStateFetcher);
// 3 threads, 1 for each stream type (GetWork, GetData, CommitWork).
this.streamStarter =
Executors.newFixedThreadPool(
Expand All @@ -105,15 +109,17 @@ static WindmillStreamSender create(
GrpcWindmillStreamFactory streamingEngineStreamFactory,
WorkItemScheduler workItemScheduler,
Function<GetDataStream, GetDataClient> getDataClientFactory,
Function<CommitWorkStream, WorkCommitter> workCommitterFactory) {
Function<CommitWorkStream, WorkCommitter> workCommitterFactory,
Function<String, Optional<ComputationState>> computationStateFetcher) {
return new WindmillStreamSender(
connection,
getWorkRequest,
new AtomicReference<>(getWorkBudget),
streamingEngineStreamFactory,
workItemScheduler,
getDataClientFactory,
workCommitterFactory);
workCommitterFactory,
computationStateFetcher);
}

private static GetWorkRequest withRequestBudget(GetWorkRequest request, GetWorkBudget budget) {
Expand Down
19 changes: 14 additions & 5 deletions ...rker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,12 +395,21 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(Work work, long bytes) {
if (keyGroupWorkQueue == null) {
return null;
}
@Nullable QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup);
if (queuedWork == null) {
return null;
while (true) {
@Nullable QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup);
if (queuedWork == null) {
return null;
}
Work work = queuedWork.getWork().work();
if (work.isFailed()) {
queuedWork.getHandle().close();

@scwhittle scwhittle Jun 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is just updating the counters, but it is not updating that this key was scheduled for processing. I think that somehow ActiveWorkState.completeWorkAndGetNextWorkForKey needs to be called so that we note that this key is completed and other work for the key can schedule.

Maybe a unit test at a higher level would help verify this behavior is working properly.

@arunpandianp arunpandianp Jun 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching it. Will fix it with a test, after #38814 is merged

work.getComputationState()

@scwhittle scwhittle Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wondering if there is someway to make the control flow clearer for these work items. One idea would be to not handle the failure here but pass the failed work items back to the caller of poll. Then StreamingWorkExecutionContext could possibly have a abortedworkitemhhandler similar to the KeyTransitionListener. Then we can have consistent error handling to StreamingWorkScheduler.handleProcessWorkFailure by passing it in instead of duplicating what it does.

Or perhaps we could collect the failed work in StreamingWorkExecutionContext and then handle it later after execution completes instead of using an injected function.

.completeWorkAndScheduleNextWorkForKey(work.getShardedKey(), work.id());
continue;
}
internalHandle.merge(queuedWork.getHandle());
return queuedWork.getWork();
}
internalHandle.merge(queuedWork.getHandle());
return queuedWork.getWork();
}

private void decrementCounters(int elements, long bytes) {
Expand Down
Loading
Loading