Refactor to use RefCountingListener

Author: JeremyDahlgren

server/src/main/java/org/elasticsearch/transport/RemoteClusterService.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.action.support.CountDownActionListener;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.RefCountingListener;
 import org.elasticsearch.action.support.RefCountingRunnable;
 import org.elasticsearch.client.internal.RemoteClusterClient;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
@@ -29,7 +30,6 @@
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
-import org.elasticsearch.common.util.concurrent.CountDown;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.TimeValue;
@@ -568,43 +568,26 @@ public void collectNodes(Set<String> clusters, ActionListener<BiFunction<String,
                 "this node does not have the " + DiscoveryNodeRole.REMOTE_CLUSTER_CLIENT_ROLE.roleName() + " role"
             );
         }
+        final var connectionsMap = new HashMap<String, RemoteClusterConnection>();
         for (String cluster : clusters) {
-            if (this.remoteClusters.containsKey(cluster) == false) {
+            final var connection = this.remoteClusters.get(cluster);
+            if (connection == null) {
                 listener.onFailure(new NoSuchRemoteClusterException(cluster));
                 return;
             }
+            connectionsMap.put(cluster, connection);
         }
 
         final Map<String, Function<String, DiscoveryNode>> clusterMap = new HashMap<>();
-        CountDown countDown = new CountDown(clusters.size());
-        Function<String, DiscoveryNode> nullFunction = s -> null;
-        for (final String cluster : clusters) {
-            RemoteClusterConnection connection = this.remoteClusters.get(cluster);
-            // Ensure the connection is not null, it could have been removed since the containsKey() call above.
-            if (connection == null) {
-                if (countDown.fastForward()) {
-                    listener.onFailure(new NoSuchRemoteClusterException(cluster));
-                }
-                break;
-            }
-            connection.collectNodes(new ActionListener<Function<String, DiscoveryNode>>() {
-                @Override
-                public void onResponse(Function<String, DiscoveryNode> nodeLookup) {
-                    synchronized (clusterMap) {
-                        clusterMap.put(cluster, nodeLookup);
-                    }
-                    if (countDown.countDown()) {
-                        listener.onResponse((clusterAlias, nodeId) -> clusterMap.getOrDefault(clusterAlias, nullFunction).apply(nodeId));
-                    }
-                }
-
-                @Override
-                public void onFailure(Exception e) {
-                    if (countDown.fastForward()) { // we need to check if it's true since we could have multiple failures
-                        listener.onFailure(e);
-                    }
+        final var finalListener = listener.<Void>safeMap(
+            ignored -> (clusterAlias, nodeId) -> clusterMap.getOrDefault(clusterAlias, s -> null).apply(nodeId)
+        );
+        try (var refs = new RefCountingListener(finalListener)) {
+            connectionsMap.forEach((cluster, connection) -> connection.collectNodes(refs.acquire(nodeLookup -> {
+                synchronized (clusterMap) {
+                    clusterMap.put(cluster, nodeLookup);
                 }
-            });
+            })));
         }
     }