Make can_match code a little easier to reuse (elastic#126588)

original-brownbear · web-flow · commit d39f4725773c · 2025-04-25T16:07:17.000+02:00
Step 1 to refactoring this with reuse in a per-datanode fashion
for batched execution.
Some obvious cleanup essentially making this a utility, removing
one weird indirection and reducing the use of the actual instance
of `CanMatchPreFilterSearchPhase` (this also results in a real
performance gain from moving work for sorting shards etc. off of
the transport_workers and closer to where its result is used).

This should by relatively trivial to review and allows for a simple follow
up that extracts the ability to run an individual round in isolation
as well as running the coordinator rewrite phase separately.
diff --git a/server/src/main/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhase.java b/server/src/main/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhase.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.util.FixedBitSet;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.CountDown;
@@ -76,7 +77,7 @@ final class CanMatchPreFilterSearchPhase {
     private int numPossibleMatches;
     private final CoordinatorRewriteContextProvider coordinatorRewriteContextProvider;
 
-    CanMatchPreFilterSearchPhase(
+    private CanMatchPreFilterSearchPhase(
         Logger logger,
         SearchTransportService searchTransportService,
         BiFunction<String, String, Transport.Connection> nodeIdToConnection,
@@ -123,6 +124,57 @@ final class CanMatchPreFilterSearchPhase {
         this.shardItIndexMap = shardItIndexMap;
     }
 
+    public static SubscribableListener<List<SearchShardIterator>> execute(
+        Logger logger,
+        SearchTransportService searchTransportService,
+        BiFunction<String, String, Transport.Connection> nodeIdToConnection,
+        Map<String, AliasFilter> aliasFilter,
+        Map<String, Float> concreteIndexBoosts,
+        Executor executor,
+        SearchRequest request,
+        List<SearchShardIterator> shardsIts,
+        TransportSearchAction.SearchTimeProvider timeProvider,
+        SearchTask task,
+        boolean requireAtLeastOneMatch,
+        CoordinatorRewriteContextProvider coordinatorRewriteContextProvider
+    ) {
+        if (shardsIts.isEmpty()) {
+            return SubscribableListener.newSucceeded(List.of());
+        }
+        final SubscribableListener<List<SearchShardIterator>> listener = new SubscribableListener<>();
+        // Note that the search is failed when this task is rejected by the executor
+        executor.execute(new AbstractRunnable() {
+            @Override
+            public void onFailure(Exception e) {
+                if (logger.isDebugEnabled()) {
+                    logger.debug(() -> format("Failed to execute [%s] while running [can_match] phase", request), e);
+                }
+                listener.onFailure(new SearchPhaseExecutionException("can_match", "start", e, ShardSearchFailure.EMPTY_ARRAY));
+            }
+
+            @Override
+            protected void doRun() {
+                assert assertSearchCoordinationThread();
+                new CanMatchPreFilterSearchPhase(
+                    logger,
+                    searchTransportService,
+                    nodeIdToConnection,
+                    aliasFilter,
+                    concreteIndexBoosts,
+                    executor,
+                    request,
+                    shardsIts,
+                    timeProvider,
+                    task,
+                    requireAtLeastOneMatch,
+                    coordinatorRewriteContextProvider,
+                    listener
+                ).runCoordinatorRewritePhase();
+            }
+        });
+        return listener;
+    }
+
     private static boolean assertSearchCoordinationThread() {
         return ThreadPool.assertCurrentThreadPool(ThreadPool.Names.SEARCH_COORDINATION);
     }
@@ -165,7 +217,7 @@ private void runCoordinatorRewritePhase() {
             }
         }
         if (matchedShardLevelRequests.isEmpty()) {
-            finishPhase();
+            listener.onResponse(getIterator(shardsIts));
         } else {
             // verify missing shards only for the shards that we hit for the query
             checkNoMissingShards(matchedShardLevelRequests);
@@ -176,18 +228,14 @@ private void runCoordinatorRewritePhase() {
     private void consumeResult(boolean canMatch, ShardSearchRequest request) {
         CanMatchShardResponse result = new CanMatchShardResponse(canMatch, null);
         result.setShardIndex(request.shardRequestIndex());
-        consumeResult(result, () -> {});
+        consumeResult(result);
     }
 
-    private void consumeResult(CanMatchShardResponse result, Runnable next) {
-        try {
-            final boolean canMatch = result.canMatch();
-            final MinAndMax<?> minAndMax = result.estimatedMinAndMax();
-            if (canMatch || minAndMax != null) {
-                consumeResult(result.getShardIndex(), canMatch, minAndMax);
-            }
-        } finally {
-            next.run();
+    private void consumeResult(CanMatchShardResponse result) {
+        final boolean canMatch = result.canMatch();
+        final MinAndMax<?> minAndMax = result.estimatedMinAndMax();
+        if (canMatch || minAndMax != null) {
+            consumeResult(result.getShardIndex(), canMatch, minAndMax);
         }
     }
 
@@ -226,7 +274,7 @@ private Map<SendingTarget, List<SearchShardIterator>> groupByNode(List<SearchSha
      * If there are failures during a round, there will be a follow-up round
      * to retry on other available shard copies.
      */
-    class Round extends AbstractRunnable {
+    private class Round extends AbstractRunnable {
         private final List<SearchShardIterator> shards;
         private final CountDown countDown;
         private final AtomicReferenceArray<Exception> failedResponses;
@@ -296,11 +344,10 @@ public void onFailure(Exception e) {
 
         private void onOperation(int idx, CanMatchShardResponse response) {
             failedResponses.set(idx, null);
-            consumeResult(response, () -> {
-                if (countDown.countDown()) {
-                    finishRound();
-                }
-            });
+            consumeResult(response);
+            if (countDown.countDown()) {
+                finishRound();
+            }
         }
 
         private void onOperationFailed(int idx, Exception e) {
@@ -322,7 +369,7 @@ private void finishRound() {
                 }
             }
             if (remainingShards.isEmpty()) {
-                finishPhase();
+                listener.onResponse(getIterator(shardsIts));
             } else {
                 // trigger another round, forcing execution
                 executor.execute(new Round(remainingShards) {
@@ -339,7 +386,7 @@ public void onFailure(Exception e) {
             if (logger.isDebugEnabled()) {
                 logger.debug(() -> format("Failed to execute [%s] while running [can_match] phase", request), e);
             }
-            onPhaseFailure("round", e);
+            listener.onFailure(new SearchPhaseExecutionException("can_match", "round", e, ShardSearchFailure.EMPTY_ARRAY));
         }
     }
 
@@ -363,13 +410,9 @@ private CanMatchNodeRequest createCanMatchRequest(Map.Entry<SendingTarget, List<
         );
     }
 
-    private void finishPhase() {
-        listener.onResponse(getIterator(shardsIts));
-    }
-
     private static final float DEFAULT_INDEX_BOOST = 1.0f;
 
-    public CanMatchNodeRequest.Shard buildShardLevelRequest(SearchShardIterator shardIt) {
+    private CanMatchNodeRequest.Shard buildShardLevelRequest(SearchShardIterator shardIt) {
         AliasFilter filter = aliasFilter.get(shardIt.shardId().getIndex().getUUID());
         assert filter != null;
         float indexBoost = concreteIndexBoosts.getOrDefault(shardIt.shardId().getIndex().getUUID(), DEFAULT_INDEX_BOOST);
@@ -386,33 +429,6 @@ public CanMatchNodeRequest.Shard buildShardLevelRequest(SearchShardIterator shar
         );
     }
 
-    public void start() {
-        if (shardsIts.isEmpty()) {
-            finishPhase();
-            return;
-        }
-        // Note that the search is failed when this task is rejected by the executor
-        executor.execute(new AbstractRunnable() {
-            @Override
-            public void onFailure(Exception e) {
-                if (logger.isDebugEnabled()) {
-                    logger.debug(() -> format("Failed to execute [%s] while running [can_match] phase", request), e);
-                }
-                onPhaseFailure("start", e);
-            }
-
-            @Override
-            protected void doRun() {
-                assert assertSearchCoordinationThread();
-                runCoordinatorRewritePhase();
-            }
-        });
-    }
-
-    private void onPhaseFailure(String msg, Exception cause) {
-        listener.onFailure(new SearchPhaseExecutionException("can_match", msg, cause, ShardSearchFailure.EMPTY_ARRAY));
-    }
-
     private synchronized List<SearchShardIterator> getIterator(List<SearchShardIterator> shardsIts) {
         // TODO: pick the local shard when possible
         if (requireAtLeastOneMatch && numPossibleMatches == 0) {
diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportOpenPointInTimeAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportOpenPointInTimeAction.java
@@ -161,7 +161,7 @@ public void runNewSearchPhase(
             // that is signaled to the local can match through the SearchShardIterator#prefiltered flag. Local shards do need to go
             // through the local can match phase.
             if (SearchService.canRewriteToMatchNone(searchRequest.source())) {
-                new CanMatchPreFilterSearchPhase(
+                CanMatchPreFilterSearchPhase.execute(
                     logger,
                     searchTransportService,
                     connectionLookup,
@@ -173,22 +173,24 @@ public void runNewSearchPhase(
                     timeProvider,
                     task,
                     false,
-                    searchService.getCoordinatorRewriteContextProvider(timeProvider::absoluteStartMillis),
-                    listener.delegateFailureAndWrap(
-                        (searchResponseActionListener, searchShardIterators) -> runOpenPointInTimePhase(
-                            task,
-                            searchRequest,
-                            executor,
-                            searchShardIterators,
-                            timeProvider,
-                            connectionLookup,
-                            clusterState,
-                            aliasFilter,
-                            concreteIndexBoosts,
-                            clusters
+                    searchService.getCoordinatorRewriteContextProvider(timeProvider::absoluteStartMillis)
+                )
+                    .addListener(
+                        listener.delegateFailureAndWrap(
+                            (searchResponseActionListener, searchShardIterators) -> runOpenPointInTimePhase(
+                                task,
+                                searchRequest,
+                                executor,
+                                searchShardIterators,
+                                timeProvider,
+                                connectionLookup,
+                                clusterState,
+                                aliasFilter,
+                                concreteIndexBoosts,
+                                clusters
+                            )
                         )
-                    )
-                ).start();
+                    );
             } else {
                 runOpenPointInTimePhase(
                     task,
diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java
@@ -1484,7 +1484,7 @@ public void runNewSearchPhase(
             if (preFilter) {
                 // only for aggs we need to contact shards even if there are no matches
                 boolean requireAtLeastOneMatch = searchRequest.source() != null && searchRequest.source().aggregations() != null;
-                new CanMatchPreFilterSearchPhase(
+                CanMatchPreFilterSearchPhase.execute(
                     logger,
                     searchTransportService,
                     connectionLookup,
@@ -1496,24 +1496,26 @@ public void runNewSearchPhase(
                     timeProvider,
                     task,
                     requireAtLeastOneMatch,
-                    searchService.getCoordinatorRewriteContextProvider(timeProvider::absoluteStartMillis),
-                    listener.delegateFailureAndWrap((l, iters) -> {
-                        runNewSearchPhase(
-                            task,
-                            searchRequest,
-                            executor,
-                            iters,
-                            timeProvider,
-                            connectionLookup,
-                            clusterState,
-                            aliasFilter,
-                            concreteIndexBoosts,
-                            false,
-                            threadPool,
-                            clusters
-                        );
-                    })
-                ).start();
+                    searchService.getCoordinatorRewriteContextProvider(timeProvider::absoluteStartMillis)
+                )
+                    .addListener(
+                        listener.delegateFailureAndWrap(
+                            (l, iters) -> runNewSearchPhase(
+                                task,
+                                searchRequest,
+                                executor,
+                                iters,
+                                timeProvider,
+                                connectionLookup,
+                                clusterState,
+                                aliasFilter,
+                                concreteIndexBoosts,
+                                false,
+                                threadPool,
+                                clusters
+                            )
+                        )
+                    );
                 return;
             }
             // for synchronous CCS minimize_roundtrips=false, use the CCSSingleCoordinatorSearchProgressListener
diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportSearchShardsAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchShardsAction.java
@@ -153,7 +153,7 @@ public void searchShards(Task task, SearchShardsRequest searchShardsRequest, Act
                         new SearchShardsResponse(toGroups(shardIts), project.cluster().nodes().getAllNodes(), aliasFilters)
                     );
                 } else {
-                    new CanMatchPreFilterSearchPhase(logger, searchTransportService, (clusterAlias, node) -> {
+                    CanMatchPreFilterSearchPhase.execute(logger, searchTransportService, (clusterAlias, node) -> {
                         assert Objects.equals(clusterAlias, searchShardsRequest.clusterAlias());
                         return transportService.getConnection(project.cluster().nodes().get(node));
                     },
@@ -165,9 +165,13 @@ public void searchShards(Task task, SearchShardsRequest searchShardsRequest, Act
                         timeProvider,
                         (SearchTask) task,
                         false,
-                        searchService.getCoordinatorRewriteContextProvider(timeProvider::absoluteStartMillis),
-                        delegate.map(its -> new SearchShardsResponse(toGroups(its), project.cluster().nodes().getAllNodes(), aliasFilters))
-                    ).start();
+                        searchService.getCoordinatorRewriteContextProvider(timeProvider::absoluteStartMillis)
+                    )
+                        .addListener(
+                            delegate.map(
+                                its -> new SearchShardsResponse(toGroups(its), project.cluster().nodes().getAllNodes(), aliasFilters)
+                            )
+                        );
                 }
             })
         );
diff --git a/server/src/test/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhaseTests.java b/server/src/test/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhaseTests.java
