[BUG] CompletionSuggestSearchIT.testSkipDuplicates is flaky (opensear…

…ch-project#8963) Signed-off-by: Andriy Redko <andriy.redko@aiven.io> Signed-off-by: Shivansh Arora <hishiv@amazon.com>
Gaurav614 · Apr 25, 2024 · b2a015f · b2a015f
1 parent 5b8811b
commit b2a015f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java
@@ -343,6 +343,10 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto
                 }
             }
         }
+
+        // Note: this is called if collection ran successfully, including the above special cases of
+        // CollectionTerminatedException and TimeExceededException, but no other exception.
+        leafCollector.finish();
     }
 
     private Weight wrapWeight(Weight weight) {

diff --git a/server/src/main/java/org/opensearch/search/suggest/completion/CompletionSuggester.java b/server/src/main/java/org/opensearch/search/suggest/completion/CompletionSuggester.java
@@ -35,6 +35,7 @@
 import org.apache.lucene.search.BulkScorer;
 import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.search.suggest.document.CompletionQuery;
 import org.apache.lucene.search.suggest.document.TopSuggestDocs;
@@ -108,15 +109,21 @@ private static void suggest(IndexSearcher searcher, CompletionQuery query, TopSu
         for (LeafReaderContext context : searcher.getIndexReader().leaves()) {
             BulkScorer scorer = weight.bulkScorer(context);
             if (scorer != null) {
+                LeafCollector leafCollector = null;
                 try {
-                    scorer.score(collector.getLeafCollector(context), context.reader().getLiveDocs());
+                    leafCollector = collector.getLeafCollector(context);
+                    scorer.score(leafCollector, context.reader().getLiveDocs());
                 } catch (CollectionTerminatedException e) {
                     // collection was terminated prematurely
                     // continue with the following leaf
                 }
+                // Note: this is called if collection ran successfully, including the above special cases of
+                // CollectionTerminatedException and TimeExceededException, but no other exception.
+                if (leafCollector != null) {
+                    leafCollector.finish();
+                }
             }
         }
-        collector.finish();
     }
 
     @Override