[8.x] Resume Driver on cancelled or early finished (#120020) (#120208)

* Resume Driver on cancelled or early finished (#120020)

Today, when a Driver is put to sleep or scheduled but the thread pool is 
busy with other tasks, canceling the Driver must wait until the Driver
is awakened and scheduled for another loop, or until the thread pool
completes the tasks ahead of it. A similar issue occurs with early
finishes.

This change enables preemptive early finishes or cancellations of the 
Driver without waiting for another run loop.

* compile

Author: dnhatn

docs/changelog/120020.yaml

@@ -0,0 +1,5 @@
+pr: 120020
+summary: Resume Driver on cancelled or early finished
+area: ES|QL
+type: enhancement
+issues: []

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/Driver.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.compute.Describable;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.exchange.ExchangeSinkOperator;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -74,10 +75,9 @@ public class Driver implements Releasable, Describable {
     private final long statusNanos;
 
     private final AtomicReference<String> cancelReason = new AtomicReference<>();
-    private final AtomicReference<SubscribableListener<Void>> blocked = new AtomicReference<>();
-
     private final AtomicBoolean started = new AtomicBoolean();
     private final SubscribableListener<Void> completionListener = new SubscribableListener<>();
+    private final DriverScheduler scheduler = new DriverScheduler();
 
     /**
      * Status reported to the tasks API. We write the status at most once every
@@ -245,31 +245,33 @@ private IsBlockedResult runSingleLoopIteration() {
         ensureNotCancelled();
         boolean movedPage = false;
 
-        for (int i = 0; i < activeOperators.size() - 1; i++) {
-            Operator op = activeOperators.get(i);
-            Operator nextOp = activeOperators.get(i + 1);
+        if (activeOperators.isEmpty() == false && activeOperators.get(activeOperators.size() - 1).isFinished() == false) {
+            for (int i = 0; i < activeOperators.size() - 1; i++) {
+                Operator op = activeOperators.get(i);
+                Operator nextOp = activeOperators.get(i + 1);
 
-            // skip blocked operator
-            if (op.isBlocked().listener().isDone() == false) {
-                continue;
-            }
+                // skip blocked operator
+                if (op.isBlocked().listener().isDone() == false) {
+                    continue;
+                }
 
-            if (op.isFinished() == false && nextOp.needsInput()) {
-                Page page = op.getOutput();
-                if (page == null) {
-                    // No result, just move to the next iteration
-                } else if (page.getPositionCount() == 0) {
-                    // Empty result, release any memory it holds immediately and move to the next iteration
-                    page.releaseBlocks();
-                } else {
-                    // Non-empty result from the previous operation, move it to the next operation
-                    nextOp.addInput(page);
-                    movedPage = true;
+                if (op.isFinished() == false && nextOp.needsInput()) {
+                    Page page = op.getOutput();
+                    if (page == null) {
+                        // No result, just move to the next iteration
+                    } else if (page.getPositionCount() == 0) {
+                        // Empty result, release any memory it holds immediately and move to the next iteration
+                        page.releaseBlocks();
+                    } else {
+                        // Non-empty result from the previous operation, move it to the next operation
+                        nextOp.addInput(page);
+                        movedPage = true;
+                    }
                 }
-            }
 
-            if (op.isFinished()) {
-                nextOp.finish();
+                if (op.isFinished()) {
+                    nextOp.finish();
+                }
             }
         }
 
@@ -312,19 +314,10 @@ private IsBlockedResult runSingleLoopIteration() {
 
     public void cancel(String reason) {
         if (cancelReason.compareAndSet(null, reason)) {
-            synchronized (this) {
-                SubscribableListener<Void> fut = this.blocked.get();
-                if (fut != null) {
-                    fut.onFailure(new TaskCancelledException(reason));
-                }
-            }
+            scheduler.runPendingTasks();
         }
     }
 
-    private boolean isCancelled() {
-        return cancelReason.get() != null;
-    }
-
     private void ensureNotCancelled() {
         String reason = cancelReason.get();
         if (reason != null) {
@@ -342,6 +335,16 @@ public static void start(
         driver.completionListener.addListener(listener);
         if (driver.started.compareAndSet(false, true)) {
             driver.updateStatus(0, 0, DriverStatus.Status.STARTING, "driver starting");
+            // Register a listener to an exchange sink to handle early completion scenarios:
+            // 1. When the query accumulates sufficient data (e.g., reaching the LIMIT).
+            // 2. When users abort the query but want to retain the current result.
+            // This allows the Driver to finish early without waiting for the scheduled task.
+            final List<Operator> operators = driver.activeOperators;
+            if (operators.isEmpty() == false) {
+                if (operators.get(operators.size() - 1) instanceof ExchangeSinkOperator sinkOperator) {
+                    sinkOperator.addCompletionListener(ActionListener.running(driver.scheduler::runPendingTasks));
+                }
+            }
             schedule(DEFAULT_TIME_BEFORE_YIELDING, maxIterations, threadContext, executor, driver, driver.completionListener);
         }
     }
@@ -371,7 +374,7 @@ private static void schedule(
         Driver driver,
         ActionListener<Void> listener
     ) {
-        executor.execute(new AbstractRunnable() {
+        final var task = new AbstractRunnable() {
 
             @Override
             protected void doRun() {
@@ -383,16 +386,12 @@ protected void doRun() {
                 if (fut.isDone()) {
                     schedule(maxTime, maxIterations, threadContext, executor, driver, listener);
                 } else {
-                    synchronized (driver) {
-                        if (driver.isCancelled() == false) {
-                            driver.blocked.set(fut);
-                        }
-                    }
                     ActionListener<Void> readyListener = ActionListener.wrap(
                         ignored -> schedule(maxTime, maxIterations, threadContext, executor, driver, listener),
                         this::onFailure
                     );
                     fut.addListener(ContextPreservingActionListener.wrapPreservingContext(readyListener, threadContext));
+                    driver.scheduler.addOrRunDelayedTask(() -> fut.onResponse(null));
                 }
             }
 
@@ -405,7 +404,8 @@ public void onFailure(Exception e) {
             void onComplete(ActionListener<Void> listener) {
                 driver.driverContext.waitForAsyncActions(ContextPreservingActionListener.wrapPreservingContext(listener, threadContext));
             }
-        });
+        };
+        driver.scheduler.scheduleOrRunTask(executor, task);
     }
 
     private static IsBlockedResult oneOf(List<IsBlockedResult> results) {

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/DriverScheduler.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.operator;
+
+import org.elasticsearch.common.util.concurrent.EsExecutors;
+
+import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * A Driver be put to sleep while its sink is full or its source is empty or be rescheduled after running several iterations.
+ * This scheduler tracks the delayed and scheduled tasks, allowing them to run without waking up the driver or waiting for
+ * the thread pool to pick up the task. This enables fast cancellation or early finishing without discarding the current result.
+ */
+final class DriverScheduler {
+    private final AtomicReference<Runnable> delayedTask = new AtomicReference<>();
+    private final AtomicReference<Runnable> scheduledTask = new AtomicReference<>();
+    private final AtomicBoolean completing = new AtomicBoolean();
+
+    void addOrRunDelayedTask(Runnable task) {
+        delayedTask.set(task);
+        if (completing.get()) {
+            final Runnable toRun = delayedTask.getAndSet(null);
+            if (toRun != null) {
+                assert task == toRun;
+                toRun.run();
+            }
+        }
+    }
+
+    void scheduleOrRunTask(Executor executor, Runnable task) {
+        final Runnable existing = scheduledTask.getAndSet(task);
+        assert existing == null : existing;
+        final Executor executorToUse = completing.get() ? EsExecutors.DIRECT_EXECUTOR_SERVICE : executor;
+        executorToUse.execute(() -> {
+            final Runnable next = scheduledTask.getAndSet(null);
+            if (next != null) {
+                assert next == task;
+                next.run();
+            }
+        });
+    }
+
+    void runPendingTasks() {
+        completing.set(true);
+        for (var taskHolder : List.of(delayedTask, scheduledTask)) {
+            final Runnable task = taskHolder.getAndSet(null);
+            if (task != null) {
+                task.run();
+            }
+        }
+    }
+}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSink.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.compute.operator.exchange;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.IsBlockedResult;
 
@@ -30,6 +31,11 @@ public interface ExchangeSink {
      */
     boolean isFinished();
 
+    /**
+     * Adds a listener that will be notified when this exchange sink is finished.
+     */
+    void addCompletionListener(ActionListener<Void> listener);
+
     /**
      * Whether the sink is blocked on adding more pages
      */

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSinkHandler.java

@@ -52,9 +52,11 @@ public ExchangeSinkHandler(BlockFactory blockFactory, int maxBufferSize, LongSup
 
     private class ExchangeSinkImpl implements ExchangeSink {
         boolean finished;
+        private final SubscribableListener<Void> onFinished = new SubscribableListener<>();
 
         ExchangeSinkImpl() {
             onChanged();
+            buffer.addCompletionListener(onFinished);
             outstandingSinks.incrementAndGet();
         }
 
@@ -68,6 +70,7 @@ public void addPage(Page page) {
         public void finish() {
             if (finished == false) {
                 finished = true;
+                onFinished.onResponse(null);
                 onChanged();
                 if (outstandingSinks.decrementAndGet() == 0) {
                     buffer.finish(false);
@@ -78,7 +81,12 @@ public void finish() {
 
         @Override
         public boolean isFinished() {
-            return finished || buffer.isFinished();
+            return onFinished.isDone();
+        }
+
+        @Override
+        public void addCompletionListener(ActionListener<Void> listener) {
+            onFinished.addListener(listener);
         }
 
         @Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSinkOperator.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -60,6 +61,10 @@ public boolean isFinished() {
         return sink.isFinished();
     }
 
+    public void addCompletionListener(ActionListener<Void> listener) {
+        sink.addCompletionListener(listener);
+    }
+
     @Override
     public void finish() {
         sink.finish();

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/DriverTests.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRunnable;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeValue;
@@ -21,6 +22,10 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.exchange.ExchangeSinkHandler;
+import org.elasticsearch.compute.operator.exchange.ExchangeSinkOperator;
+import org.elasticsearch.compute.operator.exchange.ExchangeSourceHandler;
+import org.elasticsearch.compute.operator.exchange.ExchangeSourceOperator;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.FixedExecutorBuilder;
@@ -35,8 +40,10 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.function.LongSupplier;
 
+import static org.hamcrest.Matchers.either;
 import static org.hamcrest.Matchers.equalTo;
 
 public class DriverTests extends ESTestCase {
@@ -273,6 +280,33 @@ public Page getOutput() {
         }
     }
 
+    public void testResumeOnEarlyFinish() throws Exception {
+        DriverContext driverContext = driverContext();
+        ThreadPool threadPool = threadPool();
+        try {
+            PlainActionFuture<Void> sourceFuture = new PlainActionFuture<>();
+            var sourceHandler = new ExchangeSourceHandler(between(1, 5), threadPool.executor("esql"), sourceFuture);
+            var sinkHandler = new ExchangeSinkHandler(driverContext.blockFactory(), between(1, 5), System::currentTimeMillis);
+            var sourceOperator = new ExchangeSourceOperator(sourceHandler.createExchangeSource());
+            var sinkOperator = new ExchangeSinkOperator(sinkHandler.createExchangeSink(), Function.identity());
+            Driver driver = new Driver(driverContext, sourceOperator, List.of(), sinkOperator, () -> {});
+            PlainActionFuture<Void> future = new PlainActionFuture<>();
+            Driver.start(threadPool.getThreadContext(), threadPool.executor("esql"), driver, between(1, 1000), future);
+            assertBusy(
+                () -> assertThat(
+                    driver.status().status(),
+                    either(equalTo(DriverStatus.Status.ASYNC)).or(equalTo(DriverStatus.Status.STARTING))
+                )
+            );
+            sinkHandler.fetchPageAsync(true, ActionListener.noop());
+            future.actionGet(5, TimeUnit.SECONDS);
+            assertThat(driver.status().status(), equalTo(DriverStatus.Status.DONE));
+            sourceFuture.actionGet(5, TimeUnit.SECONDS);
+        } finally {
+            terminate(threadPool);
+        }
+    }
+
     private static void assertRunningWithRegularUser(ThreadPool threadPool) {
         String user = threadPool.getThreadContext().getHeader("user");
         assertThat(user, equalTo("user1"));