Add utility for asserting listeners are completed

server/src/main/java/org/elasticsearch/action/ActionListener.java

@@ -425,6 +425,28 @@ public boolean equals(Object obj) {
         }
     }
 
+    /**
+     * @return A listener which (if assertions are enabled) wraps around the given delegate and asserts that it is called at least once.
+     */
+    static <Response> ActionListener<Response> assertAtLeastOnce(ActionListener<Response> delegate) {
+        if (Assertions.ENABLED) {
+            return new ActionListenerImplementations.AssertAtLeastOnceActionListener<>(
+                delegate,
+                // We use maybeDieOnAnotherThread here because this will get executed
+                // on the Cleaner thread, which silently swallows Throwable
+                (listener, createdAt) -> ExceptionsHelper.maybeDieOnAnotherThread(
+                    new AssertionError(
+                        "Expected listener "
+                            + delegate
+                            + " to be called at least once, but it was never called. Created:"
+                            + ExceptionsHelper.formatStackTrace(createdAt.getStackTrace())
+                    )
+                )
+            );
+        }
+        return delegate;
+    }
+
     /**
      * Execute the given action in a {@code try/catch} block which feeds all exceptions to the given listener's {@link #onFailure} method.
      */

server/src/main/java/org/elasticsearch/action/ActionListenerImplementations.java

@@ -8,12 +8,14 @@
 
 package org.elasticsearch.action;
 
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.CheckedBiConsumer;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.core.CheckedRunnable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 
+import java.lang.ref.Cleaner;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -24,6 +26,8 @@
  */
 class ActionListenerImplementations {
 
+    private static final Cleaner cleaner = Cleaner.create();
+
     private ActionListenerImplementations() {
         // no instances
     }
@@ -353,4 +357,84 @@ public String toString() {
             return "notifyOnce[" + get() + "]";
         }
     }
+
+    /**
+     * An {@link ActionListener} that wraps a delegate, registering itself with a {@link Cleaner} that will check that
+     * the delegate was called at least once prior to being garbage collected.
+     * If the listener was not completed prior to becoming unreachable it will invoke the passed {@link NotCalledListener}
+     * from the Cleaner thread.
+     */
+    static final class AssertAtLeastOnceActionListener<Response> extends DelegatingActionListener<Response, Response> {
+
+        interface NotCalledListener {
+
+            /**
+             * Called when {@link CallTrackingActionListener#run()} is called prior to the listener being
+             * called at least once.
+             *
+             * @param listener The listener that was not called at least once
+             * @param createdAt The stack trace of where the listener was created
+             */
+            void onListenerNotCalled(ActionListener<?> listener, ElasticsearchException createdAt);
+        }
+
+        /**
+         * This is the delegate that keeps track of whether {@link #onResponse(Object)} or {@link #onFailure(Exception)}
+         * were called, and invokes the {@link NotCalledListener} if neither was called.
+         * Because it is used as the "cleanup" action by the cleaner, it must be static, so it does not hold a reference
+         * to the wrapper and prevent it becoming "phantom reachable".
+         * See <a href="https://inside.java/2022/05/25/clean-cleaner/">the guidelines</a> for more details.
+         */
+        private static final class CallTrackingActionListener<R> implements Runnable, ActionListener<R> {
+
+            private final ActionListener<R> delegate;
+            private final ElasticsearchException created;
+            private final NotCalledListener listener;
+            private volatile boolean wasCalled = false;
+
+            private CallTrackingActionListener(ActionListener<R> delegate, NotCalledListener listener) {
+                this.delegate = delegate;
+                this.created = new ElasticsearchException("Listener created at...");
+                this.listener = listener;
+            }
+
+            @Override
+            public void run() {
+                if (wasCalled == false) {
+                    listener.onListenerNotCalled(delegate, created);
+                }
+            }
+
+            @Override
+            public void onResponse(R r) {
+                wasCalled = true;
+                delegate.onResponse(r);
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                wasCalled = true;
+                delegate.onFailure(e);
+            }
+        }
+
+        /**
+         * @param delegate The listener to wrap
+         * @param listener The action to perform if we detect that it wasn't called (called on the Cleaner thread)
+         */
+        AssertAtLeastOnceActionListener(ActionListener<Response> delegate, NotCalledListener listener) {
+            super(new CallTrackingActionListener<>(delegate, listener));
+            cleaner.register(this, (Runnable) this.delegate);
+        }
+
+        @Override
+        public void onResponse(Response response) {
+            delegate.onResponse(response);
+        }
+
+        @Override
+        public String toString() {
+            return delegate.toString();
+        }
+    }
 }

server/src/test/java/org/elasticsearch/action/ActionListenerTests.java

@@ -7,15 +7,18 @@
  */
 package org.elasticsearch.action;
 
+import org.apache.logging.log4j.Level;
 import org.apache.lucene.store.AlreadyClosedException;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.CheckedConsumer;
 import org.elasticsearch.core.CheckedRunnable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.MockLog;
 import org.elasticsearch.test.ReachabilityChecker;
 import org.hamcrest.Matcher;
 
@@ -28,6 +31,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
+import java.util.stream.IntStream;
 
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.endsWith;
@@ -370,6 +374,75 @@ public void onFailure(Exception e) {
         assertThat(exReference.get(), instanceOf(IllegalArgumentException.class));
     }
 
+    public void testAssertAtLeastOnceWillLogAssertionErrorWhenNotResolved() throws Exception {
+        try (MockLog mockLog = MockLog.capture(ExceptionsHelper.class)) {
+            mockLog.addExpectation(
+                new MockLog.PatternSeenEventExpectation(
+                    "action listener not called",
+                    ExceptionsHelper.class.getName(),
+                    Level.ERROR,
+                    ".*Expected listener NoopActionListener to be called at least once, but it was never called\\. Created:(.*|\\s)*"
+                )
+            );
+            final AtomicReference<ActionListener<Object>> listenerRef = new AtomicReference<>(
+                ActionListener.assertAtLeastOnce(ActionListener.noop())
+            );
+            assertBusy(() -> {
+                // Nullify reference so it becomes unreachable
+                listenerRef.set(null);
+                System.gc();
+                mockLog.assertAllExpectationsMatched();
+            });
+        }
+    }
+
+    public void testAssertAtLeastOnceWillInvokeListenerWhenNotResolved() throws Exception {
+        AtomicReference<Exception> listenerCreatedAt = new AtomicReference<>();
+        AtomicReference<ActionListener<?>> notCalledListener = new AtomicReference<>();
+        AtomicReference<ActionListener<Object>> listenerRef = new AtomicReference<>(
+            new ActionListenerImplementations.AssertAtLeastOnceActionListener<>(ActionListener.noop(), (listener, createdAt) -> {
+                logger.info("NotCalledListener called!");
+                notCalledListener.set(listener);
+                listenerCreatedAt.set(createdAt);
+            })
+        );
+        assertBusy(() -> {
+            // Nullify reference so it becomes unreachable
+            listenerRef.set(null);
+            System.gc();
+            assertNotNull(notCalledListener.get());
+            assertNotNull(listenerCreatedAt.get());
+        });
+    }
+
+    public void testAssertAtLeastOnceWillNotLogWhenResolvedOrFailed() throws Exception {
+        final ReachabilityChecker reachabilityChecker = new ReachabilityChecker();
+        final AtomicBoolean notCalledListenerCalled = new AtomicBoolean();
+        AtomicReference<ActionListener<Object>> listenerRef = new AtomicReference<>(
+            reachabilityChecker.register(
+                new ActionListenerImplementations.AssertAtLeastOnceActionListener<>(
+                    ActionListener.noop(),
+                    (listener, createdAt) -> notCalledListenerCalled.set(true)
+                )
+            )
+        );
+        // Call onResponse and/or onFailure at least once
+        IntStream.range(0, 1 + randomInt(3)).forEach(i -> {
+            if (randomBoolean()) {
+                listenerRef.get().onResponse("succeeded");
+            } else {
+                listenerRef.get().onFailure(new RuntimeException("Failed"));
+            }
+        });
+        assertBusy(() -> {
+            // Nullify reference so it becomes unreachable
+            listenerRef.set(null);
+            System.gc();
+            reachabilityChecker.ensureUnreachable();    // Only proceed once we know the object isn't reachable
+            assertFalse(notCalledListenerCalled.get());
+        });
+    }
+
     /**
      * Test that map passes the output of the function to its delegate listener and that exceptions in the function are propagated to the
      * onFailure handler. Also verify that exceptions from ActionListener.onResponse does not invoke onFailure, since it is the