Handle rejection in DriverScheduler (#122105) (#122118)

When a node is shutting down, scheduling tasks for the Driver can result 
in a rejection exception. In this case, we drain and close all
operators. However, we don't clear the pending tasks in the scheduler,
which can lead to a pending task being triggered unexpectedly, causing a
ConcurrentModificationException.

Author: dnhatn

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

@@ -7,7 +7,9 @@
 
 package org.elasticsearch.compute.operator;
 
+import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
 
 import java.util.List;
 import java.util.concurrent.Executor;
@@ -21,7 +23,7 @@
  */
 final class DriverScheduler {
     private final AtomicReference<Runnable> delayedTask = new AtomicReference<>();
-    private final AtomicReference<Runnable> scheduledTask = new AtomicReference<>();
+    private final AtomicReference<AbstractRunnable> scheduledTask = new AtomicReference<>();
     private final AtomicBoolean completing = new AtomicBoolean();
 
     void addOrRunDelayedTask(Runnable task) {
@@ -35,22 +37,32 @@ void addOrRunDelayedTask(Runnable task) {
         }
     }
 
-    void scheduleOrRunTask(Executor executor, Runnable task) {
-        final Runnable existing = scheduledTask.getAndSet(task);
+    void scheduleOrRunTask(Executor executor, AbstractRunnable task) {
+        final AbstractRunnable 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();
+        executorToUse.execute(new AbstractRunnable() {
+            @Override
+            public void onFailure(Exception e) {
+                assert e instanceof EsRejectedExecutionException : new AssertionError(e);
+                if (scheduledTask.getAndUpdate(t -> t == task ? null : t) == task) {
+                    task.onFailure(e);
+                }
+            }
+
+            @Override
+            protected void doRun() {
+                AbstractRunnable toRun = scheduledTask.getAndSet(null);
+                if (toRun == task) {
+                    task.run();
+                }
             }
         });
     }
 
     void runPendingTasks() {
         completing.set(true);
-        for (var taskHolder : List.of(delayedTask, scheduledTask)) {
+        for (var taskHolder : List.of(scheduledTask, delayedTask)) {
             final Runnable task = taskHolder.getAndSet(null);
             if (task != null) {
                 task.run();

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

@@ -0,0 +1,61 @@
+/*
+ * 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.settings.Settings;
+import org.elasticsearch.common.util.concurrent.AbstractRunnable;
+import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.threadpool.FixedExecutorBuilder;
+import org.elasticsearch.threadpool.TestThreadPool;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class DriverSchedulerTests extends ESTestCase {
+
+    public void testClearPendingTaskOnRejection() {
+        DriverScheduler scheduler = new DriverScheduler();
+        AtomicInteger counter = new AtomicInteger();
+        var threadPool = new TestThreadPool(
+            "test",
+            new FixedExecutorBuilder(Settings.EMPTY, "test", 1, 2, "test", EsExecutors.TaskTrackingConfig.DEFAULT)
+        );
+        CountDownLatch latch = new CountDownLatch(1);
+        Executor executor = threadPool.executor("test");
+        try {
+            for (int i = 0; i < 10; i++) {
+                try {
+                    executor.execute(() -> safeAwait(latch));
+                } catch (EsRejectedExecutionException e) {
+                    break;
+                }
+            }
+            scheduler.scheduleOrRunTask(executor, new AbstractRunnable() {
+                @Override
+                public void onFailure(Exception e) {
+                    counter.incrementAndGet();
+                }
+
+                @Override
+                protected void doRun() {
+                    counter.incrementAndGet();
+                }
+            });
+            scheduler.runPendingTasks();
+            assertThat(counter.get(), equalTo(1));
+        } finally {
+            latch.countDown();
+            terminate(threadPool);
+        }
+    }
+}