[BUG] CompletionSuggestSearchIT.testSkipDuplicates is flaky (#8963) (#…

…11200) (cherry picked from commit 106e83a) Signed-off-by: Andriy Redko <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Nov 15, 2023 · afcbae2 · afcbae2
1 parent 879b9af
commit afcbae2
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
@@ -350,6 +350,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