Provide the TransportRequest during validation of a search context (#24985)

This commit provides the TransportRequest that caused the retrieval of a search context to the
SearchOperationListener#validateSearchContext method so that implementers have access to the
request.
This commit is contained in:
Jay Modi 2017-06-01 07:49:58 -06:00 committed by GitHub
parent 45bdaeced5
commit 7526c29a05
3 changed files with 18 additions and 12 deletions

View File

@ -23,6 +23,7 @@ import org.apache.logging.log4j.message.ParameterizedMessage;
import org.apache.logging.log4j.util.Supplier; import org.apache.logging.log4j.util.Supplier;
import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.ExceptionsHelper;
import org.elasticsearch.search.internal.SearchContext; import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.transport.TransportRequest;
import java.util.List; import java.util.List;
@ -110,8 +111,9 @@ public interface SearchOperationListener {
* from the active contexts. If the context is deemed invalid a runtime * from the active contexts. If the context is deemed invalid a runtime
* exception can be thrown, which will prevent the context from being used. * exception can be thrown, which will prevent the context from being used.
* @param context the context retrieved from the active contexts * @param context the context retrieved from the active contexts
* @param transportRequest the request that is going to use the search context
*/ */
default void validateSearchContext(SearchContext context) {} default void validateSearchContext(SearchContext context, TransportRequest transportRequest) {}
/** /**
* A Composite listener that multiplexes calls to each of the listeners methods. * A Composite listener that multiplexes calls to each of the listeners methods.
@ -236,11 +238,11 @@ public interface SearchOperationListener {
} }
@Override @Override
public void validateSearchContext(SearchContext context) { public void validateSearchContext(SearchContext context, TransportRequest request) {
Exception exception = null; Exception exception = null;
for (SearchOperationListener listener : listeners) { for (SearchOperationListener listener : listeners) {
try { try {
listener.validateSearchContext(context); listener.validateSearchContext(context, request);
} catch (Exception e) { } catch (Exception e) {
exception = ExceptionsHelper.useOrSuppress(exception, e); exception = ExceptionsHelper.useOrSuppress(exception, e);
} }

View File

@ -85,6 +85,7 @@ import org.elasticsearch.search.suggest.completion.CompletionSuggestion;
import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.threadpool.ThreadPool.Cancellable; import org.elasticsearch.threadpool.ThreadPool.Cancellable;
import org.elasticsearch.threadpool.ThreadPool.Names; import org.elasticsearch.threadpool.ThreadPool.Names;
import org.elasticsearch.transport.TransportRequest;
import java.io.IOException; import java.io.IOException;
import java.util.Collections; import java.util.Collections;
@ -309,7 +310,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
} }
public ScrollQuerySearchResult executeQueryPhase(InternalScrollSearchRequest request, SearchTask task) { public ScrollQuerySearchResult executeQueryPhase(InternalScrollSearchRequest request, SearchTask task) {
final SearchContext context = findContext(request.id()); final SearchContext context = findContext(request.id(), request);
SearchOperationListener operationListener = context.indexShard().getSearchOperationListener(); SearchOperationListener operationListener = context.indexShard().getSearchOperationListener();
context.incRef(); context.incRef();
try { try {
@ -333,7 +334,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
} }
public QuerySearchResult executeQueryPhase(QuerySearchRequest request, SearchTask task) { public QuerySearchResult executeQueryPhase(QuerySearchRequest request, SearchTask task) {
final SearchContext context = findContext(request.id()); final SearchContext context = findContext(request.id(), request);
context.setTask(task); context.setTask(task);
IndexShard indexShard = context.indexShard(); IndexShard indexShard = context.indexShard();
SearchOperationListener operationListener = indexShard.getSearchOperationListener(); SearchOperationListener operationListener = indexShard.getSearchOperationListener();
@ -374,7 +375,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
} }
public ScrollQueryFetchSearchResult executeFetchPhase(InternalScrollSearchRequest request, SearchTask task) { public ScrollQueryFetchSearchResult executeFetchPhase(InternalScrollSearchRequest request, SearchTask task) {
final SearchContext context = findContext(request.id()); final SearchContext context = findContext(request.id(), request);
context.incRef(); context.incRef();
try { try {
context.setTask(task); context.setTask(task);
@ -405,7 +406,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
} }
public FetchSearchResult executeFetchPhase(ShardFetchRequest request, SearchTask task) { public FetchSearchResult executeFetchPhase(ShardFetchRequest request, SearchTask task) {
final SearchContext context = findContext(request.id()); final SearchContext context = findContext(request.id(), request);
final SearchOperationListener operationListener = context.indexShard().getSearchOperationListener(); final SearchOperationListener operationListener = context.indexShard().getSearchOperationListener();
context.incRef(); context.incRef();
try { try {
@ -435,7 +436,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
} }
} }
private SearchContext findContext(long id) throws SearchContextMissingException { private SearchContext findContext(long id, TransportRequest request) throws SearchContextMissingException {
SearchContext context = activeContexts.get(id); SearchContext context = activeContexts.get(id);
if (context == null) { if (context == null) {
throw new SearchContextMissingException(id); throw new SearchContextMissingException(id);
@ -443,7 +444,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
SearchOperationListener operationListener = context.indexShard().getSearchOperationListener(); SearchOperationListener operationListener = context.indexShard().getSearchOperationListener();
try { try {
operationListener.validateSearchContext(context); operationListener.validateSearchContext(context, request);
return context; return context;
} catch (Exception e) { } catch (Exception e) {
processFailure(context, e); processFailure(context, e);

View File

@ -21,6 +21,8 @@ package org.elasticsearch.index.shard;
import org.elasticsearch.search.internal.SearchContext; import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.test.TestSearchContext; import org.elasticsearch.test.TestSearchContext;
import org.elasticsearch.transport.TransportRequest;
import org.elasticsearch.transport.TransportRequest.Empty;
import java.lang.reflect.Proxy; import java.lang.reflect.Proxy;
import java.util.ArrayList; import java.util.ArrayList;
@ -112,7 +114,7 @@ public class SearchOperationListenerTests extends ESTestCase {
} }
@Override @Override
public void validateSearchContext(SearchContext context) { public void validateSearchContext(SearchContext context, TransportRequest request) {
assertNotNull(context); assertNotNull(context);
validateSearchContext.incrementAndGet(); validateSearchContext.incrementAndGet();
} }
@ -267,9 +269,10 @@ public class SearchOperationListenerTests extends ESTestCase {
assertEquals(0, validateSearchContext.get()); assertEquals(0, validateSearchContext.get());
if (throwingListeners == 0) { if (throwingListeners == 0) {
compositeListener.validateSearchContext(ctx); compositeListener.validateSearchContext(ctx, Empty.INSTANCE);
} else { } else {
RuntimeException expected = expectThrows(RuntimeException.class, () -> compositeListener.validateSearchContext(ctx)); RuntimeException expected =
expectThrows(RuntimeException.class, () -> compositeListener.validateSearchContext(ctx, Empty.INSTANCE));
assertNull(expected.getMessage()); assertNull(expected.getMessage());
assertEquals(throwingListeners - 1, expected.getSuppressed().length); assertEquals(throwingListeners - 1, expected.getSuppressed().length);
if (throwingListeners > 1) { if (throwingListeners > 1) {