Split early terminating collector into two separate collectors (elastic#96901)

We use `EarlyTerminatingCollector` for two similar, yet different purposes:
- terminate_after, for which we wrap an empty collector and we forcibly terminate collection for all collectors
- early termination of total hit count collector based on track_total_hits, which wraps `TotalHitCountCollector` and does not affect other collectors
    
Given that the two are quite separate scenarios, besides the commonality that they both do early termination of some kind, it is convenient to split the two into separate collectors which are specialized for their purposes.
    
The proposal is to have the following two collectors:
1) `TerminateAfterCollector` which forcibly terminates the collector by throwing an exception that is unknown by Lucene. It does not need to wrap another collector
2) `PartialHitCountCollector` that is an extension of `TotalHitCountCollector` that is able to early terminate total hits   tracking. This simplifies things as we no longer need to wrap and early terminate a generic collector.

Author: javanna

server/src/main/java/org/elasticsearch/search/query/EarlyTerminatingCollector.java

@@ -0,0 +1,68 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.query;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.CollectionTerminatedException;
+import org.apache.lucene.search.FilterLeafCollector;
+import org.apache.lucene.search.LeafCollector;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.TotalHitCountCollector;
+
+import java.io.IOException;
+
+/**
+ * Extension of {@link TotalHitCountCollector} that supports early termination of total hits counting based on a provided threshold.
+ * Note that the total hit count may be retrieved from {@link org.apache.lucene.search.Weight#count(LeafReaderContext)},
+ * in which case early termination is only applied to the leaves that do collect documents.
+ */
+class PartialHitCountCollector extends TotalHitCountCollector {
+
+    private final int totalHitsThreshold;
+    // we could reuse the counter that TotalHitCountCollector has and exposes through getTotalHits(),
+    // but that would make us early terminate also when retrieving count from Weight#count and would
+    // cause a behaviour that's difficult to explain and test.
+    private int numCollected = 0;
+    private boolean earlyTerminated;
+
+    PartialHitCountCollector(int totalHitsThreshold) {
+        this.totalHitsThreshold = totalHitsThreshold;
+    }
+
+    @Override
+    public ScoreMode scoreMode() {
+        // Does not need scores like TotalHitCountCollector (COMPLETE_NO_SCORES), but not exhaustive as it early terminates.
+        return ScoreMode.TOP_DOCS;
+    }
+
+    @Override
+    public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+        if (numCollected >= totalHitsThreshold) {
+            earlyTerminateIfNeeded();
+        }
+        return new FilterLeafCollector(super.getLeafCollector(context)) {
+            @Override
+            public void collect(int doc) throws IOException {
+                if (++numCollected > totalHitsThreshold) {
+                    earlyTerminateIfNeeded();
+                }
+                super.collect(doc);
+            }
+        };
+    }
+
+    private void earlyTerminateIfNeeded() {
+        earlyTerminated = true;
+        throw new CollectionTerminatedException();
+    }
+
+    boolean hasEarlyTerminated() {
+        return earlyTerminated;
+    }
+}

server/src/main/java/org/elasticsearch/search/query/PartialHitCountCollector.java

@@ -208,15 +208,11 @@ static void addCollectorsAndSearch(SearchContext searchContext) throws QueryPhas
             if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) {
                 // add terminate_after before the filter collectors
                 // it will only be applied on documents accepted by these filter collectors
-                EarlyTerminatingCollector earlyTerminatingCollector = new EarlyTerminatingCollector(
-                    EMPTY_COLLECTOR,
-                    searchContext.terminateAfter(),
-                    true
-                );
+                TerminateAfterCollector terminateAfterCollector = new TerminateAfterCollector(searchContext.terminateAfter());
                 final Collector collector = collectorManager.newCollector();
                 collectorManager = wrapWithProfilerCollectorManagerIfNeeded(
                     searchContext.getProfilers(),
-                    new SingleThreadCollectorManager(MultiCollector.wrap(earlyTerminatingCollector, collector)),
+                    new SingleThreadCollectorManager(MultiCollector.wrap(terminateAfterCollector, collector)),
                     REASON_SEARCH_TERMINATE_AFTER_COUNT,
                     collector
                 );
@@ -319,7 +315,7 @@ private static void searchWithCollectorManager(
         QuerySearchResult queryResult = searchContext.queryResult();
         try {
             searcher.search(query, collectorManager);
-        } catch (EarlyTerminatingCollector.EarlyTerminationException e) {
+        } catch (TerminateAfterCollector.EarlyTerminationException e) {
             queryResult.terminatedEarly(true);
         } catch (TimeExceededException e) {
             assert timeoutSet : "TimeExceededException thrown even though timeout wasn't set";

server/src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -0,0 +1,74 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.query;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.LeafCollector;
+import org.apache.lucene.search.Scorable;
+import org.apache.lucene.search.ScoreMode;
+
+import java.io.IOException;
+
+/**
+ * A {@link Collector} that forcibly early terminates collection after a certain number of hits have been collected.
+ * Terminates the collection across all collectors by throwing an {@link EarlyTerminationException} once the threshold is reached.
+ */
+class TerminateAfterCollector implements Collector {
+    static final class EarlyTerminationException extends RuntimeException {
+        private EarlyTerminationException(String msg) {
+            super(msg);
+        }
+
+        @Override
+        public Throwable fillInStackTrace() {
+            // never re-thrown so we can save the expensive stacktrace
+            return this;
+        }
+    }
+
+    private final int maxCountHits;
+    private int numCollected;
+
+    /**
+     *
+     * @param maxCountHits the number of hits to collect, after which the collection must be early terminated
+     */
+    TerminateAfterCollector(int maxCountHits) {
+        this.maxCountHits = maxCountHits;
+    }
+
+    @Override
+    public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+        if (numCollected >= maxCountHits) {
+            earlyTerminate();
+        }
+        return new LeafCollector() {
+            @Override
+            public void setScorer(Scorable scorer) {}
+
+            @Override
+            public void collect(int doc) {
+                if (++numCollected > maxCountHits) {
+                    earlyTerminate();
+                }
+            }
+        };
+    }
+
+    @Override
+    public ScoreMode scoreMode() {
+        // this collector is not exhaustive, as it early terminates, and never needs scores
+        return ScoreMode.TOP_DOCS;
+    }
+
+    private void earlyTerminate() {
+        throw new EarlyTerminationException("early termination [CountBased]");
+    }
+}

server/src/main/java/org/elasticsearch/search/query/TerminateAfterCollector.java

@@ -109,19 +109,19 @@ private EmptyTopDocsCollectorManagerFactory(@Nullable SortAndFormats sortAndForm
             super(REASON_SEARCH_COUNT, null);
             this.sort = sortAndFormats == null ? null : sortAndFormats.sort;
             if (trackTotalHitsUpTo == SearchContext.TRACK_TOTAL_HITS_DISABLED) {
-                this.collector = new EarlyTerminatingCollector(new TotalHitCountCollector(), 0, false);
+                this.collector = new PartialHitCountCollector(0);
                 // for bwc hit count is set to 0, it will be converted to -1 by the coordinating node
                 this.hitCountSupplier = () -> new TotalHits(0, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO);
             } else {
-                TotalHitCountCollector hitCountCollector = new TotalHitCountCollector();
                 if (trackTotalHitsUpTo == SearchContext.TRACK_TOTAL_HITS_ACCURATE) {
+                    TotalHitCountCollector hitCountCollector = new TotalHitCountCollector();
                     this.collector = hitCountCollector;
                     this.hitCountSupplier = () -> new TotalHits(hitCountCollector.getTotalHits(), TotalHits.Relation.EQUAL_TO);
                 } else {
-                    EarlyTerminatingCollector col = new EarlyTerminatingCollector(hitCountCollector, trackTotalHitsUpTo, false);
+                    PartialHitCountCollector col = new PartialHitCountCollector(trackTotalHitsUpTo);
                     this.collector = col;
                     this.hitCountSupplier = () -> new TotalHits(
-                        hitCountCollector.getTotalHits(),
+                        col.getTotalHits(),
                         col.hasEarlyTerminated() ? TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO : TotalHits.Relation.EQUAL_TO
                     );
                 }