Parallelize remotes fetching

Author: smalyshev

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

@@ -72,6 +72,7 @@
 import org.elasticsearch.action.admin.cluster.state.ClusterStateAction;
 import org.elasticsearch.action.admin.cluster.state.TransportClusterStateAction;
 import org.elasticsearch.action.admin.cluster.stats.TransportClusterStatsAction;
+import org.elasticsearch.action.admin.cluster.stats.TransportRemoteClusterStatsAction;
 import org.elasticsearch.action.admin.cluster.storedscripts.GetScriptContextAction;
 import org.elasticsearch.action.admin.cluster.storedscripts.GetScriptLanguageAction;
 import org.elasticsearch.action.admin.cluster.storedscripts.GetStoredScriptAction;
@@ -641,6 +642,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         actions.register(TransportGetDesiredBalanceAction.TYPE, TransportGetDesiredBalanceAction.class);
         actions.register(TransportDeleteDesiredBalanceAction.TYPE, TransportDeleteDesiredBalanceAction.class);
         actions.register(TransportClusterStatsAction.TYPE, TransportClusterStatsAction.class);
+        actions.register(TransportRemoteClusterStatsAction.TYPE, TransportRemoteClusterStatsAction.class);
         actions.register(ClusterStateAction.INSTANCE, TransportClusterStateAction.class);
         actions.register(TransportClusterHealthAction.TYPE, TransportClusterHealthAction.class);
         actions.register(ClusterUpdateSettingsAction.INSTANCE, TransportClusterUpdateSettingsAction.class);

server/src/main/java/org/elasticsearch/action/search/RemoteClusterActionListener.java renamed to server/src/main/java/org/elasticsearch/action/admin/cluster/stats/RemoteClusterActionListener.java

@@ -5,7 +5,7 @@
  * in compliance with, at your election, the Elastic License 2.0 or the Server
  * Side Public License, v 1.
  */
-package org.elasticsearch.action.search;
+package org.elasticsearch.action.admin.cluster.stats;
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.DelegatingActionListener;

server/src/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java

@@ -8,14 +8,13 @@
 
 package org.elasticsearch.action.admin.cluster.stats;
 
-import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.ActionFuture;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.FailedNodeException;
-import org.elasticsearch.action.search.RemoteClusterActionListener;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.ClusterSnapshotStats;
@@ -43,7 +42,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ExecutionException;
 import java.util.function.BiFunction;
 import java.util.function.BooleanSupplier;
 
@@ -89,6 +87,14 @@ public TransportClusterStatsAction(
         this.settings = settings;
     }
 
+    private ActionFuture<Map<String, RemoteClusterStatsResponse>> remoteFuture;
+
+    @Override
+    protected void doExecute(Task task, ClusterStatsRequest request, ActionListener<ClusterStatsResponse> listener) {
+        remoteFuture = getStatsFromRemotes(request);
+        super.doExecute(task, request, listener);
+    }
+
     @Override
     protected void newResponseAsync(
         final Task task,
@@ -111,8 +117,7 @@ protected void newResponseAsync(
             clusterService.threadPool().absoluteTimeInMillis()
         );
 
-        // TODO: this should not be happening here but leaving it here for now until we figure out proper
-        // threading/async model for this
+        // This will wait until remotes are done if it didn't happen yet
         var remoteClusterStats = getRemoteClusterStats(request);
 
         final ListenableFuture<MappingStats> mappingStatsStep = new ListenableFuture<>();
@@ -178,7 +183,7 @@ private Map<String, ClusterStatsResponse.RemoteClusterStats> getRemoteClusterSta
             return null;
         }
         Map<String, ClusterStatsResponse.RemoteClusterStats> remoteClustersStats = new HashMap<>();
-        Map<String, RemoteClusterStatsResponse> remoteData = getStatsFromRemotes(request);
+        Map<String, RemoteClusterStatsResponse> remoteData = resolveRemoteClusterStats();
 
         for (String clusterAlias : remoteClusterService.getRegisteredRemoteClusterNames()) {
             RemoteClusterConnection remoteConnection = remoteClusterService.getRemoteClusterConnection(clusterAlias);
@@ -196,16 +201,27 @@ private Map<String, ClusterStatsResponse.RemoteClusterStats> getRemoteClusterSta
         return remoteClustersStats;
     }
 
-    private Map<String, RemoteClusterStatsResponse> getStatsFromRemotes(ClusterStatsRequest request) {
-        // TODO: make correct pool
-        final var remoteClientResponseExecutor = transportService.getThreadPool().executor(ThreadPool.Names.MANAGEMENT);
-        if (request.doRemotes() == false) {
+    private Map<String, RemoteClusterStatsResponse> resolveRemoteClusterStats() {
+        try {
+            return remoteFuture.actionGet();
+        } catch (ElasticsearchException e) {
+            logger.warn("Failed to get remote cluster stats", e);
             return Map.of();
         }
+    }
+
+    private ActionFuture<Map<String, RemoteClusterStatsResponse>> getStatsFromRemotes(ClusterStatsRequest request) {
+        if (request.doRemotes() == false) {
+            // this will never be used since getRemoteClusterStats has the same check
+            return null;
+        }
+
+        // TODO: make correct pool
+        final var remoteClientResponseExecutor = transportService.getThreadPool().executor(ThreadPool.Names.MANAGEMENT);
         var remotes = remoteClusterService.getRegisteredRemoteClusterNames();
 
-        var remotesListener = new PlainActionFuture<Map<String, RemoteClusterStatsResponse>>();
-        var groupListener = new RemoteClusterActionListener<>(remotes.size(), remotesListener);
+        var remotesFuture = new PlainActionFuture<Map<String, RemoteClusterStatsResponse>>();
+        var groupListener = new RemoteClusterActionListener<>(remotes.size(), remotesFuture);
 
         for (String clusterAlias : remotes) {
             ClusterStatsRequest remoteRequest = request.subRequest();
@@ -222,13 +238,6 @@ private Map<String, RemoteClusterStatsResponse> getStatsFromRemotes(ClusterStats
 
         }
 
-        try {
-            // TODO: how do we report errors?
-            return remotesListener.get();
-        } catch (InterruptedException | ExecutionException e) {
-            logger.log(Level.ERROR, "Failed to get remote cluster stats: ", ExceptionsHelper.unwrapCause(e));
-            return Map.of();
-        }
+        return remotesFuture;
     }
-
 }