analytics-engine: drive reduce through stage tasks (#21723)

* analytics-engine: drive reduce through stage tasks + wire cancellation

Two related changes:

1. Drive reduce as a task. Introduces ReducingExchangeSink with a reduce/close
   split (reduce() owns the work, close() owns cleanup); collapses three
   reduce-execution variants into one ReduceStageExecution driven by the
   sink's supportsEagerScheduling(). Streaming sink anchors native cleanup on
   the reduce() caller via a SinkState machine.

2. Wire coordinator-side cancellation. Plumbs AnalyticsQueryTask.id through
   ExchangeSinkContext, registers a cancellation token in the Rust
   QUERY_REGISTRY, and has close() fire cancel_query for in-flight drains --
   coordinator reduce now participates in task-level cancellation alongside
   data-node scans.

Adds before-/after-first-batch cancellation tests in DatafusionReduceSinkTests.

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

* analytics-engine: fix race in DatafusionReduceSink.close()

  close()'s prior get()+compareAndSet pattern had a window where reduce()
  could transition READY→REDUCING between the read and the CAS, leaving
  cancel_query unfired and the parked drain hung. Replaced with a single
  compareAndExchange whose return value tells us which branch to take.

  Also collapsed the close() override into closeImpl, removed the base's
  auto-session-close (every concrete subclass already closed session, so
  this was a double-close masked by idempotency), and gated session.close
  on preparedState == null in both subclasses.

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

* spotless

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

* fix javadoc

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

* analytics-engine: fix outStream/session leak on cancel-during-REDUCING

  External close() arriving while reduce() was parked set the base class's
  `closed` flag and ran closeImpl() — which only fired cancel_query and
  deferred teardown. When reduce()'s finally then called super.close() to
  do the teardown, the base's `closed` short-circuit blocked re-entry and
  closeImpl() never ran a second time. outStream/session leaked.

  Have reduce()'s finally call closeImpl() directly, bypassing the base's
  closed-gate. Add a `torndown` AtomicBoolean so concurrent close paths
  (external + reduce-finally) can't both run the teardown body.

  Regression test: testCancelBeforeFirstBatchUnwindsDrain now asserts
  sink.torndown is set after the cancel — fails on the prior code.

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

---------

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

Author: mch2

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ReducingExchangeSink.java

@@ -0,0 +1,38 @@
+/*
+ * 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.analytics.spi;
+
+import org.opensearch.core.action.ActionListener;
+
+/**
+ * {@link ExchangeSink} that runs reduce work (execute plan, drain into downstream).
+ * The reduce-stage execution invokes {@link #reduce} from its task; {@link #close} runs
+ * separately on every terminal transition for resource cleanup — split so cancel-before-reduce
+ * paths can release resources without re-running the work.
+ *
+ * @opensearch.experimental
+ */
+public interface ReducingExchangeSink extends ExchangeSink {
+
+    /**
+     * Execute the reduce plan and drain output into the configured downstream; fire
+     * {@code listener} on completion. Single-shot. Must not release resources — that's
+     * {@link #close()}'s job.
+     */
+    void reduce(ActionListener<Void> listener);
+
+    /**
+     * {@code true} (default) — {@link #reduce} can run concurrently with producer
+     * {@link #feed} calls. Buffered sinks that need the complete input set first
+     * (e.g. memtable-backed) must override to {@code false}.
+     */
+    default boolean supportsEagerScheduling() {
+        return true;
+    }
+}

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

@@ -12,91 +12,60 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.types.pojo.Schema;
-import org.opensearch.analytics.spi.ExchangeSink;
 import org.opensearch.analytics.spi.ExchangeSinkContext;
+import org.opensearch.analytics.spi.ReducingExchangeSink;
 import org.opensearch.be.datafusion.nativelib.StreamHandle;
-import org.opensearch.core.action.ActionListener;
 
 import java.util.LinkedHashMap;
 import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.function.Consumer;
 
 /**
- * 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
- * native handoff via {@link #feedBatchUnderLock} and {@link #closeUnderLock}.
+ * Shared lifecycle skeleton for coordinator-side {@link ReducingExchangeSink}s backed by a
+ * native DataFusion local session. Subclasses customise per-batch handling, the reduce verb,
+ * and native cleanup via {@link #feedBatchUnderLock}, {@link #reduce}, {@link #closeImpl}.
  *
- * <p>Lifecycle invariants enforced by this base:
- * <ul>
- *   <li>{@link #feed} synchronises on {@link #feedLock}, short-circuits when {@link #closed},
- *       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 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>
- * </ul>
+ * <p>Invariants: {@link #feed} runs under {@link #feedLock} and always closes the batch in
+ * {@code finally}; {@link #close} flips {@link #closed} under the lock then delegates to
+ * {@link #closeImpl}. Subclasses own the full teardown sequence in {@link #closeImpl} —
+ * including {@link #session} and any sender/stream handles. The base used to auto-close
+ * {@code session} after {@code closeImpl} returned, but that double-closed (every concrete
+ * subclass already closed it) and tangled session ownership with the close-state machine
+ * the streaming subclass needs to manage.
+ * The {@link ExchangeSinkContext#downstream()} sink is NOT closed here — its lifecycle is the
+ * walker's.
  *
- * <p>Multi-input shapes (Union, future Join) are supported at this base by exposing
- * {@link #childInputs} (childStageId → producer-side plan bytes) for subclasses to register
- * one native partition per child stage. The native call returns the IPC-encoded schema the
- * session settled on after lowering; subclasses populate {@link #childSchemas} from those
- * returns so the {@code typesMatch} tripwire validates batches against the same schema the
- * native session is registered with. The {@link #INPUT_ID} constant remains as the
- * conventional name for the single-input case (childStageId=0); the per-child id is
- * computed via {@link #inputIdFor(int)}.
+ * <p>Multi-input is exposed via {@link #childInputs} (childStageId → producer plan bytes) and
+ * {@link #childSchemas} (lazily populated by subclasses from native registration replies).
+ * Use {@link #inputIdFor(int)} for the per-child table id; {@link #INPUT_ID} is the
+ * single-input shortcut.
  *
  * @opensearch.internal
  */
-abstract class AbstractDatafusionReduceSink implements ExchangeSink {
+abstract class AbstractDatafusionReduceSink implements ReducingExchangeSink {
 
-    /**
-     * Substrait/DataFusion table name used for the single-input case (childStageId=0).
-     * For multi-input shapes use {@link #inputIdFor(int)} instead.
-     */
+    /** Single-input shortcut for the per-child table id; multi-input uses {@link #inputIdFor(int)}. */
     static final String INPUT_ID = "input-0";
 
     protected final ExchangeSinkContext ctx;
     protected final NativeRuntimeHandle runtimeHandle;
     protected final DatafusionLocalSession session;
+
     /**
-     * Non-null when this sink was constructed with a pre-prepared FINAL-aggregate plan
-     * from the FinalAggregateInstructionHandler. When present, the handler already created
-     * the session, registered the input partitions, and called {@code prepareFinalPlan} on
-     * the Rust side; the sink only needs to drive {@code executeLocalPreparedPlan} and feed
-     * batches. When null, the sink falls back to the legacy path (create its own session,
-     * register its own partitions, call {@code executeLocalPlan}).
-     *
-     * <p>Close ownership: when {@code preparedState != null} the state owns session +
-     * senders and {@link #close} skips re-closing them (avoids double-close on the native
-     * side). When {@code preparedState == null} the base class closes the session itself.
+     * Non-null when constructed from a pre-prepared plan (FinalAggregateInstructionHandler):
+     * session + senders are owned by the state and {@link #close} skips re-closing them.
      */
     protected final DataFusionReduceState preparedState;
-    /**
-     * Per-child producer-side plan bytes, keyed by childStageId. Iteration order matches
-     * the order of {@code ctx.childInputs()} so subclasses get deterministic registration.
-     * Subclasses pass each entry to the native registration call, which lowers the plan
-     * and returns the schema the session settled on.
-     */
+
+    /** Per-child producer plan bytes (childStageId → bytes), iteration order = ctx.childInputs(). */
     protected final Map<Integer, byte[]> childInputs;
 
-    /**
-     * Declared Arrow {@link org.apache.arrow.vector.types.pojo.Schema} per childStageId.
-     * Populated lazily by subclasses from the IPC bytes the native registration call
-     * returns — i.e. the schema the native session itself derived from the producer plan.
-     * Used by sinks to validate incoming batches via the {@code typesMatch} tripwire.
-     */
+    /** Per-child declared schema, populated lazily from native registration replies. */
     protected final Map<Integer, Schema> childSchemas;
 
-    /** Guards {@link #closed} and serialises {@link #feed}/{@link #close} against producers. */
+    /** Serialises {@link #feed}/{@link #close} against producers. */
     protected final Object feedLock = new Object();
 
-    /** Set once in {@link #close} under {@link #feedLock}. Visible to all threads via volatile. */
+    /** Set once under {@link #feedLock} in {@link #close}. */
     protected volatile boolean closed;
 
     protected AbstractDatafusionReduceSink(ExchangeSinkContext ctx, NativeRuntimeHandle runtimeHandle) {
@@ -132,67 +101,46 @@ public void feed(VectorSchemaRoot batch) {
                 batch.close();
                 return;
             }
-            try {
+            try (batch) {
                 feedBatchUnderLock(batch);
-            } finally {
-                batch.close();
             }
         }
     }
 
     @Override
-    public final void close() {
+    public void close() {
         synchronized (feedLock) {
             if (closed) {
                 return;
             }
             closed = true;
         }
-        Throwable failure = null;
+        Exception failure = null;
         try {
-            failure = closeUnderLock();
-        } catch (Throwable t) {
+            failure = closeImpl();
+        } catch (Exception t) {
             failure = accumulate(failure, t);
-        } finally {
-            // If a preparedState owns the session/senders, let the state's close handle
-            // them (invoked by the orchestrator). Otherwise close the session we created.
-            if (preparedState == null) {
-                try {
-                    session.close();
-                } catch (Throwable t) {
-                    failure = accumulate(failure, t);
-                }
-            }
         }
         rethrow(failure);
     }
 
     /**
-     * Per-batch hook. Called inside {@code synchronized(feedLock)} after {@code closed} is
-     * verified false. Implementations export and hand off (or buffer) {@code batch} via the
-     * native bridge. Implementations MUST NOT close {@code batch} — the base class does that
-     * in {@code finally}.
+     * Per-batch hook. Called under {@link #feedLock} after the closed-check. MUST NOT
+     * close {@code batch} — the base does it in {@code finally}.
      */
     protected abstract void feedBatchUnderLock(VectorSchemaRoot batch);
 
     /**
-     * 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.
+     * Subclass shutdown. Despite the historical name, this is NOT called under {@link #feedLock} —
+     * the lock is released after {@link #closed} is flipped. Subclasses own the full teardown
+     * including {@link #session} (gate on {@link #preparedState} == null when the session is
+     * owned externally). Return the first failure, or null.
      */
-    protected abstract Throwable closeUnderLock();
+    protected abstract Exception closeImpl();
 
     /**
-     * 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.
+     * Imports each batch from {@code outStream} into a fresh {@link VectorSchemaRoot} and
+     * feeds it downstream. Caller retains ownership of {@code outStream}.
      */
     protected final void drainOutputIntoDownstream(StreamHandle outStream) {
         BufferAllocator alloc = ctx.allocator();
@@ -204,48 +152,22 @@ protected final void drainOutputIntoDownstream(StreamHandle outStream) {
         }
     }
 
-    /**
-     * Synchronously awaits the result of an async native call expressed as a
-     * {@code Consumer<ActionListener<Long>>}. Restores interrupt state on
-     * {@link InterruptedException} and unwraps {@link ExecutionException} to surface the
-     * original cause.
-     */
-    protected static long asyncCall(Consumer<ActionListener<Long>> call) {
-        CompletableFuture<Long> future = new CompletableFuture<>();
-        call.accept(ActionListener.wrap(future::complete, future::completeExceptionally));
-        try {
-            return future.get();
-        } catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
-            throw new RuntimeException(e);
-        } catch (ExecutionException e) {
-            Throwable cause = e.getCause();
-            if (cause instanceof RuntimeException re) {
-                throw re;
-            }
-            throw new RuntimeException(cause);
-        }
-    }
-
     /** Returns {@code t} if {@code acc} is null; otherwise adds {@code t} as a suppressed of {@code acc}. */
-    protected static Throwable accumulate(Throwable acc, Throwable t) {
+    protected static Exception accumulate(Exception acc, Exception t) {
         if (acc == null) {
             return t;
         }
         acc.addSuppressed(t);
         return acc;
     }
 
-    private static void rethrow(Throwable failure) {
+    private static void rethrow(Exception failure) {
         if (failure == null) {
             return;
         }
         if (failure instanceof RuntimeException re) {
             throw re;
         }
-        if (failure instanceof Error err) {
-            throw err;
-        }
         throw new RuntimeException(failure);
     }
 }

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

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

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

@@ -19,9 +19,6 @@
 
 import java.io.IOException;
 import java.util.Collection;
-import java.util.concurrent.Executor;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
 
 /**
  * Node-level service managing the DataFusion native runtime lifecycle.
@@ -42,8 +39,6 @@ public class DataFusionService extends AbstractLifecycleComponent {
     private final int cpuThreads;
     private final ClusterSettings clusterSettings;
 
-    private volatile ExecutorService drainExecutor;
-
     /** Handle to the native DataFusion global runtime (memory pool + cache). */
     private volatile NativeRuntimeHandle runtimeHandle;
 
@@ -58,14 +53,6 @@ private DataFusionService(Builder builder) {
         this.clusterSettings = builder.clusterSettings;
     }
 
-    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();
@@ -75,7 +62,6 @@ public static Builder builder() {
     protected void doStart() {
         logger.debug("Starting DataFusion service");
         NativeBridge.initTokioRuntimeManager(cpuThreads);
-        this.drainExecutor = Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("analytics-reduce-drain-", 0).factory());
         logger.debug("Tokio runtime manager initialized with {} CPU threads", cpuThreads);
 
         long cacheManagerPtr = 0L;
@@ -99,16 +85,9 @@ protected void doStop() {
         try {
             releaseRuntime();
         } finally {
-            try {
-                ExecutorService exec = drainExecutor;
-                if (exec != null) {
-                    exec.shutdown();
-                    drainExecutor = null;
-                }
-            } finally {
-                NativeBridge.shutdownTokioRuntimeManager();
-            }
+            NativeBridge.shutdownTokioRuntimeManager();
         }
+
         logger.debug("DataFusion service stopped");
     }