Improve rejection handling in ThreadedActionListener (#87042)

Today if the submission within `ThreadedActionListener#onResponse` is
rejected from its threadpool then we call `delegate#onFailure` with the
rejection exception on the calling thread. However, if the submission
within `ThreadedActionListener#onFailure` is rejected then we just drop
the listener and log an error.

In most cases completing a listener exceptionally triggers some cleanup
which is often fairly lightweight and therefore safe enough to complete
on the calling thread. In any case it's generally preferable to complete
a listener exceptionally on the wrong thread rather than just dropping
it entirely.

This commit fixes this and adds a test to verify that
`ThreadedActionListener` completes properly even in the face of
rejections.

server/src/main/java/org/elasticsearch/action/support/ThreadedActionListener.java

@@ -9,7 +9,6 @@
 package org.elasticsearch.action.support;
 
 import org.apache.logging.log4j.Logger;
-import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRunnable;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
@@ -51,6 +50,11 @@ public boolean isForceExecution() {
             protected void doRun() {
                 listener.onResponse(response);
             }
+
+            @Override
+            public String toString() {
+                return ThreadedActionListener.this + "/onResponse";
+            }
         });
     }
 
@@ -63,14 +67,36 @@ public boolean isForceExecution() {
             }
 
             @Override
-            protected void doRun() throws Exception {
+            protected void doRun() {
                 delegate.onFailure(e);
             }
 
+            @Override
+            public void onRejection(Exception e2) {
+                e.addSuppressed(e2);
+                try {
+                    delegate.onFailure(e);
+                } catch (Exception e3) {
+                    e.addSuppressed(e3);
+                    onFailure(e);
+                }
+            }
+
             @Override
             public void onFailure(Exception e) {
-                logger.warn(() -> new ParameterizedMessage("failed to execute failure callback on [{}]", delegate), e);
+                assert false : e;
+                logger.error(() -> "failed to execute failure callback on [" + delegate + "]", e);
+            }
+
+            @Override
+            public String toString() {
+                return ThreadedActionListener.this + "/onFailure";
             }
         });
     }
+
+    @Override
+    public String toString() {
+        return "ThreadedActionListener[" + executor + "/" + delegate + "]";
+    }
 }

server/src/test/java/org/elasticsearch/action/support/ThreadedActionListenerTests.java

@@ -0,0 +1,79 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.action.support;
+
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.threadpool.FixedExecutorBuilder;
+import org.elasticsearch.threadpool.ScalingExecutorBuilder;
+import org.elasticsearch.threadpool.TestThreadPool;
+
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class ThreadedActionListenerTests extends ESTestCase {
+
+    public void testRejectionHandling() throws InterruptedException {
+        final var listenerCount = between(1, 1000);
+        final var countdownLatch = new CountDownLatch(listenerCount);
+        final var threadPool = new TestThreadPool(
+            "test",
+            Settings.EMPTY,
+            new FixedExecutorBuilder(Settings.EMPTY, "fixed-bounded-queue", between(1, 3), 10, "fbq", randomBoolean()),
+            new FixedExecutorBuilder(Settings.EMPTY, "fixed-unbounded-queue", between(1, 3), -1, "fnq", randomBoolean()),
+            new ScalingExecutorBuilder("scaling-drop-if-shutdown", between(1, 3), between(3, 5), TimeValue.timeValueSeconds(1), false),
+            new ScalingExecutorBuilder("scaling-reject-if-shutdown", between(1, 3), between(3, 5), TimeValue.timeValueSeconds(1), true)
+        );
+        final var closeFlag = new AtomicBoolean();
+        try {
+            final var pools = randomNonEmptySubsetOf(
+                List.of("fixed-bounded-queue", "fixed-unbounded-queue", "scaling-drop-if-shutdown", "scaling-reject-if-shutdown")
+            );
+            final var shutdownUnsafePools = Set.of("fixed-bounded-queue", "scaling-drop-if-shutdown");
+
+            threadPool.generic().execute(() -> {
+                for (int i = 0; i < listenerCount; i++) {
+                    final var pool = randomFrom(pools);
+                    final var listener = new ThreadedActionListener<Void>(
+                        logger,
+                        threadPool,
+                        pool,
+                        ActionListener.wrap(countdownLatch::countDown),
+                        (pool.equals("fixed-bounded-queue") || pool.startsWith("scaling")) && rarely()
+                    );
+                    synchronized (closeFlag) {
+                        if (closeFlag.get() && shutdownUnsafePools.contains(pool)) {
+                            // closing, so tasks submitted to this pool may just be dropped
+                            countdownLatch.countDown();
+                        } else if (randomBoolean()) {
+                            listener.onResponse(null);
+                        } else {
+                            listener.onFailure(new ElasticsearchException("simulated"));
+                        }
+                    }
+                    Thread.yield();
+                }
+            });
+        } finally {
+            synchronized (closeFlag) {
+                assertTrue(closeFlag.compareAndSet(false, true));
+                threadPool.shutdown();
+            }
+            assertTrue(threadPool.awaitTermination(10, TimeUnit.SECONDS));
+        }
+        assertTrue(countdownLatch.await(10, TimeUnit.SECONDS));
+    }
+
+}

test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java

@@ -1208,6 +1208,13 @@ public static <T> List<T> randomSubsetOf(Collection<T> collection) {
         return randomSubsetOf(randomInt(collection.size()), collection);
     }
 
+    public static <T> List<T> randomNonEmptySubsetOf(Collection<T> collection) {
+        if (collection.isEmpty()) {
+            throw new IllegalArgumentException("Can't pick non-empty subset of an empty collection");
+        }
+        return randomSubsetOf(randomIntBetween(1, collection.size()), collection);
+    }
+
     /**
      * Returns size random values
      */