[BUG] CompletionSuggestSearchIT.testSkipDuplicates is flaky (opensearch-project#8963)

Signed-off-by: Andriy Redko <[email protected]>
Signed-off-by: Shivansh Arora <[email protected]>

Author: reta

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) {

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