Add utility for asserting listeners are completed

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

@@ -18,6 +18,7 @@
 import org.elasticsearch.core.CheckedRunnable;
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
+import org.elasticsearch.transport.LeakTracker;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -425,6 +426,16 @@ 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 ActionListener.runBefore(delegate, LeakTracker.INSTANCE.track(delegate)::close);
+        }
+        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/test/java/org/elasticsearch/action/ActionListenerTests.java

@@ -33,6 +33,8 @@
 import static org.hamcrest.Matchers.endsWith;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 public class ActionListenerTests extends ESTestCase {
 
@@ -370,6 +372,57 @@ public void onFailure(Exception e) {
         assertThat(exReference.get(), instanceOf(IllegalArgumentException.class));
     }
 
+    public void testAssertAtLeastOnceWillLogAssertionErrorWhenNotResolved() throws Exception {
+        assumeTrue("assertAtLeastOnce will be a no-op when assertions are disabled", Assertions.ENABLED);
+        final AtomicReference<ActionListener<Object>> listenerRef = new AtomicReference<>(
+            ActionListener.assertAtLeastOnce(ActionListener.running(() -> {
+                // Do nothing, but don't use ActionListener.noop() as it'll never be garbage collected
+            }))
+        );
+        // Nullify reference so it becomes unreachable
+        listenerRef.set(null);
+        assertBusy(() -> {
+            System.gc();
+            assertLeakDetected("LEAK: resource was not cleaned up before it was garbage-collected\\.(.*|\\s)*");
+        });
+    }
+
+    public void testAssertAtLeastOnceWillNotLogWhenResolvedOrFailed() {
+        assumeTrue("assertAtLeastOnce will be a no-op when assertions are disabled", Assertions.ENABLED);
+        ReachabilityChecker reachabilityChecker = new ReachabilityChecker();
+        final AtomicReference<ActionListener<Object>> listenerRef = new AtomicReference<>(
+            reachabilityChecker.register(ActionListener.assertAtLeastOnce(ActionListener.running(() -> {
+                // Do nothing, but don't use ActionListener.noop() as it'll never be garbage collected
+            })))
+        );
+        // Call onResponse or onFailure
+        if (randomBoolean()) {
+            listenerRef.get().onResponse("succeeded");
+        } else {
+            listenerRef.get().onFailure(new RuntimeException("Failed"));
+        }
+        // Nullify reference so it becomes unreachable
+        listenerRef.set(null);
+        reachabilityChecker.ensureUnreachable();
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testAssertAtLeastOnceWillDelegateResponses() {
+        ActionListener<Object> delegate = mock(ActionListener.class);
+        ActionListener<Object> listener = ActionListener.assertAtLeastOnce(delegate);
+        listener.onResponse("succeeded");
+        verify(delegate).onResponse("succeeded");
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testAssertAtLeastOnceWillDelegateFailures() {
+        ActionListener<Object> delegate = mock(ActionListener.class);
+        ActionListener<Object> listener = ActionListener.assertAtLeastOnce(delegate);
+        RuntimeException exception = new RuntimeException();
+        listener.onFailure(exception);
+        verify(delegate).onFailure(exception);
+    }
+
     /**
      * 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

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

@@ -186,6 +186,7 @@
 import java.util.function.IntFunction;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.DoubleStream;
 import java.util.stream.IntStream;
@@ -788,6 +789,18 @@ protected static void checkStaticState() throws Exception {
         }
     }
 
+    /**
+     * Assert that a leak was detected, also remove the leak from the list of detected leaks
+     * so the test won't fail for that specific leak.
+     *
+     * @param expectedPattern A pattern that matches the detected leak's exception
+     */
+    protected static void assertLeakDetected(String expectedPattern) {
+        synchronized (loggedLeaks) {
+            assertTrue(loggedLeaks.removeIf(leakText -> Pattern.matches(expectedPattern, leakText)));
+        }
+    }
+
     // this must be a separate method from other ensure checks above so suite scoped integ tests can call...TODO: fix that
     public final void ensureAllSearchContextsReleased() throws Exception {
         assertBusy(() -> MockSearchService.assertNoInFlightContext());