Streamline skip_unavailable handling (#37672)

This commit moves the collectSearchShards method out of RemoteClusterService into TransportSearchAction that currently calls it. RemoteClusterService used to be used only for cross-cluster search but is now also used in cross-cluster replication where different API are called through the RemoteClusterAwareClient. There is no reason for the collectSearchShards and fetchShards methods to be respectively in RemoteClusterService and RemoteClusterConnection. The search shards API can be called through the RemoteClusterAwareClient too, the only missing bit is a way to handle failures based on the skip_unavailable setting for each cluster (currently only supported in RemoteClusterConnection#fetchShards) which is achieved by adding a isSkipUnavailable(String clusterAlias) method to RemoteClusterService.
This change is useful for #32125 as we will very soon need to also call the search API against remote clusters, which will be done through RemoteClusterAwareClient. In that case we will also need to support skip_unavailable when calling the search API so we need some way to handle the skip_unavailable setting like we currently do for the search_shards call.

Relates to #32125

Author: javanna

server/src/main/java/org/elasticsearch/action/admin/cluster/shards/ClusterSearchShardsResponse.java

@@ -29,15 +29,11 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
 public class ClusterSearchShardsResponse extends ActionResponse implements ToXContentObject {
 
-    public static final ClusterSearchShardsResponse EMPTY = new ClusterSearchShardsResponse(new ClusterSearchShardsGroup[0],
-            new DiscoveryNode[0], Collections.emptyMap());
-
     private final ClusterSearchShardsGroup[] groups;
     private final DiscoveryNode[] nodes;
     private final Map<String, AliasFilter> indicesAndFilters;

server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java

@@ -22,10 +22,12 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.OriginalIndices;
 import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsGroup;
+import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsRequest;
 import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsResponse;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
 import org.elasticsearch.action.support.IndicesOptions;
+import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
@@ -39,6 +41,7 @@
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Setting.Property;
+import org.elasticsearch.common.util.concurrent.CountDown;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.index.shard.ShardId;
@@ -50,6 +53,7 @@
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.RemoteClusterAware;
 import org.elasticsearch.transport.RemoteClusterService;
+import org.elasticsearch.transport.RemoteTransportException;
 import org.elasticsearch.transport.Transport;
 import org.elasticsearch.transport.TransportService;
 
@@ -60,8 +64,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.LongSupplier;
@@ -195,17 +202,23 @@ protected void doExecute(Task task, SearchRequest searchRequest, ActionListener<
                 executeSearch((SearchTask)task, timeProvider, searchRequest, localIndices, Collections.emptyList(),
                     (clusterName, nodeId) -> null, clusterState, Collections.emptyMap(), listener, SearchResponse.Clusters.EMPTY);
             } else {
-                remoteClusterService.collectSearchShards(searchRequest.indicesOptions(), searchRequest.preference(),
-                    searchRequest.routing(), remoteClusterIndices, ActionListener.wrap((searchShardsResponses) -> {
-                        List<SearchShardIterator> remoteShardIterators = new ArrayList<>();
-                        Map<String, AliasFilter> remoteAliasFilters = new HashMap<>();
-                        BiFunction<String, String, DiscoveryNode> clusterNodeLookup = processRemoteShards(searchShardsResponses,
-                            remoteClusterIndices, remoteShardIterators, remoteAliasFilters);
-                        SearchResponse.Clusters clusters = buildClusters(localIndices, remoteClusterIndices, searchShardsResponses);
-                        executeSearch((SearchTask) task, timeProvider, searchRequest, localIndices,
-                            remoteShardIterators, clusterNodeLookup, clusterState, remoteAliasFilters, listener,
-                            clusters);
-                    }, listener::onFailure));
+                AtomicInteger skippedClusters = new AtomicInteger(0);
+                collectSearchShards(searchRequest.indicesOptions(), searchRequest.preference(), searchRequest.routing(), skippedClusters,
+                    remoteClusterIndices, remoteClusterService, threadPool,
+                    ActionListener.wrap(
+                        searchShardsResponses -> {
+                            List<SearchShardIterator> remoteShardIterators = new ArrayList<>();
+                            Map<String, AliasFilter> remoteAliasFilters = new HashMap<>();
+                            BiFunction<String, String, DiscoveryNode> clusterNodeLookup = processRemoteShards(
+                                searchShardsResponses, remoteClusterIndices, remoteShardIterators, remoteAliasFilters);
+                            int localClusters = localIndices == null ? 0 : 1;
+                            int totalClusters = remoteClusterIndices.size() + localClusters;
+                            int successfulClusters = searchShardsResponses.size() + localClusters;
+                            executeSearch((SearchTask) task, timeProvider, searchRequest, localIndices,
+                                remoteShardIterators, clusterNodeLookup, clusterState, remoteAliasFilters, listener,
+                                new SearchResponse.Clusters(totalClusters, successfulClusters, skippedClusters.get()));
+                        },
+                        listener::onFailure));
             }
         }, listener::onFailure);
         if (searchRequest.source() == null) {
@@ -216,18 +229,56 @@ protected void doExecute(Task task, SearchRequest searchRequest, ActionListener<
         }
     }
 
-    static SearchResponse.Clusters buildClusters(OriginalIndices localIndices, Map<String, OriginalIndices> remoteIndices,
-                                                 Map<String, ClusterSearchShardsResponse> searchShardsResponses) {
-        int localClusters = localIndices == null ? 0 : 1;
-        int totalClusters = remoteIndices.size() + localClusters;
-        int successfulClusters = localClusters;
-        for (ClusterSearchShardsResponse searchShardsResponse : searchShardsResponses.values()) {
-            if (searchShardsResponse != ClusterSearchShardsResponse.EMPTY) {
-                successfulClusters++;
-            }
+    static void collectSearchShards(IndicesOptions indicesOptions, String preference, String routing, AtomicInteger skippedClusters,
+                                    Map<String, OriginalIndices> remoteIndicesByCluster, RemoteClusterService remoteClusterService,
+                                    ThreadPool threadPool, ActionListener<Map<String, ClusterSearchShardsResponse>> listener) {
+        final CountDown responsesCountDown = new CountDown(remoteIndicesByCluster.size());
+        final Map<String, ClusterSearchShardsResponse> searchShardsResponses = new ConcurrentHashMap<>();
+        final AtomicReference<RemoteTransportException> transportException = new AtomicReference<>();
+        for (Map.Entry<String, OriginalIndices> entry : remoteIndicesByCluster.entrySet()) {
+            final String clusterAlias = entry.getKey();
+            boolean skipUnavailable = remoteClusterService.isSkipUnavailable(clusterAlias);
+            Client clusterClient = remoteClusterService.getRemoteClusterClient(threadPool, clusterAlias);
+            final String[] indices = entry.getValue().indices();
+            ClusterSearchShardsRequest searchShardsRequest = new ClusterSearchShardsRequest(indices)
+                .indicesOptions(indicesOptions).local(true).preference(preference).routing(routing);
+            clusterClient.admin().cluster().searchShards(searchShardsRequest, new ActionListener<ClusterSearchShardsResponse>() {
+                    @Override
+                    public void onResponse(ClusterSearchShardsResponse response) {
+                        searchShardsResponses.put(clusterAlias, response);
+                        maybeFinish();
+                    }
+
+                    @Override
+                    public void onFailure(Exception e) {
+                        if (skipUnavailable) {
+                            skippedClusters.incrementAndGet();
+                        } else {
+                            RemoteTransportException exception =
+                                new RemoteTransportException("error while communicating with remote cluster [" + clusterAlias + "]", e);
+                            if (transportException.compareAndSet(null, exception) == false) {
+                                transportException.accumulateAndGet(exception, (previous, current) -> {
+                                    current.addSuppressed(previous);
+                                    return current;
+                                });
+                            }
+                        }
+                        maybeFinish();
+                    }
+
+                    private void maybeFinish() {
+                        if (responsesCountDown.countDown()) {
+                            RemoteTransportException exception = transportException.get();
+                            if (exception == null) {
+                                listener.onResponse(searchShardsResponses);
+                            } else {
+                                listener.onFailure(transportException.get());
+                            }
+                        }
+                    }
+                }
+            );
         }
-        int skippedClusters = totalClusters - successfulClusters;
-        return new SearchResponse.Clusters(totalClusters, successfulClusters, skippedClusters);
     }
 
     static BiFunction<String, String, DiscoveryNode> processRemoteShards(Map<String, ClusterSearchShardsResponse> searchShardsResponses,

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

@@ -25,9 +25,6 @@
 import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsAction;
-import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsRequest;
-import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsResponse;
 import org.elasticsearch.action.admin.cluster.state.ClusterStateAction;
 import org.elasticsearch.action.admin.cluster.state.ClusterStateRequest;
 import org.elasticsearch.action.admin.cluster.state.ClusterStateResponse;
@@ -62,7 +59,6 @@
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
@@ -172,6 +168,13 @@ void updateSkipUnavailable(boolean skipUnavailable) {
         this.skipUnavailable = skipUnavailable;
     }
 
+    /**
+     * Returns whether this cluster is configured to be skipped when unavailable
+     */
+    boolean isSkipUnavailable() {
+        return skipUnavailable;
+    }
+
     @Override
     public void onNodeDisconnected(DiscoveryNode node) {
         boolean remove = connectedNodes.remove(node);
@@ -181,67 +184,18 @@ public void onNodeDisconnected(DiscoveryNode node) {
         }
     }
 
-    /**
-     * Fetches all shards for the search request from this remote connection. This is used to later run the search on the remote end.
-     */
-    public void fetchSearchShards(ClusterSearchShardsRequest searchRequest,
-                                  ActionListener<ClusterSearchShardsResponse> listener) {
-
-        final ActionListener<ClusterSearchShardsResponse> searchShardsListener;
-        final Consumer<Exception> onConnectFailure;
-        if (skipUnavailable) {
-            onConnectFailure = (exception) -> listener.onResponse(ClusterSearchShardsResponse.EMPTY);
-            searchShardsListener = ActionListener.wrap(listener::onResponse, (e) -> listener.onResponse(ClusterSearchShardsResponse.EMPTY));
-        } else {
-            onConnectFailure = listener::onFailure;
-            searchShardsListener = listener;
-        }
-        // in case we have no connected nodes we try to connect and if we fail we either notify the listener or not depending on
-        // the skip_unavailable setting
-        ensureConnected(ActionListener.wrap((x) -> fetchShardsInternal(searchRequest, searchShardsListener), onConnectFailure));
-    }
-
     /**
      * Ensures that this cluster is connected. If the cluster is connected this operation
      * will invoke the listener immediately.
      */
-    public void ensureConnected(ActionListener<Void> voidActionListener) {
+    void ensureConnected(ActionListener<Void> voidActionListener) {
         if (connectedNodes.size() == 0) {
             connectHandler.connect(voidActionListener);
         } else {
             voidActionListener.onResponse(null);
         }
     }
 
-    private void fetchShardsInternal(ClusterSearchShardsRequest searchShardsRequest,
-                                     final ActionListener<ClusterSearchShardsResponse> listener) {
-        final DiscoveryNode node = getAnyConnectedNode();
-        Transport.Connection connection = connectionManager.getConnection(node);
-        transportService.sendRequest(connection, ClusterSearchShardsAction.NAME, searchShardsRequest, TransportRequestOptions.EMPTY,
-            new TransportResponseHandler<ClusterSearchShardsResponse>() {
-
-                @Override
-                public ClusterSearchShardsResponse read(StreamInput in) throws IOException {
-                    return new ClusterSearchShardsResponse(in);
-                }
-
-                @Override
-                public void handleResponse(ClusterSearchShardsResponse clusterSearchShardsResponse) {
-                    listener.onResponse(clusterSearchShardsResponse);
-                }
-
-                @Override
-                public void handleException(TransportException e) {
-                    listener.onFailure(e);
-                }
-
-                @Override
-                public String executor() {
-                    return ThreadPool.Names.SEARCH;
-                }
-            });
-    }
-
     /**
      * Collects all nodes on the connected cluster and returns / passes a nodeID to {@link DiscoveryNode} lookup function
      * that returns <code>null</code> if the node ID is not found.