From f6f816d99d191356212a4b752e6add21eca80b9d Mon Sep 17 00:00:00 2001
From: "github-actions[bot]" <github-actions[bot]@users.noreply.github.com>
Date: Tue, 14 Nov 2023 19:21:40 +0000
Subject: [PATCH] [BUG] CompletionSuggestSearchIT.testSkipDuplicates is flaky
 (#8963)

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
(cherry picked from commit 106e83a75205e5ff5a3a760b28a53f71240e7e61)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../search/internal/ContextIndexSearcher.java         |  4 ++++
 .../suggest/completion/CompletionSuggester.java       | 11 +++++++++--
 2 files changed, 13 insertions(+), 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
index 7e2e7de2643f6..a520086de8051 100644
--- 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
index ed7bd709f94bc..fabb67d7fb841 100644
--- 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