honeymoose
/
OpenSearch
mirror of https://github.com/honeymoose/OpenSearch.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Introduce QueryPhaseSearcher extension point (SearchPlugin) (#1931) 

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
This commit is contained in:
Andriy Redko 2022-03-22 00:11:23 -04:00 committed by GitHub
parent bd2d9350b7
commit 82fb7abb71
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 259 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
server/src/main/java/org/opensearch/node/Node.java
View File

@ -174,6 +174,7 @@ import org.opensearch.search.SearchModule;
import org.opensearch.search.SearchService;
import org.opensearch.search.aggregations.support.AggregationUsageService;
import org.opensearch.search.fetch.FetchPhase;
import org.opensearch.search.query.QueryPhase;
import org.opensearch.snapshots.InternalSnapshotsInfoService;
import org.opensearch.snapshots.RestoreService;
import org.opensearch.snapshots.SnapshotShardsService;
@ -210,6 +211,7 @@ import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.function.UnaryOperator;
@ -849,9 +851,11 @@ public class Node implements Closeable {
threadPool,
scriptService,
bigArrays,
searchModule.getQueryPhase(),
searchModule.getFetchPhase(),
responseCollectorService,
circuitBreakerService
circuitBreakerService,
searchModule.getIndexSearcherExecutor(threadPool)
);
final List<PersistentTasksExecutor<?>> tasksExecutors = pluginsService.filterPlugins(PersistentTaskPlugin.class)
@ -1407,9 +1411,11 @@ public class Node implements Closeable {
ThreadPool threadPool,
ScriptService scriptService,
BigArrays bigArrays,
QueryPhase queryPhase,
FetchPhase fetchPhase,
ResponseCollectorService responseCollectorService,
CircuitBreakerService circuitBreakerService
CircuitBreakerService circuitBreakerService,
Executor indexSearcherExecutor
) {
return new SearchService(
clusterService,
@ -1417,9 +1423,11 @@ public class Node implements Closeable {
threadPool,
scriptService,
bigArrays,
queryPhase,
fetchPhase,
responseCollectorService,
circuitBreakerService
circuitBreakerService,
indexSearcherExecutor
);
}

36
server/src/main/java/org/opensearch/plugins/SearchPlugin.java
View File

@ -32,6 +32,7 @@
package org.opensearch.plugins;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.opensearch.common.CheckedFunction;
@ -40,6 +41,7 @@ import org.opensearch.common.io.stream.NamedWriteable;
import org.opensearch.common.io.stream.StreamInput;
import org.opensearch.common.io.stream.Writeable;
import org.opensearch.common.lucene.search.function.ScoreFunction;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.xcontent.ContextParser;
import org.opensearch.common.xcontent.XContent;
import org.opensearch.common.xcontent.XContentParser;
@ -61,6 +63,7 @@ import org.opensearch.search.aggregations.pipeline.PipelineAggregator;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.fetch.FetchSubPhase;
import org.opensearch.search.fetch.subphase.highlight.Highlighter;
import org.opensearch.search.query.QueryPhaseSearcher;
import org.opensearch.search.rescore.Rescorer;
import org.opensearch.search.rescore.RescorerBuilder;
import org.opensearch.search.sort.SortBuilder;
@ -68,11 +71,14 @@ import org.opensearch.search.sort.SortParser;
import org.opensearch.search.suggest.Suggest;
import org.opensearch.search.suggest.Suggester;
import org.opensearch.search.suggest.SuggestionBuilder;
import org.opensearch.threadpool.ThreadPool;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.TreeMap;
import java.util.concurrent.ExecutorService;
import java.util.function.BiFunction;
import java.util.function.Consumer;
@ -178,6 +184,36 @@ public interface SearchPlugin {
return emptyList();
}
/**
* The new {@link QueryPhaseSearcher} added by this plugin. At the moment, only one {@link QueryPhaseSearcher} is supported per
* instance, the {@link IllegalStateException} is going to be thrown if more then one plugin tries to register
* {@link QueryPhaseSearcher} implementation.
*/
default Optional<QueryPhaseSearcher> getQueryPhaseSearcher() {
return Optional.empty();
}
/**
* The executor service provider (registered through {@link Plugin#getExecutorBuilders(Settings)} to be used at search
* time by {@link IndexSearcher}. The {@link IllegalStateException} is going to be thrown if more then one
* plugin tries to register index searcher executor.
*/
default Optional<ExecutorServiceProvider> getIndexSearcherExecutorProvider() {
return Optional.empty();
}
/**
* Executor service provider
*/
interface ExecutorServiceProvider {
/**
* Provides an executor service instance
* @param threadPool thread pool
* @return executor service instance
*/
ExecutorService getExecutor(ThreadPool threadPool);
}
/**
* Specification of custom {@link ScoreFunction}.
*/

7
server/src/main/java/org/opensearch/search/DefaultSearchContext.java
View File

@ -94,6 +94,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executor;
import java.util.function.LongSupplier;
final class DefaultSearchContext extends SearchContext {
@ -177,7 +178,8 @@ final class DefaultSearchContext extends SearchContext {
FetchPhase fetchPhase,
boolean lowLevelCancellation,
Version minNodeVersion,
boolean validate
boolean validate,
Executor executor
) throws IOException {
this.readerContext = readerContext;
this.request = request;
@ -198,7 +200,8 @@ final class DefaultSearchContext extends SearchContext {
engineSearcher.getSimilarity(),
engineSearcher.getQueryCache(),
engineSearcher.getQueryCachingPolicy(),
lowLevelCancellation
lowLevelCancellation,
executor
);
this.relativeTimeSupplier = relativeTimeSupplier;
this.timeout = timeout;

52
server/src/main/java/org/opensearch/search/SearchModule.java
View File

@ -34,6 +34,7 @@ package org.opensearch.search;
import org.apache.lucene.search.BooleanQuery;
import org.opensearch.common.NamedRegistry;
import org.opensearch.common.Nullable;
import org.opensearch.common.ParseField;
import org.opensearch.common.geo.GeoShapeType;
import org.opensearch.common.geo.ShapesAvailability;
@ -273,6 +274,8 @@ import org.opensearch.search.fetch.subphase.highlight.HighlightPhase;
import org.opensearch.search.fetch.subphase.highlight.Highlighter;
import org.opensearch.search.fetch.subphase.highlight.PlainHighlighter;
import org.opensearch.search.fetch.subphase.highlight.UnifiedHighlighter;
import org.opensearch.search.query.QueryPhase;
import org.opensearch.search.query.QueryPhaseSearcher;
import org.opensearch.search.rescore.QueryRescorerBuilder;
import org.opensearch.search.rescore.RescorerBuilder;
import org.opensearch.search.sort.FieldSortBuilder;
@ -293,11 +296,14 @@ import org.opensearch.search.suggest.phrase.SmoothingModel;
import org.opensearch.search.suggest.phrase.StupidBackoff;
import org.opensearch.search.suggest.term.TermSuggestion;
import org.opensearch.search.suggest.term.TermSuggestionBuilder;
import org.opensearch.threadpool.ThreadPool;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ExecutorService;
import java.util.function.Consumer;
import java.util.function.Function;
@ -329,6 +335,8 @@ public class SearchModule {
private final List<NamedWriteableRegistry.Entry> namedWriteables = new ArrayList<>();
private final List<NamedXContentRegistry.Entry> namedXContents = new ArrayList<>();
private final ValuesSourceRegistry valuesSourceRegistry;
private final QueryPhaseSearcher queryPhaseSearcher;
private final SearchPlugin.ExecutorServiceProvider indexSearcherExecutorProvider;
/**
* Constructs a new SearchModule object
@ -355,6 +363,8 @@ public class SearchModule {
registerSearchExts(plugins);
registerShapes();
registerIntervalsSourceProviders();
queryPhaseSearcher = registerQueryPhaseSearcher(plugins);
indexSearcherExecutorProvider = registerIndexSearcherExecutorProvider(plugins);
namedWriteables.addAll(SortValue.namedWriteables());
}
@ -1282,7 +1292,49 @@ public class SearchModule {
);
}
private QueryPhaseSearcher registerQueryPhaseSearcher(List<SearchPlugin> plugins) {
QueryPhaseSearcher searcher = null;
for (SearchPlugin plugin : plugins) {
final Optional<QueryPhaseSearcher> searcherOpt = plugin.getQueryPhaseSearcher();
if (searcher == null) {
searcher = searcherOpt.orElse(null);
} else if (searcherOpt.isPresent()) {
throw new IllegalStateException("Only one QueryPhaseSearcher is allowed, but more than one are provided by the plugins");
}
}
return searcher;
}
private SearchPlugin.ExecutorServiceProvider registerIndexSearcherExecutorProvider(List<SearchPlugin> plugins) {
SearchPlugin.ExecutorServiceProvider provider = null;
for (SearchPlugin plugin : plugins) {
final Optional<SearchPlugin.ExecutorServiceProvider> providerOpt = plugin.getIndexSearcherExecutorProvider();
if (provider == null) {
provider = providerOpt.orElse(null);
} else if (providerOpt.isPresent()) {
throw new IllegalStateException(
"The index searcher executor is already assigned but more than one are provided by the plugins"
);
}
}
return provider;
}
public FetchPhase getFetchPhase() {
return new FetchPhase(fetchSubPhases);
}
public QueryPhase getQueryPhase() {
return (queryPhaseSearcher == null) ? new QueryPhase() : new QueryPhase(queryPhaseSearcher);
}
public @Nullable ExecutorService getIndexSearcherExecutor(ThreadPool pool) {
return (indexSearcherExecutorProvider == null) ? null : indexSearcherExecutorProvider.getExecutor(pool);
}
}

11
server/src/main/java/org/opensearch/search/SearchService.java
View File

@ -256,6 +256,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
private final AtomicInteger openScrollContexts = new AtomicInteger();
private final String sessionId = UUIDs.randomBase64UUID();
private final Executor indexSearcherExecutor;
public SearchService(
ClusterService clusterService,
@ -263,9 +264,11 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
ThreadPool threadPool,
ScriptService scriptService,
BigArrays bigArrays,
QueryPhase queryPhase,
FetchPhase fetchPhase,
ResponseCollectorService responseCollectorService,
CircuitBreakerService circuitBreakerService
CircuitBreakerService circuitBreakerService,
Executor indexSearcherExecutor
) {
Settings settings = clusterService.getSettings();
this.threadPool = threadPool;
@ -274,13 +277,14 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
this.scriptService = scriptService;
this.responseCollectorService = responseCollectorService;
this.bigArrays = bigArrays;
this.queryPhase = new QueryPhase();
this.queryPhase = queryPhase;
this.fetchPhase = fetchPhase;
this.multiBucketConsumerService = new MultiBucketConsumerService(
clusterService,
settings,
circuitBreakerService.getBreaker(CircuitBreaker.REQUEST)
);
this.indexSearcherExecutor = indexSearcherExecutor;
TimeValue keepAliveInterval = KEEPALIVE_INTERVAL_SETTING.get(settings);
setKeepAlives(DEFAULT_KEEPALIVE_SETTING.get(settings), MAX_KEEPALIVE_SETTING.get(settings));
@ -884,7 +888,8 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
fetchPhase,
lowLevelCancellation,
clusterService.state().nodes().getMinNodeVersion(),
validate
validate,
indexSearcherExecutor
);
// 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

70
server/src/main/java/org/opensearch/search/query/QueryPhase.java
View File

@ -72,6 +72,7 @@ import org.opensearch.threadpool.ThreadPool;
import java.io.IOException;
import java.util.LinkedList;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import static org.opensearch.search.query.QueryCollectorContext.createEarlyTerminationCollectorContext;
@ -88,12 +89,19 @@ public class QueryPhase {
private static final Logger LOGGER = LogManager.getLogger(QueryPhase.class);
// TODO: remove this property
public static final boolean SYS_PROP_REWRITE_SORT = Booleans.parseBoolean(System.getProperty("opensearch.search.rewrite_sort", "true"));
public static final QueryPhaseSearcher DEFAULT_QUERY_PHASE_SEARCHER = new DefaultQueryPhaseSearcher();
private final QueryPhaseSearcher queryPhaseSearcher;
private final AggregationPhase aggregationPhase;
private final SuggestPhase suggestPhase;
private final RescorePhase rescorePhase;
public QueryPhase() {
this(DEFAULT_QUERY_PHASE_SEARCHER);
}
public QueryPhase(QueryPhaseSearcher queryPhaseSearcher) {
this.queryPhaseSearcher = Objects.requireNonNull(queryPhaseSearcher, "QueryPhaseSearcher is required");
this.aggregationPhase = new AggregationPhase();
this.suggestPhase = new SuggestPhase();
this.rescorePhase = new RescorePhase();
@ -139,7 +147,7 @@ public class QueryPhase {
// request, preProcess is called on the DFS phase phase, this is why we pre-process them
// here to make sure it happens during the QUERY phase
aggregationPhase.preProcess(searchContext);
boolean rescore = executeInternal(searchContext);
boolean rescore = executeInternal(searchContext, queryPhaseSearcher);
if (rescore) { // only if we do a regular search
rescorePhase.execute(searchContext);
@ -162,6 +170,15 @@ public class QueryPhase {
* @return whether the rescoring phase should be executed
*/
static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExecutionException {
return executeInternal(searchContext, QueryPhase.DEFAULT_QUERY_PHASE_SEARCHER);
}
/**
* In a package-private method so that it can be tested without having to
* wire everything (mapperService, etc.)
* @return whether the rescoring phase should be executed
*/
static boolean executeInternal(SearchContext searchContext, QueryPhaseSearcher queryPhaseSearcher) throws QueryPhaseExecutionException {
final ContextIndexSearcher searcher = searchContext.searcher();
final IndexReader reader = searcher.getIndexReader();
QuerySearchResult queryResult = searchContext.queryResult();
@ -261,13 +278,22 @@ public class QueryPhase {
}
try {
boolean shouldRescore = searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, timeoutSet);
boolean shouldRescore = queryPhaseSearcher.searchWith(
searchContext,
searcher,
query,
collectors,
hasFilterCollector,
timeoutSet
);
ExecutorService executor = searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH);
if (executor instanceof QueueResizingOpenSearchThreadPoolExecutor) {
QueueResizingOpenSearchThreadPoolExecutor rExecutor = (QueueResizingOpenSearchThreadPoolExecutor) executor;
queryResult.nodeQueueSize(rExecutor.getCurrentQueueSize());
queryResult.serviceTimeEWMA((long) rExecutor.getTaskExecutionEWMA());
}
return shouldRescore;
} finally {
// Search phase has finished, no longer need to check for timeout
@ -358,5 +384,43 @@ public class QueryPhase {
return true;
}
private static class TimeExceededException extends RuntimeException {}
/**
* The exception being raised when search timeout is reached.
*/
public static class TimeExceededException extends RuntimeException {
private static final long serialVersionUID = 1L;
}
/**
* Default {@link QueryPhaseSearcher} implementation which delegates to the {@link QueryPhase}.
*/
public static class DefaultQueryPhaseSearcher implements QueryPhaseSearcher {
/**
* Please use {@link QueryPhase#DEFAULT_QUERY_PHASE_SEARCHER}
*/
protected DefaultQueryPhaseSearcher() {}
@Override
public boolean searchWith(
SearchContext searchContext,
ContextIndexSearcher searcher,
Query query,
LinkedList<QueryCollectorContext> collectors,
boolean hasFilterCollector,
boolean hasTimeout
) throws IOException {
return searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
}
protected boolean searchWithCollector(
SearchContext searchContext,
ContextIndexSearcher searcher,
Query query,
LinkedList<QueryCollectorContext> collectors,
boolean hasFilterCollector,
boolean hasTimeout
) throws IOException {
return QueryPhase.searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout);
}
}
}

41
server/src/main/java/org/opensearch/search/query/QueryPhaseSearcher.java Normal file
View File

@ -0,0 +1,41 @@
/*
* 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.query;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.Query;
import org.opensearch.search.internal.ContextIndexSearcher;
import org.opensearch.search.internal.SearchContext;
import java.io.IOException;
import java.util.LinkedList;
/**
* The extension point which allows to plug in custom search implementation to be
* used at {@link QueryPhase}.
*/
public interface QueryPhaseSearcher {
/**
* Perform search using {@link CollectorManager}
* @param searchContext search context
* @param searcher context index searcher
* @param query query
* @param hasTimeout "true" if timeout was set, "false" otherwise
* @return is rescoring required or not
* @throws IOException IOException
*/
boolean searchWith(
SearchContext searchContext,
ContextIndexSearcher searcher,
Query query,
LinkedList<QueryCollectorContext> collectors,
boolean hasFilterCollector,
boolean hasTimeout
) throws IOException;
}

18
server/src/test/java/org/opensearch/search/DefaultSearchContextTests.java
View File

@ -182,7 +182,8 @@ public class DefaultSearchContextTests extends OpenSearchTestCase {
null,
false,
Version.CURRENT,
false
false,
null
);
contextWithoutScroll.from(300);
contextWithoutScroll.close();
@ -223,7 +224,8 @@ public class DefaultSearchContextTests extends OpenSearchTestCase {
null,
false,
Version.CURRENT,
false
false,
null
);
context1.from(300);
exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
@ -292,7 +294,8 @@ public class DefaultSearchContextTests extends OpenSearchTestCase {
null,
false,
Version.CURRENT,
false
false,
null
);
SliceBuilder sliceBuilder = mock(SliceBuilder.class);
@ -330,7 +333,8 @@ public class DefaultSearchContextTests extends OpenSearchTestCase {
null,
false,
Version.CURRENT,
false
false,
null
);
ParsedQuery parsedQuery = ParsedQuery.parsedMatchAllQuery();
context3.sliceBuilder(null).parsedQuery(parsedQuery).preProcess(false);
@ -360,7 +364,8 @@ public class DefaultSearchContextTests extends OpenSearchTestCase {
null,
false,
Version.CURRENT,
false
false,
null
);
context4.sliceBuilder(new SliceBuilder(1, 2)).parsedQuery(parsedQuery).preProcess(false);
Query query1 = context4.query();
@ -440,7 +445,8 @@ public class DefaultSearchContextTests extends OpenSearchTestCase {
null,
false,
Version.CURRENT,
false
false,
null
);
assertThat(context.searcher().hasCancellations(), is(false));
context.searcher().addQueryCancellation(() -> {});

5
server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java
View File

@ -196,6 +196,7 @@ import org.opensearch.script.ScriptService;
import org.opensearch.search.SearchService;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.fetch.FetchPhase;
import org.opensearch.search.query.QueryPhase;
import org.opensearch.snapshots.mockstore.MockEventuallyConsistentRepository;
import org.opensearch.test.OpenSearchTestCase;
import org.opensearch.test.disruption.DisruptableMockTransport;
@ -1977,9 +1978,11 @@ public class SnapshotResiliencyTests extends OpenSearchTestCase {
threadPool,
scriptService,
bigArrays,
new QueryPhase(),
new FetchPhase(Collections.emptyList()),
responseCollectorService,
new NoneCircuitBreakerService()
new NoneCircuitBreakerService(),
null
);
SearchPhaseController searchPhaseController = new SearchPhaseController(
writableRegistry(),

11
test/framework/src/main/java/org/opensearch/node/MockNode.java
View File

@ -59,6 +59,7 @@ import org.opensearch.script.ScriptService;
import org.opensearch.search.MockSearchService;
import org.opensearch.search.SearchService;
import org.opensearch.search.fetch.FetchPhase;
import org.opensearch.search.query.QueryPhase;
import org.opensearch.test.MockHttpTransport;
import org.opensearch.test.transport.MockTransportService;
import org.opensearch.threadpool.ThreadPool;
@ -71,6 +72,7 @@ import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Executor;
import java.util.function.Function;
/**
@ -148,9 +150,11 @@ public class MockNode extends Node {
ThreadPool threadPool,
ScriptService scriptService,
BigArrays bigArrays,
QueryPhase queryPhase,
FetchPhase fetchPhase,
ResponseCollectorService responseCollectorService,
CircuitBreakerService circuitBreakerService
CircuitBreakerService circuitBreakerService,
Executor indexSearcherExecutor
) {
if (getPluginsService().filterPlugins(MockSearchService.TestPlugin.class).isEmpty()) {
return super.newSearchService(
@ -159,9 +163,11 @@ public class MockNode extends Node {
threadPool,
scriptService,
bigArrays,
queryPhase,
fetchPhase,
responseCollectorService,
circuitBreakerService
circuitBreakerService,
indexSearcherExecutor
);
}
return new MockSearchService(
@ -170,6 +176,7 @@ public class MockNode extends Node {
threadPool,
scriptService,
bigArrays,
queryPhase,
fetchPhase,
circuitBreakerService
);

15
test/framework/src/main/java/org/opensearch/search/MockSearchService.java
View File

@ -41,6 +41,7 @@ import org.opensearch.plugins.Plugin;
import org.opensearch.script.ScriptService;
import org.opensearch.search.fetch.FetchPhase;
import org.opensearch.search.internal.ReaderContext;
import org.opensearch.search.query.QueryPhase;
import org.opensearch.threadpool.ThreadPool;
import java.util.HashMap;
@ -91,10 +92,22 @@ public class MockSearchService extends SearchService {
ThreadPool threadPool,
ScriptService scriptService,
BigArrays bigArrays,
QueryPhase queryPhase,
FetchPhase fetchPhase,
CircuitBreakerService circuitBreakerService
) {
super(clusterService, indicesService, threadPool, scriptService, bigArrays, fetchPhase, null, circuitBreakerService);
super(
clusterService,
indicesService,
threadPool,
scriptService,
bigArrays,
queryPhase,
fetchPhase,
null,
circuitBreakerService,
null
);
}
@Override