Abort drivers on node shutting down (#101701) (#101716)

The TransportResponseHandler can be notified while the Driver is still 
running during node shutdown or the Driver hasn't started when the
parent task is canceled. In such cases, we should abort the Driver and
wait for it to finish; otherwise, multiple threads can access a Driver
at the same time

Closes #101595

Author: dnhatn

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

@@ -24,6 +24,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -56,6 +57,10 @@ public class Driver implements Releasable, Describable {
 
     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<>();
+
     /**
      * Status reported to the tasks API. We write the status at most once every
      * {@link #statusNanos}, as soon as loop has finished and after {@link #statusNanos}
@@ -149,7 +154,7 @@ private SubscribableListener<Void> run(TimeValue maxTime, int maxIterations) {
         if (isFinished()) {
             status.set(updateStatus(DriverStatus.Status.DONE));
             driverContext.finish();
-            releasable.close();
+            Releasables.close(releasable, driverContext.getSnapshot());
         } else {
             status.set(updateStatus(DriverStatus.Status.WAITING));
         }
@@ -159,7 +164,7 @@ private SubscribableListener<Void> run(TimeValue maxTime, int maxIterations) {
     /**
      * Whether the driver has run the chain of operators to completion.
      */
-    public boolean isFinished() {
+    private boolean isFinished() {
         return activeOperators.isEmpty();
     }
 
@@ -168,6 +173,19 @@ public void close() {
         drainAndCloseOperators(null);
     }
 
+    /**
+     * Abort the driver and wait for it to finish
+     */
+    public void abort(Exception reason, ActionListener<Void> listener) {
+        completionListener.addListener(listener);
+        if (started.compareAndSet(false, true)) {
+            drainAndCloseOperators(reason);
+            completionListener.onFailure(reason);
+        } else {
+            cancel(reason.getMessage());
+        }
+    }
+
     private SubscribableListener<Void> runSingleLoopIteration() {
         ensureNotCancelled();
         boolean movedPage = false;
@@ -261,8 +279,11 @@ public static void start(
         int maxIterations,
         ActionListener<Void> listener
     ) {
-        driver.status.set(driver.updateStatus(DriverStatus.Status.STARTING));
-        schedule(DEFAULT_TIME_BEFORE_YIELDING, maxIterations, threadContext, executor, driver, listener);
+        driver.completionListener.addListener(listener);
+        if (driver.started.compareAndSet(false, true)) {
+            driver.status.set(driver.updateStatus(DriverStatus.Status.STARTING));
+            schedule(DEFAULT_TIME_BEFORE_YIELDING, maxIterations, threadContext, executor, driver, driver.completionListener);
+        }
     }
 
     // Drains all active operators and closes them.
@@ -279,7 +300,7 @@ private void drainAndCloseOperators(@Nullable Exception e) {
             itr.remove();
         }
         driverContext.finish();
-        Releasables.closeWhileHandlingException(releasable);
+        Releasables.closeWhileHandlingException(releasable, driverContext.getSnapshot());
     }
 
     private static void schedule(

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
 
 import java.util.Collections;
 import java.util.IdentityHashMap;
@@ -69,7 +70,12 @@ public BlockFactory blockFactory() {
     }
 
     /** A snapshot of the driver context. */
-    public record Snapshot(Set<Releasable> releasables) {}
+    public record Snapshot(Set<Releasable> releasables) implements Releasable {
+        @Override
+        public void close() {
+            Releasables.close(releasables);
+        }
+    }
 
     /**
      * Adds a releasable to this context. Releasables are identified by Object identity.

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

@@ -12,7 +12,6 @@
 import org.elasticsearch.common.util.concurrent.AtomicArray;
 import org.elasticsearch.common.util.concurrent.CountDown;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
-import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tasks.TaskCancelledException;
 
 import java.util.HashMap;
@@ -84,13 +83,6 @@ private void done() {
                     responseHeaders.setOnce(driverIndex, threadContext.getResponseHeaders());
                     if (counter.countDown()) {
                         mergeResponseHeaders(responseHeaders);
-                        for (Driver d : drivers) {
-                            if (d.status().status() == DriverStatus.Status.QUEUED) {
-                                d.close();
-                            } else {
-                                Releasables.close(d.driverContext().getSnapshot().releasables());
-                            }
-                        }
                         Exception error = failure.get();
                         if (error != null) {
                             listener.onFailure(error);

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

@@ -52,7 +52,13 @@ protected void start(Driver driver, ActionListener<Void> driverListener) {
                     new DriverRequest(driver, executor),
                     parentTask,
                     TransportRequestOptions.EMPTY,
-                    TransportResponseHandler.empty(executor, driverListener)
+                    TransportResponseHandler.empty(
+                        executor,
+                        // The TransportResponseHandler can be notified while the Driver is still running during node shutdown
+                        // or the Driver hasn't started when the parent task is canceled. In such cases, we should abort
+                        // the Driver and wait for it to finish.
+                        ActionListener.wrap(driverListener::onResponse, e -> driver.abort(e, driverListener))
+                    )
                 );
             }
         };