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 ;
@@ -136,7 +137,7 @@ final class DefaultSearchContext extends SearchContext {
136
137
private final ShardSearchRequest request ;
137
138
private final SearchShardTarget shardTarget ;
138
139
private final LongSupplier relativeTimeSupplier ;
139
- private final Collection <ConcurrentSearchDecider > concurrentSearchDeciders ;
140
+ private final Collection <ConcurrentSearchRequestDecider . Factory > concurrentSearchDeciderFactories ;
140
141
private SearchType searchType ;
141
142
private final BigArrays bigArrays ;
142
143
private final IndexShard indexShard ;
@@ -221,7 +222,7 @@ final class DefaultSearchContext extends SearchContext {
221
222
boolean validate ,
222
223
Executor executor ,
223
224
Function <SearchSourceBuilder , InternalAggregation .ReduceContextBuilder > requestToAggReduceContextBuilder ,
224
- Collection <ConcurrentSearchDecider > concurrentSearchDeciders
225
+ Collection <ConcurrentSearchRequestDecider . Factory > concurrentSearchDeciderFactories
225
226
) throws IOException {
226
227
this .readerContext = readerContext ;
227
228
this .request = request ;
@@ -264,7 +265,7 @@ final class DefaultSearchContext extends SearchContext {
264
265
265
266
this .maxAggRewriteFilters = evaluateFilterRewriteSetting ();
266
267
this .cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold ();
267
- this .concurrentSearchDeciders = concurrentSearchDeciders ;
268
+ this .concurrentSearchDeciderFactories = concurrentSearchDeciderFactories ;
268
269
}
269
270
270
271
@ Override
@@ -928,14 +929,21 @@ public boolean shouldUseConcurrentSearch() {
928
929
929
930
private boolean evaluateAutoMode () {
930
931
931
- // filter out deciders that want to opt-out of decision-making
932
- final Set <ConcurrentSearchDecider > filteredDeciders = concurrentSearchDeciders .stream ()
933
- .filter (concurrentSearchDecider -> concurrentSearchDecider .canEvaluateForIndex (indexService .getIndexSettings ()))
934
- .collect (Collectors .toSet ());
932
+ final Set <ConcurrentSearchRequestDecider > concurrentSearchRequestDeciders = new HashSet <>();
933
+
934
+ // create the ConcurrentSearchRequestDeciders using registered factories
935
+ for (ConcurrentSearchRequestDecider .Factory deciderFactory : concurrentSearchDeciderFactories ) {
936
+ final Optional <ConcurrentSearchRequestDecider > concurrentSearchRequestDecider = deciderFactory .create (
937
+ indexService .getIndexSettings ()
938
+ );
939
+ concurrentSearchRequestDecider .ifPresent (concurrentSearchRequestDeciders ::add );
940
+
941
+ }
942
+
935
943
// evaluate based on concurrent search query visitor
936
- if (filteredDeciders .size () > 0 ) {
944
+ if (concurrentSearchRequestDeciders .size () > 0 ) {
937
945
ConcurrentSearchVisitor concurrentSearchVisitor = new ConcurrentSearchVisitor (
938
- filteredDeciders ,
946
+ concurrentSearchRequestDeciders ,
939
947
indexService .getIndexSettings ()
940
948
);
941
949
if (request ().source () != null && request ().source ().query () != null ) {
@@ -945,7 +953,7 @@ private boolean evaluateAutoMode() {
945
953
}
946
954
947
955
final List <ConcurrentSearchDecision > decisions = new ArrayList <>();
948
- for (ConcurrentSearchDecider decider : filteredDeciders ) {
956
+ for (ConcurrentSearchRequestDecider decider : concurrentSearchRequestDeciders ) {
949
957
ConcurrentSearchDecision decision = decider .getConcurrentSearchDecision ();
950
958
if (decision != null ) {
951
959
if (logger .isDebugEnabled ()) {
0 commit comments