Merge branch 'main' into unified-allocator-followup

Author: gaurav-amz

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LocalTaskRunner.java

@@ -9,16 +9,16 @@
 package org.opensearch.analytics.exec.stage.coordinator;
 
 import org.opensearch.analytics.exec.task.TaskRunner;
+import org.opensearch.common.util.concurrent.AbstractRunnable;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.NotifyOnceListener;
 
 import java.util.concurrent.Executor;
 
 /**
- * LOCAL-kind task runner: submits the {@link LocalStageTask#body()} to a per-query
- * virtual-thread executor. The body owns the listener — wrapped in a
+ * LOCAL-kind task runner: submits the {@link LocalStageTask} to the search
+ * executor. The body owns the listener — wrapped in a
  * {@link NotifyOnceListener} so a body that both fires and throws can't double-notify.
- * Virtual threads keep blocking reduce drains off SEARCH workers.
  *
  * @opensearch.internal
  */
@@ -43,10 +43,21 @@ protected void innerOnFailure(Exception cause) {
                 listener.onFailure(cause);
             }
         };
-        executor.execute(() -> {
-            try {
+        executor.execute(new AbstractRunnable() {
+            @Override
+            protected void doRun() {
                 task.body().accept(once);
-            } catch (Exception e) {
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                once.onFailure(e);
+            }
+
+            @Override
+            public void onRejection(Exception e) {
+                // Bounded pool queue full — fail the query gracefully instead of letting the
+                // rejection escape run() and hang the query with no terminal callback.
                 once.onFailure(e);
             }
         });

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/PassThroughStageExecution.java

@@ -37,7 +37,7 @@ public PassThroughStageExecution(Stage stage, QueryContext config, ExchangeSink
             throw new IllegalArgumentException("PassThroughStageExecution requires a RowProducingSink");
         }
         this.ownedSink = (RowProducingSink) sink;
-        this.runner = new LocalTaskRunner(config.localTaskExecutor());
+        this.runner = new LocalTaskRunner(config.searchExecutor());
     }
 
     @Override

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java

@@ -38,7 +38,7 @@ public ReduceStageExecution(Stage stage, QueryContext config, ReducingExchangeSi
         super(stage, config.queryId(), config.operationListeners(), config.parentTask());
         this.backendSink = backendSink;
         this.downstream = downstream;
-        this.runner = new LocalTaskRunner(config.localTaskExecutor());
+        this.runner = new LocalTaskRunner(config.searchExecutor());
     }
 
     @Override

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/LocalTaskRunnerTests.java

@@ -10,9 +10,12 @@
 
 import org.opensearch.analytics.exec.stage.coordinator.LocalStageTask;
 import org.opensearch.analytics.exec.stage.coordinator.LocalTaskRunner;
+import org.opensearch.common.util.concurrent.AbstractRunnable;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -74,6 +77,29 @@ public void testDispatchSubmitsToExecutor() {
         assertTrue("onCompleted fires once the executor runs the body", completed.get());
     }
 
+    public void testRejectionSignalsFailedGracefully() {
+        // A bounded pool rejects when its queue is full. OpenSearchThreadPoolExecutor invokes
+        // onRejection on an AbstractRunnable rather than rethrowing — the runner must route that
+        // to the listener so the query fails gracefully instead of hanging with no callback.
+        OpenSearchRejectedExecutionException rejected = new OpenSearchRejectedExecutionException("queue full");
+        AtomicBoolean ran = new AtomicBoolean();
+        AtomicBoolean completed = new AtomicBoolean();
+        AtomicReference<Exception> failure = new AtomicReference<>();
+
+        Executor rejectingExecutor = r -> ((AbstractRunnable) r).onRejection(rejected);
+        LocalTaskRunner dispatcher = new LocalTaskRunner(rejectingExecutor);
+        LocalStageTask task = new LocalStageTask(new StageTaskId(0, 0), listener -> {
+            ran.set(true);
+            listener.onResponse(null);
+        });
+
+        dispatcher.run(task, handle(completed, failure));
+
+        assertFalse("body must not run when the task is rejected", ran.get());
+        assertFalse("onCompleted must not fire when rejected", completed.get());
+        assertSame("rejection is routed to listener.onFailure", rejected, failure.get());
+    }
+
     private static ActionListener<Void> handle(AtomicBoolean completed, AtomicReference<Exception> failure) {
         return new ActionListener<Void>() {
             @Override