Move duplicate connection lookup logic to AbstractSearchAsyncAction (#117055)

We found this duplication today when working on batching query phase requests.
For batching it would be nice to have the connection already available at a higher
level in the AbstractSearchAsyncAction and this is a worthwhile cleanup in general,
given how many issues we had around connection lookup recently.

Author: original-brownbear

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

@@ -299,7 +299,7 @@ protected void performPhaseOnShard(final int shardIndex, final SearchShardIterat
     }
 
     private void doPerformPhaseOnShard(int shardIndex, SearchShardIterator shardIt, SearchShardTarget shard, Releasable releasable) {
-        executePhaseOnShard(shardIt, shard, new SearchActionListener<>(shard, shardIndex) {
+        var shardListener = new SearchActionListener<Result>(shard, shardIndex) {
             @Override
             public void innerOnResponse(Result result) {
                 try {
@@ -315,7 +315,15 @@ public void onFailure(Exception e) {
                 releasable.close();
                 onShardFailure(shardIndex, shard, shardIt, e);
             }
-        });
+        };
+        final Transport.Connection connection;
+        try {
+            connection = getConnection(shard.getClusterAlias(), shard.getNodeId());
+        } catch (Exception e) {
+            shardListener.onFailure(e);
+            return;
+        }
+        executePhaseOnShard(shardIt, connection, shardListener);
     }
 
     private void failOnUnavailable(int shardIndex, SearchShardIterator shardIt) {
@@ -327,12 +335,12 @@ private void failOnUnavailable(int shardIndex, SearchShardIterator shardIt) {
     /**
      * Sends the request to the actual shard.
      * @param shardIt the shards iterator
-     * @param shard the shard routing to send the request for
+     * @param connection to node that the shard is located on
      * @param listener the listener to notify on response
      */
     protected abstract void executePhaseOnShard(
         SearchShardIterator shardIt,
-        SearchShardTarget shard,
+        Transport.Connection connection,
         SearchActionListener<Result> listener
     );
 

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

@@ -84,16 +84,9 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
     @Override
     protected void executePhaseOnShard(
         final SearchShardIterator shardIt,
-        final SearchShardTarget shard,
+        final Transport.Connection connection,
         final SearchActionListener<DfsSearchResult> listener
     ) {
-        final Transport.Connection connection;
-        try {
-            connection = getConnection(shard.getClusterAlias(), shard.getNodeId());
-        } catch (Exception e) {
-            listener.onFailure(e);
-            return;
-        }
         getSearchTransport().sendExecuteDfs(connection, buildShardSearchRequest(shardIt, listener.requestIndex), getTask(), listener);
     }
 

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

@@ -79,7 +79,7 @@ protected static void doCheckNoMissingShards(
     /**
      * Releases shard targets that are not used in the docsIdsToLoad.
      */
-    protected void releaseIrrelevantSearchContext(SearchPhaseResult searchPhaseResult, AbstractSearchAsyncAction<?> context) {
+    protected static void releaseIrrelevantSearchContext(SearchPhaseResult searchPhaseResult, AbstractSearchAsyncAction<?> context) {
         // we only release search context that we did not fetch from, if we are not scrolling
         // or using a PIT and if it has at least one hit that didn't make it to the global topDocs
         if (searchPhaseResult == null) {

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

@@ -91,16 +91,9 @@ class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<SearchPh
 
     protected void executePhaseOnShard(
         final SearchShardIterator shardIt,
-        final SearchShardTarget shard,
+        final Transport.Connection connection,
         final SearchActionListener<SearchPhaseResult> listener
     ) {
-        final Transport.Connection connection;
-        try {
-            connection = getConnection(shard.getClusterAlias(), shard.getNodeId());
-        } catch (Exception e) {
-            listener.onFailure(e);
-            return;
-        }
         ShardSearchRequest request = rewriteShardSearchRequest(super.buildShardSearchRequest(shardIt, listener.requestIndex));
         getSearchTransport().sendExecuteQuery(connection, request, getTask(), listener);
     }

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

@@ -35,7 +35,6 @@
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchService;
-import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.search.internal.ShardSearchContextId;
@@ -252,16 +251,9 @@ protected String missingShardsErrorMessage(StringBuilder missingShards) {
                 @Override
                 protected void executePhaseOnShard(
                     SearchShardIterator shardIt,
-                    SearchShardTarget shard,
+                    Transport.Connection connection,
                     SearchActionListener<SearchPhaseResult> phaseListener
                 ) {
-                    final Transport.Connection connection;
-                    try {
-                        connection = connectionLookup.apply(shardIt.getClusterAlias(), shard.getNodeId());
-                    } catch (Exception e) {
-                        phaseListener.onFailure(e);
-                        return;
-                    }
                     transportService.sendChildRequest(
                         connection,
                         OPEN_SHARD_READER_CONTEXT_NAME,

server/src/test/java/org/elasticsearch/action/search/AbstractSearchAsyncActionTests.java

@@ -101,7 +101,7 @@ protected SearchPhase getNextPhase() {
             @Override
             protected void executePhaseOnShard(
                 final SearchShardIterator shardIt,
-                final SearchShardTarget shard,
+                final Transport.Connection shard,
                 final SearchActionListener<SearchPhaseResult> listener
             ) {}
 

server/src/test/java/org/elasticsearch/action/search/MockSearchPhaseContext.java

@@ -147,7 +147,7 @@ public void executeNextPhase(SearchPhase currentPhase, Supplier<SearchPhase> nex
     @Override
     protected void executePhaseOnShard(
         SearchShardIterator shardIt,
-        SearchShardTarget shard,
+        Transport.Connection shard,
         SearchActionListener<SearchPhaseResult> listener
     ) {
         onShardResult(new SearchPhaseResult() {

server/src/test/java/org/elasticsearch/action/search/SearchAsyncActionTests.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.search.SearchHits;
 import org.elasticsearch.search.SearchPhaseResult;
-import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.search.internal.ShardSearchContextId;
 import org.elasticsearch.test.ESTestCase;
@@ -119,16 +118,15 @@ public void testSkipSearchShards() throws InterruptedException {
             @Override
             protected void executePhaseOnShard(
                 SearchShardIterator shardIt,
-                SearchShardTarget shard,
+                Transport.Connection connection,
                 SearchActionListener<TestSearchPhaseResult> listener
             ) {
-                seenShard.computeIfAbsent(shard.getShardId(), (i) -> {
+                seenShard.computeIfAbsent(shardIt.shardId(), (i) -> {
                     numRequests.incrementAndGet(); // only count this once per replica
                     return Boolean.TRUE;
                 });
 
                 new Thread(() -> {
-                    Transport.Connection connection = getConnection(null, shard.getNodeId());
                     TestSearchPhaseResult testSearchPhaseResult = new TestSearchPhaseResult(
                         new ShardSearchContextId(UUIDs.randomBase64UUID(), contextIdGenerator.incrementAndGet()),
                         connection.getNode()
@@ -227,23 +225,22 @@ public void testLimitConcurrentShardRequests() throws InterruptedException {
                 @Override
                 protected void executePhaseOnShard(
                     SearchShardIterator shardIt,
-                    SearchShardTarget shard,
+                    Transport.Connection connection,
                     SearchActionListener<TestSearchPhaseResult> listener
                 ) {
-                    seenShard.computeIfAbsent(shard.getShardId(), (i) -> {
+                    seenShard.computeIfAbsent(shardIt.shardId(), (i) -> {
                         numRequests.incrementAndGet(); // only count this once per shard copy
                         return Boolean.TRUE;
                     });
 
                     new Thread(() -> {
                         safeAwait(awaitInitialRequests);
-                        Transport.Connection connection = getConnection(null, shard.getNodeId());
                         TestSearchPhaseResult testSearchPhaseResult = new TestSearchPhaseResult(
                             new ShardSearchContextId(UUIDs.randomBase64UUID(), contextIdGenerator.incrementAndGet()),
                             connection.getNode()
                         );
                         try {
-                            if (shardFailures[shard.getShardId().id()]) {
+                            if (shardFailures[shardIt.shardId().id()]) {
                                 listener.onFailure(new RuntimeException());
                             } else {
                                 listener.onResponse(testSearchPhaseResult);
@@ -340,11 +337,11 @@ public void sendFreeContext(Transport.Connection connection, ShardSearchContextI
                 @Override
                 protected void executePhaseOnShard(
                     SearchShardIterator shardIt,
-                    SearchShardTarget shard,
+                    Transport.Connection connection,
                     SearchActionListener<TestSearchPhaseResult> listener
                 ) {
-                    assertTrue("shard: " + shard.getShardId() + " has been queried twice", testResponse.queried.add(shard.getShardId()));
-                    Transport.Connection connection = getConnection(null, shard.getNodeId());
+                    var shardId = shardIt.shardId();
+                    assertTrue("shard: " + shardId + " has been queried twice", testResponse.queried.add(shardId));
                     TestSearchPhaseResult testSearchPhaseResult = new TestSearchPhaseResult(
                         new ShardSearchContextId(UUIDs.randomBase64UUID(), contextIdGenerator.incrementAndGet()),
                         connection.getNode()
@@ -464,13 +461,13 @@ public void sendFreeContext(Transport.Connection connection, ShardSearchContextI
                 @Override
                 protected void executePhaseOnShard(
                     SearchShardIterator shardIt,
-                    SearchShardTarget shard,
+                    Transport.Connection connection,
                     SearchActionListener<TestSearchPhaseResult> listener
                 ) {
-                    assertTrue("shard: " + shard.getShardId() + " has been queried twice", response.queried.add(shard.getShardId()));
-                    Transport.Connection connection = getConnection(null, shard.getNodeId());
+                    var shardId = shardIt.shardId();
+                    assertTrue("shard: " + shardId + " has been queried twice", response.queried.add(shardId));
                     final TestSearchPhaseResult testSearchPhaseResult;
-                    if (shard.getShardId().id() == 0) {
+                    if (shardId.id() == 0) {
                         testSearchPhaseResult = new TestSearchPhaseResult(null, connection.getNode());
                     } else {
                         testSearchPhaseResult = new TestSearchPhaseResult(
@@ -573,15 +570,14 @@ public void testAllowPartialResults() throws InterruptedException {
                 @Override
                 protected void executePhaseOnShard(
                     SearchShardIterator shardIt,
-                    SearchShardTarget shard,
+                    Transport.Connection connection,
                     SearchActionListener<TestSearchPhaseResult> listener
                 ) {
-                    seenShard.computeIfAbsent(shard.getShardId(), (i) -> {
+                    seenShard.computeIfAbsent(shardIt.shardId(), (i) -> {
                         numRequests.incrementAndGet(); // only count this once per shard copy
                         return Boolean.TRUE;
                     });
                     new Thread(() -> {
-                        Transport.Connection connection = getConnection(null, shard.getNodeId());
                         TestSearchPhaseResult testSearchPhaseResult = new TestSearchPhaseResult(
                             new ShardSearchContextId(UUIDs.randomBase64UUID(), contextIdGenerator.incrementAndGet()),
                             connection.getNode()
@@ -673,7 +669,7 @@ public void testSkipUnavailableSearchShards() throws InterruptedException {
             @Override
             protected void executePhaseOnShard(
                 SearchShardIterator shardIt,
-                SearchShardTarget shard,
+                Transport.Connection connection,
                 SearchActionListener<TestSearchPhaseResult> listener
             ) {
                 assert false : "Expected to skip all shards";