Add support for local cluster alias to SearchRequest (#36997)

With the upcoming cross-cluster search alternate execution mode, the CCS
 node will be able to split a CCS request into multiple search requests,
 one per remote cluster involved. In order to do that, the CCS node has
to be able to signal to each remote cluster that such sub-requests are
part of a CCS request. Each cluster does not know about the other
clusters involved, and does not know either what alias it is given in
the CCS node, hence the CCS coordinating node needs to be able to provide
the alias as part of the search request so that it is used as index prefix
in the returned search hits.

The cluster alias is a notion that's already supported in the search shards
iterator and search shard target, but it is currently used in CCS as both
 index prefix and connection lookup key when fanning out to all the shards.
With CCS alternate execution mode the provided cluster alias needs to be
used only as index prefix, as shards are local to each cluster hence no
cluster alias should be used for connection lookups.

The local cluster alias can be set to the SearchRequest at the transport layer
only, and its constructor/getter methods are package private.

Relates to #32125

Author: javanna

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

@@ -318,16 +318,16 @@ public final void onFailure(Exception e) {
         listener.onFailure(e);
     }
 
+    @Override
     public final ShardSearchTransportRequest buildShardSearchRequest(SearchShardIterator shardIt) {
-        String clusterAlias = shardIt.getClusterAlias();
         AliasFilter filter = aliasFilter.get(shardIt.shardId().getIndex().getUUID());
         assert filter != null;
         float indexBoost = concreteIndexBoosts.getOrDefault(shardIt.shardId().getIndex().getUUID(), DEFAULT_INDEX_BOOST);
         String indexName = shardIt.shardId().getIndex().getName();
         final String[] routings = indexRoutings.getOrDefault(indexName, Collections.emptySet())
             .toArray(new String[0]);
         return new ShardSearchTransportRequest(shardIt.getOriginalIndices(), request, shardIt.shardId(), getNumShards(),
-            filter, indexBoost, timeProvider.getAbsoluteStartMillis(), clusterAlias, routings);
+            filter, indexBoost, timeProvider.getAbsoluteStartMillis(), shardIt.getClusterAlias(), routings);
     }
 
     /**

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

@@ -85,8 +85,7 @@ private void onShardFailure(final int shardIndex, @Nullable ShardRouting shard,
                                 final SearchShardIterator shardIt, Exception e) {
         // we always add the shard failure for a specific shard instance
         // we do make sure to clean it on a successful response from a shard
-        SearchShardTarget shardTarget = new SearchShardTarget(nodeId, shardIt.shardId(), shardIt.getClusterAlias(),
-                shardIt.getOriginalIndices());
+        SearchShardTarget shardTarget = shardIt.newSearchShardTarget(nodeId);
         onShardFailure(shardIndex, shardTarget, e);
 
         if (totalOps.incrementAndGet() == expectedTotalOps) {
@@ -209,8 +208,8 @@ private void performPhaseOnShard(final int shardIndex, final SearchShardIterator
             fork(() -> onShardFailure(shardIndex, null, null, shardIt, new NoShardAvailableActionException(shardIt.shardId())));
         } else {
             try {
-                executePhaseOnShard(shardIt, shard, new SearchActionListener<FirstResult>(new SearchShardTarget(shard.currentNodeId(),
-                    shardIt.shardId(), shardIt.getClusterAlias(), shardIt.getOriginalIndices()), shardIndex) {
+                executePhaseOnShard(shardIt, shard, new SearchActionListener<FirstResult>(
+                    shardIt.newSearchShardTarget(shard.currentNodeId()), shardIndex) {
                     @Override
                     public void innerOnResponse(FirstResult result) {
                         maybeFork(thread, () -> onShardResult(result, shardIt));

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

@@ -36,6 +36,7 @@ public String getNode() {
         return node;
     }
 
+    @Nullable
     public String getClusterAlias() {
         return clusterAlias;
     }

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

@@ -65,6 +65,8 @@ public final class SearchRequest extends ActionRequest implements IndicesRequest
     public static final int DEFAULT_PRE_FILTER_SHARD_SIZE = 128;
     public static final int DEFAULT_BATCHED_REDUCE_SIZE = 512;
 
+    private String localClusterAlias;
+
     private SearchType searchType = SearchType.DEFAULT;
 
     private String[] indices = Strings.EMPTY_ARRAY;
@@ -95,6 +97,7 @@ public final class SearchRequest extends ActionRequest implements IndicesRequest
     private IndicesOptions indicesOptions = DEFAULT_INDICES_OPTIONS;
 
     public SearchRequest() {
+        this.localClusterAlias = null;
     }
 
     /**
@@ -114,6 +117,7 @@ public SearchRequest(SearchRequest searchRequest) {
         this.searchType = searchRequest.searchType;
         this.source = searchRequest.source;
         this.types = searchRequest.types;
+        this.localClusterAlias = searchRequest.localClusterAlias;
     }
 
     /**
@@ -128,13 +132,22 @@ public SearchRequest(String... indices) {
      * Constructs a new search request against the provided indices with the given search source.
      */
     public SearchRequest(String[] indices, SearchSourceBuilder source) {
+        this();
         if (source == null) {
             throw new IllegalArgumentException("source must not be null");
         }
         indices(indices);
         this.source = source;
     }
 
+    /**
+     * Creates a new search request by providing the alias of the cluster where it will be executed. Used when a {@link SearchRequest}
+     * is created and executed as part of a cross-cluster search request performing local reduction on each cluster.
+     * The coordinating CCS node provides the alias to prefix index names with in the returned search results.
+     */
+    SearchRequest(String localClusterAlias) {
+        this.localClusterAlias = Objects.requireNonNull(localClusterAlias, "cluster alias must not be null");
+    }
 
     @Override
     public ActionRequestValidationException validate() {
@@ -162,6 +175,16 @@ public ActionRequestValidationException validate() {
         return validationException;
     }
 
+    /**
+     * Returns the alias of the cluster that this search request is being executed on. A non-null value indicates that this search request
+     * is being executed as part of a locally reduced cross-cluster search request. The cluster alias is used to prefix index names
+     * returned as part of search hits with the alias of the cluster they came from.
+     */
+    @Nullable
+    String getLocalClusterAlias() {
+        return localClusterAlias;
+    }
+
     /**
      * Sets the indices the search will be executed on.
      */
@@ -470,7 +493,12 @@ public void readFrom(StreamInput in) throws IOException {
         }
         if (in.getVersion().onOrAfter(Version.V_6_3_0)) {
             allowPartialSearchResults = in.readOptionalBoolean();
-        }        
+        }
+        if (in.getVersion().onOrAfter(Version.V_6_7_0)) {
+            localClusterAlias = in.readOptionalString();
+        } else {
+            localClusterAlias = null;
+        }
     }
 
     @Override
@@ -495,7 +523,10 @@ public void writeTo(StreamOutput out) throws IOException {
         }
         if (out.getVersion().onOrAfter(Version.V_6_3_0)) {
             out.writeOptionalBoolean(allowPartialSearchResults);
-        }        
+        }
+        if (out.getVersion().onOrAfter(Version.V_6_7_0)) {
+            out.writeOptionalString(localClusterAlias);
+        }
     }
 
     @Override
@@ -519,14 +550,15 @@ public boolean equals(Object o) {
                 Objects.equals(maxConcurrentShardRequests, that.maxConcurrentShardRequests) &&
                 Objects.equals(preFilterShardSize, that.preFilterShardSize) &&
                 Objects.equals(indicesOptions, that.indicesOptions) &&
-                Objects.equals(allowPartialSearchResults, that.allowPartialSearchResults);
+                Objects.equals(allowPartialSearchResults, that.allowPartialSearchResults) &&
+                Objects.equals(localClusterAlias, that.localClusterAlias);
     }
 
     @Override
     public int hashCode() {
         return Objects.hash(searchType, Arrays.hashCode(indices), routing, preference, source, requestCache,
-                scroll, Arrays.hashCode(types), indicesOptions, batchedReduceSize, maxConcurrentShardRequests, preFilterShardSize, 
-                allowPartialSearchResults);
+                scroll, Arrays.hashCode(types), indicesOptions, batchedReduceSize, maxConcurrentShardRequests, preFilterShardSize,
+                allowPartialSearchResults, localClusterAlias);
     }
 
     @Override
@@ -544,6 +576,7 @@ public String toString() {
                 ", batchedReduceSize=" + batchedReduceSize +
                 ", preFilterShardSize=" + preFilterShardSize +
                 ", allowPartialSearchResults=" + allowPartialSearchResults +
+                ", localClusterAlias=" + localClusterAlias +
                 ", source=" + source + '}';
     }
 }

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

@@ -22,28 +22,34 @@
 import org.elasticsearch.action.OriginalIndices;
 import org.elasticsearch.cluster.routing.PlainShardIterator;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.search.SearchShardTarget;
 
 import java.util.List;
 
 /**
  * Extension of {@link PlainShardIterator} used in the search api, which also holds the {@link OriginalIndices}
- * of the search request. Useful especially with cross cluster search, as each cluster has its own set of original indices.
+ * of the search request (useful especially with cross-cluster search, as each cluster has its own set of original indices) as well as
+ * the cluster alias.
+ * @see OriginalIndices
  */
 public final class SearchShardIterator extends PlainShardIterator {
 
     private final OriginalIndices originalIndices;
-    private String clusterAlias;
+    private final String clusterAlias;
     private boolean skip = false;
 
     /**
      * Creates a {@link PlainShardIterator} instance that iterates over a subset of the given shards
      * this the a given <code>shardId</code>.
      *
+     * @param clusterAlias the alias of the cluster where the shard is located
      * @param shardId shard id of the group
      * @param shards  shards to iterate
+     * @param originalIndices the indices that the search request originally related to (before any rewriting happened)
      */
-    public SearchShardIterator(String clusterAlias, ShardId shardId, List<ShardRouting> shards, OriginalIndices originalIndices) {
+    public SearchShardIterator(@Nullable String clusterAlias, ShardId shardId, List<ShardRouting> shards, OriginalIndices originalIndices) {
         super(shardId, shards);
         this.originalIndices = originalIndices;
         this.clusterAlias = clusterAlias;
@@ -56,10 +62,22 @@ public OriginalIndices getOriginalIndices() {
         return originalIndices;
     }
 
+    /**
+     * Returns the alias of the cluster where the shard is located.
+     */
+    @Nullable
     public String getClusterAlias() {
         return clusterAlias;
     }
 
+    /**
+     * Creates a new shard target from this iterator, pointing at the node identified by the provided identifier.
+     * @see SearchShardTarget
+     */
+    SearchShardTarget newSearchShardTarget(String nodeId) {
+        return new SearchShardTarget(nodeId, shardId(), clusterAlias, originalIndices);
+    }
+
     /**
      * Reset the iterator and mark it as skippable
      * @see #skip()

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

@@ -98,8 +98,8 @@ public static ShardSearchFailure readShardSearchFailure(StreamInput in) throws I
 
     @Override
     public void readFrom(StreamInput in) throws IOException {
-        if (in.readBoolean()) {
-            shardTarget = new SearchShardTarget(in);
+        shardTarget = in.readOptionalWriteable(SearchShardTarget::new);
+        if (shardTarget != null) {
             index = shardTarget.getFullyQualifiedIndexName();
             shardId = shardTarget.getShardId().getId();
         }
@@ -110,12 +110,7 @@ public void readFrom(StreamInput in) throws IOException {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        if (shardTarget == null) {
-            out.writeBoolean(false);
-        } else {
-            out.writeBoolean(true);
-            shardTarget.writeTo(out);
-        }
+        out.writeOptionalWriteable(shardTarget);
         out.writeString(reason);
         RestStatus.writeTo(out, status);
         out.writeException(cause);
@@ -175,7 +170,7 @@ public static ShardSearchFailure fromXContent(XContentParser parser) throws IOEx
         SearchShardTarget searchShardTarget = null;
         if (nodeId != null) {
             searchShardTarget = new SearchShardTarget(nodeId,
-                    new ShardId(new Index(indexName, IndexMetaData.INDEX_UUID_NA_VALUE), shardId), clusterAlias, OriginalIndices.NONE);
+                new ShardId(new Index(indexName, IndexMetaData.INDEX_UUID_NA_VALUE), shardId), clusterAlias, OriginalIndices.NONE);
         }
         return new ShardSearchFailure(exception, searchShardTarget);
     }

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

@@ -35,6 +35,7 @@
 import org.elasticsearch.cluster.routing.GroupShardsIterator;
 import org.elasticsearch.cluster.routing.ShardIterator;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Setting.Property;
@@ -61,6 +62,7 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.LongSupplier;
 
 import static org.elasticsearch.action.search.SearchType.QUERY_THEN_FETCH;
@@ -311,7 +313,7 @@ private void executeSearch(SearchTask task, SearchTimeProvider timeProvider, Sea
         GroupShardsIterator<ShardIterator> localShardsIterator = clusterService.operationRouting().searchShards(clusterState,
                 concreteIndices, routingMap, searchRequest.preference(), searchService.getResponseCollectorService(), nodeSearchCounts);
         GroupShardsIterator<SearchShardIterator> shardIterators = mergeShardsIterators(localShardsIterator, localIndices,
-            remoteShardIterators);
+            searchRequest.getLocalClusterAlias(), remoteShardIterators);
 
         failIfOverShardCountLimit(clusterService, shardIterators.size());
 
@@ -338,13 +340,8 @@ private void executeSearch(SearchTask task, SearchTimeProvider timeProvider, Sea
         }
 
         final DiscoveryNodes nodes = clusterState.nodes();
-        BiFunction<String, String, Transport.Connection> connectionLookup = (clusterName, nodeId) -> {
-            final DiscoveryNode discoveryNode = clusterName == null ? nodes.get(nodeId) : remoteConnections.apply(clusterName, nodeId);
-            if (discoveryNode == null) {
-                throw new IllegalStateException("no node found for id: " + nodeId);
-            }
-            return searchTransportService.getConnection(clusterName, discoveryNode);
-        };
+        BiFunction<String, String, Transport.Connection> connectionLookup = buildConnectionLookup(searchRequest.getLocalClusterAlias(),
+            nodes::get, remoteConnections, searchTransportService::getConnection);
         if (searchRequest.isMaxConcurrentShardRequestsSet() == false) {
             // we try to set a default of max concurrent shard requests based on
             // the node count but upper-bound it by 256 by default to keep it sane. A single
@@ -359,7 +356,27 @@ private void executeSearch(SearchTask task, SearchTimeProvider timeProvider, Sea
             Collections.unmodifiableMap(aliasFilter), concreteIndexBoosts, routingMap, listener, preFilterSearchShards, clusters).start();
     }
 
-    private boolean shouldPreFilterSearchShards(SearchRequest searchRequest, GroupShardsIterator<SearchShardIterator> shardIterators) {
+    static BiFunction<String, String, Transport.Connection> buildConnectionLookup(String requestClusterAlias,
+                                                              Function<String, DiscoveryNode> localNodes,
+                                                              BiFunction<String, String, DiscoveryNode> remoteNodes,
+                                                              BiFunction<String, DiscoveryNode, Transport.Connection> nodeToConnection) {
+        return (clusterAlias, nodeId) -> {
+            final DiscoveryNode discoveryNode;
+            if (clusterAlias == null || requestClusterAlias != null) {
+                assert requestClusterAlias == null || requestClusterAlias.equals(clusterAlias);
+                discoveryNode = localNodes.apply(nodeId);
+            } else {
+                discoveryNode = remoteNodes.apply(clusterAlias, nodeId);
+            }
+            if (discoveryNode == null) {
+                throw new IllegalStateException("no node found for id: " + nodeId);
+            }
+            return nodeToConnection.apply(clusterAlias, discoveryNode);
+        };
+    }
+
+    private static boolean shouldPreFilterSearchShards(SearchRequest searchRequest,
+                                                       GroupShardsIterator<SearchShardIterator> shardIterators) {
         SearchSourceBuilder source = searchRequest.source();
         return searchRequest.searchType() == QUERY_THEN_FETCH && // we can't do this for DFS it needs to fan out to all shards all the time
                 SearchService.canRewriteToMatchNone(source) &&
@@ -368,10 +385,11 @@ private boolean shouldPreFilterSearchShards(SearchRequest searchRequest, GroupSh
 
     static GroupShardsIterator<SearchShardIterator> mergeShardsIterators(GroupShardsIterator<ShardIterator> localShardsIterator,
                                                              OriginalIndices localIndices,
+                                                             @Nullable String localClusterAlias,
                                                              List<SearchShardIterator> remoteShardIterators) {
         List<SearchShardIterator> shards = new ArrayList<>(remoteShardIterators);
         for (ShardIterator shardIterator : localShardsIterator) {
-            shards.add(new SearchShardIterator(null, shardIterator.shardId(), shardIterator.getShardRoutings(), localIndices));
+            shards.add(new SearchShardIterator(localClusterAlias, shardIterator.shardId(), shardIterator.getShardRoutings(), localIndices));
         }
         return new GroupShardsIterator<>(shards);
     }

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

@@ -44,8 +44,8 @@ static String buildScrollId(AtomicArray<? extends SearchPhaseResult> searchPhase
                 out.writeLong(searchPhaseResult.getRequestId());
                 SearchShardTarget searchShardTarget = searchPhaseResult.getSearchShardTarget();
                 if (searchShardTarget.getClusterAlias() != null) {
-                    out.writeString(RemoteClusterAware.buildRemoteIndexName(searchShardTarget.getClusterAlias(),
-                        searchShardTarget.getNodeId()));
+                    out.writeString(
+                        RemoteClusterAware.buildRemoteIndexName(searchShardTarget.getClusterAlias(), searchShardTarget.getNodeId()));
                 } else {
                     out.writeString(searchShardTarget.getNodeId());
                 }

server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java

@@ -159,7 +159,7 @@ final class DefaultSearchContext extends SearchContext {
     DefaultSearchContext(long id, ShardSearchRequest request, SearchShardTarget shardTarget,
                          Engine.Searcher engineSearcher, ClusterService clusterService, IndexService indexService,
                          IndexShard indexShard, BigArrays bigArrays, Counter timeEstimateCounter, TimeValue timeout,
-                         FetchPhase fetchPhase, String clusterAlias, Version minNodeVersion) {
+                         FetchPhase fetchPhase, Version minNodeVersion) {
         this.id = id;
         this.request = request;
         this.fetchPhase = fetchPhase;
@@ -179,7 +179,7 @@ final class DefaultSearchContext extends SearchContext {
         this.timeout = timeout;
         this.minNodeVersion = minNodeVersion;
         queryShardContext = indexService.newQueryShardContext(request.shardId().id(), searcher.getIndexReader(), request::nowInMillis,
-            clusterAlias);
+            shardTarget.getClusterAlias());
         queryShardContext.setTypes(request.types());
         queryBoost = request.indexBoost();
     }