72
72
import org .opensearch .search .aggregations .SearchContextAggregations ;
73
73
import org .opensearch .search .builder .SearchSourceBuilder ;
74
74
import org .opensearch .search .collapse .CollapseContext ;
75
- import org .opensearch .search .deciders .ConcurrentSearchDecider ;
76
75
import org .opensearch .search .deciders .ConcurrentSearchDecision ;
76
+ import org .opensearch .search .deciders .ConcurrentSearchRequestDecider ;
77
77
import org .opensearch .search .deciders .ConcurrentSearchVisitor ;
78
78
import org .opensearch .search .dfs .DfsSearchResult ;
79
79
import org .opensearch .search .fetch .FetchPhase ;
106
106
import java .util .Collection ;
107
107
import java .util .Collections ;
108
108
import java .util .HashMap ;
109
+ import java .util .HashSet ;
109
110
import java .util .List ;
110
111
import java .util .Map ;
112
+ import java .util .Optional ;
111
113
import java .util .Set ;
112
114
import java .util .concurrent .Executor ;
113
115
import java .util .function .Function ;
114
116
import java .util .function .LongSupplier ;
115
- import java .util .stream .Collectors ;
116
117
117
118
import static org .opensearch .search .SearchService .CARDINALITY_AGGREGATION_PRUNING_THRESHOLD ;
118
119
import static org .opensearch .search .SearchService .CLUSTER_CONCURRENT_SEGMENT_SEARCH_MODE ;
@@ -137,7 +138,7 @@ final class DefaultSearchContext extends SearchContext {
137
138
private final ShardSearchRequest request ;
138
139
private final SearchShardTarget shardTarget ;
139
140
private final LongSupplier relativeTimeSupplier ;
140
- private final Collection <ConcurrentSearchDecider > concurrentSearchDeciders ;
141
+ private final Collection <ConcurrentSearchRequestDecider . Factory > concurrentSearchDeciderFactories ;
141
142
private SearchType searchType ;
142
143
private final BigArrays bigArrays ;
143
144
private final IndexShard indexShard ;
@@ -223,7 +224,7 @@ final class DefaultSearchContext extends SearchContext {
223
224
boolean validate ,
224
225
Executor executor ,
225
226
Function <SearchSourceBuilder , InternalAggregation .ReduceContextBuilder > requestToAggReduceContextBuilder ,
226
- Collection <ConcurrentSearchDecider > concurrentSearchDeciders
227
+ Collection <ConcurrentSearchRequestDecider . Factory > concurrentSearchDeciderFactories
227
228
) throws IOException {
228
229
this .readerContext = readerContext ;
229
230
this .request = request ;
@@ -267,7 +268,7 @@ final class DefaultSearchContext extends SearchContext {
267
268
268
269
this .maxAggRewriteFilters = evaluateFilterRewriteSetting ();
269
270
this .cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold ();
270
- this .concurrentSearchDeciders = concurrentSearchDeciders ;
271
+ this .concurrentSearchDeciderFactories = concurrentSearchDeciderFactories ;
271
272
this .keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled ();
272
273
}
273
274
@@ -932,14 +933,21 @@ public boolean shouldUseConcurrentSearch() {
932
933
933
934
private boolean evaluateAutoMode () {
934
935
935
- // filter out deciders that want to opt-out of decision-making
936
- final Set <ConcurrentSearchDecider > filteredDeciders = concurrentSearchDeciders .stream ()
937
- .filter (concurrentSearchDecider -> concurrentSearchDecider .canEvaluateForIndex (indexService .getIndexSettings ()))
938
- .collect (Collectors .toSet ());
936
+ final Set <ConcurrentSearchRequestDecider > concurrentSearchRequestDeciders = new HashSet <>();
937
+
938
+ // create the ConcurrentSearchRequestDeciders using registered factories
939
+ for (ConcurrentSearchRequestDecider .Factory deciderFactory : concurrentSearchDeciderFactories ) {
940
+ final Optional <ConcurrentSearchRequestDecider > concurrentSearchRequestDecider = deciderFactory .create (
941
+ indexService .getIndexSettings ()
942
+ );
943
+ concurrentSearchRequestDecider .ifPresent (concurrentSearchRequestDeciders ::add );
944
+
945
+ }
946
+
939
947
// evaluate based on concurrent search query visitor
940
- if (filteredDeciders .size () > 0 ) {
948
+ if (concurrentSearchRequestDeciders .size () > 0 ) {
941
949
ConcurrentSearchVisitor concurrentSearchVisitor = new ConcurrentSearchVisitor (
942
- filteredDeciders ,
950
+ concurrentSearchRequestDeciders ,
943
951
indexService .getIndexSettings ()
944
952
);
945
953
if (request ().source () != null && request ().source ().query () != null ) {
@@ -949,7 +957,7 @@ private boolean evaluateAutoMode() {
949
957
}
950
958
951
959
final List <ConcurrentSearchDecision > decisions = new ArrayList <>();
952
- for (ConcurrentSearchDecider decider : filteredDeciders ) {
960
+ for (ConcurrentSearchRequestDecider decider : concurrentSearchRequestDeciders ) {
953
961
ConcurrentSearchDecision decision = decider .getConcurrentSearchDecision ();
954
962
if (decision != null ) {
955
963
if (logger .isDebugEnabled ()) {
0 commit comments