From b318c6a47a0adbfcdedf7227ff65e388426e8959 Mon Sep 17 00:00:00 2001 From: Ganesh Ramadurai Date: Wed, 4 Sep 2024 22:37:21 -0700 Subject: [PATCH] Provide factory for pluggable deciders Signed-off-by: Ganesh Ramadurai --- .../main/java/org/opensearch/node/Node.java | 8 +-- .../org/opensearch/plugins/SearchPlugin.java | 10 +-- .../search/DefaultSearchContext.java | 32 +++++---- .../org/opensearch/search/SearchModule.java | 24 +++---- .../org/opensearch/search/SearchService.java | 10 +-- .../deciders/ConcurrentSearchDecision.java | 2 +- ...va => ConcurrentSearchRequestDecider.java} | 36 ++++++---- .../deciders/ConcurrentSearchVisitor.java | 6 +- .../search/DefaultSearchContextTests.java | 70 +++++++++++-------- .../opensearch/search/SearchModuleTests.java | 49 ++++++------- .../java/org/opensearch/node/MockNode.java | 6 +- 11 files changed, 136 insertions(+), 117 deletions(-) rename server/src/main/java/org/opensearch/search/deciders/{ConcurrentSearchDecider.java => ConcurrentSearchRequestDecider.java} (50%) diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index c87f4288cefe7..72531b98ec102 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -233,7 +233,7 @@ import org.opensearch.search.aggregations.support.AggregationUsageService; import org.opensearch.search.backpressure.SearchBackpressureService; import org.opensearch.search.backpressure.settings.SearchBackpressureSettings; -import org.opensearch.search.deciders.ConcurrentSearchDecider; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.fetch.FetchPhase; import org.opensearch.search.pipeline.SearchPipelineService; import org.opensearch.search.query.QueryPhase; @@ -1337,7 +1337,7 @@ protected Node( circuitBreakerService, searchModule.getIndexSearcherExecutor(threadPool), taskResourceTrackingService, - searchModule.getConcurrentSearchDeciders() + searchModule.getConcurrentSearchRequestDeciderFactories() ); final List> tasksExecutors = pluginsService.filterPlugins(PersistentTaskPlugin.class) @@ -1987,7 +1987,7 @@ protected SearchService newSearchService( CircuitBreakerService circuitBreakerService, Executor indexSearcherExecutor, TaskResourceTrackingService taskResourceTrackingService, - Collection concurrentSearchDecidersList + Collection concurrentSearchDeciderFactories ) { return new SearchService( clusterService, @@ -2001,7 +2001,7 @@ protected SearchService newSearchService( circuitBreakerService, indexSearcherExecutor, taskResourceTrackingService, - concurrentSearchDecidersList + concurrentSearchDeciderFactories ); } diff --git a/server/src/main/java/org/opensearch/plugins/SearchPlugin.java b/server/src/main/java/org/opensearch/plugins/SearchPlugin.java index 895e6ed2971d8..60cb2184b5ab5 100644 --- a/server/src/main/java/org/opensearch/plugins/SearchPlugin.java +++ b/server/src/main/java/org/opensearch/plugins/SearchPlugin.java @@ -65,7 +65,7 @@ import org.opensearch.search.aggregations.pipeline.MovAvgPipelineAggregator; import org.opensearch.search.aggregations.pipeline.PipelineAggregator; import org.opensearch.search.aggregations.support.ValuesSourceRegistry; -import org.opensearch.search.deciders.ConcurrentSearchDecider; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.fetch.FetchSubPhase; import org.opensearch.search.fetch.subphase.highlight.Highlighter; import org.opensearch.search.query.QueryPhaseSearcher; @@ -141,12 +141,12 @@ default Map getHighlighters() { } /** - * Allows plugins to register custom decider for concurrent search - * @return A {@link ConcurrentSearchDecider} + * Allows plugins to register a factory to create custom decider for concurrent search + * @return A {@link ConcurrentSearchRequestDecider.Factory} */ @ExperimentalApi - default ConcurrentSearchDecider getConcurrentSearchDecider() { - return null; + default Optional getConcurrentSearchRequestDeciderFactory() { + return Optional.empty(); } /** diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java index 4576921b8426e..f148bfb164d7f 100644 --- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java +++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java @@ -72,8 +72,8 @@ import org.opensearch.search.aggregations.SearchContextAggregations; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.collapse.CollapseContext; -import org.opensearch.search.deciders.ConcurrentSearchDecider; import org.opensearch.search.deciders.ConcurrentSearchDecision; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.deciders.ConcurrentSearchVisitor; import org.opensearch.search.dfs.DfsSearchResult; import org.opensearch.search.fetch.FetchPhase; @@ -106,13 +106,14 @@ import java.util.Collection; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.Executor; import java.util.function.Function; import java.util.function.LongSupplier; -import java.util.stream.Collectors; import static org.opensearch.search.SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD; import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_MODE; @@ -136,7 +137,7 @@ final class DefaultSearchContext extends SearchContext { private final ShardSearchRequest request; private final SearchShardTarget shardTarget; private final LongSupplier relativeTimeSupplier; - private final Collection concurrentSearchDeciders; + private final Collection concurrentSearchDeciderFactories; private SearchType searchType; private final BigArrays bigArrays; private final IndexShard indexShard; @@ -221,7 +222,7 @@ final class DefaultSearchContext extends SearchContext { boolean validate, Executor executor, Function requestToAggReduceContextBuilder, - Collection concurrentSearchDeciders + Collection concurrentSearchDeciderFactories ) throws IOException { this.readerContext = readerContext; this.request = request; @@ -264,7 +265,7 @@ final class DefaultSearchContext extends SearchContext { this.maxAggRewriteFilters = evaluateFilterRewriteSetting(); this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold(); - this.concurrentSearchDeciders = concurrentSearchDeciders; + this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories; } @Override @@ -928,14 +929,21 @@ public boolean shouldUseConcurrentSearch() { private boolean evaluateAutoMode() { - // filter out deciders that want to opt-out of decision-making - final Set filteredDeciders = concurrentSearchDeciders.stream() - .filter(concurrentSearchDecider -> concurrentSearchDecider.canEvaluateForIndex(indexService.getIndexSettings())) - .collect(Collectors.toSet()); + final Set concurrentSearchRequestDeciders = new HashSet<>(); + + // create the ConcurrentSearchRequestDeciders using registered factories + for (ConcurrentSearchRequestDecider.Factory deciderFactory : concurrentSearchDeciderFactories) { + final Optional concurrentSearchRequestDecider = deciderFactory.create( + indexService.getIndexSettings() + ); + concurrentSearchRequestDecider.ifPresent(concurrentSearchRequestDeciders::add); + + } + // evaluate based on concurrent search query visitor - if (filteredDeciders.size() > 0) { + if (concurrentSearchRequestDeciders.size() > 0) { ConcurrentSearchVisitor concurrentSearchVisitor = new ConcurrentSearchVisitor( - filteredDeciders, + concurrentSearchRequestDeciders, indexService.getIndexSettings() ); if (request().source() != null && request().source().query() != null) { @@ -945,7 +953,7 @@ private boolean evaluateAutoMode() { } final List decisions = new ArrayList<>(); - for (ConcurrentSearchDecider decider : filteredDeciders) { + for (ConcurrentSearchRequestDecider decider : concurrentSearchRequestDeciders) { ConcurrentSearchDecision decision = decider.getConcurrentSearchDecision(); if (decision != null) { if (logger.isDebugEnabled()) { diff --git a/server/src/main/java/org/opensearch/search/SearchModule.java b/server/src/main/java/org/opensearch/search/SearchModule.java index e9ed02828b971..b8d3a13e0df20 100644 --- a/server/src/main/java/org/opensearch/search/SearchModule.java +++ b/server/src/main/java/org/opensearch/search/SearchModule.java @@ -255,7 +255,7 @@ import org.opensearch.search.aggregations.pipeline.SumBucketPipelineAggregationBuilder; import org.opensearch.search.aggregations.pipeline.SumBucketPipelineAggregator; import org.opensearch.search.aggregations.support.ValuesSourceRegistry; -import org.opensearch.search.deciders.ConcurrentSearchDecider; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.fetch.FetchPhase; import org.opensearch.search.fetch.FetchSubPhase; import org.opensearch.search.fetch.subphase.ExplainPhase; @@ -334,7 +334,7 @@ public class SearchModule { private final QueryPhaseSearcher queryPhaseSearcher; private final SearchPlugin.ExecutorServiceProvider indexSearcherExecutorProvider; - private final Collection concurrentSearchDeciders; + private final Collection concurrentSearchDeciderFactories; /** * Constructs a new SearchModule object @@ -364,25 +364,23 @@ public SearchModule(Settings settings, List plugins) { queryPhaseSearcher = registerQueryPhaseSearcher(plugins); indexSearcherExecutorProvider = registerIndexSearcherExecutorProvider(plugins); namedWriteables.addAll(SortValue.namedWriteables()); - concurrentSearchDeciders = registerConcurrentSearchDeciders(plugins); + concurrentSearchDeciderFactories = registerConcurrentSearchDeciderFactories(plugins); } - private Collection registerConcurrentSearchDeciders(List plugins) { - List concurrentSearchDeciders = new ArrayList<>(); + private Collection registerConcurrentSearchDeciderFactories(List plugins) { + List concurrentSearchDeciderFactories = new ArrayList<>(); for (SearchPlugin plugin : plugins) { - ConcurrentSearchDecider decider = plugin.getConcurrentSearchDecider(); - if (decider != null) { - concurrentSearchDeciders.add(decider); - } + final Optional deciderFactory = plugin.getConcurrentSearchRequestDeciderFactory(); + deciderFactory.ifPresent(concurrentSearchDeciderFactories::add); } - return concurrentSearchDeciders; + return concurrentSearchDeciderFactories; } /** - * Returns the concurrent search deciders that the plugins have registered + * Returns the concurrent search decider factories that the plugins have registered */ - public Collection getConcurrentSearchDeciders() { - return concurrentSearchDeciders; + public Collection getConcurrentSearchRequestDeciderFactories() { + return concurrentSearchDeciderFactories; } public List getNamedWriteables() { diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index e33a47fe8e178..626f1b574871e 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -105,7 +105,7 @@ import org.opensearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.collapse.CollapseContext; -import org.opensearch.search.deciders.ConcurrentSearchDecider; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.dfs.DfsPhase; import org.opensearch.search.dfs.DfsSearchResult; import org.opensearch.search.fetch.FetchPhase; @@ -358,7 +358,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv private final QueryPhase queryPhase; private final FetchPhase fetchPhase; - private final Collection concurrentSearchDeciders; + private final Collection concurrentSearchDeciderFactories; private volatile long defaultKeepAlive; @@ -404,7 +404,7 @@ public SearchService( CircuitBreakerService circuitBreakerService, Executor indexSearcherExecutor, TaskResourceTrackingService taskResourceTrackingService, - Collection concurrentSearchDeciders + Collection concurrentSearchDeciderFactories ) { Settings settings = clusterService.getSettings(); this.threadPool = threadPool; @@ -460,7 +460,7 @@ public SearchService( allowDerivedField = CLUSTER_ALLOW_DERIVED_FIELD_SETTING.get(settings); clusterService.getClusterSettings().addSettingsUpdateConsumer(CLUSTER_ALLOW_DERIVED_FIELD_SETTING, this::setAllowDerivedField); - this.concurrentSearchDeciders = concurrentSearchDeciders; + this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories; } private void validateKeepAlives(TimeValue defaultKeepAlive, TimeValue maxKeepAlive) { @@ -1161,7 +1161,7 @@ private DefaultSearchContext createSearchContext(ReaderContext reader, ShardSear validate, indexSearcherExecutor, this::aggReduceContextBuilder, - concurrentSearchDeciders + concurrentSearchDeciderFactories ); // we clone the query shard context here just for rewriting otherwise we // might end up with incorrect state since we are using now() or script services diff --git a/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchDecision.java b/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchDecision.java index 2a30413eff9c8..4ac47221856d1 100644 --- a/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchDecision.java +++ b/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchDecision.java @@ -13,7 +13,7 @@ import java.util.Collection; /** - * This Class defines the decisions that a {@link ConcurrentSearchDecider#getConcurrentSearchDecision} can return. + * This Class defines the decisions that a {@link ConcurrentSearchRequestDecider#getConcurrentSearchDecision} can return. * */ @ExperimentalApi diff --git a/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchDecider.java b/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchRequestDecider.java similarity index 50% rename from server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchDecider.java rename to server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchRequestDecider.java index 9c588bb45b4ec..ec40527314454 100644 --- a/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchDecider.java +++ b/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchRequestDecider.java @@ -12,17 +12,21 @@ import org.opensearch.index.IndexSettings; import org.opensearch.index.query.QueryBuilder; +import java.util.Optional; + /** - * {@link ConcurrentSearchDecider} allows pluggable way to evaluate if a query in the search request + * {@link ConcurrentSearchRequestDecider} allows pluggable way to evaluate if a query in the search request * can use concurrent segment search using the passed in queryBuilders from query tree and index settings * on a per shard request basis. - * Implementations can also opt out of the evaluation process for certain indices based on the index settings. - * For all the deciders which can evaluate query tree for an index, its evaluateForQuery method - * will be called for each node in the query tree. After traversing of the query tree is completed, the final - * decision from the deciders will be obtained using {@link ConcurrentSearchDecider#getConcurrentSearchDecision} + * Implementations will need to implement the Factory interface that can be used to create the ConcurrentSearchRequestDecider + * This factory will be called on each shard search request to create the ConcurrentSearchRequestDecider and get the + * concurrent search decision from the created decider on a per-request basis. + * For all the deciders the evaluateForQuery method will be called for each node in the query tree. + * After traversing of the query tree is completed, the final decision from the deciders will be + * obtained using {@link ConcurrentSearchRequestDecider#getConcurrentSearchDecision} */ @ExperimentalApi -public abstract class ConcurrentSearchDecider { +public abstract class ConcurrentSearchRequestDecider { /** * Evaluate for the passed in queryBuilder node in the query tree of the search request @@ -31,14 +35,6 @@ public abstract class ConcurrentSearchDecider { */ public abstract void evaluateForQuery(QueryBuilder queryBuilder, IndexSettings indexSettings); - /** - * Provides a way for deciders to opt out of decision-making process for certain requests based on - * index settings. - * Return true if interested in decision making for index, - * false, otherwise - */ - public abstract boolean canEvaluateForIndex(IndexSettings indexSettings); - /** * Provide the final decision for concurrent search based on all evaluations * Plugins may need to maintain internal state of evaluations to provide a final decision @@ -47,4 +43,16 @@ public abstract class ConcurrentSearchDecider { */ public abstract ConcurrentSearchDecision getConcurrentSearchDecision(); + /** + * Factory interface that can be implemented to create the ConcurrentSearchRequestDecider object. + * Implementations can use the passed in indexSettings to decide whether to create the decider object or + * return {@link Optional#empty()}. + */ + @ExperimentalApi + public interface Factory { + default Optional create(IndexSettings indexSettings) { + return Optional.empty(); + } + } + } diff --git a/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchVisitor.java b/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchVisitor.java index 12ba1b2a9cc5f..d1a4fa982dc7e 100644 --- a/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchVisitor.java +++ b/server/src/main/java/org/opensearch/search/deciders/ConcurrentSearchVisitor.java @@ -19,15 +19,15 @@ /** * Class to traverse the QueryBuilder tree and invoke the - * {@link ConcurrentSearchDecider#evaluateForQuery} at each node of the query tree + * {@link ConcurrentSearchRequestDecider#evaluateForQuery} at each node of the query tree */ @ExperimentalApi public class ConcurrentSearchVisitor implements QueryBuilderVisitor { - private final Set deciders; + private final Set deciders; private final IndexSettings indexSettings; - public ConcurrentSearchVisitor(Set concurrentSearchVisitorDeciders, IndexSettings idxSettings) { + public ConcurrentSearchVisitor(Set concurrentSearchVisitorDeciders, IndexSettings idxSettings) { Objects.requireNonNull(concurrentSearchVisitorDeciders, "Concurrent search deciders cannot be null"); deciders = concurrentSearchVisitorDeciders; indexSettings = idxSettings; diff --git a/server/src/test/java/org/opensearch/search/DefaultSearchContextTests.java b/server/src/test/java/org/opensearch/search/DefaultSearchContextTests.java index 491a0377ab32e..a56b9c860c85c 100644 --- a/server/src/test/java/org/opensearch/search/DefaultSearchContextTests.java +++ b/server/src/test/java/org/opensearch/search/DefaultSearchContextTests.java @@ -76,8 +76,8 @@ import org.opensearch.search.aggregations.MultiBucketConsumerService; import org.opensearch.search.aggregations.SearchContextAggregations; import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.search.deciders.ConcurrentSearchDecider; import org.opensearch.search.deciders.ConcurrentSearchDecision; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.internal.AliasFilter; import org.opensearch.search.internal.LegacyReaderContext; import org.opensearch.search.internal.PitReaderContext; @@ -96,6 +96,8 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.List; +import java.util.Optional; import java.util.UUID; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -988,14 +990,34 @@ protected Engine.Searcher acquireSearcherInternal(String source) { // Case4: multiple deciders are registered and all of them opt out of decision-making // with supported agg query so concurrent path is used - ConcurrentSearchDecider decider1 = mock(ConcurrentSearchDecider.class); - when(decider1.canEvaluateForIndex(any())).thenReturn(false); - ConcurrentSearchDecider decider2 = mock(ConcurrentSearchDecider.class); - when(decider2.canEvaluateForIndex(any())).thenReturn(false); + ConcurrentSearchRequestDecider decider1 = mock(ConcurrentSearchRequestDecider.class); - Collection concurrentSearchDeciders = new ArrayList<>(); - concurrentSearchDeciders.add(decider1); - concurrentSearchDeciders.add(decider2); + ConcurrentSearchRequestDecider decider2 = mock(ConcurrentSearchRequestDecider.class); + + ConcurrentSearchRequestDecider.Factory factory1 = new ConcurrentSearchRequestDecider.Factory() { + @Override + public Optional create(IndexSettings indexSettings) { + return Optional.ofNullable(decider1); + } + }; + + ConcurrentSearchRequestDecider.Factory factory2 = new ConcurrentSearchRequestDecider.Factory() { + @Override + public Optional create(IndexSettings indexSettings) { + return Optional.ofNullable(decider2); + } + }; + ConcurrentSearchRequestDecider.Factory factory3 = new ConcurrentSearchRequestDecider.Factory() { + @Override + public Optional create(IndexSettings indexSettings) { + return Optional.empty(); + } + }; + + List concurrentSearchRequestDeciders = new ArrayList<>(); + concurrentSearchRequestDeciders.add(factory1); + concurrentSearchRequestDeciders.add(factory2); + concurrentSearchRequestDeciders.add(factory3); context = new DefaultSearchContext( readerContext, @@ -1011,7 +1033,7 @@ protected Engine.Searcher acquireSearcherInternal(String source) { false, executor, null, - concurrentSearchDeciders + concurrentSearchRequestDeciders ); // create a supported agg operation context.aggregations(mockAggregations); @@ -1025,15 +1047,9 @@ protected Engine.Searcher acquireSearcherInternal(String source) { // Case5: multiple deciders are registered and one of them returns ConcurrentSearchDecision.DecisionStatus.NO // use non-concurrent path even if query contains supported agg - when(decider1.canEvaluateForIndex(any())).thenReturn(true); when(decider1.getConcurrentSearchDecision()).thenReturn( new ConcurrentSearchDecision(ConcurrentSearchDecision.DecisionStatus.NO, "disable concurrent search") ); - when(decider2.canEvaluateForIndex(any())).thenReturn(false); - - concurrentSearchDeciders.clear(); - concurrentSearchDeciders.add(decider1); - concurrentSearchDeciders.add(decider2); // create a source so that query tree is parsed by visitor SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); @@ -1055,7 +1071,7 @@ protected Engine.Searcher acquireSearcherInternal(String source) { false, executor, null, - concurrentSearchDeciders + concurrentSearchRequestDeciders ); // create a supported agg operation @@ -1071,20 +1087,17 @@ protected Engine.Searcher acquireSearcherInternal(String source) { // Case6: multiple deciders are registered and first decider returns ConcurrentSearchDecision.DecisionStatus.YES // while second decider returns ConcurrentSearchDecision.DecisionStatus.NO // use non-concurrent path even if query contains supported agg - when(decider1.canEvaluateForIndex(any())).thenReturn(true); + when(decider1.getConcurrentSearchDecision()).thenReturn( new ConcurrentSearchDecision(ConcurrentSearchDecision.DecisionStatus.YES, "enable concurrent search") ); - when(decider2.canEvaluateForIndex(any())).thenReturn(true); + when(decider2.getConcurrentSearchDecision()).thenReturn( new ConcurrentSearchDecision(ConcurrentSearchDecision.DecisionStatus.NO, "disable concurrent search") ); - concurrentSearchDeciders.clear(); - concurrentSearchDeciders.add(decider1); - concurrentSearchDeciders.add(decider2); - // create a source so that query tree is parsed by visitor + when(shardSearchRequest.source()).thenReturn(sourceBuilder); context = new DefaultSearchContext( readerContext, @@ -1100,7 +1113,7 @@ protected Engine.Searcher acquireSearcherInternal(String source) { false, executor, null, - concurrentSearchDeciders + concurrentSearchRequestDeciders ); // create a supported agg operation @@ -1115,22 +1128,19 @@ protected Engine.Searcher acquireSearcherInternal(String source) { // Case7: multiple deciders are registered and all return ConcurrentSearchDecision.DecisionStatus.NO_OP // but un-supported agg query is present, use non-concurrent path - when(decider1.canEvaluateForIndex(any())).thenReturn(true); + when(decider1.getConcurrentSearchDecision()).thenReturn( new ConcurrentSearchDecision(ConcurrentSearchDecision.DecisionStatus.NO_OP, "noop") ); - when(decider2.canEvaluateForIndex(any())).thenReturn(true); + when(decider2.getConcurrentSearchDecision()).thenReturn( new ConcurrentSearchDecision(ConcurrentSearchDecision.DecisionStatus.NO_OP, "noop") ); when(mockAggregations.factories().allFactoriesSupportConcurrentSearch()).thenReturn(false); - concurrentSearchDeciders.clear(); - concurrentSearchDeciders.add(decider1); - concurrentSearchDeciders.add(decider2); - // create a source so that query tree is parsed by visitor + when(shardSearchRequest.source()).thenReturn(sourceBuilder); context = new DefaultSearchContext( readerContext, @@ -1146,7 +1156,7 @@ protected Engine.Searcher acquireSearcherInternal(String source) { false, executor, null, - concurrentSearchDeciders + concurrentSearchRequestDeciders ); // create a supported agg operation diff --git a/server/src/test/java/org/opensearch/search/SearchModuleTests.java b/server/src/test/java/org/opensearch/search/SearchModuleTests.java index b3483b76dee1c..81b7ca8aef30b 100644 --- a/server/src/test/java/org/opensearch/search/SearchModuleTests.java +++ b/server/src/test/java/org/opensearch/search/SearchModuleTests.java @@ -41,6 +41,7 @@ import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.index.IndexSettings; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryRewriteContext; import org.opensearch.index.query.QueryShardContext; @@ -70,7 +71,7 @@ import org.opensearch.search.aggregations.support.ValuesSourceConfig; import org.opensearch.search.aggregations.support.ValuesSourceRegistry; import org.opensearch.search.aggregations.support.ValuesSourceType; -import org.opensearch.search.deciders.ConcurrentSearchDecider; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.fetch.FetchSubPhase; import org.opensearch.search.fetch.subphase.ExplainPhase; import org.opensearch.search.fetch.subphase.highlight.CustomHighlighter; @@ -508,12 +509,12 @@ public Optional getIndexSearcherExecutorProvider() { expectThrows(IllegalStateException.class, () -> new SearchModule(Settings.EMPTY, searchPlugins)); } - public void testRegisterConcurrentSearchDecidersNoExternalPlugins() { + public void testRegisterConcurrentSearchRequestDecidersNoExternalPlugins() { SearchModule searchModule = new SearchModule(Settings.EMPTY, Collections.emptyList()); - assertEquals(searchModule.getConcurrentSearchDeciders().size(), 0); + assertEquals(searchModule.getConcurrentSearchRequestDeciderFactories().size(), 0); } - public void testRegisterConcurrentSearchDecidersExternalPluginsWithNoDeciders() { + public void testRegisterConcurrentSearchRequestDecidersExternalPluginsWithNoDeciders() { SearchPlugin plugin1 = new SearchPlugin() { @Override public Optional getIndexSearcherExecutorProvider() { @@ -528,10 +529,10 @@ public Optional getIndexSearcherExecutorProvider() { searchPlugins.add(plugin2); SearchModule searchModule = new SearchModule(Settings.EMPTY, searchPlugins); - assertEquals(searchModule.getConcurrentSearchDeciders().size(), 0); + assertEquals(searchModule.getConcurrentSearchRequestDeciderFactories().size(), 0); } - public void testRegisterConcurrentSearchDecidersExternalPluginsWithDeciders() { + public void testRegisterConcurrentSearchRequestDecidersExternalPluginsWithDeciders() { SearchPlugin pluginDecider1 = new SearchPlugin() { @Override public Optional getIndexSearcherExecutorProvider() { @@ -539,15 +540,25 @@ public Optional getIndexSearcherExecutorProvider() { } @Override - public ConcurrentSearchDecider getConcurrentSearchDecider() { - return mock(ConcurrentSearchDecider.class); + public Optional getConcurrentSearchRequestDeciderFactory() { + return Optional.of(new ConcurrentSearchRequestDecider.Factory() { + @Override + public Optional create(IndexSettings indexSettings) { + return Optional.of(mock(ConcurrentSearchRequestDecider.class)); + } + }); } }; SearchPlugin pluginDecider2 = new SearchPlugin() { @Override - public ConcurrentSearchDecider getConcurrentSearchDecider() { - return mock(ConcurrentSearchDecider.class); + public Optional getConcurrentSearchRequestDeciderFactory() { + return Optional.of(new ConcurrentSearchRequestDecider.Factory() { + @Override + public Optional create(IndexSettings indexSettings) { + return Optional.of(mock(ConcurrentSearchRequestDecider.class)); + } + }); } }; @@ -556,23 +567,7 @@ public ConcurrentSearchDecider getConcurrentSearchDecider() { searchPlugins.add(pluginDecider2); SearchModule searchModule = new SearchModule(Settings.EMPTY, searchPlugins); - assertEquals(searchModule.getConcurrentSearchDeciders().size(), 2); - } - - public void testRegisterConcurrentSearchDecidersPluginWithNullDecider() { - SearchPlugin pluginWithNullDecider = new SearchPlugin() { - @Override - public ConcurrentSearchDecider getConcurrentSearchDecider() { - return null; - } - }; - - List searchPlugins = new ArrayList<>(); - searchPlugins.add(pluginWithNullDecider); - SearchModule searchModule = new SearchModule(Settings.EMPTY, searchPlugins); - // null decider is filtered out, so 0 deciders - assertEquals(searchModule.getConcurrentSearchDeciders().size(), 0); - + assertEquals(searchModule.getConcurrentSearchRequestDeciderFactories().size(), 2); } private static final String[] NON_DEPRECATED_QUERIES = new String[] { diff --git a/test/framework/src/main/java/org/opensearch/node/MockNode.java b/test/framework/src/main/java/org/opensearch/node/MockNode.java index 09df9b85320f0..97c06962ca2e7 100644 --- a/test/framework/src/main/java/org/opensearch/node/MockNode.java +++ b/test/framework/src/main/java/org/opensearch/node/MockNode.java @@ -57,7 +57,7 @@ import org.opensearch.script.ScriptService; import org.opensearch.search.MockSearchService; import org.opensearch.search.SearchService; -import org.opensearch.search.deciders.ConcurrentSearchDecider; +import org.opensearch.search.deciders.ConcurrentSearchRequestDecider; import org.opensearch.search.fetch.FetchPhase; import org.opensearch.search.query.QueryPhase; import org.opensearch.tasks.TaskResourceTrackingService; @@ -158,7 +158,7 @@ protected SearchService newSearchService( CircuitBreakerService circuitBreakerService, Executor indexSearcherExecutor, TaskResourceTrackingService taskResourceTrackingService, - Collection concurrentSearchDecidersList + Collection concurrentSearchDeciderFactories ) { if (getPluginsService().filterPlugins(MockSearchService.TestPlugin.class).isEmpty()) { return super.newSearchService( @@ -173,7 +173,7 @@ protected SearchService newSearchService( circuitBreakerService, indexSearcherExecutor, taskResourceTrackingService, - concurrentSearchDecidersList + concurrentSearchDeciderFactories ); } return new MockSearchService(