Allow early termination in Driver

Author: dnhatn

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

@@ -186,7 +186,14 @@ SubscribableListener<Void> run(TimeValue maxTime, int maxIterations, LongSupplie
         long nextStatus = startTime + statusNanos;
         int iter = 0;
         while (true) {
-            IsBlockedResult isBlocked = runSingleLoopIteration();
+            IsBlockedResult isBlocked;
+            try {
+                isBlocked = runSingleLoopIteration();
+            } catch (DriverEarlyTerminationException e) {
+                drainAndCloseOperators(null);
+                updateStatus(finishNanos - startTime, iter, DriverStatus.Status.DONE, "early termination");
+                return Operator.NOT_BLOCKED.listener();
+            }
             iter++;
             if (isBlocked.listener().isDone() == false) {
                 updateStatus(nowSupplier.getAsLong() - startTime, iter, DriverStatus.Status.ASYNC, isBlocked.reason());
@@ -227,6 +234,17 @@ public void close() {
         drainAndCloseOperators(null);
     }
 
+    private void checkForEarlyTermination() throws DriverEarlyTerminationException {
+        if (activeOperators.size() >= 2 && activeOperators.getLast().isFinished()) {
+            for (int i = activeOperators.size() - 2; i >= 0; i--) {
+                Operator op = activeOperators.get(i);
+                if (op.isFinished() == false) {
+                    throw new DriverEarlyTerminationException();
+                }
+            }
+        }
+    }
+
     /**
      * Abort the driver and wait for it to finish
      */
@@ -253,6 +271,7 @@ private IsBlockedResult runSingleLoopIteration() {
             if (op.isBlocked().listener().isDone() == false) {
                 continue;
             }
+            checkForEarlyTermination();
 
             if (op.isFinished() == false && nextOp.needsInput()) {
                 Page page = op.getOutput();
@@ -263,6 +282,15 @@ private IsBlockedResult runSingleLoopIteration() {
                     page.releaseBlocks();
                 } else {
                     // Non-empty result from the previous operation, move it to the next operation
+                    boolean terminated = true;
+                    try {
+                        checkForEarlyTermination();
+                        terminated = false;
+                    } finally {
+                        if (terminated) {
+                            page.releaseBlocks();
+                        }
+                    }
                     nextOp.addInput(page);
                     movedPage = true;
                 }
@@ -290,6 +318,7 @@ private IsBlockedResult runSingleLoopIteration() {
                     itr.remove();
                 }
 
+                checkForEarlyTermination();
                 // Finish the next operator, which is now the first operator.
                 if (activeOperators.isEmpty() == false) {
                     Operator newRootOperator = activeOperators.get(0);

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

@@ -0,0 +1,12 @@
+/*
+ * 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;
+
+final class DriverEarlyTerminationException extends RuntimeException {
+    DriverEarlyTerminationException() {}
+}

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java

@@ -345,7 +345,8 @@ private void assertCancelled(ActionFuture<EsqlQueryResponse> response) throws Ex
                     "task cancelled",
                     "request cancelled test cancel",
                     "parent task was cancelled [test cancel]",
-                    "cancelled on failure"
+                    "cancelled on failure",
+                    "task cancelled [cancelled on failure]"
                 )
             )
         );