Return an empty suggestion when suggest phase times out

We recently fixed timeout handling in the suggest phase. A test failure on SearchTimeoutIT
surfaced an issue with the current approach. In case partial results are allowed, it may
happen that some shards time out while executing the suggest phase and some don't.

SearchPhaseController assumes that if one shard has suggest results, all of the other shards
will have suggest results too. We could address that assertion and check is the search
timed out, instead this commit changes timeout handling in the suggest phase to return
an empty suggestion instead of null. This seems appropriate in terms of providing some
results and makes the assertion about non null suggestions in SearchPhaseController happy.

Relates to elastic#122357

Author: javanna

server/src/internalClusterTest/java/org/elasticsearch/search/SearchTimeoutIT.java

@@ -484,8 +484,7 @@ protected TermSuggestion innerExecute(
             CharsRefBuilder spare
         ) {
             contextIndexSearcher.throwTimeExceededException();
-            assert false;
-            return new TermSuggestion(name, suggestion.getSize(), SortBy.SCORE);
+            throw new AssertionError("should have thrown TimeExceededException");
         }
 
         @Override

server/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java

@@ -475,11 +475,13 @@ static ReducedQueryPhase reducedQueryPhase(
             }
 
             if (hasSuggest) {
-                assert result.suggest() != null;
-                for (Suggestion<? extends Suggestion.Entry<? extends Suggestion.Entry.Option>> suggestion : result.suggest()) {
-                    groupedSuggestions.computeIfAbsent(suggestion.getName(), s -> new ArrayList<>()).add(suggestion);
-                    if (suggestion instanceof CompletionSuggestion completionSuggestion) {
-                        completionSuggestion.setShardIndex(result.getShardIndex());
+                assert result.suggest() != null || result.searchTimedOut();
+                if (result.suggest() != null) {
+                    for (Suggestion<? extends Suggestion.Entry<? extends Suggestion.Entry.Option>> suggestion : result.suggest()) {
+                        groupedSuggestions.computeIfAbsent(suggestion.getName(), s -> new ArrayList<>()).add(suggestion);
+                        if (suggestion instanceof CompletionSuggestion completionSuggestion) {
+                            completionSuggestion.setShardIndex(result.getShardIndex());
+                        }
                     }
                 }
             }

server/src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -126,15 +126,7 @@ static void executeRank(SearchContext searchContext) throws QueryPhaseExecutionE
 
     static void executeQuery(SearchContext searchContext) throws QueryPhaseExecutionException {
         if (searchContext.hasOnlySuggest()) {
-            try {
-                SuggestPhase.execute(searchContext);
-            } catch (ContextIndexSearcher.TimeExceededException timeExceededException) {
-                SearchTimeoutException.handleTimeout(
-                    searchContext.request().allowPartialSearchResults(),
-                    searchContext.shardTarget(),
-                    searchContext.queryResult()
-                );
-            }
+            SuggestPhase.execute(searchContext);
             searchContext.queryResult().topDocs(new TopDocsAndMaxScore(Lucene.EMPTY_TOP_DOCS, Float.NaN), new DocValueFormat[0]);
             return;
         }
@@ -150,18 +142,10 @@ static void executeQuery(SearchContext searchContext) throws QueryPhaseExecution
 
         addCollectorsAndSearch(searchContext);
 
-        try {
-            RescorePhase.execute(searchContext);
-            SuggestPhase.execute(searchContext);
-            if (searchContext.getProfilers() != null) {
-                searchContext.queryResult().profileResults(searchContext.getProfilers().buildQueryPhaseResults());
-            }
-        } catch (ContextIndexSearcher.TimeExceededException timeExceededException) {
-            SearchTimeoutException.handleTimeout(
-                searchContext.request().allowPartialSearchResults(),
-                searchContext.shardTarget(),
-                searchContext.queryResult()
-            );
+        RescorePhase.execute(searchContext);
+        SuggestPhase.execute(searchContext);
+        if (searchContext.getProfilers() != null) {
+            searchContext.queryResult().profileResults(searchContext.getProfilers().buildQueryPhaseResults());
         }
     }
 

server/src/main/java/org/elasticsearch/search/rescore/RescorePhase.java

@@ -18,7 +18,9 @@
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.lucene.grouping.TopFieldGroups;
+import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.query.SearchTimeoutException;
 import org.elasticsearch.search.sort.ShardDocSortField;
 import org.elasticsearch.search.sort.SortAndFormats;
 
@@ -84,6 +86,13 @@ public static void execute(SearchContext context) {
                 .topDocs(new TopDocsAndMaxScore(topDocs, topDocs.scoreDocs[0].score), context.queryResult().sortValueFormats());
         } catch (IOException e) {
             throw new ElasticsearchException("Rescore Phase Failed", e);
+        } catch (ContextIndexSearcher.TimeExceededException timeExceededException) {
+            SearchTimeoutException.handleTimeout(
+                context.request().allowPartialSearchResults(),
+                context.shardTarget(),
+                context.queryResult()
+            );
+            //if the rescore phase times out and partial results are allowed, the returned top docs from this shard won't be rescored
         }
     }
 

server/src/main/java/org/elasticsearch/search/suggest/SuggestPhase.java

@@ -10,7 +10,9 @@
 
 import org.apache.lucene.util.CharsRefBuilder;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.query.SearchTimeoutException;
 import org.elasticsearch.search.suggest.Suggest.Suggestion;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry.Option;
@@ -40,12 +42,22 @@ public static void execute(SearchContext context) {
             for (Map.Entry<String, SuggestionSearchContext.SuggestionContext> entry : suggest.suggestions().entrySet()) {
                 SuggestionSearchContext.SuggestionContext suggestion = entry.getValue();
                 Suggester<SuggestionContext> suggester = suggestion.getSuggester();
-                Suggestion<? extends Entry<? extends Option>> result = suggester.execute(
-                    entry.getKey(),
-                    suggestion,
-                    context.searcher(),
-                    spare
-                );
+                Suggestion<? extends Entry<? extends Option>> result;
+                try {
+                    result = suggester.execute(
+                        entry.getKey(),
+                        suggestion,
+                        context.searcher(),
+                        spare
+                    );
+                } catch(ContextIndexSearcher.TimeExceededException timeExceededException) {
+                    SearchTimeoutException.handleTimeout(
+                        context.request().allowPartialSearchResults(),
+                        context.shardTarget(),
+                        context.queryResult()
+                    );
+                    result = suggester.emptySuggestion(entry.getKey(), suggestion, spare);
+                }
                 if (result != null) {
                     assert entry.getKey().equals(result.name);
                     suggestions.add(result);