Add IT for num_reduced_phases with batched query execution

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search/120_batch_reduce_size.yml

@@ -1,7 +1,4 @@
 setup:
-  - skip:
-      awaits_fix: "TODO fix this test, the response with batched execution is not deterministic enough for the available matchers"
-
   - do:
       indices.create:
           index: test_1

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

@@ -0,0 +1,107 @@
+/*
+ * 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.action.search;
+
+import org.elasticsearch.action.search.SearchQueryThenFetchAsyncAction.NodeQueryRequest;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.plugins.NetworkPlugin;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.test.ESIntegTestCase;
+import org.elasticsearch.transport.TransportInterceptor;
+import org.elasticsearch.transport.TransportRequest;
+import org.elasticsearch.transport.TransportRequestHandler;
+import org.junit.Before;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.Executor;
+
+import static org.elasticsearch.action.search.SearchQueryThenFetchAsyncAction.NODE_SEARCH_ACTION_NAME;
+import static org.elasticsearch.action.search.SearchType.QUERY_THEN_FETCH;
+import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
+import static org.hamcrest.Matchers.equalTo;
+
+public class BatchedQueryPhaseIT extends ESIntegTestCase {
+
+    // All the batched query requests that were made in each test
+    private static final List<NodeQueryRequest> batchedQueryRequests = new CopyOnWriteArrayList<>();
+
+    @Before
+    public void clear() {
+        batchedQueryRequests.clear();
+    }
+
+    public static class BatchedQueryCapturePlugin extends Plugin implements NetworkPlugin {
+        @Override
+        public List<TransportInterceptor> getTransportInterceptors(
+            NamedWriteableRegistry namedWriteableRegistry,
+            ThreadContext threadContext
+        ) {
+            return List.of(new TransportInterceptor() {
+                @Override
+                public <T extends TransportRequest> TransportRequestHandler<T> interceptHandler(
+                    String action,
+                    Executor executor,
+                    boolean forceExecution,
+                    TransportRequestHandler<T> actualHandler
+                ) {
+                    if (NODE_SEARCH_ACTION_NAME.equals(action)) {
+                        return (request, channel, task) -> {
+                            batchedQueryRequests.add((NodeQueryRequest) request);
+                            actualHandler.messageReceived(request, channel, task);
+                        };
+                    }
+                    return actualHandler;
+                }
+            });
+        }
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return List.of(BatchedQueryCapturePlugin.class);
+    }
+
+    /**
+     * num_reduce_phases tracks the number of times a partial reduction occurs on the coordinating node.
+     * This test must be aware of how batched queries are executed because reductions on the data nodes are
+     * not counted.
+     */
+    public void testNumReducePhases() {
+        assertAcked(prepareCreate("test-idx").setMapping("title", "type=keyword"));
+        for (int i = 0; i < 100; i++) {
+            prepareIndex("test-idx").setId(Integer.toString(i)).setSource("title", "testing" + i).get();
+        }
+        refresh();
+
+        assertNoFailuresAndResponse(
+            prepareSearch("test-idx").setBatchedReduceSize(2).addAggregation(terms("terms").field("title")).setSearchType(QUERY_THEN_FETCH),
+            response -> {
+                final int totalShards = response.getTotalShards();
+                final List<Integer> shardsPerDataNode = batchedQueryRequests.stream()
+                    .map(nodeQueryRequest -> nodeQueryRequest.shards().size())
+                    .toList();
+                final int coordNodeShards = totalShards - shardsPerDataNode.stream().mapToInt(Integer::intValue).sum();
+
+                // Because batched_reduce_size = 2, whenever two or more shard results exist on the coordinating node, they will be
+                // partially reduced. This reduction happens when either:
+                // - The search fans out directly to shards on the coordinating node
+                // - A batched result comes back
+                // Hence the formula: (# of shards on the coordinating node) + (# of batched requests) - 1
+                final int expectedNumReducePhases = Math.max(1, coordNodeShards + shardsPerDataNode.size() - 1);
+                assertThat(response.getNumReducePhases(), equalTo(expectedNumReducePhases));
+            }
+        );
+    }
+}

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

@@ -350,6 +350,10 @@ public String[] indices() {
         public IndicesOptions indicesOptions() {
             return searchRequest.indicesOptions();
         }
+
+        public List<ShardToQuery> shards() {
+            return shards;
+        }
     }
 
     private record ShardToQuery(float boost, String[] originalIndices, int shardIndex, ShardId shardId, ShardSearchContextId contextId)