From 83ec94bc006576951ab22e88085e37437ffc3650 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Fri, 21 Aug 2026 07:58:46 +0000 Subject: [PATCH 1/3] Use one stream-scoped Arrow import staging allocator instead of closing one per batch Signed-off-by: Lantao Jin --- .../be/datafusion/DatafusionResultStream.java | 85 +++++--- ...afusionStagingAllocatorLifecycleTests.java | 182 ++++++++++++++++++ .../be/lucene/LuceneResultStream.java | 63 +++--- 3 files changed, 277 insertions(+), 53 deletions(-) create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java index c07574b76abed..8eac3cd317bef 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java @@ -27,7 +27,6 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.action.ActionListener; -import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; @@ -76,7 +75,7 @@ public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.reclaimDrainedStaging(); + iteratorInstance.closeStagingAllocator(); } } finally { try { @@ -99,9 +98,22 @@ static class BatchIterator implements Iterator { private Boolean nextAvailable; private boolean batchEmitted; private boolean nativeStreamExhausted; - // Per-batch staging allocators used by {@link #importBatch}. Each is reclaimed once its batch's - // buffers have been released by the consumer (see {@link #reclaimDrainedStaging}). - private final List stagingAllocators = new ArrayList<>(); + /** + * ONE staging allocator for the whole stream, created lazily on first import (see + * {@link #importBatch}). Deliberately NOT per-batch: the Flight transport builds its reused stream + * root on {@code fieldVectors.getFirst().getAllocator()}, i.e. the FIRST batch's staging allocator + * ({@code FlightServerChannel#transferIntoStreamRoot}, whose comment states "The producer's + * allocator must be long-lived (not closed per-request)"), and then charges that same allocator for + * every later batch via {@code BaseFixedWidthVector#transferTo}. Closing a staging allocator + * per-batch therefore closed an allocator the transport was still using, and + * {@code BaseAllocator#close} threw {@code IllegalStateException: Memory was leaked by query} + * whenever a transfer landed in the window between the drained-check and the close. + * + *

Unboundedness is the only property the original per-batch design actually needed (so an import + * cannot OOM part-way through a C Data array — see {@link #importBatch}); a single unbounded child + * preserves that while removing the check-then-close race entirely. + */ + private BufferAllocator stagingAllocator; BatchIterator(StreamHandle streamHandle, BufferAllocator allocator, CDataDictionaryProvider dictionaryProvider) { this.streamHandle = streamHandle; @@ -159,36 +171,53 @@ private boolean loadNextBatch() { * unbounded staging child can't OOM mid-array, so the release callback always fires. * *

The batch is returned as-is (zero-copy); its buffers are released by the existing consumer close - * paths, which drives the C Data reference count to zero. Each staging allocator is reclaimed once - * drained (see {@link #reclaimDrainedStaging}); on import failure it is closed immediately. + * paths, which drives the C Data reference count to zero. The staging allocator is stream-scoped and + * outlives every batch (see {@link #stagingAllocator}), so nothing is closed per-batch; on import + * failure the partially-imported root is closed by {@link #importOntoStaging} but the allocator + * itself stays open for subsequent batches. */ 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; + if (stagingAllocator == null) { + stagingAllocator = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); } + return importOntoStaging(stagingAllocator, schema, arrowArray, dictionaryProvider); } /** - * 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. + * Closes the stream-scoped staging allocator, if it can be closed. Called only from + * {@link DatafusionResultStream#close()} — never per-batch. + * + *

A non-zero balance here means the transport still holds buffers charged to this allocator: the + * Flight channel's reused stream root is freed in its own {@code close()}, which may run after ours. + * Closing anyway would throw {@code IllegalStateException} from {@code BaseAllocator#close} AND — worse + * — leave the allocator permanently half-closed, because {@code close()} sets {@code isClosed = true} + * BEFORE its leak check, so its bytes would never be returned to the parent and a later retry would + * early-return as a no-op. Leaving it open hands ownership to the root allocator, which is the same + * outcome the previous per-batch code produced for any still-in-flight batch. */ - private void reclaimDrainedStaging() { - stagingAllocators.removeIf(a -> { - if (a.getAllocatedMemory() == 0) { - a.close(); - return true; - } - return false; - }); + void closeStagingAllocator() { + if (stagingAllocator == null) { + return; + } + if (stagingAllocator.getAllocatedMemory() == 0) { + stagingAllocator.close(); + stagingAllocator = null; + } + } + + /** The stream-scoped staging allocator, or null before the first import. For tests. */ + BufferAllocator stagingAllocator() { + return stagingAllocator; + } + + /** + * Drives the exact production {@link #importBatch} path with an explicitly supplied schema, so the + * staging-allocator lifetime regression test can import several batches without a native stream. + * {@code schema} is normally set by {@link #ensureSchema()} from the native handle. + */ + VectorSchemaRoot importBatchForTest(Schema batchSchema, ArrowArray arrowArray) { + this.schema = batchSchema; + return importBatch(arrowArray); } /** diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java new file mode 100644 index 0000000000000..b00ed74409a5b --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java @@ -0,0 +1,182 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion; + +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.CDataDictionaryProvider; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.List; + +/** + * Regression tests for the staging-allocator LIFETIME contract, i.e. the bug behind + * {@code IllegalStateException: Memory was leaked by query. Memory leaked: (1024)} thrown from + * {@code BaseAllocator#close} via the old {@code reclaimDrainedStaging}. + * + *

Why this asserts the contract rather than the race. The failure was a data race between the + * producer thread (sweeping drained staging allocators) and the Flight executor thread (charging the + * reused stream root's allocator in {@code BaseFixedWidthVector#transferTo}). The window is a couple of + * statements wide with no injection point, so a thread-racing test would reproduce only probabilistically. + * What IS deterministic is the invariant the transport depends on and the old code broke: + * {@code FlightServerChannel#transferIntoStreamRoot} builds its long-lived stream root on the FIRST + * emitted batch's vector allocator, and its comment states "The producer's allocator must be long-lived + * (not closed per-request)". So: + *

    + *
  • every batch of one stream must be imported on the SAME allocator, and
    • + *
  • that allocator must still be OPEN after later batches are imported and earlier ones released.
    • + *
+ * The old per-batch design violated both: batch N+1 got a fresh allocator, and batch N's was closed as + * soon as its own buffers drained — while the transport was still charging it. + */ +public class DatafusionStagingAllocatorLifecycleTests extends OpenSearchTestCase { + + private static final int ROWS = 8192; // DataFusion's default batch size + private RootAllocator root; + // Resources the production DatafusionResultStream#close would own; the test must release them itself. + private final List providers = new ArrayList<>(); + private final List producers = new ArrayList<>(); + private final List iterators = new ArrayList<>(); + + @Override + public void setUp() throws Exception { + super.setUp(); + root = new RootAllocator(Long.MAX_VALUE); + } + + @Override + public void tearDown() throws Exception { + // Release in the order production does: staging allocator, then the dictionary provider, then the + // stand-in native producer allocators. Only then can the root close cleanly. + for (DatafusionResultStream.BatchIterator it : iterators) { + it.closeStagingAllocator(); + } + for (CDataDictionaryProvider dp : providers) { + dp.close(); + } + for (BufferAllocator p : producers) { + p.close(); + } + root.close(); + super.tearDown(); + } + + /** + * Imports three batches through the production path and asserts all three land on the SAME staging + * allocator, which is still open at the end. + * + *

Fails on the pre-fix code: each {@code importBatch} minted a new + * {@code datafusion-import-staging} child, so the allocators differ between batches. + */ + public void testAllBatchesOfAStreamShareOneOpenStagingAllocator() throws Exception { + DatafusionResultStream.BatchIterator it = newIterator(); + + BufferAllocator first = null; + for (int i = 0; i < 3; i++) { + VectorSchemaRoot imported = importOneBatch(it); + BufferAllocator batchAllocator = imported.getFieldVectors().getFirst().getAllocator(); + if (first == null) { + first = batchAllocator; + } else { + assertSame( + "every batch of a stream must import onto the SAME staging allocator — the Flight " + + "stream root is built on the first batch's allocator and reused for all later batches", + first, + batchAllocator + ); + } + // The consumer releases this batch. Under the old code this drained the allocator to zero and + // made it eligible for the per-batch sweep on the NEXT import. + imported.close(); + } + + assertNotNull("staging allocator must exist after importing", it.stagingAllocator()); + assertSame("the tracked staging allocator is the one batches were imported on", first, it.stagingAllocator()); + // The decisive assertion: still usable by the transport after later imports + earlier releases. + assertEquals("staging allocator must NOT have been closed per-batch", 0L, it.stagingAllocator().getAllocatedMemory()); + it.stagingAllocator().assertOpen(); + } + + /** + * Pins the exact sequence that threw in production: import batch 1, release it (allocator drains to + * zero), then import batch 2 — under the old code the sweep at the head of {@code importBatch} closed + * batch 1's allocator, which is the very allocator the Flight stream root was built on. + */ + public void testReleasingABatchDoesNotCloseTheAllocatorTheTransportHolds() throws Exception { + DatafusionResultStream.BatchIterator it = newIterator(); + + VectorSchemaRoot batch1 = importOneBatch(it); + BufferAllocator transportAllocator = batch1.getFieldVectors().getFirst().getAllocator(); + batch1.close(); // consumer drains it to zero + assertEquals( + "precondition: allocator is drained, i.e. sweep-eligible under the old design", + 0L, + transportAllocator.getAllocatedMemory() + ); + + VectorSchemaRoot batch2 = importOneBatch(it); // old code swept batch1's allocator here + try { + // Simulates FlightServerChannel#transferIntoStreamRoot charging the long-lived allocator for a + // LATER batch. On a closed allocator this throws (IllegalStateException under -ea, which tests + // run with) — that is the production failure, reproduced deterministically. + try (VectorSchemaRoot streamRoot = VectorSchemaRoot.create(batch2.getSchema(), transportAllocator)) { + streamRoot.allocateNew(); + streamRoot.setRowCount(1); + } + } finally { + batch2.close(); + } + } + + // ── helpers ──────────────────────────────────────────────────────────────────────────────────── + + private DatafusionResultStream.BatchIterator newIterator() { + // streamHandle is never touched by importBatch, so a null handle is safe and keeps the test free of + // the native runtime (no .so needed). + CDataDictionaryProvider dp = new CDataDictionaryProvider(); + providers.add(dp); + DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(null, root, dp); + iterators.add(it); + return it; + } + + /** Exports a fresh batch across the C Data Interface and imports it via the production path. */ + private VectorSchemaRoot importOneBatch(DatafusionResultStream.BatchIterator it) { + BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); + producers.add(producer); + try (ArrowArray array = ArrowArray.allocateNew(producer); ArrowSchema cSchema = ArrowSchema.allocateNew(producer)) { + try (VectorSchemaRoot source = VectorSchemaRoot.create(intSchema(producer), producer)) { + IntVector v = (IntVector) source.getVector(0); + v.allocateNew(ROWS); + for (int i = 0; i < ROWS; i++) { + v.set(i, i); + } + source.setRowCount(ROWS); + Data.exportVectorSchemaRoot(producer, source, null, array, cSchema); + } + // Consume the exported C schema (not source.getSchema()) so its C-side release callback fires; + // an unconsumed ArrowSchema strands its exported children in the producer allocator. + Schema schema = Data.importSchema(root, cSchema, providers.getLast()); + return it.importBatchForTest(schema, array); + } + } + + private static Schema intSchema(BufferAllocator alloc) { + try (IntVector probe = new IntVector("n", alloc)) { + return new Schema(List.of(probe.getField())); + } + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java index 357b40fcaea76..4cf95456582f1 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java @@ -23,7 +23,6 @@ import org.opensearch.analytics.exec.ArrowValues; import org.opensearch.common.annotation.ExperimentalApi; -import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; @@ -80,7 +79,7 @@ public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.reclaimDrainedStaging(); + iteratorInstance.closeStagingAllocator(); } } finally { try { @@ -112,9 +111,18 @@ static class BatchIterator implements Iterator { private Boolean nextAvailable; private boolean batchEmitted; private boolean exhausted; - // Per-batch staging allocators used by {@link #importBatch}. Each is reclaimed once its batch's - // buffers have been released by the consumer (see {@link #reclaimDrainedStaging}). - private final List stagingAllocators = new ArrayList<>(); + /** + * ONE staging allocator for the whole stream, created lazily on first import (see + * {@link #importBatch}). Deliberately NOT per-batch: the Flight transport builds its reused stream + * root on {@code fieldVectors.getFirst().getAllocator()}, i.e. the FIRST batch's staging allocator + * ({@code FlightServerChannel#transferIntoStreamRoot}, whose comment states "The producer's + * allocator must be long-lived (not closed per-request)"), then charges that same allocator for every + * later batch. Closing one per-batch closed an allocator the transport was still using, so + * {@code BaseAllocator#close} threw {@code IllegalStateException: Memory was leaked by query} when a + * transfer landed between the drained-check and the close. Unboundedness is the only property the + * per-batch design needed; one unbounded child keeps it without the race. + */ + private BufferAllocator stagingAllocator; BatchIterator( ArrowArray arrowArray, @@ -159,38 +167,43 @@ private boolean loadNextBatch() { * mid-array, so the release callback always fires. * *

The batch is returned as-is (zero-copy); its buffers are released by the existing consumer close - * paths, which drives the C Data reference count to zero. Each staging allocator is reclaimed once - * drained (see {@link #reclaimDrainedStaging}); on import failure it is closed immediately. + * paths, which drives the C Data reference count to zero. The staging allocator is stream-scoped and + * outlives every batch (see {@link #stagingAllocator}), so nothing is closed per-batch. */ private VectorSchemaRoot importBatch() { - reclaimDrainedStaging(); - BufferAllocator staging = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); - VectorSchemaRoot root = VectorSchemaRoot.create(schema, staging); + if (stagingAllocator == null) { + stagingAllocator = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); + } + VectorSchemaRoot root = VectorSchemaRoot.create(schema, stagingAllocator); try { - Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); + Data.importIntoVectorSchemaRoot(stagingAllocator, arrowArray, root, dictionaryProvider); } catch (RuntimeException e) { + // Close the partially-imported root (fires the native release) but KEEP the allocator: it is + // stream-scoped and a later batch may still use it. root.close(); - staging.close(); throw e; } - stagingAllocators.add(staging); return root; } /** - * 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. + * Closes the stream-scoped staging allocator, if it can be closed. Called only from + * {@link LuceneResultStream#close()} — never per-batch. + * + *

A non-zero balance means the transport still holds buffers charged here (its reused stream root + * is freed in its own {@code close()}, which may run after ours). Closing anyway would throw AND leave + * the allocator permanently half-closed, because {@code BaseAllocator#close} sets {@code isClosed} + * BEFORE its leak check — so its bytes would never return to the parent. Leaving it open hands + * ownership to the root, matching what the old per-batch code did for an in-flight batch. */ - private void reclaimDrainedStaging() { - stagingAllocators.removeIf(a -> { - if (a.getAllocatedMemory() == 0) { - a.close(); - return true; - } - return false; - }); + void closeStagingAllocator() { + if (stagingAllocator == null) { + return; + } + if (stagingAllocator.getAllocatedMemory() == 0) { + stagingAllocator.close(); + stagingAllocator = null; + } } @Override From 70f32a2ca5ffcef4b9b8be59bc4f215450312283 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Fri, 21 Aug 2026 08:43:13 +0000 Subject: [PATCH 2/3] address comment Signed-off-by: Lantao Jin --- .../be/datafusion/DatafusionResultStream.java | 24 +++++++++++++++++- .../be/lucene/LuceneResultStream.java | 25 +++++++++++++++++-- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java index 8eac3cd317bef..d5e375cab6335 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java @@ -18,6 +18,8 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.analytics.backend.EngineResultBatch; import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.exec.ArrowValues; @@ -45,6 +47,8 @@ @ExperimentalApi public class DatafusionResultStream implements EngineResultStream, FragmentResources.MetricsCapable { + private static final Logger LOGGER = LogManager.getLogger(DatafusionResultStream.class); + private final StreamHandle streamHandle; private final BufferAllocator allocator; private final CDataDictionaryProvider dictionaryProvider; @@ -194,6 +198,10 @@ private VectorSchemaRoot importBatch(ArrowArray arrowArray) { * BEFORE its leak check, so its bytes would never be returned to the parent and a later retry would * early-return as a no-op. Leaving it open hands ownership to the root allocator, which is the same * outcome the previous per-batch code produced for any still-in-flight batch. + * + *

The deferral is logged at DEBUG, not WARN: the transport frees its stream root asynchronously + * (posted to the flight executor by {@code FlightServerChannel#close}), so a non-zero balance here is + * the expected outcome of every streaming query, not a signal of a leak. */ void closeStagingAllocator() { if (stagingAllocator == null) { @@ -202,6 +210,13 @@ void closeStagingAllocator() { if (stagingAllocator.getAllocatedMemory() == 0) { stagingAllocator.close(); stagingAllocator = null; + } else { + LOGGER.debug( + "Deferring close of staging allocator [{}] with {} bytes outstanding; the transport still " + + "holds them and frees them with its stream root", + stagingAllocator.getName(), + stagingAllocator.getAllocatedMemory() + ); } } @@ -237,7 +252,14 @@ static VectorSchemaRoot importOntoStaging( try { Data.importIntoVectorSchemaRoot(staging, arrowArray, root, dictionaryProvider); } catch (RuntimeException e) { - root.close(); + // Releasing a partially-imported root can itself throw (VectorSchemaRoot#close rethrows any + // RuntimeException from the vectors' release). Attach it rather than let it mask the import + // failure that is the real diagnosis. + try { + root.close(); + } catch (RuntimeException releaseFailure) { + e.addSuppressed(releaseFailure); + } throw e; } return root; diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java index 4cf95456582f1..bf780b40acf11 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java @@ -18,6 +18,8 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.analytics.backend.EngineResultBatch; import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.exec.ArrowValues; @@ -47,6 +49,8 @@ @ExperimentalApi public class LuceneResultStream implements EngineResultStream { + private static final Logger LOGGER = LogManager.getLogger(LuceneResultStream.class); + /** C-Data array carrying the populated batch. Owned by this stream until {@link #close}. */ private final ArrowArray arrowArray; /** C-Data schema describing {@link #arrowArray}. */ @@ -179,8 +183,14 @@ private VectorSchemaRoot importBatch() { Data.importIntoVectorSchemaRoot(stagingAllocator, arrowArray, root, dictionaryProvider); } catch (RuntimeException e) { // Close the partially-imported root (fires the native release) but KEEP the allocator: it is - // stream-scoped and a later batch may still use it. - root.close(); + // stream-scoped and a later batch may still use it. The release can itself throw + // (VectorSchemaRoot#close rethrows any RuntimeException from the vectors); attach it rather + // than let it mask the import failure that is the real diagnosis. + try { + root.close(); + } catch (RuntimeException releaseFailure) { + e.addSuppressed(releaseFailure); + } throw e; } return root; @@ -195,6 +205,10 @@ private VectorSchemaRoot importBatch() { * the allocator permanently half-closed, because {@code BaseAllocator#close} sets {@code isClosed} * BEFORE its leak check — so its bytes would never return to the parent. Leaving it open hands * ownership to the root, matching what the old per-batch code did for an in-flight batch. + * + *

The deferral is logged at DEBUG, not WARN: the transport frees its stream root asynchronously + * (posted to the flight executor by {@code FlightServerChannel#close}), so a non-zero balance here is + * the expected outcome of every streaming query, not a signal of a leak. */ void closeStagingAllocator() { if (stagingAllocator == null) { @@ -203,6 +217,13 @@ void closeStagingAllocator() { if (stagingAllocator.getAllocatedMemory() == 0) { stagingAllocator.close(); stagingAllocator = null; + } else { + LOGGER.debug( + "Deferring close of staging allocator [{}] with {} bytes outstanding; the transport still " + + "holds them and frees them with its stream root", + stagingAllocator.getName(), + stagingAllocator.getAllocatedMemory() + ); } } From 1deaf216527283c296db32bbc1380b088b095687 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Tue, 25 Aug 2026 06:24:49 +0000 Subject: [PATCH 3/3] Make the Arrow import staging allocator node-scoped instead of per-stream Signed-off-by: Lantao Jin --- .../backend/ShardScanExecutionContext.java | 26 ++++ .../spi/AnalyticsSearchBackendPlugin.java | 5 +- .../analytics/spi/ExchangeSinkContext.java | 27 +++- .../AbstractDatafusionReduceSink.java | 10 +- .../DataFusionAnalyticsBackendPlugin.java | 5 +- .../be/datafusion/DatafusionResultStream.java | 110 ++++++---------- .../DatafusionSearchExecEngine.java | 8 +- .../opensearch/be/datafusion/GetService.java | 16 ++- .../DatafusionResultStreamTests.java | 13 +- .../DatafusionSearchExecEngineTests.java | 7 + ...afusionStagingAllocatorLifecycleTests.java | 120 ++++++++++-------- .../be/lucene/LuceneResultStream.java | 98 +++++--------- .../be/lucene/LuceneSearchExecEngine.java | 8 +- .../opensearch/analytics/AnalyticsPlugin.java | 5 +- .../exec/AnalyticsSearchService.java | 57 ++++++++- .../exec/CoordinatorAllocatorHandle.java | 21 ++- .../analytics/exec/DefaultPlanExecutor.java | 6 + .../analytics/exec/QueryContext.java | 21 +++ .../LocalComputeStageExecutionFactory.java | 3 +- .../ReduceStageExecutionFactory.java | 3 +- 20 files changed, 364 insertions(+), 205 deletions(-) diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java index 6681cd65a0cbc..714295b9342d5 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java @@ -33,6 +33,7 @@ public class ShardScanExecutionContext implements CommonExecutionContext { private final Task task; private byte[] fragmentBytes; private BufferAllocator allocator; + private BufferAllocator importStagingAllocator; private MapperService mapperService; private IndexSettings indexSettings; private NamedWriteableRegistry namedWriteableRegistry; @@ -89,6 +90,31 @@ public void setAllocator(BufferAllocator allocator) { this.allocator = allocator; } + /** + * Returns the allocator Arrow C Data Interface imports are staged on, or null if the caller did not + * set one. Distinct from {@link #getAllocator()} on two properties a batch import depends on: + *

    + *
  • Unbounded, parented at the root. {@code Data#importIntoVectorSchemaRoot} charges each + * buffer as it walks the array; a target that fills part-way through throws, and arrow-java + * (≤ 18.1.0) retains the imported array before the throwing {@code wrapForeignAllocation} + * without rolling back, so the C Data release callback never fires and the whole batch leaks in + * the producer's native allocator — invisible to the JVM heap and to Java Arrow accounting.
    • + *
  • Node-scoped, never closed per request. The Flight transport builds its reused stream + * root on the FIRST emitted batch's vector allocator and charges that same allocator for every + * later batch ({@code FlightServerChannel#transferIntoStreamRoot}: "The producer's allocator must + * be long-lived (not closed per-request)"), and frees it asynchronously with its own channel.
    • + *
+ * Caller-owned: the engine imports onto it and must never close it. + */ + public BufferAllocator getImportStagingAllocator() { + return importStagingAllocator; + } + + /** Sets the node-scoped import staging allocator. The caller owns its lifecycle. */ + public void setImportStagingAllocator(BufferAllocator importStagingAllocator) { + this.importStagingAllocator = importStagingAllocator; + } + /** Returns the shard's mapper service for field type resolution. */ public MapperService getMapperService() { return mapperService; diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java index 071d1578561f1..1fc98dc7e7da7 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java @@ -220,6 +220,8 @@ default CanMatchResult canMatchWithBounds(IndexShard shard, byte[] filterBytes, * @param rowIdVector Arrow BigIntVector containing global row IDs * @param columns column names to read * @param allocator Arrow buffer allocator for result import + * @param importStagingAllocator node-scoped allocator to stage Arrow C Data imports on; see + * {@link org.opensearch.analytics.backend.ShardScanExecutionContext#getImportStagingAllocator()} * @return a result stream containing the requested rows */ default EngineResultStream fetchByRowIds( @@ -227,7 +229,8 @@ default EngineResultStream fetchByRowIds( BigIntVector rowIdVector, String[] columns, BufferAllocator allocator, - long contextId + long contextId, + BufferAllocator importStagingAllocator ) { throw new UnsupportedOperationException("fetchByRowIds not implemented for [" + name() + "]"); } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java index f374c689aa722..ae14fbde36538 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java @@ -43,12 +43,37 @@ *
  • {@code downstream} — sink the backend drains its reduced output * into. The backend owns {@code downstream}'s lifecycle: it must * feed every produced batch and close it when draining is complete.
    • + *
  • {@code importStagingAllocator} — node-scoped allocator the backend stages + * Arrow C Data Interface imports on. Unbounded and parented at the root so an + * import cannot fail part-way through an array (which strands the whole native + * batch — see {@link org.opensearch.analytics.backend.ShardScanExecutionContext#getImportStagingAllocator()}), + * and long-lived because the Flight transport keeps charging it after the + * importing stream closes. Caller-owned: the backend must never close it, and + * must never derive a per-stream child to import onto — an un-closed child stays + * registered in the root's {@code childAllocators} map until node restart.
    • * * * @opensearch.internal */ public record ExchangeSinkContext(String queryId, int stageId, long taskId, byte[] fragmentBytes, BufferAllocator allocator, List< - ChildInput> childInputs, ExchangeSink downstream) implements CommonExecutionContext { + ChildInput> childInputs, ExchangeSink downstream, BufferAllocator importStagingAllocator) implements CommonExecutionContext { + + /** + * Stages imports on {@code allocator} itself. For callers whose allocator is already an unbounded + * root — chiefly tests: that is the pre-staging behaviour, correct but without the mid-import-OOM + * mitigation described above. Production paths pass a dedicated staging allocator. + */ + public ExchangeSinkContext( + String queryId, + int stageId, + long taskId, + byte[] fragmentBytes, + BufferAllocator allocator, + List childInputs, + ExchangeSink downstream + ) { + this(queryId, stageId, taskId, fragmentBytes, allocator, childInputs, downstream, allocator); + } /** * Per-child input descriptor: the child stage id and the producer-side plan bytes the diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java index eb2c3855f498f..5cbb68415954e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java @@ -168,7 +168,15 @@ public byte[] getExecutionMetrics() { protected final void drainOutputIntoDownstream(StreamHandle outStream) { BufferAllocator alloc = ctx.allocator(); try (CDataDictionaryProvider dictProvider = new CDataDictionaryProvider()) { - DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(outStream, alloc, dictProvider); + // Imports go onto the caller-owned, node-scoped staging allocator — never a child minted here: + // downstream may hand the batch to the Flight transport, which keeps charging that allocator + // long after this drain returns. + DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator( + outStream, + alloc, + ctx.importStagingAllocator(), + dictProvider + ); while (it.hasNext()) { // next() transfers ownership of the imported VSR to us. feed() takes ownership only // on success; if it throws (e.g. the downstream sink was torn down on a concurrent diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index 1e3f75b5d9800..d4e24eb555176 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -1079,7 +1079,8 @@ public EngineResultStream fetchByRowIds( BigIntVector rowIdVector, String[] columns, BufferAllocator allocator, - long contextId + long contextId, + BufferAllocator importStagingAllocator ) { DataFusionService dataFusionService = plugin.getDataFusionService(); if (dataFusionService == null) { @@ -1115,7 +1116,7 @@ public EngineResultStream fetchByRowIds( throw new IllegalStateException("BigIntVector buffer address is 0 or count is 0"); } StreamHandle streamHandle = new StreamHandle(streamPtr, dataFusionService.getNativeRuntime()); - return new DatafusionResultStream(streamHandle, allocator); + return new DatafusionResultStream(streamHandle, allocator, importStagingAllocator); } @Override diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java index d5e375cab6335..be6b68ac08306 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java @@ -18,8 +18,6 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; import org.opensearch.analytics.backend.EngineResultBatch; import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.exec.ArrowValues; @@ -32,6 +30,7 @@ import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; +import java.util.Objects; import java.util.concurrent.CompletableFuture; import static org.apache.arrow.c.Data.importField; @@ -47,24 +46,31 @@ @ExperimentalApi public class DatafusionResultStream implements EngineResultStream, FragmentResources.MetricsCapable { - private static final Logger LOGGER = LogManager.getLogger(DatafusionResultStream.class); - private final StreamHandle streamHandle; private final BufferAllocator allocator; + private final BufferAllocator stagingAllocator; private final CDataDictionaryProvider dictionaryProvider; private volatile BatchIterator iteratorInstance; - // Allocator is caller-owned; this stream imports into it but never closes it. - public DatafusionResultStream(StreamHandle streamHandle, BufferAllocator allocator) { + /** + * Both allocators are caller-owned; this stream imports onto them but never closes either. + * + * @param stagingAllocator node-scoped, unbounded allocator every batch of every stream is imported onto + * (see {@link BatchIterator#importBatch}). MUST outlive this stream: the Flight transport builds + * its reused stream root on the first batch's vector allocator and frees it asynchronously with + * its own channel. + */ + public DatafusionResultStream(StreamHandle streamHandle, BufferAllocator allocator, BufferAllocator stagingAllocator) { this.streamHandle = streamHandle; this.allocator = allocator; + this.stagingAllocator = Objects.requireNonNull(stagingAllocator, "stagingAllocator"); this.dictionaryProvider = new CDataDictionaryProvider(); } @Override public Iterator iterator() { if (iteratorInstance == null) { - iteratorInstance = new BatchIterator(streamHandle, allocator, dictionaryProvider); + iteratorInstance = new BatchIterator(streamHandle, allocator, stagingAllocator, dictionaryProvider); } return iteratorInstance; } @@ -79,7 +85,6 @@ public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.closeStagingAllocator(); } } finally { try { @@ -96,32 +101,34 @@ static class BatchIterator implements Iterator { private final StreamHandle streamHandle; private final BufferAllocator allocator; + /** + * Caller-owned, node-scoped, unbounded staging allocator every batch is imported onto (see + * {@link #importBatch}). Owned by {@code AnalyticsSearchService}, never by this stream: the Flight + * transport builds its reused stream root on {@code fieldVectors.getFirst().getAllocator()}, i.e. the + * FIRST batch's staging allocator ({@code FlightServerChannel#transferIntoStreamRoot}, whose comment + * states "The producer's allocator must be long-lived (not closed per-request)"), then charges that + * same allocator for every later batch via {@code BaseFixedWidthVector#transferTo} and frees the + * stream root asynchronously in its own {@code close()}. So this stream can neither close it at batch + * boundaries (the transport is still using it) nor at stream close (the transport's free may run + * after ours) — it does not own it at all. + */ + private final BufferAllocator stagingAllocator; private final CDataDictionaryProvider dictionaryProvider; private Schema schema; private VectorSchemaRoot nextBatch; private Boolean nextAvailable; private boolean batchEmitted; private boolean nativeStreamExhausted; - /** - * ONE staging allocator for the whole stream, created lazily on first import (see - * {@link #importBatch}). Deliberately NOT per-batch: the Flight transport builds its reused stream - * root on {@code fieldVectors.getFirst().getAllocator()}, i.e. the FIRST batch's staging allocator - * ({@code FlightServerChannel#transferIntoStreamRoot}, whose comment states "The producer's - * allocator must be long-lived (not closed per-request)"), and then charges that same allocator for - * every later batch via {@code BaseFixedWidthVector#transferTo}. Closing a staging allocator - * per-batch therefore closed an allocator the transport was still using, and - * {@code BaseAllocator#close} threw {@code IllegalStateException: Memory was leaked by query} - * whenever a transfer landed in the window between the drained-check and the close. - * - *

    Unboundedness is the only property the original per-batch design actually needed (so an import - * cannot OOM part-way through a C Data array — see {@link #importBatch}); a single unbounded child - * preserves that while removing the check-then-close race entirely. - */ - private BufferAllocator stagingAllocator; - BatchIterator(StreamHandle streamHandle, BufferAllocator allocator, CDataDictionaryProvider dictionaryProvider) { + BatchIterator( + StreamHandle streamHandle, + BufferAllocator allocator, + BufferAllocator stagingAllocator, + CDataDictionaryProvider dictionaryProvider + ) { this.streamHandle = streamHandle; this.allocator = allocator; + this.stagingAllocator = stagingAllocator; this.dictionaryProvider = dictionaryProvider; } @@ -163,64 +170,27 @@ private boolean loadNextBatch() { } /** - * Imports one native batch across the Arrow C Data Interface into a per-batch staging allocator - * (an unbounded child of the root) rather than directly into {@code allocator}. + * Imports one native batch across the Arrow C Data Interface onto the caller-supplied + * {@link #stagingAllocator} rather than directly into {@code allocator}. * *

    {@link Data#importIntoVectorSchemaRoot} charges each buffer against the target allocator as it * walks the array. Against a bounded target that fills part-way through a wide batch the import * throws, and arrow-java's {@code ReferenceCountedArrowArray#unsafeAssociateAllocation} retains the * imported array before the throwing {@code wrapForeignAllocation} without rolling back, so * the C Data release callback never fires and the whole native batch leaks in the producer's native - * allocator — invisible to the JVM heap and the Java Arrow allocator (arrow-java ≤ 18.1.0). An - * unbounded staging child can't OOM mid-array, so the release callback always fires. + * allocator — invisible to the JVM heap and the Java Arrow allocator (arrow-java ≤ 18.1.0). The + * staging allocator is unbounded and parented at the root, so it can't OOM mid-array before the pool + * itself is exhausted and the release callback always fires. * *

    The batch is returned as-is (zero-copy); its buffers are released by the existing consumer close - * paths, which drives the C Data reference count to zero. The staging allocator is stream-scoped and - * outlives every batch (see {@link #stagingAllocator}), so nothing is closed per-batch; on import - * failure the partially-imported root is closed by {@link #importOntoStaging} but the allocator - * itself stays open for subsequent batches. + * paths, which drives the C Data reference count to zero. On import failure the partially-imported + * root is closed by {@link #importOntoStaging}; the allocator is untouched either way. */ private VectorSchemaRoot importBatch(ArrowArray arrowArray) { - if (stagingAllocator == null) { - stagingAllocator = allocator.getRoot().newChildAllocator("datafusion-import-staging", 0, Long.MAX_VALUE); - } return importOntoStaging(stagingAllocator, schema, arrowArray, dictionaryProvider); } - /** - * Closes the stream-scoped staging allocator, if it can be closed. Called only from - * {@link DatafusionResultStream#close()} — never per-batch. - * - *

    A non-zero balance here means the transport still holds buffers charged to this allocator: the - * Flight channel's reused stream root is freed in its own {@code close()}, which may run after ours. - * Closing anyway would throw {@code IllegalStateException} from {@code BaseAllocator#close} AND — worse - * — leave the allocator permanently half-closed, because {@code close()} sets {@code isClosed = true} - * BEFORE its leak check, so its bytes would never be returned to the parent and a later retry would - * early-return as a no-op. Leaving it open hands ownership to the root allocator, which is the same - * outcome the previous per-batch code produced for any still-in-flight batch. - * - *

    The deferral is logged at DEBUG, not WARN: the transport frees its stream root asynchronously - * (posted to the flight executor by {@code FlightServerChannel#close}), so a non-zero balance here is - * the expected outcome of every streaming query, not a signal of a leak. - */ - void closeStagingAllocator() { - if (stagingAllocator == null) { - return; - } - if (stagingAllocator.getAllocatedMemory() == 0) { - stagingAllocator.close(); - stagingAllocator = null; - } else { - LOGGER.debug( - "Deferring close of staging allocator [{}] with {} bytes outstanding; the transport still " - + "holds them and frees them with its stream root", - stagingAllocator.getName(), - stagingAllocator.getAllocatedMemory() - ); - } - } - - /** The stream-scoped staging allocator, or null before the first import. For tests. */ + /** The caller-supplied staging allocator batches are imported onto. For tests. */ BufferAllocator stagingAllocator() { return stagingAllocator; } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java index cd0721d212552..72b37f0c95ec7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearchExecEngine.java @@ -48,6 +48,12 @@ public EngineResultStream execute(ShardScanExecutionContext requestContext) thro if (allocator == null) { throw new IllegalStateException("ExecutionContext.allocator must be set by the caller before execute()"); } + // Node-scoped and unbounded — batches are imported onto it and the Flight transport keeps charging it + // after this stream closes, so it must be supplied rather than minted per stream. + BufferAllocator stagingAllocator = requestContext.getImportStagingAllocator(); + if (stagingAllocator == null) { + throw new IllegalStateException("ExecutionContext.importStagingAllocator must be set by the caller before execute()"); + } // Register cancellation hook so HTTP disconnect / _tasks/_cancel / timeout // immediately fires the Rust CancellationToken. @@ -63,7 +69,7 @@ public EngineResultStream execute(ShardScanExecutionContext requestContext) thro DatafusionSearcher searcher = datafusionContext.getSearcher(); searcher.search(datafusionContext); StreamHandle handle = datafusionContext.takeStreamHandle(); - return new DatafusionResultStream(handle, allocator); + return new DatafusionResultStream(handle, allocator, stagingAllocator); } @Override diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java index 52a7b82f40de1..bdc3bc1455707 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java @@ -89,6 +89,17 @@ static final class NativeBridgeExecutor implements DocumentRowReader, Closeable private final DataFusionPlugin dfPlugin; private final BufferAllocator sharedAllocator = new RootAllocator(64 * 1024 * 1024); + /** + * Unbounded staging child every batch of every stream this executor opens is imported onto — never + * one per stream, see {@code DatafusionResultStream.BatchIterator#stagingAllocator}. These streams are + * drained inline (not handed to the Flight transport), so it is always drained by the time + * {@link #close()} runs. + */ + private final BufferAllocator importStagingAllocator = sharedAllocator.newChildAllocator( + "datafusion-get-import-staging", + 0, + Long.MAX_VALUE + ); NativeBridgeExecutor(DataFusionPlugin dfPlugin) { this.dfPlugin = dfPlugin; @@ -96,6 +107,7 @@ static final class NativeBridgeExecutor implements DocumentRowReader, Closeable @Override public void close() { + importStagingAllocator.close(); sharedAllocator.close(); } @@ -165,7 +177,7 @@ private List> readAllRows(long streamPtr) { List> results = new ArrayList<>(); try ( StreamHandle streamHandle = new StreamHandle(streamPtr, dfPlugin.getDataFusionService().getNativeRuntime()); - DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator) + DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator, importStagingAllocator) ) { var iter = stream.iterator(); while (iter.hasNext()) { @@ -188,7 +200,7 @@ private List> readAllRows(long streamPtr) { private Map readSingleRow(long streamPtr) { try ( StreamHandle streamHandle = new StreamHandle(streamPtr, dfPlugin.getDataFusionService().getNativeRuntime()); - DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator) + DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator, importStagingAllocator) ) { var iter = stream.iterator(); if (!iter.hasNext()) return null; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java index 1136afa93b98b..8138100a2c93f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java @@ -40,6 +40,8 @@ public class DatafusionResultStreamTests extends OpenSearchTestCase { private long queryConfigPtr; private long tieredStorePtr; private NativeStoreHandle storeHandle; + /** Stands in for the node-scoped staging allocator AnalyticsSearchService hands to every stream. */ + private BufferAllocator stagingAllocator; private final java.util.List allocatorsToClose = new java.util.ArrayList<>(); @Override @@ -50,6 +52,8 @@ public void setUp() throws Exception { long ptr = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024); runtimeHandle = new NativeRuntimeHandle(ptr); testRootAllocator = new RootAllocator(Long.MAX_VALUE); + stagingAllocator = testRootAllocator.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); + allocatorsToClose.add(stagingAllocator); // Create a real TieredObjectStore (local-only) and wrap in NativeStoreHandle tieredStorePtr = NativeStoreTestHelper.createTieredObjectStore(0L, 0L); @@ -272,7 +276,8 @@ public void onFailure(Exception e) { allocatorsToClose.add(failureAlloc); DatafusionResultStream stream = new DatafusionResultStream( new org.opensearch.be.datafusion.nativelib.StreamHandle(streamPtr, tempRuntime), - failureAlloc + failureAlloc, + stagingAllocator ); // Close runtime — streamNext should now fail with IllegalStateException from NativeRuntimeHandle.get() @@ -323,7 +328,8 @@ private DatafusionResultStream createStreamWithLimit(String sql, long limitBytes allocatorsToClose.add(childAllocator); return new DatafusionResultStream( new org.opensearch.be.datafusion.nativelib.StreamHandle(streamPtr, runtimeHandle), - childAllocator + childAllocator, + stagingAllocator ); } @@ -358,7 +364,8 @@ private DatafusionResultStream createStream(String sql) { allocatorsToClose.add(childAllocator); return new DatafusionResultStream( new org.opensearch.be.datafusion.nativelib.StreamHandle(streamPtr, runtimeHandle), - childAllocator + childAllocator, + stagingAllocator ); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java index 48b380ea44056..ebee77c1022ba 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java @@ -8,6 +8,7 @@ package org.opensearch.be.datafusion; +import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.opensearch.analytics.backend.EngineResultBatch; import org.opensearch.analytics.backend.EngineResultStream; @@ -87,10 +88,12 @@ public void testEngineExecuteSelectAll() throws Exception { try ( RootAllocator alloc = new RootAllocator(Long.MAX_VALUE); + BufferAllocator staging = alloc.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context) ) { ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context); execCtx.setAllocator(alloc); + execCtx.setImportStagingAllocator(staging); engine.prepare(execCtx); try (EngineResultStream stream = engine.execute(execCtx)) { List rows = collectRows(stream); @@ -116,10 +119,12 @@ public void testEngineExecuteAggregation() throws Exception { try ( RootAllocator alloc = new RootAllocator(Long.MAX_VALUE); + BufferAllocator staging = alloc.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context) ) { ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context); execCtx.setAllocator(alloc); + execCtx.setImportStagingAllocator(staging); engine.prepare(execCtx); try (EngineResultStream stream = engine.execute(execCtx)) { List rows = collectRows(stream); @@ -142,10 +147,12 @@ public void testEngineExecuteFilter() throws Exception { try ( RootAllocator alloc = new RootAllocator(Long.MAX_VALUE); + BufferAllocator staging = alloc.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context) ) { ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context); execCtx.setAllocator(alloc); + execCtx.setImportStagingAllocator(staging); engine.prepare(execCtx); try (EngineResultStream stream = engine.execute(execCtx)) { List rows = collectRows(stream); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java index b00ed74409a5b..fad807a96ac67 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionStagingAllocatorLifecycleTests.java @@ -20,94 +20,89 @@ import org.opensearch.test.OpenSearchTestCase; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Set; /** - * Regression tests for the staging-allocator LIFETIME contract, i.e. the bug behind - * {@code IllegalStateException: Memory was leaked by query. Memory leaked: (1024)} thrown from - * {@code BaseAllocator#close} via the old {@code reclaimDrainedStaging}. + * Regression tests for the staging-allocator OWNERSHIP contract: the allocator batches are imported onto is + * supplied by the caller ({@code AnalyticsSearchService} in production, via + * {@code ShardScanExecutionContext#getImportStagingAllocator}), node-scoped, and never created or closed by + * the stream. Two separate bugs are pinned here, both rooted in a stream owning its staging allocator: * - *

    Why this asserts the contract rather than the race. The failure was a data race between the - * producer thread (sweeping drained staging allocators) and the Flight executor thread (charging the - * reused stream root's allocator in {@code BaseFixedWidthVector#transferTo}). The window is a couple of - * statements wide with no injection point, so a thread-racing test would reproduce only probabilistically. - * What IS deterministic is the invariant the transport depends on and the old code broke: - * {@code FlightServerChannel#transferIntoStreamRoot} builds its long-lived stream root on the FIRST - * emitted batch's vector allocator, and its comment states "The producer's allocator must be long-lived - * (not closed per-request)". So: *

      - *
    • every batch of one stream must be imported on the SAME allocator, and
      • - *
    • that allocator must still be OPEN after later batches are imported and earlier ones released.
      • + *
    • Use-after-close. {@code FlightServerChannel#transferIntoStreamRoot} builds its long-lived + * stream root on the FIRST emitted batch's vector allocator — its comment states "The producer's + * allocator must be long-lived (not closed per-request)" — and charges that same allocator for every + * later batch. A per-batch staging allocator closed as soon as its own buffers drained was therefore + * closed while the transport was still using it, throwing {@code IllegalStateException: Memory was + * leaked by query} from {@code BaseAllocator#close}. The window is a couple of statements wide with no + * injection point, so what these tests assert is the deterministic invariant behind it: every batch of a + * stream imports onto the SAME allocator, which is still OPEN afterwards.
      • + *
    • Stranded child allocators. Keeping a per-stream allocator open instead (to dodge the race) + * leaks it until node reboot: {@code BaseAllocator#newChildAllocator} registers every child in the + * parent's {@code childAllocators} map unconditionally, and only the child's own {@code close()} + * deregisters it — so an un-closed staging child stays strongly referenced by the node-lifetime root, + * one per query. {@link #testStreamsDoNotMintOrStrandChildAllocators} pins that the stream mints no + * child allocator at all.
      • *
    - * The old per-batch design violated both: batch N+1 got a fresh allocator, and batch N's was closed as - * soon as its own buffers drained — while the transport was still charging it. */ public class DatafusionStagingAllocatorLifecycleTests extends OpenSearchTestCase { private static final int ROWS = 8192; // DataFusion's default batch size private RootAllocator root; + /** Stands in for the node-scoped allocator AnalyticsSearchService owns and hands to every stream. */ + private BufferAllocator staging; + /** Stands in for the native (Rust) allocator owning the exported buffers; long-lived like the real one. */ + private BufferAllocator producer; // Resources the production DatafusionResultStream#close would own; the test must release them itself. private final List providers = new ArrayList<>(); - private final List producers = new ArrayList<>(); - private final List iterators = new ArrayList<>(); @Override public void setUp() throws Exception { super.setUp(); root = new RootAllocator(Long.MAX_VALUE); + staging = root.newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); + producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); } @Override public void tearDown() throws Exception { - // Release in the order production does: staging allocator, then the dictionary provider, then the - // stand-in native producer allocators. Only then can the root close cleanly. - for (DatafusionResultStream.BatchIterator it : iterators) { - it.closeStagingAllocator(); - } + // No stream closes the staging allocator — the caller does, exactly as AnalyticsSearchService#close + // does in production. for (CDataDictionaryProvider dp : providers) { dp.close(); } - for (BufferAllocator p : producers) { - p.close(); - } + staging.close(); + producer.close(); root.close(); super.tearDown(); } /** * Imports three batches through the production path and asserts all three land on the SAME staging - * allocator, which is still open at the end. - * - *

    Fails on the pre-fix code: each {@code importBatch} minted a new - * {@code datafusion-import-staging} child, so the allocators differ between batches. + * allocator — the caller-supplied one — which is still open at the end. */ public void testAllBatchesOfAStreamShareOneOpenStagingAllocator() throws Exception { DatafusionResultStream.BatchIterator it = newIterator(); - BufferAllocator first = null; for (int i = 0; i < 3; i++) { VectorSchemaRoot imported = importOneBatch(it); - BufferAllocator batchAllocator = imported.getFieldVectors().getFirst().getAllocator(); - if (first == null) { - first = batchAllocator; - } else { - assertSame( - "every batch of a stream must import onto the SAME staging allocator — the Flight " - + "stream root is built on the first batch's allocator and reused for all later batches", - first, - batchAllocator - ); - } - // The consumer releases this batch. Under the old code this drained the allocator to zero and - // made it eligible for the per-batch sweep on the NEXT import. + assertSame( + "every batch of a stream must import onto the caller-supplied staging allocator — the Flight " + + "stream root is built on the first batch's allocator and reused for all later batches", + staging, + imported.getFieldVectors().getFirst().getAllocator() + ); + // The consumer releases this batch. Under the old per-batch code this drained the allocator to + // zero and made it eligible for the sweep on the NEXT import. imported.close(); } - assertNotNull("staging allocator must exist after importing", it.stagingAllocator()); - assertSame("the tracked staging allocator is the one batches were imported on", first, it.stagingAllocator()); + assertSame("the iterator must not substitute an allocator of its own", staging, it.stagingAllocator()); // The decisive assertion: still usable by the transport after later imports + earlier releases. - assertEquals("staging allocator must NOT have been closed per-batch", 0L, it.stagingAllocator().getAllocatedMemory()); - it.stagingAllocator().assertOpen(); + assertEquals("staging allocator must NOT have been closed per-batch", 0L, staging.getAllocatedMemory()); + staging.assertOpen(); } /** @@ -141,6 +136,31 @@ public void testReleasingABatchDoesNotCloseTheAllocatorTheTransportHolds() throw } } + /** + * The leak regression: streaming N results must not add a single child allocator to the root. A stream + * that mints its own staging child either has to close it (racing the transport — see the tests above) or + * leave it open, in which case the root's {@code childAllocators} map retains it for the node's lifetime. + * Neither is acceptable, so the stream mints nothing. + */ + public void testStreamsDoNotMintOrStrandChildAllocators() throws Exception { + Set before = new HashSet<>(root.getChildAllocators()); + + for (int stream = 0; stream < 3; stream++) { + DatafusionResultStream.BatchIterator it = newIterator(); + VectorSchemaRoot imported = importOneBatch(it); + imported.close(); // consumer releases the batch + it.closeLastBatch(); // what DatafusionResultStream#close does to the iterator + } + + assertEquals( + "importing on a caller-supplied staging allocator must add no child allocator to the root — an " + + "un-closed child stays registered in the parent's childAllocators map until node reboot", + before, + new HashSet<>(root.getChildAllocators()) + ); + staging.assertOpen(); + } + // ── helpers ──────────────────────────────────────────────────────────────────────────────────── private DatafusionResultStream.BatchIterator newIterator() { @@ -148,15 +168,11 @@ private DatafusionResultStream.BatchIterator newIterator() { // the native runtime (no .so needed). CDataDictionaryProvider dp = new CDataDictionaryProvider(); providers.add(dp); - DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(null, root, dp); - iterators.add(it); - return it; + return new DatafusionResultStream.BatchIterator(null, root, staging, dp); } /** Exports a fresh batch across the C Data Interface and imports it via the production path. */ private VectorSchemaRoot importOneBatch(DatafusionResultStream.BatchIterator it) { - BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); - producers.add(producer); try (ArrowArray array = ArrowArray.allocateNew(producer); ArrowSchema cSchema = ArrowSchema.allocateNew(producer)) { try (VectorSchemaRoot source = VectorSchemaRoot.create(intSchema(producer), producer)) { IntVector v = (IntVector) source.getVector(0); diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java index bf780b40acf11..5495d8aac3fa6 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneResultStream.java @@ -18,8 +18,6 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; import org.opensearch.analytics.backend.EngineResultBatch; import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.exec.ArrowValues; @@ -28,6 +26,7 @@ import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; +import java.util.Objects; import static org.apache.arrow.c.Data.importField; @@ -49,31 +48,36 @@ @ExperimentalApi public class LuceneResultStream implements EngineResultStream { - private static final Logger LOGGER = LogManager.getLogger(LuceneResultStream.class); - /** C-Data array carrying the populated batch. Owned by this stream until {@link #close}. */ private final ArrowArray arrowArray; /** C-Data schema describing {@link #arrowArray}. */ private final ArrowSchema arrowSchema; private final BufferAllocator allocator; + private final BufferAllocator stagingAllocator; private final CDataDictionaryProvider dictionaryProvider; private volatile BatchIterator iteratorInstance; /** * Caller hands over ownership of {@code arrowArray} and {@code arrowSchema}; this stream - * closes them in {@link #close}. + * closes them in {@link #close}. Both allocators stay caller-owned. + * + * @param stagingAllocator node-scoped, unbounded allocator the batch is imported onto (see + * {@link BatchIterator#importBatch}). MUST outlive this stream: the Flight transport builds its + * reused stream root on the first batch's vector allocator and frees it asynchronously with its + * own channel. */ - public LuceneResultStream(ArrowArray arrowArray, ArrowSchema arrowSchema, BufferAllocator allocator) { + public LuceneResultStream(ArrowArray arrowArray, ArrowSchema arrowSchema, BufferAllocator allocator, BufferAllocator stagingAllocator) { this.arrowArray = arrowArray; this.arrowSchema = arrowSchema; this.allocator = allocator; + this.stagingAllocator = Objects.requireNonNull(stagingAllocator, "stagingAllocator"); this.dictionaryProvider = new CDataDictionaryProvider(); } @Override public Iterator iterator() { if (iteratorInstance == null) { - iteratorInstance = new BatchIterator(arrowArray, arrowSchema, allocator, dictionaryProvider); + iteratorInstance = new BatchIterator(arrowArray, arrowSchema, allocator, stagingAllocator, dictionaryProvider); } return iteratorInstance; } @@ -83,7 +87,6 @@ public void close() { try { if (iteratorInstance != null) { iteratorInstance.closeLastBatch(); - iteratorInstance.closeStagingAllocator(); } } finally { try { @@ -109,34 +112,34 @@ static class BatchIterator implements Iterator { private final ArrowArray arrowArray; private final ArrowSchema arrowSchema; private final BufferAllocator allocator; + /** + * Caller-owned, node-scoped, unbounded staging allocator the batch is imported onto (see + * {@link #importBatch}). Owned by {@code AnalyticsSearchService}, never by this stream: the Flight + * transport builds its reused stream root on {@code fieldVectors.getFirst().getAllocator()}, i.e. the + * FIRST batch's staging allocator ({@code FlightServerChannel#transferIntoStreamRoot}, whose comment + * states "The producer's allocator must be long-lived (not closed per-request)"), then charges that + * same allocator for every later batch and frees the stream root asynchronously in its own + * {@code close()} — which may run after ours. + */ + private final BufferAllocator stagingAllocator; private final CDataDictionaryProvider dictionaryProvider; private Schema schema; private VectorSchemaRoot nextBatch; private Boolean nextAvailable; private boolean batchEmitted; private boolean exhausted; - /** - * ONE staging allocator for the whole stream, created lazily on first import (see - * {@link #importBatch}). Deliberately NOT per-batch: the Flight transport builds its reused stream - * root on {@code fieldVectors.getFirst().getAllocator()}, i.e. the FIRST batch's staging allocator - * ({@code FlightServerChannel#transferIntoStreamRoot}, whose comment states "The producer's - * allocator must be long-lived (not closed per-request)"), then charges that same allocator for every - * later batch. Closing one per-batch closed an allocator the transport was still using, so - * {@code BaseAllocator#close} threw {@code IllegalStateException: Memory was leaked by query} when a - * transfer landed between the drained-check and the close. Unboundedness is the only property the - * per-batch design needed; one unbounded child keeps it without the race. - */ - private BufferAllocator stagingAllocator; BatchIterator( ArrowArray arrowArray, ArrowSchema arrowSchema, BufferAllocator allocator, + BufferAllocator stagingAllocator, CDataDictionaryProvider dictionaryProvider ) { this.arrowArray = arrowArray; this.arrowSchema = arrowSchema; this.allocator = allocator; + this.stagingAllocator = stagingAllocator; this.dictionaryProvider = dictionaryProvider; } @@ -159,33 +162,31 @@ private boolean loadNextBatch() { } /** - * Imports the batch across the Arrow C Data Interface into a per-batch staging allocator (an - * unbounded child of the root) rather than directly into {@code allocator}. + * Imports the batch across the Arrow C Data Interface onto the caller-supplied + * {@link #stagingAllocator} rather than directly into {@code allocator}. * *

    {@link Data#importIntoVectorSchemaRoot} charges each buffer against the target allocator as it * walks the array. Against a bounded target that fills part-way through a wide batch the import * throws, and arrow-java's {@code ReferenceCountedArrowArray#unsafeAssociateAllocation} retains the * imported array before the throwing {@code wrapForeignAllocation} without rolling back, so * the C Data release callback never fires and the whole native batch leaks — invisible to the JVM - * heap and the Java Arrow allocator (arrow-java ≤ 18.1.0). An unbounded staging child can't OOM - * mid-array, so the release callback always fires. + * heap and the Java Arrow allocator (arrow-java ≤ 18.1.0). The staging allocator is unbounded and + * parented at the root, so it can't OOM mid-array before the pool itself is exhausted and the release + * callback always fires. * *

    The batch is returned as-is (zero-copy); its buffers are released by the existing consumer close - * paths, which drives the C Data reference count to zero. The staging allocator is stream-scoped and - * outlives every batch (see {@link #stagingAllocator}), so nothing is closed per-batch. + * paths, which drives the C Data reference count to zero. The staging allocator is caller-owned and + * outlives this stream (see {@link #stagingAllocator}), so nothing here closes it. */ private VectorSchemaRoot importBatch() { - if (stagingAllocator == null) { - stagingAllocator = allocator.getRoot().newChildAllocator("lucene-import-staging", 0, Long.MAX_VALUE); - } VectorSchemaRoot root = VectorSchemaRoot.create(schema, stagingAllocator); try { Data.importIntoVectorSchemaRoot(stagingAllocator, arrowArray, root, dictionaryProvider); } catch (RuntimeException e) { - // Close the partially-imported root (fires the native release) but KEEP the allocator: it is - // stream-scoped and a later batch may still use it. The release can itself throw - // (VectorSchemaRoot#close rethrows any RuntimeException from the vectors); attach it rather - // than let it mask the import failure that is the real diagnosis. + // Close the partially-imported root (fires the native release); the allocator is caller-owned + // and untouched. The release can itself throw (VectorSchemaRoot#close rethrows any + // RuntimeException from the vectors); attach it rather than let it mask the import failure + // that is the real diagnosis. try { root.close(); } catch (RuntimeException releaseFailure) { @@ -196,37 +197,6 @@ private VectorSchemaRoot importBatch() { return root; } - /** - * Closes the stream-scoped staging allocator, if it can be closed. Called only from - * {@link LuceneResultStream#close()} — never per-batch. - * - *

    A non-zero balance means the transport still holds buffers charged here (its reused stream root - * is freed in its own {@code close()}, which may run after ours). Closing anyway would throw AND leave - * the allocator permanently half-closed, because {@code BaseAllocator#close} sets {@code isClosed} - * BEFORE its leak check — so its bytes would never return to the parent. Leaving it open hands - * ownership to the root, matching what the old per-batch code did for an in-flight batch. - * - *

    The deferral is logged at DEBUG, not WARN: the transport frees its stream root asynchronously - * (posted to the flight executor by {@code FlightServerChannel#close}), so a non-zero balance here is - * the expected outcome of every streaming query, not a signal of a leak. - */ - void closeStagingAllocator() { - if (stagingAllocator == null) { - return; - } - if (stagingAllocator.getAllocatedMemory() == 0) { - stagingAllocator.close(); - stagingAllocator = null; - } else { - LOGGER.debug( - "Deferring close of staging allocator [{}] with {} bytes outstanding; the transport still " - + "holds them and frees them with its stream root", - stagingAllocator.getName(), - stagingAllocator.getAllocatedMemory() - ); - } - } - @Override public boolean hasNext() { if (nextAvailable == null) { diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java index 701840299c8b1..bd8e4dd6c28e1 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java @@ -77,13 +77,19 @@ public EngineResultStream execute(ShardScanExecutionContext context) throws IOEx state.outputColumnNames() ); BufferAllocator allocator = context.getAllocator(); + // Node-scoped and unbounded — the batch is imported onto it and the Flight transport keeps charging it + // after this stream closes, so it must be supplied rather than minted per stream. + BufferAllocator stagingAllocator = context.getImportStagingAllocator(); + if (stagingAllocator == null) { + throw new IllegalStateException("ExecutionContext.importStagingAllocator must be set by the caller before execute()"); + } Schema schema = buildSchema(state.outputColumnNames()); ArrowArray array = ArrowArray.allocateNew(allocator); ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); boolean transferred = false; try { populateBatchToCData(allocator, schema, state.outputColumnNames(), count, array, arrowSchema); - LuceneResultStream stream = new LuceneResultStream(array, arrowSchema, allocator); + LuceneResultStream stream = new LuceneResultStream(array, arrowSchema, allocator, stagingAllocator); transferred = true; return stream; } finally { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java index 39f6d54c871e2..2c30ea396038a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java @@ -225,8 +225,11 @@ public Collection createComponents( // plugin's lifecycle owns its lifetime. The Guice-bound DefaultPlanExecutor consumes // it via the handle without taking on close responsibility — mirroring how // AnalyticsSearchService's allocator is owned and closed by this plugin. + // The import staging allocator is created and closed by the search service; the handle only + // carries it so the coordinator-reduce path stages its imports on the same node-scoped allocator. coordinatorAllocatorHandle = new CoordinatorAllocatorHandle( - nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).newChildAllocator("coordinator", 0, Long.MAX_VALUE) + nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).newChildAllocator("coordinator", 0, Long.MAX_VALUE), + searchService.getImportStagingAllocator() ); // Returned as components so Guice can inject them into DefaultPlanExecutor diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java index 1046cbaa873f1..520647468289f 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java @@ -97,6 +97,19 @@ public class AnalyticsSearchService implements AutoCloseable { private org.opensearch.threadpool.ThreadPool threadPool; private org.opensearch.cluster.service.ClusterService clusterService; private final BufferAllocator allocator; + /** + * ONE node-scoped allocator that every backend stages its Arrow C Data Interface imports on, handed to + * the backends via {@link ShardScanExecutionContext#setImportStagingAllocator} and + * {@link AnalyticsSearchBackendPlugin#fetchByRowIds}. Deliberately not per-stream: the Flight transport + * builds its reused stream root on the first emitted batch's vector allocator and charges that same + * allocator for every later batch ({@code FlightServerChannel#transferIntoStreamRoot}: "The producer's + * allocator must be long-lived (not closed per-request)"), then frees it asynchronously with its own + * channel. A per-stream staging child therefore cannot be closed at stream close without racing the + * transport, and one that is left open instead stays registered in the root's {@code childAllocators} + * map for the node's lifetime ({@code BaseAllocator#newChildAllocator} registers unconditionally; only + * the child's own {@code close()} deregisters) — one stranded allocator per query. + */ + private final BufferAllocator importStagingAllocator; private final ArrowNativeAllocator nativeAllocator; public AnalyticsSearchService(Map backends, ArrowNativeAllocator nativeAllocator) { @@ -131,13 +144,47 @@ public AnalyticsSearchService( // effect immediately via Arrow's parent-cap check at allocateBytes — no listener needed. BufferAllocator queryPool = nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); this.allocator = queryPool.newChildAllocator("analytics-search-service", 0, Long.MAX_VALUE); + // Parented at the ROOT rather than the query pool, and unbounded: a C Data import that fails + // part-way through an array strands the whole native batch (see + // ShardScanExecutionContext#getImportStagingAllocator), so the staging target must not be able to + // fill up before the pool itself is exhausted. + this.importStagingAllocator = allocator.getRoot().newChildAllocator("arrow-import-staging", 0, Long.MAX_VALUE); this.namedWriteableRegistry = namedWriteableRegistry; this.readerContextStore = readerContextStore; } @Override public void close() { - allocator.close(); + // finally, so a leak report from the service allocator cannot skip the staging allocator's release. + try { + allocator.close(); + } finally { + // A stream still in flight at shutdown keeps its batches charged to the staging allocator until + // the Flight channel frees its stream root — and that free is posted to the flight executor, + // which may already refuse tasks (FlightServerChannel#close logs "root reclaimed at process + // exit"). Closing on a non-zero balance would throw from BaseAllocator#close AND strand the + // bytes, because close() sets isClosed before its leak check. Report and let process exit + // reclaim them instead. + long outstanding = importStagingAllocator.getAllocatedMemory(); + if (outstanding == 0) { + importStagingAllocator.close(); + } else { + LOGGER.warn( + "Arrow import staging allocator [{}] still holds {} bytes at shutdown; reclaimed at process exit", + importStagingAllocator.getName(), + outstanding + ); + } + } + } + + /** + * The node-scoped allocator Arrow C Data imports are staged on. Owned (and closed) here; the + * coordinator-reduce path borrows the same one via {@code CoordinatorAllocatorHandle} so a node has + * exactly one, whatever the execution path. + */ + public BufferAllocator getImportStagingAllocator() { + return importStagingAllocator; } public void setTaskResourceTrackingService(TaskResourceTrackingService service) { @@ -488,6 +535,7 @@ public void executeWorkerFragmentStreamingAsync( ShardScanExecutionContext ctx = new ShardScanExecutionContext(/* tableName */ "", task, /* reader */ null); ctx.setFragmentBytes(plan.getFragmentBytes()); ctx.setAllocator(allocator); + ctx.setImportStagingAllocator(importStagingAllocator); ctx.setNamedWriteableRegistry(namedWriteableRegistry); ctx.setShuffleBufferRegistry(shuffleBufferRegistry); @@ -707,7 +755,8 @@ private void drainFetchByRowIds( rowIdVector, columns, allocator, - task.getNativeTaskId() + task.getNativeTaskId(), + importStagingAllocator ); // FragmentResources keeps the rowIdVector alive until the stream drains — closing // it earlier would pull off-heap memory out from under the native FFM call. @@ -914,7 +963,8 @@ private ExchangeSink buildPartitionedSink( new byte[0], ctx.getAllocator(), List.of(), - /* downstream */ null + /* downstream */ null, + importStagingAllocator ); return backend.getExchangeSinkProvider() .createPartitionedSink( @@ -1061,6 +1111,7 @@ private ShardScanExecutionContext buildContext( ShardScanExecutionContext ctx = new ShardScanExecutionContext(tableName, task, reader); ctx.setFragmentBytes(plan.getFragmentBytes()); ctx.setAllocator(allocator); + ctx.setImportStagingAllocator(importStagingAllocator); ctx.setMapperService(shard.mapperService()); ctx.setIndexSettings(shard.indexSettings()); ctx.setNamedWriteableRegistry(namedWriteableRegistry); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java index ab6f4a6d786b4..b53122422f07f 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/CoordinatorAllocatorHandle.java @@ -24,18 +24,27 @@ * *

    The plugin owns the lifetime; consumers ({@code DefaultPlanExecutor}) only * read the underlying allocator via {@link #getAllocator()}. + * + *

    Also carries the node-scoped Arrow import staging allocator so the coordinator-reduce path can + * reach it. That one is borrowed, not owned: {@code AnalyticsSearchService} creates and closes + * it, so a node has exactly one whatever the execution path. */ public final class CoordinatorAllocatorHandle implements AutoCloseable { private final BufferAllocator allocator; + private final BufferAllocator importStagingAllocator; /** * Wraps an already-constructed coordinator allocator (typically a child of * {@code POOL_QUERY}). The caller (the plugin) retains responsibility for * having created the allocator under the right parent. + * + * @param importStagingAllocator the node-scoped import staging allocator owned by + * {@code AnalyticsSearchService}; borrowed here, never closed by this handle */ - public CoordinatorAllocatorHandle(BufferAllocator allocator) { + public CoordinatorAllocatorHandle(BufferAllocator allocator, BufferAllocator importStagingAllocator) { this.allocator = allocator; + this.importStagingAllocator = importStagingAllocator; } /** Returns the wrapped allocator. Consumers must not close it directly. */ @@ -43,6 +52,16 @@ public BufferAllocator getAllocator() { return allocator; } + /** + * Returns the node-scoped allocator Arrow C Data imports are staged on — unbounded and parented at + * the root so an import cannot fail part-way through an array, and long-lived because the Flight + * transport keeps charging it after the importing stream closes. Owned by + * {@code AnalyticsSearchService}; consumers must not close it. + */ + public BufferAllocator getImportStagingAllocator() { + return importStagingAllocator; + } + @Override public void close() { allocator.close(); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java index ad4f94be09dcc..6861772cd55d4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java @@ -125,6 +125,8 @@ public class DefaultPlanExecutor extends HandledTransportAction perQueryBufferLimit = v); @@ -530,6 +533,9 @@ private void executeInternal( ownsAllocator, profile ); + // Node-scoped and unbounded: coordinator-reduce batches are imported onto it and the Flight + // transport keeps charging it after the importing stream closes, so it must not be per-query. + context.setImportStagingAllocator(importStagingAllocator); } catch (Exception e) { if (ownsAllocator) queryAllocator.close(); throw e; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java index 113e1e855f7a3..2a52fd813a52d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java @@ -53,6 +53,8 @@ public class QueryContext { private final List operationListeners; private final BufferAllocator allocator; private final boolean ownsAllocator; + /** Caller-owned; see {@link #importStagingAllocator()}. Never closed by this context. */ + private BufferAllocator importStagingAllocator; /** Whether profiling is enabled for this query (data nodes should collect and return metrics). */ private final boolean profile; /** @@ -311,6 +313,25 @@ public BufferAllocator bufferAllocator() { return allocator; } + /** + * Returns the node-scoped allocator coordinator-side Arrow C Data imports are staged on — unbounded + * and parented at the root so an import cannot fail part-way through an array (which strands the whole + * native batch, see {@code ShardScanExecutionContext#getImportStagingAllocator}), and long-lived + * because the Flight transport keeps charging it after the importing stream closes. + * + *

    Falls back to {@link #bufferAllocator()} when unset — that is the pre-staging behaviour, correct + * but without the mid-import-OOM mitigation, and it keeps test contexts (whose allocators are unbounded + * roots anyway) working without wiring one. + */ + public BufferAllocator importStagingAllocator() { + return importStagingAllocator != null ? importStagingAllocator : allocator; + } + + /** Set once by {@code DefaultPlanExecutor} after construction. The caller owns the allocator. */ + public void setImportStagingAllocator(BufferAllocator importStagingAllocator) { + this.importStagingAllocator = importStagingAllocator; + } + /** Lazy per-query virtual-thread executor for LOCAL tasks. Shared across phased contexts. */ public ExecutorService localTaskExecutor() { ExecutorService exec = sharedState.localTaskExecutor; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java index 60a1803750064..292fd258967b0 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalComputeStageExecutionFactory.java @@ -50,7 +50,8 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte chosenBytes(stage), config.bufferAllocator(), List.of(), - sink + sink, + config.importStagingAllocator() ); ExchangeSink backendSink = provider.createSink(context, null); if (backendSink instanceof ReducingExchangeSink reducing) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java index 0e5aad1cb61d1..494086f93bfa2 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java @@ -51,7 +51,8 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte chosenBytes(stage), config.bufferAllocator(), buildChildInputs(stage), - sink + sink, + config.importStagingAllocator() ); // Apply instruction handlers for the reduce stage.