Keep progress renewal heartbeats active through bundle teardown - #40020
Draft
kennknowles wants to merge 4 commits into
Draft
Keep progress renewal heartbeats active through bundle teardown#40020kennknowles wants to merge 4 commits into
kennknowles wants to merge 4 commits into
Conversation
BatchDataflowWorker shut down WorkProgressUpdater in the finally block of executeWork() before calling workItemStatusClient.reportSuccess(). If the worker thread was delayed in closing connections, flushing shuffle state, or serializing counters and metrics, no lease renewals could be sent, allowing the 180s work item lease to expire on the Dataflow service. Furthermore, reportSuccess() and reportUpdate() in WorkItemStatusClient both synchronized on the same monitor lock, blocking any concurrent lease renewals while the worker thread was assembling final metrics. To fix this: 1. Extend progress reporting in BatchDataflowWorker so progressUpdater stays active until reportSuccess() completes, while reporting any unreported dynamic split before success. 2. Decouple monitor locking in WorkItemStatusClient so counter and metric extraction run outside the RPC lock, using a dedicated metricsLock to serialize metric extraction and commits. 3. Support lightweight lease renewal pings without heavyweight metric serialization to keep leases alive under high CPU/memory utilization. 4. Gracefully handle in-flight non-final status updates in execute() if a final completion state was already sent.
kennknowles
commented
Sep 4, 2026
| * WorkItemStatusClient#reportSuccess} completes, lease renewal heartbeats continue even if bundle | ||
| * completion or status reporting encounters delays, preventing lease expiration. | ||
| */ | ||
| void executeWork( |
Member
Author
There was a problem hiding this comment.
Don't use method overrides. They are a human-oriented stylistic choice that obscures readability. Give the method a meaningful name that indicates why it takes more parameters.
| */ | ||
| public @Nullable WorkItemServiceState reportLeasePing(Duration requestedLeaseDuration) | ||
| throws Exception { | ||
| checkState(worker != null, "setWorker should be called before reportLeasePing"); |
Member
Author
There was a problem hiding this comment.
Use checkStateNotNull
| if (finalStateSent) { | ||
| return null; | ||
| } | ||
| checkArgument(requestedLeaseDuration != null, "requestLeaseDuration must be non-null"); |
Member
Author
There was a problem hiding this comment.
Use checkArgumentNotNull
|
|
||
| private synchronized @Nullable WorkItemServiceState execute(WorkItemStatus status) | ||
| throws IOException { | ||
| if (finalStateSent && !Boolean.TRUE.equals(status.getCompleted())) { |
Member
Author
There was a problem hiding this comment.
Do not use Boolean.TRUE.equals. getCompleted() is already a boolean.
- Rename executeWork overload taking WorkItemStatusClient to executeWorkAndReportSuccess in BatchDataflowWorker, clarifying that it executes work and completes reporting while progress heartbeats are active. - Use Preconditions.checkStateNotNull and checkArgumentNotNull in WorkItemStatusClient#reportLeasePing. - Avoid Boolean.TRUE.equals when checking status.getCompleted() in WorkItemStatusClient#execute. - Add unit tests for reportLeasePing precondition validations.
The previous change attempted to split synchronization in WorkItemStatusClient via a dedicated metricsLock to allow progress heartbeats during success reporting. However, splitting the lock violated the single-reader invariant of MetricsContainerImpl and BatchModeExecutionContext, introduced a data race on lastReportedMillis, dropped destructively-drained metric updates if progress updates raced reportSuccess(), and caused TOCTOU IllegalStateExceptions. Furthermore, reportLeasePing was dead code never invoked by the worker lifecycle. To address this while fulfilling the requirement to keep leases alive until reportSuccess completes: 1. Remove metricsLock and restore synchronization on WorkItemStatusClient's intrinsic monitor across extract, RPC, and commit. Because the Dataflow service strictly sequences reportIndex per WorkItem, concurrent status RPCs are invalid and monitor serialization is required. 2. Hold the monitor across the extract-RPC-commit sequence so at most one thread extracts or commits metrics, preserving single-reader invariants and preventing dropped counter and metric updates. 3. Wire reportLeasePing into the worker lifecycle: once BatchDataflowWorker completes compute and dynamic splitting, it sets leaseRenewalOnly(true) on the progress updater. Any heartbeat occurring during final bundle teardown and reportSuccess preparation sends a lightweight lease ping without metric extraction overhead. 4. Treat finalStateSent gracefully in reportUpdate, reportLeasePing, and execute, returning null without throwing IllegalStateException or logging spurious errors when a heartbeat races reportSuccess. 5. Add multi-threaded concurrency tests in LeaseRenewalRaceTest to verify that races between reportSuccess and progress updates or lease pings do not throw exceptions or lose metrics.
Revert redundant internal synchronization in DataflowWorkProgressUpdater and restore executor visibility to private in WorkProgressUpdater. Because WorkProgressUpdater's caller methods already hold the monitor on executor, internal synchronized blocks were reentrant no-ops that caused SpotBugs warnings. Additionally: - Use a volatile boolean for leaseRenewalOnly so status updates can be switched to lightweight lease renewal pings after execute() finishes. - In reportProgressHelper(), do not bypass reporting an uncommitted dynamic split if dynamicSplitResultToReport is present when in lease renewal mode. - Reduce reportLeasePing visibility to package-private. - Rename concurrency test to drainedMetricUpdatesSurviveConcurrentReportSuccess and remove the vestigial latch wait that added artificial test suite delay.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
BatchDataflowWorker shut down WorkProgressUpdater in the finally block of executeWork() before calling workItemStatusClient.reportSuccess(). If the worker thread was delayed in closing connections, flushing shuffle state, or serializing counters and metrics, no lease renewals could be sent, allowing the 180s work item lease to expire on the Dataflow service.
Furthermore, reportSuccess() and reportUpdate() in WorkItemStatusClient both synchronized on the same monitor lock, blocking any concurrent lease renewals while the worker thread was assembling final metrics.
To fix this:
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.