Handle search timeout in SuggestPhase (elastic#122357)

Whenever a search timeout is set to a search request, the timeout may be triggered
by the suggest phase via exitable directory reader. In that case, the exception that
is thrown by the timeout check needs to be handled, instead of returned back to the user.

Instead of handling the timeout in each phase, this commit handles it as part of QueryPhase
for both SuggestPhase and RescorePhase.

For rescore phase, one integration test that is time dependent is also rewritten to remove
the time dependency and moved from QueryRescorerIT to SearchTimeoutIT.

Closes elastic#122186

Author: javanna

docs/changelog/122357.yaml

@@ -0,0 +1,6 @@
+pr: 122357
+summary: Handle search timeout in `SuggestPhase`
+area: Search
+type: bug
+issues:
+ - 122186

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

@@ -14,15 +14,18 @@
 import org.apache.lucene.search.ConstantScoreScorer;
 import org.apache.lucene.search.ConstantScoreWeight;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.CharsRefBuilder;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.search.SearchRequestBuilder;
@@ -32,12 +35,23 @@
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.SearchPlugin;
 import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
 import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
 import org.elasticsearch.search.internal.ContextIndexSearcher;
+import org.elasticsearch.search.rescore.RescoreContext;
+import org.elasticsearch.search.rescore.Rescorer;
+import org.elasticsearch.search.rescore.RescorerBuilder;
+import org.elasticsearch.search.suggest.SortBy;
+import org.elasticsearch.search.suggest.SuggestBuilder;
+import org.elasticsearch.search.suggest.Suggester;
+import org.elasticsearch.search.suggest.SuggestionSearchContext;
+import org.elasticsearch.search.suggest.term.TermSuggestion;
+import org.elasticsearch.search.suggest.term.TermSuggestionBuilder;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.hamcrest.ElasticsearchAssertions;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -57,7 +71,7 @@ public class SearchTimeoutIT extends ESIntegTestCase {
 
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
-        return Collections.singleton(BulkScorerTimeoutQueryPlugin.class);
+        return Collections.singleton(SearchTimeoutPlugin.class);
     }
 
     @Override
@@ -71,6 +85,9 @@ protected void setupSuiteScopeCluster() throws Exception {
         indexRandom(true, "test", randomIntBetween(20, 50));
     }
 
+    /**
+     * Test the scenario where the query times out before starting to collect documents, verify that partial hits are not returned
+     */
     public void testTopHitsTimeoutBeforeCollecting() {
         // setting the timeout is necessary only because we check that if a TimeExceededException is thrown, a timeout was set
         SearchRequestBuilder searchRequestBuilder = prepareSearch("test").setTimeout(new TimeValue(10, TimeUnit.SECONDS))
@@ -87,6 +104,9 @@ public void testTopHitsTimeoutBeforeCollecting() {
         });
     }
 
+    /**
+     * Test the scenario where the query times out while collecting documents, verify that partial hits results are returned
+     */
     public void testTopHitsTimeoutWhileCollecting() {
         // setting the timeout is necessary only because we check that if a TimeExceededException is thrown, a timeout was set
         SearchRequestBuilder searchRequestBuilder = prepareSearch("test").setTimeout(new TimeValue(10, TimeUnit.SECONDS))
@@ -102,6 +122,9 @@ public void testTopHitsTimeoutWhileCollecting() {
         });
     }
 
+    /**
+     * Test the scenario where the query times out before starting to collect documents, verify that partial aggs results are not returned
+     */
     public void testAggsTimeoutBeforeCollecting() {
         SearchRequestBuilder searchRequestBuilder = prepareSearch("test").setSize(0)
             // setting the timeout is necessary only because we check that if a TimeExceededException is thrown, a timeout was set
@@ -122,6 +145,9 @@ public void testAggsTimeoutBeforeCollecting() {
         });
     }
 
+    /**
+     * Test the scenario where the query times out while collecting documents, verify that partial aggs results are returned
+     */
     public void testAggsTimeoutWhileCollecting() {
         SearchRequestBuilder searchRequestBuilder = prepareSearch("test").setSize(0)
             // setting the timeout is necessary only because we check that if a TimeExceededException is thrown, a timeout was set
@@ -144,6 +170,56 @@ public void testAggsTimeoutWhileCollecting() {
         });
     }
 
+    /**
+     * Test the scenario where the suggest phase (part of the query phase) times out, yet there are results
+     * available coming from executing the query and aggs on each shard.
+     */
+    public void testSuggestTimeoutWithPartialResults() {
+        SuggestBuilder suggestBuilder = new SuggestBuilder();
+        suggestBuilder.setGlobalText("text");
+        TimeoutSuggestionBuilder timeoutSuggestionBuilder = new TimeoutSuggestionBuilder();
+        suggestBuilder.addSuggestion("suggest", timeoutSuggestionBuilder);
+        SearchRequestBuilder searchRequestBuilder = prepareSearch("test").suggest(suggestBuilder)
+            .addAggregation(new TermsAggregationBuilder("terms").field("field.keyword"));
+        ElasticsearchAssertions.assertResponse(searchRequestBuilder, searchResponse -> {
+            assertThat(searchResponse.isTimedOut(), equalTo(true));
+            assertEquals(0, searchResponse.getShardFailures().length);
+            assertEquals(0, searchResponse.getFailedShards());
+            assertThat(searchResponse.getSuccessfulShards(), greaterThan(0));
+            assertEquals(searchResponse.getSuccessfulShards(), searchResponse.getTotalShards());
+            assertThat(searchResponse.getHits().getTotalHits().value, greaterThan(0L));
+            assertThat(searchResponse.getHits().getHits().length, greaterThan(0));
+            StringTerms terms = searchResponse.getAggregations().get("terms");
+            assertEquals(1, terms.getBuckets().size());
+            StringTerms.Bucket bucket = terms.getBuckets().get(0);
+            assertEquals("value", bucket.getKeyAsString());
+            assertThat(bucket.getDocCount(), greaterThan(0L));
+        });
+    }
+
+    /**
+     * Test the scenario where the rescore phase (part of the query phase) times out, yet there are results
+     * available coming from executing the query and aggs on each shard.
+     */
+    public void testRescoreTimeoutWithPartialResults() {
+        SearchRequestBuilder searchRequestBuilder = prepareSearch("test").setRescorer(new TimeoutRescorerBuilder())
+            .addAggregation(new TermsAggregationBuilder("terms").field("field.keyword"));
+        ElasticsearchAssertions.assertResponse(searchRequestBuilder, searchResponse -> {
+            assertThat(searchResponse.isTimedOut(), equalTo(true));
+            assertEquals(0, searchResponse.getShardFailures().length);
+            assertEquals(0, searchResponse.getFailedShards());
+            assertThat(searchResponse.getSuccessfulShards(), greaterThan(0));
+            assertEquals(searchResponse.getSuccessfulShards(), searchResponse.getTotalShards());
+            assertThat(searchResponse.getHits().getTotalHits().value, greaterThan(0L));
+            assertThat(searchResponse.getHits().getHits().length, greaterThan(0));
+            StringTerms terms = searchResponse.getAggregations().get("terms");
+            assertEquals(1, terms.getBuckets().size());
+            StringTerms.Bucket bucket = terms.getBuckets().get(0);
+            assertEquals("value", bucket.getKeyAsString());
+            assertThat(bucket.getDocCount(), greaterThan(0L));
+        });
+    }
+
     public void testPartialResultsIntolerantTimeoutBeforeCollecting() {
         ElasticsearchException ex = expectThrows(
             ElasticsearchException.class,
@@ -170,13 +246,67 @@ public void testPartialResultsIntolerantTimeoutWhileCollecting() {
         assertEquals(504, ex.status().getStatus());
     }
 
-    public static final class BulkScorerTimeoutQueryPlugin extends Plugin implements SearchPlugin {
+    public void testPartialResultsIntolerantTimeoutWhileSuggestingOnly() {
+        SuggestBuilder suggestBuilder = new SuggestBuilder();
+        suggestBuilder.setGlobalText("text");
+        TimeoutSuggestionBuilder timeoutSuggestionBuilder = new TimeoutSuggestionBuilder();
+        suggestBuilder.addSuggestion("suggest", timeoutSuggestionBuilder);
+        ElasticsearchException ex = expectThrows(
+            ElasticsearchException.class,
+            prepareSearch("test").suggest(suggestBuilder).setAllowPartialSearchResults(false) // this line causes timeouts to report
+                                                                                              // failures
+        );
+        assertTrue(ex.toString().contains("Time exceeded"));
+        assertEquals(504, ex.status().getStatus());
+    }
+
+    public void testPartialResultsIntolerantTimeoutWhileSuggesting() {
+        SuggestBuilder suggestBuilder = new SuggestBuilder();
+        suggestBuilder.setGlobalText("text");
+        TimeoutSuggestionBuilder timeoutSuggestionBuilder = new TimeoutSuggestionBuilder();
+        suggestBuilder.addSuggestion("suggest", timeoutSuggestionBuilder);
+        ElasticsearchException ex = expectThrows(
+            ElasticsearchException.class,
+            prepareSearch("test").setQuery(new TermQueryBuilder("field", "value"))
+                .suggest(suggestBuilder)
+                .setAllowPartialSearchResults(false) // this line causes timeouts to report failures
+        );
+        assertTrue(ex.toString().contains("Time exceeded"));
+        assertEquals(504, ex.status().getStatus());
+    }
+
+    public void testPartialResultsIntolerantTimeoutWhileRescoring() {
+        ElasticsearchException ex = expectThrows(
+            ElasticsearchException.class,
+            prepareSearch("test").setQuery(new TermQueryBuilder("field", "value"))
+                .setRescorer(new TimeoutRescorerBuilder())
+                .setAllowPartialSearchResults(false) // this line causes timeouts to report failures
+        );
+        assertTrue(ex.toString().contains("Time exceeded"));
+        assertEquals(504, ex.status().getStatus());
+    }
+
+    public static final class SearchTimeoutPlugin extends Plugin implements SearchPlugin {
         @Override
         public List<QuerySpec<?>> getQueries() {
             return Collections.singletonList(new QuerySpec<QueryBuilder>("timeout", BulkScorerTimeoutQuery::new, parser -> {
                 throw new UnsupportedOperationException();
             }));
         }
+
+        @Override
+        public List<SuggesterSpec<?>> getSuggesters() {
+            return Collections.singletonList(new SuggesterSpec<>("timeout", TimeoutSuggestionBuilder::new, parser -> {
+                throw new UnsupportedOperationException();
+            }, TermSuggestion::new));
+        }
+
+        @Override
+        public List<RescorerSpec<?>> getRescorers() {
+            return Collections.singletonList(new RescorerSpec<>("timeout", TimeoutRescorerBuilder::new, parser -> {
+                throw new UnsupportedOperationException();
+            }));
+        }
     }
 
     /**
@@ -314,4 +444,111 @@ public TransportVersion getMinimalSupportedVersion() {
             return null;
         }
     }
+
+    /**
+     * Suggestion builder that triggers a timeout as part of its execution
+     */
+    private static final class TimeoutSuggestionBuilder extends TermSuggestionBuilder {
+        TimeoutSuggestionBuilder() {
+            super("field");
+        }
+
+        TimeoutSuggestionBuilder(StreamInput in) throws IOException {
+            super(in);
+        }
+
+        @Override
+        public String getWriteableName() {
+            return "timeout";
+        }
+
+        @Override
+        public SuggestionSearchContext.SuggestionContext build(SearchExecutionContext context) {
+            return new TimeoutSuggestionContext(new TimeoutSuggester((ContextIndexSearcher) context.searcher()), context);
+        }
+    }
+
+    private static final class TimeoutSuggester extends Suggester<TimeoutSuggestionContext> {
+        private final ContextIndexSearcher contextIndexSearcher;
+
+        TimeoutSuggester(ContextIndexSearcher contextIndexSearcher) {
+            this.contextIndexSearcher = contextIndexSearcher;
+        }
+
+        @Override
+        protected TermSuggestion innerExecute(
+            String name,
+            TimeoutSuggestionContext suggestion,
+            IndexSearcher searcher,
+            CharsRefBuilder spare
+        ) {
+            contextIndexSearcher.throwTimeExceededException();
+            assert false;
+            return new TermSuggestion(name, suggestion.getSize(), SortBy.SCORE);
+        }
+
+        @Override
+        protected TermSuggestion emptySuggestion(String name, TimeoutSuggestionContext suggestion, CharsRefBuilder spare) {
+            return new TermSuggestion(name, suggestion.getSize(), SortBy.SCORE);
+        }
+    }
+
+    private static final class TimeoutSuggestionContext extends SuggestionSearchContext.SuggestionContext {
+        TimeoutSuggestionContext(Suggester<?> suggester, SearchExecutionContext searchExecutionContext) {
+            super(suggester, searchExecutionContext);
+        }
+    }
+
+    private static final class TimeoutRescorerBuilder extends RescorerBuilder<TimeoutRescorerBuilder> {
+        TimeoutRescorerBuilder() {
+            super();
+        }
+
+        TimeoutRescorerBuilder(StreamInput in) throws IOException {
+            super(in);
+        }
+
+        @Override
+        protected void doWriteTo(StreamOutput out) {}
+
+        @Override
+        protected void doXContent(XContentBuilder builder, Params params) {}
+
+        @Override
+        protected RescoreContext innerBuildContext(int windowSize, SearchExecutionContext context) throws IOException {
+            return new RescoreContext(10, new Rescorer() {
+                @Override
+                public TopDocs rescore(TopDocs topDocs, IndexSearcher searcher, RescoreContext rescoreContext) {
+                    ((ContextIndexSearcher) context.searcher()).throwTimeExceededException();
+                    assert false;
+                    return null;
+                }
+
+                @Override
+                public Explanation explain(
+                    int topLevelDocId,
+                    IndexSearcher searcher,
+                    RescoreContext rescoreContext,
+                    Explanation sourceExplanation
+                ) {
+                    throw new UnsupportedOperationException();
+                }
+            });
+        }
+
+        @Override
+        public String getWriteableName() {
+            return "timeout";
+        }
+
+        @Override
+        public TransportVersion getMinimalSupportedVersion() {
+            return null;
+        }
+
+        @Override
+        public RescorerBuilder<TimeoutRescorerBuilder> rewrite(QueryRewriteContext ctx) {
+            return this;
+        }
+    }
 }

server/src/internalClusterTest/java/org/elasticsearch/search/functionscore/QueryRescorerIT.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.common.lucene.search.function.ScoreFunction;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.settings.Settings.Builder;
-import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.query.Operator;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
@@ -1016,22 +1015,6 @@ public void testRescoreAfterCollapseRandom() throws Exception {
         });
     }
 
-    public void testRescoreWithTimeout() throws Exception {
-        // no dummy docs since merges can change scores while we run queries.
-        int numDocs = indexRandomNumbers("whitespace", -1, false);
-
-        String intToEnglish = English.intToEnglish(between(0, numDocs - 1));
-        String query = intToEnglish.split(" ")[0];
-        assertResponse(
-            prepareSearch().setSearchType(SearchType.QUERY_THEN_FETCH)
-                .setQuery(QueryBuilders.matchQuery("field1", query).operator(Operator.OR))
-                .setSize(10)
-                .addRescorer(new QueryRescorerBuilder(functionScoreQuery(new TestTimedScoreFunctionBuilder())).windowSize(100))
-                .setTimeout(TimeValue.timeValueMillis(10)),
-            r -> assertTrue(r.isTimedOut())
-        );
-    }
-
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return List.of(TestTimedQueryPlugin.class);