Fix topDocs.totalHits assignment on scroll queries (#37180)

This change fixes an unreleased bug that assigns the wrong totalHits to scroll
queries.

Closes #37179

Author: jimczi

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

@@ -166,7 +166,6 @@ static boolean execute(SearchContext searchContext,
                         }
                         // ... and stop collecting after ${size} matches
                         searchContext.terminateAfter(searchContext.size());
-                        searchContext.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_DISABLED);
                     } else if (canEarlyTerminate(reader, searchContext.sort())) {
                         // now this gets interesting: since the search sort is a prefix of the index sort, we can directly
                         // skip to the desired doc
@@ -177,7 +176,6 @@ static boolean execute(SearchContext searchContext,
                                 .build();
                             query = bq;
                         }
-                        searchContext.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_DISABLED);
                     }
                 }
             }

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

@@ -217,15 +217,15 @@ private SimpleTopDocsCollectorContext(IndexReader reader,
             super(REASON_SEARCH_TOP_HITS, numHits);
             this.sortAndFormats = sortAndFormats;
 
-            // implicit total hit counts are valid only when there is no filter collector in the chain
-            final int hitCount = hasFilterCollector ? -1 : shortcutTotalHitCount(reader, query);
             final TopDocsCollector<?> topDocsCollector;
             if (trackTotalHitsUpTo == SearchContext.TRACK_TOTAL_HITS_DISABLED) {
                 // don't compute hit counts via the collector
                 topDocsCollector = createCollector(sortAndFormats, numHits, searchAfter, 1);
                 topDocsSupplier = new CachedSupplier<>(topDocsCollector::topDocs);
                 totalHitsSupplier = () -> new TotalHits(0, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO);
             } else {
+                // implicit total hit counts are valid only when there is no filter collector in the chain
+                final int hitCount = hasFilterCollector ? -1 : shortcutTotalHitCount(reader, query);
                 if (hitCount == -1) {
                     topDocsCollector = createCollector(sortAndFormats, numHits, searchAfter, trackTotalHitsUpTo);
                     topDocsSupplier = new CachedSupplier<>(topDocsCollector::topDocs);
@@ -293,12 +293,11 @@ private ScrollingTopDocsCollectorContext(IndexReader reader,
         @Override
         void postProcess(QuerySearchResult result) throws IOException {
             final TopDocs topDocs = topDocsSupplier.get();
-            topDocs.totalHits = totalHitsSupplier.get();
-            float maxScore = maxScoreSupplier.get();
+            final float maxScore;
             if (scrollContext.totalHits == null) {
                 // first round
-                scrollContext.totalHits = topDocs.totalHits;
-                scrollContext.maxScore = maxScore;
+                topDocs.totalHits = scrollContext.totalHits = totalHitsSupplier.get();
+                maxScore = scrollContext.maxScore = maxScoreSupplier.get();
             } else {
                 // subsequent round: the total number of hits and
                 // the maximum score were computed on the first round
@@ -367,7 +366,7 @@ static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searc
             // we can disable the tracking of total hits after the initial scroll query
             // since the total hits is preserved in the scroll context.
             int trackTotalHitsUpTo = searchContext.scrollContext().totalHits != null ?
-                SearchContext.TRACK_TOTAL_HITS_DISABLED : searchContext.trackTotalHitsUpTo();
+                SearchContext.TRACK_TOTAL_HITS_DISABLED : SearchContext.TRACK_TOTAL_HITS_ACCURATE;
             // no matter what the value of from is
             int numDocs = Math.min(searchContext.size(), totalNumDocs);
             return new ScrollingTopDocsCollectorContext(reader, query, searchContext.scrollContext(),

server/src/test/java/org/elasticsearch/search/stats/SearchStatsIT.java

@@ -178,7 +178,6 @@ private Set<String> nodeIdsWithIndex(String... indices) {
         return nodes;
     }
 
-    @AwaitsFix(bugUrl="https://github.com/elastic/elasticsearch/issues/37179")
     public void testOpenContexts() {
         String index = "test1";
         createIndex(index);