Skip to content

[Dataflow Streaming] Activate SourceState Finalizers before submitting workitem to harness threads #38921

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

Merged
scwhittle merged 3 commits into from
Jun 12, 2026
+2 −3
Merged

[Dataflow Streaming] Activate SourceState Finalizers before submitting workitem to harness threads #38921

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c21a7a6
Activate SourceState Finalizers before submitting workitem to harness…
arunpandianp Jun 11, 2026
4238c18
fix top level finalizer calls
arunpandianp Jun 11, 2026
6852ecb
add comments
arunpandianp Jun 11, 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
2 changes: 1 addition & 1 deletion ...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 @@ -418,14 +418,14 @@ private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHar
.ifPresent(
computationState -> {
memoryMonitor.waitForResources("GetWork");
streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
streamingWorkScheduler.scheduleWork(
computationState,
workItem,
serializedWorkItemSize,
watermarks,
processingContext,
drainMode,
appliedFinalizeIds,
getWorkStreamLatencies);
}),
ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache),
Expand Down
8 changes: 2 additions & 6 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 @@ -165,9 +165,6 @@ private void streamingEngineDispatchLoop(
.ifPresent(
computationState -> {
waitForResources.run();
if (!appliedFinalizeIds.isEmpty()) {
streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
}
streamingWorkScheduler.scheduleWork(
computationState,
workItem,
Expand All @@ -184,6 +181,7 @@ private void streamingEngineDispatchLoop(
workCommitter::commit,
heartbeatSender),
drainMode,
appliedFinalizeIds,
getWorkStreamLatencies);
}));
try {
Expand Down Expand Up @@ -221,9 +219,6 @@ private void applianceDispatchLoop(Supplier<Windmill.GetWorkResponse> getWorkFn)
ImmutableList<Long> appliedFinalizeIds =
ImmutableList.copyOf(
Preconditions.checkNotNull(workResponse).getAppliedFinalizeIdsList());
if (!appliedFinalizeIds.isEmpty()) {
streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
}
for (Windmill.ComputationWorkItems computationWork :
Preconditions.checkNotNull(workResponse).getWorkList()) {
String computationId = computationWork.getComputationId();
Expand Down Expand Up @@ -252,6 +247,7 @@ private void applianceDispatchLoop(Supplier<Windmill.GetWorkResponse> getWorkFn)
Work.createProcessingContext(
computationId, getDataClient, workCommitter::commit, heartbeatSender),
computationWork.getDrainMode(),
appliedFinalizeIds,
/* getWorkStreamLatencies= */ ImmutableList.of());
}
}
Expand Down
12 changes: 4 additions & 8 deletions .../apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,19 +215,18 @@ public void scheduleWork(
Watermarks watermarks,
Work.ProcessingContext processingContext,
boolean drainMode,
ImmutableList<Long> appliedFinalizeIds,
ImmutableList<LatencyAttribution> getWorkStreamLatencies) {
// Before any processing starts, call any pending OnCommit callbacks.
commitFinalizer.finalizeCommits(workItem.getSourceState().getFinalizeIdsList());
Comment thread
arunpandianp marked this conversation as resolved.
commitFinalizer.finalizeCommits(appliedFinalizeIds);
Comment thread
arunpandianp marked this conversation as resolved.
Outdated
computationState.activateWork(
ExecutableWork.create(
Work.create(
workItem, serializedWorkItemSize, watermarks, processingContext, drainMode, clock),
(work, handle) -> processWork(computationState, work, getWorkStreamLatencies, handle)));
}

/** Adds any applied finalize ids to the commit finalizer to have their callbacks executed. */
public void queueAppliedFinalizeIds(ImmutableList<Long> appliedFinalizeIds) {
commitFinalizer.finalizeCommits(appliedFinalizeIds);
}

/**
* Executes the user DoFns processing {@link Work} then queues the {@link Commit}(s) to be sent to
* backing persistent store to mark that the {@link Work} has finished processing. May retry
Expand Down Expand Up @@ -255,9 +254,6 @@ private void processWork(
setUpWorkLoggingContext(work.getLatencyTrackingId(), computationId);
LOG.debug("Starting processing for {}:\n{}", computationId, work);

// Before any processing starts, call any pending OnCommit callbacks. Nothing that requires
// cleanup should be done before this, since we might exit early here.
commitFinalizer.finalizeCommits(workItem.getSourceState().getFinalizeIdsList());
if (workItem.getSourceState().getOnlyFinalize()) {
Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem);
outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true));
Expand Down
Loading