[Analytics Engine] Fix Arrow import staging allocator lifetime (transport race and leaked child allocators)

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;
@@ -27,7 +29,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;
@@ -46,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;
@@ -76,7 +79,7 @@ public void close() {
         try {
             if (iteratorInstance != null) {
                 iteratorInstance.closeLastBatch();
-                iteratorInstance.reclaimDrainedStaging();
+                iteratorInstance.closeStagingAllocator();
             }
         } finally {
             try {
@@ -99,9 +102,22 @@ static class BatchIterator implements Iterator<EngineResultBatch> {
         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<BufferAllocator> 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.
+         *
+         * <p>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 +175,64 @@ private boolean loadNextBatch() {
          * unbounded staging child can't OOM mid-array, so the release callback always fires.
          *
          * <p>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.
+         *
+         * <p>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.
+         *
+         * <p>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.
          */
-        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;
+            } 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. */
+        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);
         }
 
         /**
@@ -208,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;

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}.
+ *
+ * <p><b>Why this asserts the contract rather than the race.</b> 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:
+ * <ul>
+ *   <li>every batch of one stream must be imported on the SAME allocator, and</li>
+ *   <li>that allocator must still be OPEN after later batches are imported and earlier ones released.</li>
+ * </ul>
+ * 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<CDataDictionaryProvider> providers = new ArrayList<>();
+    private final List<BufferAllocator> producers = new ArrayList<>();
+    private final List<DatafusionResultStream.BatchIterator> 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.
+     *
+     * <p>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()));
+        }
+    }
+}