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,53 @@ 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);
+        ActionListener<Object> listenerRef = 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 = 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();
+        ActionListener<Object> listenerRef = 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.onResponse("succeeded");
+        } else {
+            listenerRef.onFailure(new RuntimeException("Failed"));
+        }
+        // Nullify reference so it becomes unreachable
+        listenerRef = 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

server/src/test/java/org/elasticsearch/transport/LeakTrackerTests.java

@@ -0,0 +1,193 @@
+/*
+ * 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.transport;
+
+import com.carrotsearch.randomizedtesting.annotations.Name;
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.common.Randomness;
+import org.elasticsearch.core.AbstractRefCounted;
+import org.elasticsearch.core.Assertions;
+import org.elasticsearch.core.RefCounted;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.ReachabilityChecker;
+import org.junit.Before;
+
+import java.util.stream.Stream;
+
+public class LeakTrackerTests extends ESTestCase {
+
+    private static final Logger logger = LogManager.getLogger();
+
+    private final TrackedObjectTestCase trackedObjectTestCase;
+    private ReachabilityChecker reachabilityChecker;
+
+    @Before
+    public void createReachabilityTracker() {
+        reachabilityChecker = new ReachabilityChecker();
+    }
+
+    public LeakTrackerTests(@Name("trackingMethod") TrackedObjectTestCase trackedObjectTestCase) {
+        this.trackedObjectTestCase = trackedObjectTestCase;
+    }
+
+    @ParametersFactory(shuffle = false)
+    public static Iterable<TrackedObjectTestCase[]> parameters() {
+        return Stream.of(
+            new PojoTrackedObjectTestCase(),
+            new ReleasableTrackedObjectTestCase(),
+            new ReferenceCountedTrackedObjectTestCase()
+        ).map(i -> new TrackedObjectTestCase[] { i }).toList();
+    }
+
+    public void testWillLogErrorWhenTrackedObjectIsNotClosed() throws Exception {
+        trackedObjectTestCase.createAndTrack(reachabilityChecker);
+        // Do not close leak before nullifying
+        trackedObjectTestCase.nullifyReference();
+        reachabilityChecker.ensureUnreachable();
+        assertBusy(() -> assertLeakDetected("LEAK: resource was not cleaned up before it was garbage-collected\\.(.*|\\s)*"));
+    }
+
+    public void testWillNotLogErrorWhenTrackedObjectIsClosed() {
+        assumeTrue("assertAtLeastOnce will be a no-op when assertions are disabled", Assertions.ENABLED);
+        trackedObjectTestCase.createAndTrack(reachabilityChecker);
+        trackedObjectTestCase.closeLeak();
+        trackedObjectTestCase.nullifyReference();
+        // Should not detect a leak
+        reachabilityChecker.ensureUnreachable();
+    }
+
+    /**
+     * Encapsulates the lifecycle for a particular type of tracked object
+     */
+    public interface TrackedObjectTestCase {
+
+        /**
+         * Create the tracked object, implementations must
+         * - track it with the {@link LeakTracker}
+         * - register it with the passed reachability checker
+         * - retain a reference to it
+         * @param reachabilityChecker The reachability checker
+         */
+        void createAndTrack(ReachabilityChecker reachabilityChecker);
+
+        /**
+         * Nullify the retained reference
+         */
+        void nullifyReference();
+
+        /**
+         * Close the {@link LeakTracker.Leak}
+         */
+        void closeLeak();
+    }
+
+    private static class PojoTrackedObjectTestCase implements TrackedObjectTestCase {
+
+        private Object object;
+        private LeakTracker.Leak leak;
+
+        @Override
+        public void createAndTrack(ReachabilityChecker reachabilityChecker) {
+            object = reachabilityChecker.register(new Object());
+            leak = LeakTracker.INSTANCE.track(object);
+        }
+
+        @Override
+        public void nullifyReference() {
+            object = null;
+        }
+
+        @Override
+        public void closeLeak() {
+            leak.close();
+        }
+
+        @Override
+        public String toString() {
+            return "LeakTracker.INSTANCE.track(Object)";
+        }
+    }
+
+    private static class ReferenceCountedTrackedObjectTestCase implements TrackedObjectTestCase {
+
+        private RefCounted refCounted;
+
+        @Override
+        public void createAndTrack(ReachabilityChecker reachabilityChecker) {
+            RefCounted refCounted = reachabilityChecker.register(createRefCounted());
+            this.refCounted = LeakTracker.wrap(refCounted);
+            this.refCounted.incRef();
+            this.refCounted.tryIncRef();
+        }
+
+        @Override
+        public void nullifyReference() {
+            refCounted = null;
+        }
+
+        @Override
+        public void closeLeak() {
+            refCounted.decRef();
+            refCounted.decRef();
+            refCounted.decRef();
+        }
+
+        @Override
+        public String toString() {
+            return "LeakTracker.wrap(RefCounted)";
+        }
+
+        private static RefCounted createRefCounted() {
+            int number = Randomness.get().nextInt();
+            return new AbstractRefCounted() {
+
+                @Override
+                protected void closeInternal() {
+                    // Do nothing
+                    logger.info("Prevents this returning a non-collectible constant {}", number);
+                }
+            };
+        }
+    }
+
+    private static class ReleasableTrackedObjectTestCase implements TrackedObjectTestCase {
+
+        private Releasable releasable;
+
+        @Override
+        public void createAndTrack(ReachabilityChecker reachabilityChecker) {
+            Releasable releasable = reachabilityChecker.register(createReleasable());
+            this.releasable = LeakTracker.wrap(releasable);
+        }
+
+        @Override
+        public void nullifyReference() {
+            releasable = null;
+        }
+
+        @Override
+        public void closeLeak() {
+            releasable.close();
+        }
+
+        @Override
+        public String toString() {
+            return "LeakTracker.wrap(Releasable)";
+        }
+
+        private static Releasable createReleasable() {
+            int number = Randomness.get().nextInt();
+            return () -> logger.info("Prevents this returning a non-collectible constant {}", number);
+        }
+    }
+}

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());