perf(flink): preempt inactive write buckets on memory exhaustion - #19728
perf(flink): preempt inactive write buckets on memory exhaustion#19728fhan688 wants to merge 3 commits into
Conversation
hudi-agent
left a comment
hudi-agent
left a comment
There was a problem hiding this comment.
Thanks for the contribution! This PR adds a PreemptiveMemorySegmentPool that lets a Flink write bucket reclaim managed memory from the largest inactive bucket (flushing and disposing it) before falling back to the diverged-record path introduced in #19692, with owner tracking, a bounded single retry, and a recursion guard. I traced the preemption mechanism, the owner-tracking lifecycle, reentrancy of the mid-write flush, the divergence interaction, tracer accounting, and subclass coverage, and did not flag any correctness issues from this automated pass — a Hudi committer or PMC member can take it from here for a final review. One minor naming nit on the new boolean field; otherwise the code is clean and well-structured.
cc @yihua
|
|
||
| @Nullable | ||
| private String currentOwnerId; | ||
| private boolean preempting; |
There was a problem hiding this comment.
🤖 nit: could you rename preempting to isPreempting? Boolean fields typically read as predicates in this codebase (e.g. isDiverged), and the is prefix makes the guard in nextSegment() a bit easier to parse at a glance.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #19728 +/- ##
============================================
- Coverage 77.88% 73.79% -4.10%
+ Complexity 33264 31759 -1505
============================================
Files 2533 2540 +7
Lines 140342 141003 +661
Branches 16912 17162 +250
============================================
- Hits 109310 104048 -5262
- Misses 23401 29766 +6365
+ Partials 7631 7189 -442
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
cshuo
left a comment
cshuo
left a comment
There was a problem hiding this comment.
Thks for the optimization, left some comments.
|
|
||
| @Nullable | ||
| private String currentOwnerId; | ||
| private boolean preempting; |
There was a problem hiding this comment.
Is preempting flag necessary, Could you clarify the purpose of the flag or add some comments.
There was a problem hiding this comment.
Is
preemptingflag necessary, Could you clarify the purpose of the flag or add some comments.
Yes, the flag is needed as a re-entrancy guard.
The reclamation callback flushes and disposes a victim while that victim may still be registered. If the callback path transitively requests another page, re-entering the reclaimer could select and flush the same victim recursively.
I kept the flag and added comments to clarify both its purpose and the guarded branch in nextSegment(). This behavior is also covered by testRetriesAllocationOnceWithoutNestedPreemption.
| @Override | ||
| public MemorySegment nextSegment() { | ||
| MemorySegment segment = delegate.nextSegment(); | ||
| if (segment != null || currentOwnerId == null || preempting) { |
There was a problem hiding this comment.
Could you clarify when will currentOwnerId be null?
There was a problem hiding this comment.
Could you clarify when will
currentOwnerIdbe null?
currentOwnerId is non-null only while bucket.writeRow() is serializing a row. It is null before and after writeRow(), including while a new buffer is being created.
For a buffer-creation allocation failure, owner-aware preemption is intentionally skipped and the existing creation-failure recovery in StreamWriteFunction flushes the largest non-empty bucket and retries the creation.
I expanded the condition and added comments to make these cases explicit.
There was a problem hiding this comment.
Thanks for the clarification. I don't think we need to expand the condition here. The existing behavior looks fine; adding some comments to explain should be sufficient.
| @@ -149,6 +150,8 @@ public class StreamWriteFunction extends AbstractStreamWriteFunction<HoodieFlink | |||
|
|
|||
| protected transient MemorySegmentPool memorySegmentPool; | |||
There was a problem hiding this comment.
Do we still need to keep memorySegmentPool, seems we can use preemptiveMemorySegmentPool throughout.
There was a problem hiding this comment.
Do we still need to keep
memorySegmentPool, seems we can usepreemptiveMemorySegmentPoolthroughout.
Good point. Both fields referenced the same wrapped pool after initialization.
I removed memorySegmentPool and now use preemptiveMemorySegmentPool throughout, including buffer creation, owner tracking, cleanup, and the test accessor.
There was a problem hiding this comment.
As discussed in #19692 (comment), Now that memory exhaustion during writeRow() already triggers MemoryReclaimer#reclaim to flush the largest inactive bucket, we don't need to eagerly flush bucketToFlush when failedBucket is diverged?
There was a problem hiding this comment.
As discussed in #19692 (comment), Now that memory exhaustion during
writeRow()already triggersMemoryReclaimer#reclaimto flush the largest inactive bucket, we don't need to eagerly flushbucketToFlushwhen failedBucket is diverged?
Agreed. With owner-aware preemption, an allocation failure during writeRow() already attempts to flush the largest inactive bucket before the current bucket can diverge. Eagerly flushing another bucket after divergence is therefore redundant.
I simplified the diverged path to flush and dispose only the failed bucket before retrying. The largest-bucket fallback is still retained for buffer-creation failures, where there is no active owner and preemptive reclamation is not triggered.
hudi-agent
left a comment
hudi-agent
left a comment
There was a problem hiding this comment.
Thanks for the contribution! This PR adds a PreemptiveMemorySegmentPool that lets a Flink write bucket reclaim managed memory from the largest inactive non-empty bucket (flush + dispose) and retry allocation once before falling back to the existing diverged-record path. I traced the divergence/preemption interaction, the re-entrancy guard (preempting), the currentOwnerId lifecycle, page/tracer accounting, and the reclaimMemoryAfterFailedWrite simplification against the PR-head sources, and the logic holds up — the removed dual-flush is redundant once a bucket diverges (divergence only occurs after preemptMemory already found no inactive non-empty bucket), and a preemption flush failure propagates out of nextSegment rather than being silently swallowed by BinaryInMemorySortBuffer.write (which only catches EOFException). No new issues flagged from this automated pass beyond the points already under discussion in the inline threads — a Hudi committer or PMC member can take it from here for a final review.
cc @yihua
cshuo
left a comment
cshuo
left a comment
There was a problem hiding this comment.
Thks for the updating, some minor comments.
| .orElse(null); | ||
|
|
||
| if (failedBucket == null) { | ||
| RowDataBucket bucketToFlush = findLargestNonEmptyBucketExcluding(bucketID); |
There was a problem hiding this comment.
Since this branch performs the same largest-non-empty-bucket selection and flush as preemptMemory(bucketID), we can reuse it here:
if (!preemptMemory(bucketID)) {
throw new HoodieException(
"Not enough memory pages to create a RowData buffer and no non-empty bucket can be flushed");
}
return;
| @Override | ||
| public MemorySegment nextSegment() { | ||
| MemorySegment segment = delegate.nextSegment(); | ||
| if (segment != null || currentOwnerId == null || preempting) { |
There was a problem hiding this comment.
Thanks for the clarification. I don't think we need to expand the condition here. The existing behavior looks fine; adding some comments to explain should be sufficient.
| int initialFreePages = pipeline.freePages(); | ||
| try { | ||
| boolean reclaimedOtherBucketBeforeDivergedBucket = false; | ||
| boolean preemptedInactiveBucket = false; |
There was a problem hiding this comment.
The updated test verifies successful preemption of an inactive bucket, but it no longer explicitly verifies the fallback when no eligible victim exists.
Could we retain or add a scenario where only the current bucket holds buffered rows, the next record exhausts the remaining pages, and the record fits after the diverged bucket is flushed and disposed? This would cover reclaim() returning false, followed by flushing the diverged bucket and successfully retrying the record, while also verifying that no records or memory pages are lost.
hudi-agent
left a comment
hudi-agent
left a comment
There was a problem hiding this comment.
Thanks for working on this! This PR adds a PreemptiveMemorySegmentPool that lets a Flink write bucket reclaim managed memory by flushing the largest inactive bucket before falling back to the existing diverged-record path. I traced the reclaim/retry flow, the re-entrancy guard, owner lifecycle, tracer accounting, and the simplified diverged-bucket branch, and did not find a correctness issue in this pass. No correctness issues found. A few style/readability suggestions in the inline comments. Please take a look, and this should be ready for a Hudi committer or PMC member to take it from here.; a Hudi committer or PMC member can take it from here for a final review.
. One minor naming-consistency suggestion below; otherwise the code is clean and well-documented.
cc @yihua
| boolean reclaim(String excludedOwnerId); | ||
| } | ||
|
|
||
| private final MemorySegmentPool delegate; |
There was a problem hiding this comment.
🤖 nit: the code mixes "preempt" (class name, StreamWriteFunction#preemptMemory) and "reclaim" (MemoryReclaimer#reclaim, memoryReclaimer) for the same concept — could you settle on one verb so the vocabulary stays consistent?
5 participants
Describe the issue this Pull Request addresses
Closes #19664.
This is a follow-up to #19692, which made the Flink writer gracefully fall back to a diverged-record bucket when a write bucket cannot acquire additional memory segments.
However, the current bucket may fail to allocate a memory segment while reclaimable segments are still retained by other inactive write buckets. In that case, directly using the diverged-record fallback can cause unnecessary bucket divergence and may leave reclaimable managed memory unused.
This PR allows the writer to preemptively flush and dispose an inactive write bucket before falling back to the existing divergence path.
Summary and Changelog
When a write bucket cannot allocate another memory segment, the writer now attempts to reclaim memory from the largest non-empty inactive bucket and retries the allocation once.
The changes include:
PreemptiveMemorySegmentPool, which wraps an existingMemorySegmentPool.writeRowas the active memory owner.No code was copied from another project.
Impact
There are no changes to the storage format, public APIs, configuration options, or default configuration values.
For normal writes with sufficient memory, behavior remains unchanged apart from lightweight active-bucket tracking.
Under memory pressure, an inactive bucket may be flushed earlier so that its memory segments can be reused by the active bucket. This reduces unnecessary bucket divergence and improves progress when memory is distributed across multiple buffered buckets.
Earlier flushing under severe memory pressure may produce smaller write batches, but this only occurs after a segment allocation has already failed.
Risk Level
Medium.
The change affects the Flink writer's bucket flushing and managed-memory lifecycle under memory exhaustion. The risk is mitigated by:
The following validations passed:
TestPreemptiveMemorySegmentPool: 5 tests passed.TestBucketStreamWriteMemoryExhaustion: 2 parameterized scenarios passed for the default and LSM Tree layouts.git diff --checkpassed.Documentation Update
None.
This is an internal memory-reclamation improvement and does not introduce or modify any user-facing configuration or API.
Contributor's checklist