Fix concurrent search NPE with track_total_hits, size=0, and terminat…

…e_after Signed-off-by: Jay Deng <jayd0104@gmail.com>
opensearch-project · Sep 20, 2023 · 1c84cec · 1c84cec
1 parent 0c4e950
commit 1c84cec
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -104,7 +104,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Fixed
 - Fix ignore_missing parameter has no effect when using template snippet in rename ingest processor ([#9725](https://github.com/opensearch-project/OpenSearch/pull/9725))
-- Fix broken backward compatibility from 2.7 for IndexSorted field indices ([#10045](https://github.com/opensearch-project/OpenSearch/pull/9725))
+- Fix broken backward compatibility from 2.7 for IndexSorted field indices ([#10045](https://github.com/opensearch-project/OpenSearch/pull/10045))
+- Fix concurrent search NPE when track_total_hits, terminate_after and size=0 are used ([#10082](https://github.com/opensearch-project/OpenSearch/pull/10082))
 
 ### Security
 

diff --git a/server/src/internalClusterTest/java/org/opensearch/search/simple/SimpleSearchIT.java b/server/src/internalClusterTest/java/org/opensearch/search/simple/SimpleSearchIT.java
@@ -292,6 +292,45 @@ public void testSimpleTerminateAfterCount() throws Exception {
         assertFalse(searchResponse.isTerminatedEarly());
     }
 
+    public void testSimpleTerminateAfterCountWithSizeAndTrackHits() throws Exception {
+        prepareCreate("test").setSettings(Settings.builder().put(SETTING_NUMBER_OF_SHARDS, 1).put(SETTING_NUMBER_OF_REPLICAS, 0)).get();
+        ensureGreen();
+        int max = randomIntBetween(3, 29);
+        List<IndexRequestBuilder> docbuilders = new ArrayList<>(max);
+
+        for (int i = 1; i <= max; i++) {
+            String id = String.valueOf(i);
+            docbuilders.add(client().prepareIndex("test").setId(id).setSource("field", i));
+        }
+
+        indexRandom(true, docbuilders);
+        ensureGreen();
+        refresh();
+
+        SearchResponse searchResponse;
+        for (int i = 1; i < max; i++) {
+            searchResponse = client().prepareSearch("test")
+                .setQuery(QueryBuilders.matchAllQuery())
+                .setTerminateAfter(i)
+                .setSize(0)
+                .setTrackTotalHits(true)
+                .get();
+            assertEquals(0, searchResponse.getFailedShards());
+            assertHitCount(searchResponse, i);
+            assertTrue(searchResponse.isTerminatedEarly());
+
+            searchResponse = client().prepareSearch("test")
+                .setQuery(QueryBuilders.matchAllQuery())
+                .setTerminateAfter(i)
+                .setSize(randomIntBetween(1, max))
+                .setTrackTotalHits(true)
+                .get();
+            assertEquals(0, searchResponse.getFailedShards());
+            assertHitCount(searchResponse, i);
+            assertTrue(searchResponse.isTerminatedEarly());
+        }
+    }
+
     public void testSimpleIndexSortEarlyTerminate() throws Exception {
         prepareCreate("test").setSettings(
             Settings.builder().put(SETTING_NUMBER_OF_SHARDS, 1).put(SETTING_NUMBER_OF_REPLICAS, 0).put("index.sort.field", "rank")

diff --git a/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java b/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java
@@ -130,15 +130,17 @@ static class EmptyTopDocsCollectorContext extends TopDocsCollectorContext {
          * @param query The query to execute
          * @param trackTotalHitsUpTo True if the total number of hits should be tracked
          * @param hasFilterCollector True if the collector chain contains a filter
+         * @param numDocs The number of top hits to retrieve
          */
         private EmptyTopDocsCollectorContext(
             IndexReader reader,
             Query query,
             @Nullable SortAndFormats sortAndFormats,
             int trackTotalHitsUpTo,
-            boolean hasFilterCollector
+            boolean hasFilterCollector,
+            int numDocs
         ) throws IOException {
-            super(REASON_SEARCH_COUNT, 0);
+            super(REASON_SEARCH_COUNT, numDocs);
             this.sort = sortAndFormats == null ? null : sortAndFormats.sort;
             this.trackTotalHitsUpTo = trackTotalHitsUpTo;
             if (this.trackTotalHitsUpTo == SearchContext.TRACK_TOTAL_HITS_DISABLED) {
@@ -189,6 +191,12 @@ CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, Re
                             trackTotalHitsUpTo,
                             false
                         );
+                    } else {
+                        manager = new EarlyTerminatingCollectorManager<>(
+                            new TotalHitCountCollectorManager.Empty(new TotalHits(numHits, TotalHits.Relation.EQUAL_TO), sort),
+                            numHits,
+                            false
+                        );
                     }
                 } else {
                     manager = new EarlyTerminatingCollectorManager<>(
@@ -778,7 +786,8 @@ public static TopDocsCollectorContext createTopDocsCollectorContext(SearchContex
                 query,
                 searchContext.sort(),
                 searchContext.trackTotalHitsUpTo(),
-                hasFilterCollector
+                hasFilterCollector,
+                totalNumDocs
             );
         } else if (searchContext.scrollContext() != null) {
             // we can disable the tracking of total hits after the initial scroll query