Fix bogus assertion tripped by force-executed tasks (#104581)

Today if you submit a force-executing task to a fixed executor with an
unbounded queue after the executor has been shut down then the
`EsAbortPolicy` will incorrectly throw an `IllegalStateException` which
ends up tripping an assertion. This is a legitimate thing to happen, so
this commit introduces a different rejection handler to deal with it. It
also introduces a check for a shut-down bounded-queue executor on
rejection to make sure we don't force the task onto a queue that's never
going to be processed.

Closes #104580

docs/changelog/104581.yaml

@@ -0,0 +1,6 @@
+pr: 104581
+summary: Fix bogus assertion tripped by force-executed tasks
+area: Infra/Core
+type: bug
+issues:
+ - 104580

server/src/main/java/org/elasticsearch/common/util/concurrent/EsAbortPolicy.java

@@ -8,26 +8,27 @@
 
 package org.elasticsearch.common.util.concurrent;
 
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 
 public class EsAbortPolicy extends EsRejectedExecutionHandler {
 
     @Override
     public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
-        if (r instanceof AbstractRunnable) {
-            if (((AbstractRunnable) r).isForceExecution()) {
-                BlockingQueue<Runnable> queue = executor.getQueue();
-                if ((queue instanceof SizeBlockingQueue) == false) {
-                    throw new IllegalStateException("forced execution, but expected a size queue");
+        if (r instanceof AbstractRunnable abstractRunnable) {
+            if (abstractRunnable.isForceExecution()) {
+                if (executor.getQueue() instanceof SizeBlockingQueue<Runnable> sizeBlockingQueue) {
+                    try {
+                        sizeBlockingQueue.forcePut(r);
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        throw new IllegalStateException("forced execution, but got interrupted", e);
+                    }
+                    if ((executor.isShutdown() && sizeBlockingQueue.remove(r)) == false) {
+                        return;
+                    } // else fall through and reject the task since the executor is shut down
+                } else {
+                    throw new IllegalStateException("expected but did not find SizeBlockingQueue: " + executor);
                 }
-                try {
-                    ((SizeBlockingQueue<Runnable>) queue).forcePut(r);
-                } catch (InterruptedException e) {
-                    Thread.currentThread().interrupt();
-                    throw new IllegalStateException("forced execution, but got interrupted", e);
-                }
-                return;
             }
         }
         incrementRejections();

server/src/main/java/org/elasticsearch/common/util/concurrent/EsExecutors.java

@@ -126,11 +126,14 @@ public static EsThreadPoolExecutor newFixed(
         ThreadContext contextHolder,
         TaskTrackingConfig config
     ) {
-        BlockingQueue<Runnable> queue;
+        final BlockingQueue<Runnable> queue;
+        final EsRejectedExecutionHandler rejectedExecutionHandler;
         if (queueCapacity < 0) {
             queue = ConcurrentCollections.newBlockingQueue();
+            rejectedExecutionHandler = new RejectOnShutdownOnlyPolicy();
         } else {
             queue = new SizeBlockingQueue<>(ConcurrentCollections.<Runnable>newBlockingQueue(), queueCapacity);
+            rejectedExecutionHandler = new EsAbortPolicy();
         }
         if (config.trackExecutionTime()) {
             return new TaskExecutionTimeTrackingEsThreadPoolExecutor(
@@ -142,7 +145,7 @@ public static EsThreadPoolExecutor newFixed(
                 queue,
                 TimedRunnable::new,
                 threadFactory,
-                new EsAbortPolicy(),
+                rejectedExecutionHandler,
                 contextHolder,
                 config
             );
@@ -155,7 +158,7 @@ public static EsThreadPoolExecutor newFixed(
                 TimeUnit.MILLISECONDS,
                 queue,
                 threadFactory,
-                new EsAbortPolicy(),
+                rejectedExecutionHandler,
                 contextHolder
             );
         }
@@ -411,6 +414,15 @@ private void reject(ThreadPoolExecutor executor, Runnable task) {
         }
     }
 
+    static class RejectOnShutdownOnlyPolicy extends EsRejectedExecutionHandler {
+        @Override
+        public void rejectedExecution(Runnable task, ThreadPoolExecutor executor) {
+            assert executor.isShutdown() : executor;
+            incrementRejections();
+            throw newRejectedException(task, executor, true);
+        }
+    }
+
     public static class TaskTrackingConfig {
         // This is a random starting point alpha. TODO: revisit this with actual testing and/or make it configurable
         public static double DEFAULT_EWMA_ALPHA = 0.3;

server/src/test/java/org/elasticsearch/common/util/concurrent/EsExecutorsTests.java

@@ -8,6 +8,8 @@
 
 package org.elasticsearch.common.util.concurrent;
 
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.ActionRunnable;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.Processors;
@@ -18,6 +20,8 @@
 import java.util.Locale;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -539,4 +543,129 @@ public void onFailure(Exception e) {
             ThreadPool.terminate(executorService, 10, TimeUnit.SECONDS);
         }
     }
+
+    public void testScalingDropOnShutdown() {
+        final var executor = EsExecutors.newScaling(
+            getName(),
+            0,
+            between(1, 5),
+            60,
+            TimeUnit.SECONDS,
+            false,
+            EsExecutors.daemonThreadFactory(getName()),
+            new ThreadContext(Settings.EMPTY)
+        );
+        ThreadPool.terminate(executor, 10, TimeUnit.SECONDS);
+        executor.execute(() -> fail("should not run")); // no-op
+        executor.execute(new AbstractRunnable() {
+            @Override
+            public void onFailure(Exception e) {
+                fail("should not call onFailure");
+            }
+
+            @Override
+            protected void doRun() {
+                fail("should not call doRun");
+            }
+
+            @Override
+            public boolean isForceExecution() {
+                return randomBoolean();
+            }
+
+            @Override
+            public void onRejection(Exception e) {
+                fail("should not call onRejection");
+            }
+
+            @Override
+            public void onAfter() {
+                fail("should not call onAfter");
+            }
+        });
+    }
+
+    public void testScalingRejectOnShutdown() {
+        runRejectOnShutdownTest(
+            EsExecutors.newScaling(
+                getName(),
+                0,
+                between(1, 5),
+                60,
+                TimeUnit.SECONDS,
+                true,
+                EsExecutors.daemonThreadFactory(getName()),
+                new ThreadContext(Settings.EMPTY)
+            )
+        );
+    }
+
+    public void testFixedBoundedRejectOnShutdown() {
+        runRejectOnShutdownTest(
+            EsExecutors.newFixed(
+                getName(),
+                between(1, 5),
+                between(1, 5),
+                EsExecutors.daemonThreadFactory(getName()),
+                threadContext,
+                randomFrom(DEFAULT, DO_NOT_TRACK)
+            )
+        );
+    }
+
+    public void testFixedUnboundedRejectOnShutdown() {
+        runRejectOnShutdownTest(
+            EsExecutors.newFixed(
+                getName(),
+                between(1, 5),
+                -1,
+                EsExecutors.daemonThreadFactory(getName()),
+                threadContext,
+                randomFrom(DEFAULT, DO_NOT_TRACK)
+            )
+        );
+    }
+
+    private static void runRejectOnShutdownTest(ExecutorService executor) {
+        for (int i = between(0, 10); i > 0; i--) {
+            final var delayMillis = between(0, 100);
+            executor.execute(ActionRunnable.wrap(ActionListener.noop(), l -> safeSleep(delayMillis)));
+        }
+        try {
+            executor.shutdown();
+            assertShutdownAndRejectingTasks(executor);
+        } finally {
+            ThreadPool.terminate(executor, 10, TimeUnit.SECONDS);
+        }
+        assertShutdownAndRejectingTasks(executor);
+    }
+
+    private static void assertShutdownAndRejectingTasks(Executor executor) {
+        final var rejected = new AtomicBoolean();
+        final var shouldBeRejected = new AbstractRunnable() {
+            @Override
+            public void onFailure(Exception e) {
+                fail("should not call onFailure");
+            }
+
+            @Override
+            protected void doRun() {
+                fail("should not call doRun");
+            }
+
+            @Override
+            public boolean isForceExecution() {
+                return randomBoolean();
+            }
+
+            @Override
+            public void onRejection(Exception e) {
+                assertTrue(asInstanceOf(EsRejectedExecutionException.class, e).isExecutorShutdown());
+                assertTrue(rejected.compareAndSet(false, true));
+            }
+        };
+        assertTrue(expectThrows(EsRejectedExecutionException.class, () -> executor.execute(shouldBeRejected::doRun)).isExecutorShutdown());
+        executor.execute(shouldBeRejected);
+        assertTrue(rejected.get());
+    }
 }