fix(datafusion): tighten reduce-sink drain + sender lifecycle (opensearch-project#21671)

Lifecycle and correctness fixes on the coordinator reduce-sink path:
locked sender_close, deterministic FFI release on drain/send races,
shared drain + session teardown in AbstractDatafusionReduceSink, and
moves the drain off SEARCH onto a virtual-thread executor.

Signed-off-by: Marc Handalian <marc.handalian@gmail.com>

Author: mch2

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java

@@ -8,17 +8,12 @@
 
 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.vector.VectorSchemaRoot;
-import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.opensearch.analytics.spi.ExchangeSink;
 import org.opensearch.analytics.spi.ExchangeSinkContext;
-import org.opensearch.be.datafusion.nativelib.NativeBridge;
 import org.opensearch.be.datafusion.nativelib.StreamHandle;
 import org.opensearch.core.action.ActionListener;
 
@@ -28,8 +23,6 @@
 import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
 
-import static org.apache.arrow.c.Data.importField;
-
 /**
  * Shared lifecycle skeleton for coordinator-side {@link ExchangeSink}s backed by a native
  * DataFusion local session. Subclasses customise per-batch handling and the close-time
@@ -41,8 +34,9 @@
  *       and always closes the supplied {@link VectorSchemaRoot} in {@code finally} regardless
  *       of whether {@link #feedBatchUnderLock} succeeds.</li>
  *   <li>{@link #close} flips {@link #closed} once under {@link #feedLock}, runs the
- *       subclass-specific {@link #closeUnderLock} hook, and unconditionally closes
- *       {@link #session} in {@code finally}, accumulating any failures and rethrowing.</li>
+ *       subclass-specific {@link #closeUnderLock} hook, and rethrows any accumulated
+ *       failure. Subclasses must close {@link #session} themselves inside
+ *       {@link #closeUnderLock} (typically last, after any owned native streams).</li>
  *   <li>The downstream from {@link ExchangeSinkContext#downstream()} is intentionally NOT
  *       closed here — it accumulates drained results consumed by the walker after the
  *       sink is done.</li>
@@ -179,38 +173,30 @@ public final void close() {
     protected abstract void feedBatchUnderLock(VectorSchemaRoot batch);
 
     /**
-     * Subclass-specific shutdown. Runs after {@link #closed} is set and before
-     * {@link #session} is closed. Implementations should close their owned native resources
-     * (sender, output stream, accumulated FFI structs, …) and drain any pending output.
+     * Subclass-specific shutdown. Runs after {@link #closed} is set. Implementations must
+     * close all owned native resources including {@link #session} — close owned streams
+     * before the session.
      *
      * @return the first failure encountered (use {@link #accumulate(Throwable, Throwable)}
      *         when multiple steps may fail), or {@code null} on clean shutdown.
      */
     protected abstract Throwable closeUnderLock();
 
     /**
-     * Drains a native output stream into {@link ExchangeSinkContext#downstream()}, importing
-     * each {@link ArrowArray} into a fresh {@link VectorSchemaRoot} on the Java side.
+     * Drains a native output stream into {@link ExchangeSinkContext#downstream()},
+     * importing each native batch into a fresh {@link VectorSchemaRoot}.
+     *
+     * <p>Uses {@link DatafusionResultStream.BatchIterator} directly (instead of
+     * {@link DatafusionResultStream}) so the caller retains ownership of {@code outStream} —
+     * the iterator manages schema, dictionary provider, and per-batch allocation, but
+     * does not close the underlying stream handle.
      */
     protected final void drainOutputIntoDownstream(StreamHandle outStream) {
         BufferAllocator alloc = ctx.allocator();
         try (CDataDictionaryProvider dictProvider = new CDataDictionaryProvider()) {
-            long schemaAddr = asyncCall(listener -> NativeBridge.streamGetSchema(outStream.getPointer(), listener));
-            Schema outSchema;
-            try (ArrowSchema arrowSchema = ArrowSchema.wrap(schemaAddr)) {
-                Field structField = importField(alloc, arrowSchema, dictProvider);
-                outSchema = new Schema(structField.getChildren(), structField.getMetadata());
-            }
-            while (true) {
-                long arrayAddr = asyncCall(listener -> NativeBridge.streamNext(runtimeHandle.get(), outStream.getPointer(), listener));
-                if (arrayAddr == 0) {
-                    break;
-                }
-                VectorSchemaRoot vsr = VectorSchemaRoot.create(outSchema, alloc);
-                try (ArrowArray arrowArray = ArrowArray.wrap(arrayAddr)) {
-                    Data.importIntoVectorSchemaRoot(alloc, arrowArray, vsr, dictProvider);
-                }
-                ctx.downstream().feed(vsr);
+            DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(outStream, alloc, dictProvider);
+            while (it.hasNext()) {
+                ctx.downstream().feed(it.next().getArrowRoot());
             }
         }
     }

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java

@@ -613,7 +613,7 @@ public ExchangeSinkProvider getExchangeSinkProvider() {
             if ("memtable".equals(mode) && ctx.childInputs().size() == 1 && preparedState == null) {
                 return new DatafusionMemtableReduceSink(ctx, svc.getNativeRuntime());
             }
-            return new DatafusionReduceSink(ctx, svc.getNativeRuntime(), preparedState);
+            return new DatafusionReduceSink(ctx, svc.getNativeRuntime(), svc.getDrainExecutor(), preparedState);
         };
     }
 

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionService.java

@@ -21,6 +21,9 @@
 
 import java.io.IOException;
 import java.util.Collection;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -42,6 +45,15 @@ public class DataFusionService extends AbstractLifecycleComponent {
     private final int cpuThreads;
     private final ClusterSettings clusterSettings;
 
+    /**
+     * Virtual-thread-per-task executor used for {@link DatafusionReduceSink}'s drain
+     * loop. Drain tasks park on {@code streamNext().join()} while the native side
+     * produces the next batch; on a virtual thread the park unmounts from its carrier
+     * so concurrent queries don't consume fixed platform-pool slots for the full
+     * query lifetime.
+     */
+    private volatile ExecutorService drainExecutor;
+
     /** Handle to the native DataFusion global runtime (memory pool + cache). */
     private volatile NativeRuntimeHandle runtimeHandle;
 
@@ -62,6 +74,20 @@ private DataFusionService(Builder builder) {
         this.clusterSettings = builder.clusterSettings;
     }
 
+    /**
+     * Returns the virtual-thread-per-task executor the coordinator-reduce sink uses for
+     * its drain loop. Tasks spawned here park on native futures; virtual threads unmount
+     * from their carriers during the park so concurrent queries don't consume fixed
+     * platform-pool slots for the query's full lifetime.
+     */
+    public Executor getDrainExecutor() {
+        ExecutorService exec = drainExecutor;
+        if (exec == null) {
+            throw new IllegalStateException("DataFusionService has not been started");
+        }
+        return exec;
+    }
+
     /** Creates a new builder. */
     public static Builder builder() {
         return new Builder();
@@ -73,6 +99,8 @@ protected void doStart() {
         NativeBridge.initTokioRuntimeManager(cpuThreads);
         logger.debug("Tokio runtime manager initialized with {} CPU threads", cpuThreads);
 
+        this.drainExecutor = Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("analytics-reduce-drain-", 0).factory());
+
         long cacheManagerPtr = 0L;
         if (clusterSettings != null) {
             cacheManagerPtr = CacheUtils.createCacheConfig(clusterSettings);
@@ -101,7 +129,15 @@ protected void doStop() {
                     rootAllocator = null;
                 }
             } finally {
-                NativeBridge.shutdownTokioRuntimeManager();
+                try {
+                    ExecutorService exec = drainExecutor;
+                    if (exec != null) {
+                        exec.shutdown();
+                        drainExecutor = null;
+                    }
+                } finally {
+                    NativeBridge.shutdownTokioRuntimeManager();
+                }
             }
         }
         logger.debug("DataFusion service stopped");

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSink.java

@@ -154,6 +154,7 @@ protected Throwable closeUnderLock() {
             if (streamPtr != 0) {
                 NativeBridge.streamClose(streamPtr);
             }
+            session.close();
         }
         return failure;
     }

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java

@@ -43,6 +43,7 @@ public void send(long arrayAddr, long schemaAddr) {
     public void close() {
         lifecycle.writeLock().lock();
         try {
+            assert lifecycle.isWriteLockedByCurrentThread() : "close must hold the write lock across super.close()";
             super.close();
         } finally {
             lifecycle.writeLock().unlock();