[analytics-engine] Fix native memory leak on Arrow C Data batch import under allocator pressure

[sandeshkr419]

Description

DatafusionResultStream and LuceneResultStream imported each native Arrow batch directly into the caller's BufferAllocator via Data.importIntoVectorSchemaRoot, which charges every imported buffer against that allocator as it walks the array. When the allocator is bounded and its limit is hit part-way through a wide batch (e.g. chained-concat Utf8View columns under concurrent load), the import throws OutOfMemoryException.

arrow-java's ReferenceCountedArrowArray.unsafeAssociateAllocation calls retain() before the throwing wrapForeignAllocation and does not roll it back, so the imported array's reference count stays elevated, the C Data release callback never fires, and the entire native batch leaks in the producer's native allocator — invisible to the JVM heap and the Java Arrow allocator. Under a concurrent chained-concat storm this accumulates multiple GB of native memory that never drains, eventually wedging the node into HTTP 429s until restart. It is visible only in _plugins/arrow_base/stats → memory_pools.runtime.allocated_bytes; the DataFusion pool stats and circuit breakers read 0.

Fix: import each batch into a per-batch staging allocator that is a fresh unbounded child of the root allocator, so a single import can never OOM part-way through the array and the release callback always fires. The batch is returned as-is (zero-copy, no buffer transfer) and released by the existing consumer close paths, which drives the C Data reference count to zero. Staging allocators are reclaimed once drained (per import and at stream close); on import failure the staging allocator is closed immediately, firing the native release for the whole batch. The same fix is applied to both result-stream import paths.

This is independent of whether POOL_QUERY is bounded: making the shared pool unbounded does not fix it, because many concurrent imports still pile onto the one shared pool and hit the node's overall memory ceiling, so an import still OOMs mid-batch and strands. A per-import staging allocator is what prevents the mid-array OOM. The underlying arrow-java bug is being reported upstream separately; this change works around it for the OpenSearch import paths.

Testing

• DatafusionImportLeakTests (new, self-contained — no native runtime): exports a multi-buffer batch across the C Data Interface and imports it back through the production staging path under a caller allocator far too small to hold it; asserts the producer allocator (standing in for the native producer) fully drains. This test fails on the unfixed direct-import path (Memory was leaked ... (8521000)) and passes with the fix, so it is a genuine regression guard. A companion test pins the arrow-java mid-import-OOM leak deterministically.
• DatafusionResultStreamTests.testStreamConsumeAndCloseDrainsAllocator verifies the staging import path end-to-end against a real native stream under a small caller allocator and asserts the shared root allocator drains after consume + close.
• Validated on a live 1-node/3-shard m8g.large cluster (textbench_50m) under a 16-worker chained-concat storm: stock leaks ~3.1 GB that stays flat at idle forever; with the fix native allocated_bytes drains from ~2.2 GB back to ~13 MB after every storm, in both bounded and unbounded POOL_QUERY configurations.

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 1b0cd41)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 Multiple PR themes Sub-PR theme: Fix native leak on DataFusion result stream import Relevant files: sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionImportLeakTests.java sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java Sub-PR theme: Fix native leak on Lucene result stream import Relevant files: sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java

⚡ Recommended focus areas for review Unbounded staging allocator list stagingAllocators accumulates one entry per imported batch, and reclaimDrainedStaging() is only invoked at the start of each importBatch and on close. If a consumer holds onto multiple batches concurrently (or releases them out of order), drained allocators for earlier batches are only reclaimed the next time importBatch runs. In a long-running stream where the consumer eagerly buffers batches, this list can grow, and each open child allocator keeps parent bookkeeping alive. Consider reclaiming after each consumer close or bounding the list. private VectorSchemaRoot importBatch(ArrowArray arrowArray) { reclaimDrainedStaging(); BufferAllocator staging = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); try { VectorSchemaRoot root = importOntoStaging(staging, schema, arrowArray, dictionaryProvider); stagingAllocators.add(staging); return root; } catch (RuntimeException e) { staging.close(); throw e; } } /** * Closes staging allocators whose batches have been fully released (drained to zero). A batch still * in flight keeps its staging allocator open so the eventual release callback frees the small C Data * bookkeeping allocation against a live allocator; that allocator is a leaf child of the root and * holds no batch data once drained. */ private void reclaimDrainedStaging() { stagingAllocators.removeIf(a -> { if (a.getAllocatedMemory() == 0) { a.close(); return true; } return false; }); } Unbounded native memory via staging Staging allocators are created as unbounded children of the root (Long.MAX_VALUE). This intentionally avoids mid-import OOM, but it also bypasses the caller's bounded allocator limit entirely: every batch is now charged against the root rather than the caller-supplied bound. Under concurrent chained-concat load the node's overall memory ceiling could still be exceeded without ever triggering the caller's circuit breaker. Confirm that root-level accounting/limits are enforced elsewhere; otherwise the fix trades a mid-import leak for uncapped growth. BufferAllocator staging = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE);

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 1b0cd41

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Ensure staging allocators close on stream close reclaimDrainedStaging() only closes staging allocators that have drained to zero; any allocator still holding memory (e.g., a batch the consumer failed to close, or a partial reclamation) will be silently leaked on stream close. Force-close all remaining staging allocators on stream close (or at least log/track them) so they are not orphaned when the stream terminates. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [75-84] public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.reclaimDrainedStaging(); + iteratorInstance.closeAllStaging(); } } finally { try { streamHandle.close(); Suggestion importance[1-10]: 6 __ Why: Valid concern: reclaimDrainedStaging() on close only reclaims allocators drained to zero, so any still-holding staging allocator would be leaked. However, per the PR's design, batches must be released by the consumer before stream close, so this is more of a defensive robustness improvement.	Low
Possible issue	Guard staging allocator against creation-time failures VectorSchemaRoot.create may itself throw (e.g., OOM allocating initial buffers) before the try block, which would leak staging. Wrap the VectorSchemaRoot.create call within the try/catch, or use a try/catch that guarantees staging.close() on any failure prior to registration. sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java [165-178] private VectorSchemaRoot importBatch() { reclaimDrainedStaging(); BufferAllocator staging = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); - VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); try { - Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); + try { + Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); + } catch (RuntimeException e) { + root.close(); + throw e; + } + stagingAllocators.add(staging); + return root; } catch (RuntimeException e) { - root.close(); staging.close(); throw e; } - stagingAllocators.add(staging); - return root; } Suggestion importance[1-10]: 5 __ Why: Correctly identifies that VectorSchemaRoot.create sits outside the try/catch, so a failure there would leak the staging allocator. This is a minor robustness issue since create is unlikely to fail against an unbounded staging allocator, but the fix aligns Lucene with the Datafusion version.	Low

---

Previous suggestions

Suggestions up to commit 29cd8ce

Category	Suggestion	Impact
Possible issue	Guard staging allocator against create failure VectorSchemaRoot.create itself can throw (e.g., OOM allocating validity/offset buffers on the staging allocator's parent chain, or schema issues), which would leak the freshly opened staging allocator since it's not yet in the try/catch. Wrap the VectorSchemaRoot.create call in the try block or a separate try-catch that closes staging on failure. sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java [165-178] private VectorSchemaRoot importBatch() { reclaimDrainedStaging(); BufferAllocator staging = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); - VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); + VectorSchemaRoot root; try { + root = VectorSchemaRoot.create(schema, staging); Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); } catch (RuntimeException e) { - root.close(); staging.close(); throw e; } stagingAllocators.add(staging); return root; } Suggestion importance[1-10]: 6 __ Why: Valid observation: if VectorSchemaRoot.create throws before the try block, the freshly opened staging allocator would leak. Moving the create into the try/catch is a reasonable defensive improvement, though such failures are unlikely on an unbounded staging allocator.	Low
General	Force-close leftover staging allocators on close reclaimDrainedStaging only closes staging allocators whose allocated memory is zero, leaving any non-drained staging allocators leaked when the stream is closed. On close(), after the consumer has been given a chance to release, remaining staging allocators should be force-closed (or logged) to avoid a slow leak when a consumer forgets to close a batch. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [75-84] public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); iteratorInstance.reclaimDrainedStaging(); + iteratorInstance.forceCloseStagingAllocators(); } } finally { try { streamHandle.close(); Suggestion importance[1-10]: 4 __ Why: The suggestion identifies a real potential leak if consumers forget to close batches, but forcibly closing non-drained staging allocators could cause issues since in-flight batches still reference them for the C Data release callback. The suggestion references a method forceCloseStagingAllocators that doesn't exist and doesn't fully address the correctness concerns.	Low

Suggestions up to commit 83a65b1

Category	Suggestion	Impact
Possible issue	Close staging allocator on any construction failure VectorSchemaRoot.create itself can throw (e.g. if the schema causes allocation failure), in which case staging would leak because it is only closed inside the catch that wraps the import call. Wrap the VectorSchemaRoot.create call in the same try/catch (or in a separate try that closes staging on failure) to guarantee the staging allocator is closed on any construction failure, mirroring the Datafusion implementation's stricter ordering. sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java [167-177] BufferAllocator staging = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); -VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); try { - Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); + try { + Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); + } catch (RuntimeException e) { + root.close(); + throw e; + } + stagingAllocators.add(staging); + return root; } catch (RuntimeException e) { - root.close(); staging.close(); throw e; } -stagingAllocators.add(staging); -return root; Suggestion importance[1-10]: 6 __ Why: Correct observation: if VectorSchemaRoot.create throws, the staging allocator leaks. The Datafusion counterpart handles this properly via importOntoStaging wrapped in try/catch. This aligns the Lucene path with the stricter Datafusion ordering.	Low
General	Ensure staging allocators cannot leak on close reclaimDrainedStaging only closes staging allocators that are already drained to zero; any still-open staging allocators (e.g. those whose consumers never released the batch, or on abnormal shutdown) will silently leak. Consider force-closing remaining staging allocators on stream close, or logging a warning, so the leak surface at shutdown is bounded. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java [75-84] public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); iteratorInstance.reclaimDrainedStaging(); + iteratorInstance.forceCloseRemainingStaging(); } } finally { try { streamHandle.close(); Suggestion importance[1-10]: 5 __ Why: Valid concern: staging allocators still in-flight at close won't be reclaimed by reclaimDrainedStaging, potentially masking leaks. However, force-closing allocators with live buffers could cause errors, so the suggested fix would need care.	Low

[github-actions]

✅ Gradle check result for 83a65b1: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.48%. Comparing base (d69286a) to head (1b0cd41).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22448      +/-   ##
============================================
+ Coverage     73.44%   73.48%   +0.03%     
+ Complexity    76488    76479       -9     
============================================
  Files          6104     6104              
  Lines        346576   346578       +2     
  Branches      49885    49886       +1     
============================================
+ Hits         254558   254676     +118     
+ Misses        71736    71611     -125     
- Partials      20282    20291       +9     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 29cd8ce

[github-actions]

❕ Gradle check result for 29cd8ce: UNSTABLE

Please review all flaky tests that succeeded after retry and create an issue if one does not already exist to track the flaky failure.

[mch2]

this doesn't feel right - we are storing allocators for each batch and only clearing when last batch is consumed? We should be releasing each batch allocator as its read?

[sandeshkr419]

Reclaim actually runs at the start of every importBatch(), not just at last-batch — each drained staging allocator is closed as soon as its batch's buffers are released (getAllocatedMemory() == 0), and the call in close() is only a final sweep. So it's incremental, not stored-till-the-end.

[github-actions]

Persistent review updated to latest commit 1b0cd41

[github-actions]

✅ Gradle check result for 1b0cd41: SUCCESS