Unique Remote Async Actions (elastic#139652)

Author: Mikep86

server/src/main/java/org/elasticsearch/index/query/QueryRewriteAsyncAction.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.client.internal.Client;
 
 import java.util.List;
+import java.util.Objects;
 import java.util.function.Consumer;
 
 /**
@@ -24,7 +25,7 @@
  * Since we need to determine whether an action has already been registered, we require implementors to provide implementations for
  * {@link #hashCode()} and {@link #equals(Object)},
  */
-public abstract class QueryRewriteAsyncAction<T> {
+public abstract class QueryRewriteAsyncAction<T, U extends QueryRewriteAsyncAction<T, U>> {
     /**
      * The execute method will:
      * - Execute the action using {@link #execute(Client, ActionListener)}
@@ -36,7 +37,7 @@ public abstract class QueryRewriteAsyncAction<T> {
      * @param consumers A list of consumer that expect the result of the action.
      */
     @SuppressWarnings("unchecked")
-    public void execute(Client client, ActionListener<?> listener, List<Consumer<?>> consumers) {
+    public final void execute(Client client, ActionListener<?> listener, List<Consumer<?>> consumers) {
         execute(client, listener.delegateFailureAndWrap((l, result) -> {
             consumers.forEach(consumer -> ((Consumer<T>) consumer).accept(result));
             l.onResponse(null);
@@ -52,8 +53,25 @@ public void execute(Client client, ActionListener<?> listener, List<Consumer<?>>
     protected abstract void execute(Client client, ActionListener<T> listener);
 
     @Override
-    public abstract int hashCode();
+    public int hashCode() {
+        return Objects.hash(getClass(), doHashCode());
+    };
+
+    public abstract int doHashCode();
 
     @Override
-    public abstract boolean equals(Object obj);
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        if (obj == null || getClass() != obj.getClass()) {
+            return false;
+        }
+
+        @SuppressWarnings("unchecked")
+        U other = (U) obj;
+        return doEquals(other);
+    }
+
+    public abstract boolean doEquals(U other);
 }

server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java

@@ -12,18 +12,15 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ResolvedIndices;
 import org.elasticsearch.client.internal.Client;
-import org.elasticsearch.client.internal.RemoteClusterClient;
 import org.elasticsearch.cluster.metadata.DataStream;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.routing.allocation.DataTier;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.TriConsumer;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.CountDown;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexSettings;
@@ -39,7 +36,6 @@
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.builder.PointInTimeBuilder;
 import org.elasticsearch.transport.RemoteClusterAware;
-import org.elasticsearch.transport.RemoteClusterService;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 
@@ -59,7 +55,6 @@
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.search.SearchService.DEFAULT_ALLOW_PARTIAL_SEARCH_RESULTS;
-import static org.elasticsearch.threadpool.ThreadPool.Names.SEARCH_COORDINATION;
 
 /**
  * Context object used to rewrite {@link QueryBuilder} instances into simplified version.
@@ -81,9 +76,7 @@ public class QueryRewriteContext {
     protected final Client client;
     protected final LongSupplier nowInMillis;
     private final List<BiConsumer<Client, ActionListener<?>>> asyncActions = new ArrayList<>();
-    private final Map<String, List<TriConsumer<RemoteClusterClient, ThreadContext, ActionListener<?>>>> remoteAsyncActions =
-        new HashMap<>();
-    private final Map<QueryRewriteAsyncAction<?>, List<Consumer<?>>> uniqueRewriteActions = new HashMap<>();
+    private final Map<QueryRewriteAsyncAction<?, ?>, List<Consumer<?>>> uniqueAsyncActions = new HashMap<>();
     protected boolean allowUnmappedFields;
     protected boolean mapUnmappedFieldAsString;
     protected Predicate<String> allowedFields;
@@ -467,22 +460,11 @@ public void registerAsyncAction(BiConsumer<Client, ActionListener<?>> asyncActio
         asyncActions.add(asyncAction);
     }
 
-    public void registerRemoteAsyncAction(
-        String clusterAlias,
-        TriConsumer<RemoteClusterClient, ThreadContext, ActionListener<?>> asyncAction
-    ) {
-        List<TriConsumer<RemoteClusterClient, ThreadContext, ActionListener<?>>> asyncActions = remoteAsyncActions.computeIfAbsent(
-            clusterAlias,
-            k -> new ArrayList<>()
-        );
-        asyncActions.add(asyncAction);
-    }
-
     /**
      * Returns <code>true</code> if there are any registered async actions.
      */
     public boolean hasAsyncActions() {
-        return asyncActions.isEmpty() == false || uniqueRewriteActions.isEmpty() == false || remoteAsyncActions.isEmpty() == false;
+        return asyncActions.isEmpty() == false || uniqueAsyncActions.isEmpty() == false;
     }
 
     /**
@@ -493,10 +475,7 @@ public void executeAsyncActions(ActionListener<Void> listener) {
         if (hasAsyncActions() == false) {
             listener.onResponse(null);
         } else {
-            int actionCount = asyncActions.size() + uniqueRewriteActions.size();
-            for (var actionList : remoteAsyncActions.values()) {
-                actionCount += actionList.size();
-            }
+            final int actionCount = asyncActions.size() + uniqueAsyncActions.size();
 
             CountDown countDown = new CountDown(actionCount);
             ActionListener<?> internalListener = new ActionListener<>() {
@@ -522,27 +501,10 @@ public void onFailure(Exception e) {
                 action.accept(client, internalListener);
             }
 
-            var copyUniqueRewriteActions = new HashMap<>(uniqueRewriteActions);
-            uniqueRewriteActions.clear();
-            for (var entry : copyUniqueRewriteActions.keySet()) {
-                entry.execute(client, internalListener, copyUniqueRewriteActions.get(entry));
-            }
-
-            var remoteAsyncActionsCopy = new HashMap<>(remoteAsyncActions);
-            remoteAsyncActions.clear();
-            for (var entry : remoteAsyncActionsCopy.entrySet()) {
-                String clusterAlias = entry.getKey();
-                List<TriConsumer<RemoteClusterClient, ThreadContext, ActionListener<?>>> remoteTriConsumers = entry.getValue();
-
-                RemoteClusterClient remoteClient = client.getRemoteClusterClient(
-                    clusterAlias,
-                    client.threadPool().executor(SEARCH_COORDINATION),
-                    RemoteClusterService.DisconnectedStrategy.RECONNECT_UNLESS_SKIP_UNAVAILABLE
-                );
-                ThreadContext threadContext = client.threadPool().getThreadContext();
-                for (var action : remoteTriConsumers) {
-                    action.apply(remoteClient, threadContext, internalListener);
-                }
+            var copyUniqueAsyncActions = new HashMap<>(uniqueAsyncActions);
+            uniqueAsyncActions.clear();
+            for (var entry : copyUniqueAsyncActions.keySet()) {
+                entry.execute(client, internalListener, copyUniqueAsyncActions.get(entry));
             }
         }
     }
@@ -733,7 +695,10 @@ public void setTrackTimeRangeFilterFrom(boolean trackTimeRangeFilterFrom) {
      * After the async action is executed, all consumers associated with it will be executed and receive as argument
      * the result of the async action.
      */
-    public <T> void registerUniqueAsyncAction(QueryRewriteAsyncAction<T> action, Consumer<T> consumer) {
-        uniqueRewriteActions.computeIfAbsent(action, k -> new ArrayList<>()).add(consumer);
+    public <T, U extends QueryRewriteAsyncAction<T, U>> void registerUniqueAsyncAction(
+        QueryRewriteAsyncAction<T, U> action,
+        Consumer<T> consumer
+    ) {
+        uniqueAsyncActions.computeIfAbsent(action, k -> new ArrayList<>()).add(consumer);
     }
 }

server/src/main/java/org/elasticsearch/index/query/QueryRewriteRemoteAsyncAction.java

@@ -0,0 +1,66 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.query;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.client.internal.RemoteClusterClient;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.RemoteClusterService;
+
+import java.util.Objects;
+
+import static org.elasticsearch.threadpool.ThreadPool.Names.SEARCH_COORDINATION;
+
+public abstract class QueryRewriteRemoteAsyncAction<T, U extends QueryRewriteRemoteAsyncAction<T, U>> extends QueryRewriteAsyncAction<
+    T,
+    U> {
+
+    private final String clusterAlias;
+
+    public QueryRewriteRemoteAsyncAction(String clusterAlias) {
+        this.clusterAlias = Objects.requireNonNull(clusterAlias);
+    }
+
+    public String getClusterAlias() {
+        return clusterAlias;
+    }
+
+    @Override
+    protected final void execute(Client client, ActionListener<T> listener) {
+        ThreadPool threadPool = client.threadPool();
+        RemoteClusterClient remoteClient = client.getRemoteClusterClient(
+            clusterAlias,
+            threadPool.executor(SEARCH_COORDINATION),
+            RemoteClusterService.DisconnectedStrategy.RECONNECT_UNLESS_SKIP_UNAVAILABLE
+        );
+        ThreadContext threadContext = threadPool.getThreadContext();
+        execute(remoteClient, threadContext, listener);
+    }
+
+    protected abstract void execute(RemoteClusterClient client, ThreadContext threadContext, ActionListener<T> listener);
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), clusterAlias);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        boolean equals = super.equals(obj);
+        if (equals) {
+            QueryRewriteRemoteAsyncAction<?, ?> other = (QueryRewriteRemoteAsyncAction<?, ?>) obj;
+            equals = Objects.equals(clusterAlias, other.clusterAlias);
+        }
+
+        return equals;
+    }
+}