From 668acd3481d881e6050ab04e9d063755ae6aff12 Mon Sep 17 00:00:00 2001 From: jaymode Date: Wed, 31 May 2017 10:34:52 -0600 Subject: [PATCH] Provide the TransportRequest during validation of a search context 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. --- .../index/shard/SearchOperationListener.java | 8 +++++--- .../org/elasticsearch/search/SearchService.java | 13 +++++++------ .../index/shard/SearchOperationListenerTests.java | 9 ++++++--- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/core/src/main/java/org/elasticsearch/index/shard/SearchOperationListener.java b/core/src/main/java/org/elasticsearch/index/shard/SearchOperationListener.java index 583bcbc561d6a..153a985ab0892 100644 --- a/core/src/main/java/org/elasticsearch/index/shard/SearchOperationListener.java +++ b/core/src/main/java/org/elasticsearch/index/shard/SearchOperationListener.java @@ -23,6 +23,7 @@ import org.apache.logging.log4j.util.Supplier; import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.search.internal.SearchContext; +import org.elasticsearch.transport.TransportRequest; import java.util.List; @@ -110,8 +111,9 @@ public interface SearchOperationListener { * from the active contexts. If the context is deemed invalid a runtime * exception can be thrown, which will prevent the context from being used. * @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. @@ -236,11 +238,11 @@ public void onFreeScrollContext(SearchContext context) { } @Override - public void validateSearchContext(SearchContext context) { + public void validateSearchContext(SearchContext context, TransportRequest request) { Exception exception = null; for (SearchOperationListener listener : listeners) { try { - listener.validateSearchContext(context); + listener.validateSearchContext(context, request); } catch (Exception e) { exception = ExceptionsHelper.useOrSuppress(exception, e); } diff --git a/core/src/main/java/org/elasticsearch/search/SearchService.java b/core/src/main/java/org/elasticsearch/search/SearchService.java index 93217819ab2a3..f960d8eb0bae4 100644 --- a/core/src/main/java/org/elasticsearch/search/SearchService.java +++ b/core/src/main/java/org/elasticsearch/search/SearchService.java @@ -85,6 +85,7 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.threadpool.ThreadPool.Cancellable; import org.elasticsearch.threadpool.ThreadPool.Names; +import org.elasticsearch.transport.TransportRequest; import java.io.IOException; import java.util.Collections; @@ -309,7 +310,7 @@ private QueryFetchSearchResult executeFetchPhase(SearchContext context, SearchOp } 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(); context.incRef(); try { @@ -333,7 +334,7 @@ public ScrollQuerySearchResult executeQueryPhase(InternalScrollSearchRequest req } public QuerySearchResult executeQueryPhase(QuerySearchRequest request, SearchTask task) { - final SearchContext context = findContext(request.id()); + final SearchContext context = findContext(request.id(), request); context.setTask(task); IndexShard indexShard = context.indexShard(); SearchOperationListener operationListener = indexShard.getSearchOperationListener(); @@ -374,7 +375,7 @@ private boolean fetchPhaseShouldFreeContext(SearchContext context) { } public ScrollQueryFetchSearchResult executeFetchPhase(InternalScrollSearchRequest request, SearchTask task) { - final SearchContext context = findContext(request.id()); + final SearchContext context = findContext(request.id(), request); context.incRef(); try { context.setTask(task); @@ -405,7 +406,7 @@ public ScrollQueryFetchSearchResult executeFetchPhase(InternalScrollSearchReques } 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(); context.incRef(); try { @@ -435,7 +436,7 @@ public FetchSearchResult executeFetchPhase(ShardFetchRequest request, SearchTask } } - private SearchContext findContext(long id) throws SearchContextMissingException { + private SearchContext findContext(long id, TransportRequest request) throws SearchContextMissingException { SearchContext context = activeContexts.get(id); if (context == null) { throw new SearchContextMissingException(id); @@ -443,7 +444,7 @@ private SearchContext findContext(long id) throws SearchContextMissingException SearchOperationListener operationListener = context.indexShard().getSearchOperationListener(); try { - operationListener.validateSearchContext(context); + operationListener.validateSearchContext(context, request); return context; } catch (Exception e) { processFailure(context, e); diff --git a/core/src/test/java/org/elasticsearch/index/shard/SearchOperationListenerTests.java b/core/src/test/java/org/elasticsearch/index/shard/SearchOperationListenerTests.java index fafdbe6755bc3..28bab8da0fdfb 100644 --- a/core/src/test/java/org/elasticsearch/index/shard/SearchOperationListenerTests.java +++ b/core/src/test/java/org/elasticsearch/index/shard/SearchOperationListenerTests.java @@ -21,6 +21,8 @@ import org.elasticsearch.search.internal.SearchContext; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.TestSearchContext; +import org.elasticsearch.transport.TransportRequest; +import org.elasticsearch.transport.TransportRequest.Empty; import java.lang.reflect.Proxy; import java.util.ArrayList; @@ -112,7 +114,7 @@ public void onFreeScrollContext(SearchContext context) { } @Override - public void validateSearchContext(SearchContext context) { + public void validateSearchContext(SearchContext context, TransportRequest request) { assertNotNull(context); validateSearchContext.incrementAndGet(); } @@ -267,9 +269,10 @@ public void validateSearchContext(SearchContext context) { assertEquals(0, validateSearchContext.get()); if (throwingListeners == 0) { - compositeListener.validateSearchContext(ctx); + compositeListener.validateSearchContext(ctx, Empty.INSTANCE); } else { - RuntimeException expected = expectThrows(RuntimeException.class, () -> compositeListener.validateSearchContext(ctx)); + RuntimeException expected = + expectThrows(RuntimeException.class, () -> compositeListener.validateSearchContext(ctx, Empty.INSTANCE)); assertNull(expected.getMessage()); assertEquals(throwingListeners - 1, expected.getSuppressed().length); if (throwingListeners > 1) {