GH-1239: Fix memory leak when C Data import hits allocator limit mid-array#1240
Open
sandeshkr419 wants to merge 2 commits into
Open
GH-1239: Fix memory leak when C Data import hits allocator limit mid-array#1240sandeshkr419 wants to merge 2 commits into
sandeshkr419 wants to merge 2 commits into
Conversation
sandeshkr419
requested review from
jbonofre,
laurentgo,
lidavidm and
wgtmac
as code owners
July 20, 2026 22:22
This comment has been minimized.
This comment has been minimized.
sandeshkr419
force-pushed
the
fix/cdata-import-oom-leak
branch
from
July 21, 2026 00:43
b0472b2 to
5607f1f
Compare
There was a problem hiding this comment.
Pull request overview
This PR addresses a native-memory leak in the C Data Interface import path when the target BufferAllocator hits its allocation limit mid-import, preventing the C Data release callback from firing and leaving the producer’s exported memory stranded.
Changes:
- Adjusts
ReferenceCountedArrowArray.unsafeAssociateAllocationto avoid leaking a retain whenwrapForeignAllocationthrows due to allocator limits. - Adds a regression test that exports a batch from a producer allocator, forces an OOM during import into a tiny consumer allocator, and asserts the producer allocator drains to zero.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java |
Changes retain/wrap ordering during buffer association to prevent leaking retains on allocator-limit failures. |
c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java |
Adds a regression test to verify producer-owned exported memory is released even when import fails with OOM mid-array. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
ReferenceCountedArrowArray.unsafeAssociateAllocation retained the imported array's shared reference count *before* calling wrapForeignAllocation, which can throw OutOfMemoryException when the tracking (importing) allocator is over its limit. The retain was not rolled back on that throw, so the reference count stayed permanently elevated: release() never reached zero, the producer's C Data release callback was never invoked, and the entire imported array's buffers leaked in the producer's memory (invisible to the JVM heap). This strands whole batches when importing a wide/variable-width array into a bounded allocator that fills part-way through the array. Retain only after wrapForeignAllocation succeeds, so each retain() is paired with exactly one ForeignAllocation.release0(). If it throws, no ForeignAllocation was created, so skipping the retain is correct: the initial reference held by ArrayImporter.importArray (released in its finally) then drives the count to zero and fires the release callback, freeing the array. Adds ImportOutOfMemoryTest, which exports a multi-buffer batch from a dedicated producer allocator and imports it into an allocator too small to hold it; it asserts the import throws an allocator OOM and that the producer drains to zero (release callback fired). The test fails on the previous retain-before-wrap code (producer strands ~1 MB) and passes with this change.
…with-resources Expand the comment in unsafeAssociateAllocation to explain why retain-after is correct for both throw paths in BaseAllocator.wrapForeignAllocation: the OOM path throws before the ForeignAllocation is associated (release0 is not called), and the unexpected-error path has BaseAllocator's own catch call release0 before rethrowing (a retain-before + catch pattern would double-release on that path). Wrap all child allocators in ImportOutOfMemoryTest in try-with-resources so an early assertion failure does not leave them open for tearDown to find. Signed-off-by: Sandesh Kumar <kusandes@amazon.com>
sandeshkr419
force-pushed
the
fix/cdata-import-oom-leak
branch
from
July 21, 2026 18:24
5607f1f to
9a7c5c3
Compare
sandeshkr419
added a commit
to sandeshkr419/OpenSearch
that referenced
this pull request
Jul 21, 2026
The arrow-java C-Data importer retains a native batch reference before calling allocateBytes. If the target allocator is bounded and throws OutOfMemoryException mid-array, the refcount stays elevated, the C-Data release callback never fires, and the native batch leaks permanently. POOL_QUERY is now registered as a special unmanaged pool (Long.MAX_VALUE, excluded from the rebalancer and budget validation) so imports can never OOM mid-array. Real memory enforcement lives Rust-side (DataFusion MemoryPool + RSS-pressure gate). The pool.query.max and pool.query.min settings are deprecated — they had no effect once the pool is unmanaged. Reverts the per-batch staging workaround from opensearch-project#22448 (DatafusionResultStream, LuceneResultStream) and removes DatafusionImportLeakTests — those were a temporary mitigation. With POOL_QUERY permanently unbounded the trigger cannot fire, making the staging path unnecessary. The definitive fix for the underlying arrow-java bug is apache/arrow-java#1240. Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
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
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.
ReferenceCountedArrowArray.unsafeAssociateAllocationcallsretain()beforewrapForeignAllocation. IfwrapForeignAllocationthrows (allocator over its limit), the retain is never balanced, the reference count stays elevated, and the C Data release callback never fires — leaking the producer's native memory.Fix: call
retain()afterwrapForeignAllocationreturns. Same behavior on the success path; on failure the existingfinallyinArrayImporter.importArraydrives the count to zero and fires the release callback.Test:
ImportOutOfMemoryTest— exports a batch from a "producer" allocator, tries to import it into a too-small consumer, and asserts the producer drains to zero after the OOM. Fails on the original code, passes with the fix.Fixes: #1239