Record stats for "adaptive replica selection" in batched queries (#144563)

Batched queries currently do not collect any of the stats needed for the
ARS system to function, which leads to hot spotting when many queries
are batched.

This records the S (service time) and q (queue size) values from the
last QuerySearchResult in the batched response. Shard queries execute in
parallel but are invoked in order, so the last result gives the most
recent queue size picture. The service time may vary across shards so
the last result is simply an arbitrary sample.

Note: we still collect R (total response time) as a side effect of using
SearchExecutionStatsCollector, but it is not particularly useful for
batched requests in the ARS formula, which only uses (R - S) to estimate
network overhead and queue time.

Author: benchaplin

docs/changelog/144563.yaml

@@ -0,0 +1,6 @@
+area: Search
+issues:
+ - 144146
+pr: 144563
+summary: Record stats for "adaptive replica selection" in batched queries
+type: bug

server/src/internalClusterTest/java/org/elasticsearch/action/search/BatchedQueryPhaseIT.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.node.ResponseCollectorService;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.test.ESIntegTestCase;
 
@@ -28,6 +29,8 @@
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.notNullValue;
 
 public class BatchedQueryPhaseIT extends ESIntegTestCase {
 
@@ -73,6 +76,69 @@ public void testNumReducePhases() {
         );
     }
 
+    public void testAdaptiveReplicaSelectionStatsWithBatchedQueryPhase() {
+        assumeTrue(
+            "test skipped because batched query execution disabled by feature flag",
+            SearchService.BATCHED_QUERY_PHASE_FEATURE_FLAG.isEnabled()
+        );
+        internalCluster().ensureAtLeastNumDataNodes(3);
+
+        String indexName = "test-ars-stats";
+        int numDataNodes = internalCluster().numDataNodes();
+        int numShards = numDataNodes * 4;
+        assertAcked(
+            prepareCreate(indexName).setMapping("title", "type=keyword")
+                .setSettings(
+                    Settings.builder()
+                        .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, numShards)
+                        .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+                )
+        );
+        for (int i = 0; i < 100; i++) {
+            prepareIndex(indexName).setId(Integer.toString(i)).setSource("title", "testing" + i).get();
+        }
+        refresh();
+        ensureGreen(indexName);
+
+        String coordinatorNode = internalCluster().getRandomNodeName();
+        String coordinatorNodeId = getNodeId(coordinatorNode);
+
+        for (int i = 0; i < 20; i++) {
+            assertNoFailuresAndResponse(client(coordinatorNode).prepareSearch(indexName).setSearchType(QUERY_THEN_FETCH), response -> {});
+        }
+
+        ResponseCollectorService responseCollectorService = internalCluster().getInstance(ResponseCollectorService.class, coordinatorNode);
+        Map<String, ResponseCollectorService.ComputedNodeStats> allStats = responseCollectorService.getAllNodeStatistics();
+        Map<String, Integer> shardsPerNode = getNodeToShardCountMap(indexName);
+
+        long batchedRemoteNodes = shardsPerNode.entrySet()
+            .stream()
+            .filter(e -> e.getKey().equals(coordinatorNodeId) == false && e.getValue() >= 2)
+            .count();
+        assertThat(
+            "Expected at least one remote node with 2+ shards to exercise the batched query path",
+            batchedRemoteNodes,
+            greaterThanOrEqualTo(1L)
+        );
+
+        for (Map.Entry<String, Integer> entry : shardsPerNode.entrySet()) {
+            String nodeId = entry.getKey();
+            if (nodeId.equals(coordinatorNodeId)) {
+                continue;
+            }
+            assertThat(
+                "Expected adaptive replica selection stats for remote node ["
+                    + nodeId
+                    + "] with ["
+                    + entry.getValue()
+                    + "] shards, but found none. Stats were only collected for nodes: "
+                    + allStats.keySet(),
+                allStats.get(nodeId),
+                notNullValue()
+            );
+        }
+    }
+
     private Map<String, Integer> getNodeToShardCountMap(String indexName) {
         ClusterState clusterState = clusterAdmin().prepareState(TEST_REQUEST_TIMEOUT).get().getState();
         IndexRoutingTable indexRoutingTable = clusterState.routingTable(ProjectId.DEFAULT).index(indexName);

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

@@ -534,6 +534,7 @@ protected void doRun(Map<SearchShardIterator, Integer> shardIndexMap) {
                 executeWithoutBatching(routing, request);
                 return;
             }
+            final ActionListener<? super SearchPhaseResult> statsCollector = searchTransportService.newStatsCollector(connection);
             searchTransportService.transportService()
                 .sendChildRequest(connection, NODE_SEARCH_ACTION_NAME, request, task, new TransportResponseHandler<NodeQueryResponse>() {
                     @Override
@@ -556,13 +557,15 @@ public void handleResponse(NodeQueryResponse response) {
                                 queryPhaseResultConsumer.addBatchedPartialResult(response.topDocsStats, response.mergeResult);
                             }
                         }
+                        SearchPhaseResult lastResult = null;
                         for (int i = 0; i < response.results.length; i++) {
                             var s = request.shards.get(i);
                             int shardIdx = s.shardIndex;
                             final SearchShardTarget target = new SearchShardTarget(routing.nodeId(), s.shardId, routing.clusterAlias());
                             switch (response.results[i]) {
                                 case Exception e -> onShardFailure(shardIdx, target, shardIterators[shardIdx], e);
                                 case SearchPhaseResult q -> {
+                                    lastResult = q;
                                     q.setShardIndex(shardIdx);
                                     q.setSearchShardTarget(target);
                                     onShardResult(q);
@@ -572,6 +575,12 @@ public void handleResponse(NodeQueryResponse response) {
                                 }
                             }
                         }
+                        if (statsCollector != null && lastResult != null) {
+                            // The last result is the most likely to give us the most recent picture of queue size, as shards queries are
+                            // invoked in the same order as the response's results array.
+                            // The service time may vary across shards so the last result is simply an arbitrary choice.
+                            statsCollector.onResponse(lastResult);
+                        }
                     }
 
                     @Override

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

@@ -211,6 +211,13 @@ public void sendExecuteQuery(
         );
     }
 
+    ActionListener<? super SearchPhaseResult> newStatsCollector(Transport.Connection connection) {
+        if (responseWrapper == null) {
+            return null;
+        }
+        return responseWrapper.apply(connection, ActionListener.noop());
+    }
+
     public void sendExecuteQuery(
         Transport.Connection connection,
         final QuerySearchRequest request,