Fix RestCancellableNodeClientTests (#128085) (#128087)

These tests set `closeListener` racily, such that it's possible to see
it still be `null` in `TestHttpChannel#close()`. This commit replaces it
with a `SubscribableListener` so that we always complete the subscribed
listener regardless of order.

It also fixes a separate race in `testCompletedTasks` where (very
rarely) all of the `client.execute` calls had completed their listener
without getting to the point where they'd subscribed a `closeListener`.

Closes #128057

Author: DaveCTurner

server/src/test/java/org/elasticsearch/rest/action/RestCancellableNodeClientTests.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.TransportSearchAction;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.set.Sets;
@@ -44,7 +45,6 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.LongSupplier;
 
 public class RestCancellableNodeClientTests extends ESTestCase {
@@ -79,7 +79,9 @@ public void testCompletedTasks() throws Exception {
             for (int j = 0; j < numTasks; j++) {
                 PlainActionFuture<SearchResponse> actionFuture = new PlainActionFuture<>();
                 RestCancellableNodeClient client = new RestCancellableNodeClient(testClient, channel);
-                threadPool.generic().submit(() -> client.execute(TransportSearchAction.TYPE, new SearchRequest(), actionFuture));
+                futures.add(
+                    threadPool.generic().submit(() -> client.execute(TransportSearchAction.TYPE, new SearchRequest(), actionFuture))
+                );
                 futures.add(actionFuture);
             }
         }
@@ -150,7 +152,7 @@ public void testChannelAlreadyClosed() {
         assertEquals(totalSearches, testClient.cancelledTasks.size());
     }
 
-    public void testConcurrentExecuteAndClose() throws Exception {
+    public void testConcurrentExecuteAndClose() {
         final var testClient = new TestClient(Settings.EMPTY, threadPool, true);
         int initialHttpChannels = RestCancellableNodeClient.getNumChannels();
         int numTasks = randomIntBetween(1, 30);
@@ -254,7 +256,7 @@ public String getLocalNodeId() {
 
     private class TestHttpChannel implements HttpChannel {
         private final AtomicBoolean open = new AtomicBoolean(true);
-        private final AtomicReference<ActionListener<Void>> closeListener = new AtomicReference<>();
+        private final SubscribableListener<ActionListener<Void>> closeListener = new SubscribableListener<>();
         private final CountDownLatch closeLatch = new CountDownLatch(1);
 
         @Override
@@ -273,8 +275,7 @@ public InetSocketAddress getRemoteAddress() {
         @Override
         public void close() {
             assertTrue("HttpChannel is already closed", open.compareAndSet(true, false));
-            ActionListener<Void> listener = closeListener.get();
-            if (listener != null) {
+            closeListener.andThenAccept(listener -> {
                 boolean failure = randomBoolean();
                 threadPool.generic().submit(() -> {
                     if (failure) {
@@ -284,11 +285,10 @@ public void close() {
                     }
                     closeLatch.countDown();
                 });
-            }
+            });
         }
 
         private void awaitClose() throws InterruptedException {
-            assertNotNull("must set closeListener before calling awaitClose", closeListener.get());
             close();
             closeLatch.await();
         }
@@ -304,9 +304,8 @@ public void addCloseListener(ActionListener<Void> listener) {
             if (open.get() == false) {
                 listener.onResponse(null);
             } else {
-                if (closeListener.compareAndSet(null, listener) == false) {
-                    throw new AssertionError("close listener already set, only one is allowed!");
-                }
+                assertFalse("close listener already set, only one is allowed!", closeListener.isDone());
+                closeListener.onResponse(ActionListener.assertOnce(listener));
             }
         }
     }