From 31e67c17ecc6eb70d28e991b13085411e961b5bd Mon Sep 17 00:00:00 2001 From: Jay Deng Date: Fri, 7 Jul 2023 13:53:50 -0700 Subject: [PATCH] Add early termination support for concurrent segment search (#8306) Signed-off-by: Jay Deng --- CHANGELOG.md | 1 + ...ConcurrentSegmentSearchCancellationIT.java | 26 +++++++++++ .../ConcurrentSegmentSearchTimeoutIT.java | 27 ++++++++++++ .../opensearch/search/SearchTimeoutIT.java | 29 ++++++++++--- .../search/internal/ContextIndexSearcher.java | 43 ++++++++----------- .../search/internal/SearchContext.java | 10 +++++ .../query/ConcurrentQueryPhaseSearcher.java | 32 +++++++++++--- .../opensearch/search/query/QueryPhase.java | 4 +- .../search/SearchCancellationTests.java | 16 +++++-- .../internal/ContextIndexSearcherTests.java | 9 +++- .../profile/query/QueryProfilerTests.java | 11 ++++- .../search/query/QueryPhaseTests.java | 12 +++++- .../search/query/QueryProfilePhaseTests.java | 12 +++++- .../aggregations/AggregatorTestCase.java | 5 ++- 14 files changed, 188 insertions(+), 49 deletions(-) create mode 100644 server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchCancellationIT.java create mode 100644 server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchTimeoutIT.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 843529bc1ffd1..e0146af42ee3b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Allow mmap to use new JDK-19 preview APIs in Apache Lucene 9.4+ ([#5151](https://github.com/opensearch-project/OpenSearch/pull/5151)) - Add events correlation engine plugin ([#6854](https://github.com/opensearch-project/OpenSearch/issues/6854)) - Add support for ignoring missing Javadoc on generated code using annotation ([#7604](https://github.com/opensearch-project/OpenSearch/pull/7604)) +- Add partial results support for concurrent segment search ([#8306](https://github.com/opensearch-project/OpenSearch/pull/8306)) ### Dependencies - Bump `log4j-core` from 2.18.0 to 2.19.0 diff --git a/server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchCancellationIT.java b/server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchCancellationIT.java new file mode 100644 index 0000000000000..3c50627e342dd --- /dev/null +++ b/server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchCancellationIT.java @@ -0,0 +1,26 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search; + +import org.opensearch.common.settings.FeatureFlagSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; + +public class ConcurrentSegmentSearchCancellationIT extends SearchCancellationIT { + @Override + protected Settings featureFlagSettings() { + Settings.Builder featureSettings = Settings.builder(); + for (Setting builtInFlag : FeatureFlagSettings.BUILT_IN_FEATURE_FLAGS) { + featureSettings.put(builtInFlag.getKey(), builtInFlag.getDefaultRaw(Settings.EMPTY)); + } + featureSettings.put(FeatureFlags.CONCURRENT_SEGMENT_SEARCH, "true"); + return featureSettings.build(); + } +} diff --git a/server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchTimeoutIT.java b/server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchTimeoutIT.java new file mode 100644 index 0000000000000..c19f762679fb0 --- /dev/null +++ b/server/src/internalClusterTest/java/org/opensearch/search/ConcurrentSegmentSearchTimeoutIT.java @@ -0,0 +1,27 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search; + +import org.opensearch.common.settings.FeatureFlagSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; + +public class ConcurrentSegmentSearchTimeoutIT extends SearchTimeoutIT { + + @Override + protected Settings featureFlagSettings() { + Settings.Builder featureSettings = Settings.builder(); + for (Setting builtInFlag : FeatureFlagSettings.BUILT_IN_FEATURE_FLAGS) { + featureSettings.put(builtInFlag.getKey(), builtInFlag.getDefaultRaw(Settings.EMPTY)); + } + featureSettings.put(FeatureFlags.CONCURRENT_SEGMENT_SEARCH, "true"); + return featureSettings.build(); + } +} diff --git a/server/src/internalClusterTest/java/org/opensearch/search/SearchTimeoutIT.java b/server/src/internalClusterTest/java/org/opensearch/search/SearchTimeoutIT.java index 049dcb50024ba..c7392b260319a 100644 --- a/server/src/internalClusterTest/java/org/opensearch/search/SearchTimeoutIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/search/SearchTimeoutIT.java @@ -51,7 +51,6 @@ import static org.opensearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE; import static org.opensearch.index.query.QueryBuilders.scriptQuery; import static org.opensearch.search.SearchTimeoutIT.ScriptedTimeoutPlugin.SCRIPT_NAME; -import static org.hamcrest.Matchers.equalTo; @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE) public class SearchTimeoutIT extends OpenSearchIntegTestCase { @@ -67,17 +66,37 @@ protected Settings nodeSettings(int nodeOrdinal) { } public void testSimpleTimeout() throws Exception { - for (int i = 0; i < 32; i++) { + final int numDocs = 1000; + for (int i = 0; i < numDocs; i++) { client().prepareIndex("test").setId(Integer.toString(i)).setSource("field", "value").get(); } refresh("test"); SearchResponse searchResponse = client().prepareSearch("test") - .setTimeout(new TimeValue(10, TimeUnit.MILLISECONDS)) + .setTimeout(new TimeValue(5, TimeUnit.MILLISECONDS)) .setQuery(scriptQuery(new Script(ScriptType.INLINE, "mockscript", SCRIPT_NAME, Collections.emptyMap()))) .setAllowPartialSearchResults(true) .get(); - assertThat(searchResponse.isTimedOut(), equalTo(true)); + assertTrue(searchResponse.isTimedOut()); + assertEquals(0, searchResponse.getFailedShards()); + assertTrue(numDocs > searchResponse.getHits().getTotalHits().value); + } + + public void testSimpleDoesNotTimeout() throws Exception { + final int numDocs = 10; + for (int i = 0; i < numDocs; i++) { + client().prepareIndex("test").setId(Integer.toString(i)).setSource("field", "value").get(); + } + refresh("test"); + + SearchResponse searchResponse = client().prepareSearch("test") + .setTimeout(new TimeValue(10000, TimeUnit.SECONDS)) + .setQuery(scriptQuery(new Script(ScriptType.INLINE, "mockscript", SCRIPT_NAME, Collections.emptyMap()))) + .setAllowPartialSearchResults(true) + .get(); + assertFalse(searchResponse.isTimedOut()); + assertEquals(0, searchResponse.getFailedShards()); + assertEquals(numDocs, searchResponse.getHits().getTotalHits().value); } public void testPartialResultsIntolerantTimeout() throws Exception { @@ -91,7 +110,7 @@ public void testPartialResultsIntolerantTimeout() throws Exception { .setAllowPartialSearchResults(false) // this line causes timeouts to report failures .get() ); - assertTrue(ex.toString().contains("Time exceeded")); + assertTrue(ex.toString().contains("QueryPhaseExecutionException[Time exceeded]")); } public static class ScriptedTimeoutPlugin extends MockScriptPlugin { 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 f1650b686f070..2f1c5475c9cf6 100644 --- a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java +++ b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java @@ -73,6 +73,7 @@ import org.opensearch.search.profile.query.ProfileWeight; import org.opensearch.search.profile.query.QueryProfiler; import org.opensearch.search.profile.query.QueryTimingType; +import org.opensearch.search.query.QueryPhase; import org.opensearch.search.query.QuerySearchResult; import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.MinAndMax; @@ -103,26 +104,6 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable { private MutableQueryTimeout cancellable; private SearchContext searchContext; - public ContextIndexSearcher( - IndexReader reader, - Similarity similarity, - QueryCache queryCache, - QueryCachingPolicy queryCachingPolicy, - boolean wrapWithExitableDirectoryReader, - Executor executor - ) throws IOException { - this( - reader, - similarity, - queryCache, - queryCachingPolicy, - new MutableQueryTimeout(), - wrapWithExitableDirectoryReader, - executor, - null - ); - } - public ContextIndexSearcher( IndexReader reader, Similarity similarity, @@ -310,18 +291,22 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto return; } - cancellable.checkCancelled(); - weight = wrapWeight(weight); - // See please https://github.com/apache/lucene/pull/964 - collector.setWeight(weight); final LeafCollector leafCollector; try { + cancellable.checkCancelled(); + weight = wrapWeight(weight); + // See please https://github.com/apache/lucene/pull/964 + collector.setWeight(weight); leafCollector = collector.getLeafCollector(ctx); } catch (CollectionTerminatedException e) { // there is no doc of interest in this reader context // continue with the following leaf return; + } catch (QueryPhase.TimeExceededException e) { + searchContext.setSearchTimedOut(true); + return; } + // catch early terminated exception and rethrow? Bits liveDocs = ctx.reader().getLiveDocs(); BitSet liveDocsBitSet = getSparseBitSetOrNull(liveDocs); if (liveDocsBitSet == null) { @@ -332,6 +317,9 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto } catch (CollectionTerminatedException e) { // collection was terminated prematurely // continue with the following leaf + } catch (QueryPhase.TimeExceededException e) { + searchContext.setSearchTimedOut(true); + return; } } } else { @@ -348,6 +336,9 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto } catch (CollectionTerminatedException e) { // collection was terminated prematurely // continue with the following leaf + } catch (QueryPhase.TimeExceededException e) { + searchContext.setSearchTimedOut(true); + return; } } } @@ -492,7 +483,7 @@ private boolean canMatch(LeafReaderContext ctx) throws IOException { } private boolean canMatchSearchAfter(LeafReaderContext ctx) throws IOException { - if (searchContext != null && searchContext.request() != null && searchContext.request().source() != null) { + if (searchContext.request() != null && searchContext.request().source() != null) { // Only applied on primary sort field and primary search_after. FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(searchContext.request().source()); if (primarySortField != null) { @@ -512,7 +503,7 @@ private boolean shouldReverseLeafReaderContexts() { // This is actually beneficial for search queries to start search on latest segments first for time series workload. // That can slow down ASC order queries on timestamp workload. So to avoid that slowdown, we will reverse leaf // reader order here. - if (searchContext != null && searchContext.indexShard().isTimeSeriesDescSortOptimizationEnabled()) { + if (searchContext.indexShard().isTimeSeriesDescSortOptimizationEnabled()) { // Only reverse order for asc order sort queries if (searchContext.sort() != null && searchContext.sort().sort != null diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java index 9c6a0498ec848..319a5624bbf56 100644 --- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java +++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java @@ -98,6 +98,8 @@ public abstract class SearchContext implements Releasable { private final AtomicBoolean closed = new AtomicBoolean(false); private InnerHitsContext innerHitsContext; + private volatile boolean searchTimedOut; + protected SearchContext() {} public abstract void setTask(SearchShardTask task); @@ -106,6 +108,14 @@ protected SearchContext() {} public abstract boolean isCancelled(); + public boolean isSearchTimedOut() { + return this.searchTimedOut; + } + + public void setSearchTimedOut(boolean searchTimedOut) { + this.searchTimedOut = searchTimedOut; + } + @Override public final void close() { if (closed.compareAndSet(false, true)) { diff --git a/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java b/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java index e3ba0eda4af55..4d3b2a124f64f 100644 --- a/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java +++ b/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java @@ -19,10 +19,10 @@ import org.opensearch.search.internal.SearchContext; import org.opensearch.search.profile.query.ProfileCollectorManager; import org.opensearch.search.query.QueryPhase.DefaultQueryPhaseSearcher; -import org.opensearch.search.query.QueryPhase.TimeExceededException; import java.io.IOException; import java.util.LinkedList; +import java.util.concurrent.ExecutionException; import static org.opensearch.search.query.TopDocsCollectorContext.createTopDocsCollectorContext; @@ -80,12 +80,12 @@ private static boolean searchWithCollectorManager( try { final ReduceableSearchResult result = searcher.search(query, collectorManager); result.reduce(queryResult); - } catch (EarlyTerminatingCollector.EarlyTerminationException e) { - queryResult.terminatedEarly(true); - } catch (TimeExceededException e) { + } catch (RuntimeException re) { + rethrowCauseIfPossible(re, searchContext); + } + if (searchContext.isSearchTimedOut()) { assert timeoutSet : "TimeExceededException thrown even though timeout wasn't set"; if (searchContext.request().allowPartialSearchResults() == false) { - // Can't rethrow TimeExceededException because not serializable throw new QueryPhaseExecutionException(searchContext.shardTarget(), "Time exceeded"); } queryResult.searchTimedOut(true); @@ -101,4 +101,26 @@ private static boolean searchWithCollectorManager( public AggregationProcessor aggregationProcessor(SearchContext searchContext) { return aggregationProcessor; } + + private static void rethrowCauseIfPossible(RuntimeException re, SearchContext searchContext) throws T { + // Rethrow exception if cause is null + if (re.getCause() == null) { + throw re; + } + + // Unwrap the RuntimeException and ExecutionException from Lucene concurrent search method and rethrow + if (re.getCause() instanceof ExecutionException || re.getCause() instanceof InterruptedException) { + Throwable t = re.getCause(); + if (t.getCause() != null) { + throw (T) t.getCause(); + } + } + + // Rethrow any unexpected exception types + throw new QueryPhaseExecutionException( + searchContext.shardTarget(), + "Failed to execute concurrent segment search thread", + re.getCause() + ); + } } diff --git a/server/src/main/java/org/opensearch/search/query/QueryPhase.java b/server/src/main/java/org/opensearch/search/query/QueryPhase.java index 069b410280d63..8418fdca2f777 100644 --- a/server/src/main/java/org/opensearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/opensearch/search/query/QueryPhase.java @@ -354,10 +354,10 @@ private static boolean searchWithCollector( searcher.search(query, queryCollector); } catch (EarlyTerminatingCollector.EarlyTerminationException e) { queryResult.terminatedEarly(true); - } catch (TimeExceededException e) { + } + if (searchContext.isSearchTimedOut()) { assert timeoutSet : "TimeExceededException thrown even though timeout wasn't set"; if (searchContext.request().allowPartialSearchResults() == false) { - // Can't rethrow TimeExceededException because not serializable throw new QueryPhaseExecutionException(searchContext.shardTarget(), "Time exceeded"); } queryResult.searchTimedOut(true); diff --git a/server/src/test/java/org/opensearch/search/SearchCancellationTests.java b/server/src/test/java/org/opensearch/search/SearchCancellationTests.java index db807e31d0238..e67123bf2c51e 100644 --- a/server/src/test/java/org/opensearch/search/SearchCancellationTests.java +++ b/server/src/test/java/org/opensearch/search/SearchCancellationTests.java @@ -49,7 +49,9 @@ import org.apache.lucene.util.automaton.CompiledAutomaton; import org.apache.lucene.util.automaton.RegExp; import org.opensearch.common.util.io.IOUtils; +import org.opensearch.index.shard.IndexShard; import org.opensearch.search.internal.ContextIndexSearcher; +import org.opensearch.search.internal.SearchContext; import org.opensearch.tasks.TaskCancelledException; import org.opensearch.test.OpenSearchTestCase; import org.junit.AfterClass; @@ -59,6 +61,8 @@ import java.util.concurrent.atomic.AtomicBoolean; import static org.hamcrest.Matchers.equalTo; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class SearchCancellationTests extends OpenSearchTestCase { @@ -109,7 +113,8 @@ public void testAddingCancellationActions() throws IOException { IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + null, + mock(SearchContext.class) ); NullPointerException npe = expectThrows(NullPointerException.class, () -> searcher.addQueryCancellation(null)); assertEquals("cancellation runnable should not be null", npe.getMessage()); @@ -123,13 +128,17 @@ public void testAddingCancellationActions() throws IOException { public void testCancellableCollector() throws IOException { TotalHitCountCollector collector1 = new TotalHitCountCollector(); Runnable cancellation = () -> { throw new TaskCancelledException("cancelled"); }; + SearchContext searchContext = mock(SearchContext.class); + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); ContextIndexSearcher searcher = new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + null, + searchContext ); searcher.search(new MatchAllDocsQuery(), collector1); @@ -157,7 +166,8 @@ public void testExitableDirectoryReader() throws IOException { IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + null, + mock(SearchContext.class) ); searcher.addQueryCancellation(cancellation); CompiledAutomaton automaton = new CompiledAutomaton(new RegExp("a.*").toAutomaton()); diff --git a/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java b/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java index 82dcbd6f69aa4..c9755cec2a33d 100644 --- a/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java +++ b/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java @@ -80,6 +80,7 @@ import org.opensearch.common.util.io.IOUtils; import org.opensearch.index.IndexSettings; import org.opensearch.index.cache.bitset.BitsetFilterCache; +import org.opensearch.index.shard.IndexShard; import org.opensearch.index.shard.ShardId; import org.opensearch.search.aggregations.LeafBucketCollector; import org.opensearch.test.OpenSearchTestCase; @@ -91,6 +92,8 @@ import java.util.IdentityHashMap; import java.util.Set; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.opensearch.search.internal.ContextIndexSearcher.intersectScorerAndBitSet; import static org.opensearch.search.internal.ExitableDirectoryReader.ExitableLeafReader; import static org.opensearch.search.internal.ExitableDirectoryReader.ExitablePointValues; @@ -253,13 +256,17 @@ public void onRemoval(ShardId shardId, Accountable accountable) { DocumentSubsetDirectoryReader filteredReader = new DocumentSubsetDirectoryReader(reader, cache, roleQuery); + SearchContext searchContext = mock(SearchContext.class); + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); ContextIndexSearcher searcher = new ContextIndexSearcher( filteredReader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + null, + searchContext ); for (LeafReaderContext context : searcher.getIndexReader().leaves()) { diff --git a/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java b/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java index 8a8700d853ad7..f0f692e15f066 100644 --- a/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java +++ b/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java @@ -60,7 +60,9 @@ import org.apache.lucene.store.Directory; import org.apache.lucene.tests.util.TestUtil; import org.opensearch.common.util.io.IOUtils; +import org.opensearch.index.shard.IndexShard; import org.opensearch.search.internal.ContextIndexSearcher; +import org.opensearch.search.internal.SearchContext; import org.opensearch.search.profile.ProfileResult; import org.opensearch.test.OpenSearchTestCase; import org.junit.After; @@ -78,6 +80,8 @@ import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThan; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class QueryProfilerTests extends OpenSearchTestCase { private Directory dir; @@ -112,13 +116,18 @@ public void setUp() throws Exception { } reader = w.getReader(); w.close(); + + SearchContext searchContext = mock(SearchContext.class); + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); searcher = new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), ALWAYS_CACHE_POLICY, true, - executor + executor, + searchContext ); } diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java index dcec4842fc81e..cf4d7798056c4 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java @@ -1204,25 +1204,33 @@ public boolean lowLevelCancellation() { } private static ContextIndexSearcher newContextSearcher(IndexReader reader, ExecutorService executor) throws IOException { + SearchContext searchContext = mock(SearchContext.class); + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - executor + executor, + searchContext ); } private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size, ExecutorService executor) throws IOException { + SearchContext searchContext = mock(SearchContext.class); + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - executor + executor, + searchContext ) { @Override diff --git a/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java index b1fc5f1354039..9c27783263c5a 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java @@ -1424,25 +1424,33 @@ private final void assertProfileData( } private static ContextIndexSearcher newContextSearcher(IndexReader reader, ExecutorService executor) throws IOException { + SearchContext searchContext = mock(SearchContext.class); + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - executor + executor, + searchContext ); } private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size, ExecutorService executor) throws IOException { + SearchContext searchContext = mock(SearchContext.class); + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - executor + executor, + searchContext ) { @Override diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java index 48eea430aca76..cf0f9dbb04f7f 100644 --- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java +++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java @@ -329,16 +329,17 @@ public boolean shouldCache(Query query) { return false; } }; + SearchContext searchContext = mock(SearchContext.class); ContextIndexSearcher contextIndexSearcher = new ContextIndexSearcher( indexSearcher.getIndexReader(), indexSearcher.getSimilarity(), queryCache, queryCachingPolicy, false, - null + null, + searchContext ); - SearchContext searchContext = mock(SearchContext.class); when(searchContext.numberOfShards()).thenReturn(1); when(searchContext.searcher()).thenReturn(contextIndexSearcher); when(searchContext.fetchPhase()).thenReturn(new FetchPhase(Arrays.asList(new FetchSourcePhase(), new FetchDocValuesPhase())));