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.

Author: jaymode

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);
                 }

core/src/main/java/org/elasticsearch/search/SearchService.java

@@ -86,6 +86,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;
@@ -310,7 +311,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 {
@@ -334,7 +335,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();
@@ -375,7 +376,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);
@@ -406,7 +407,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 {
@@ -436,15 +437,15 @@ 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);
         }
 
         SearchOperationListener operationListener = context.indexShard().getSearchOperationListener();
         try {
-            operationListener.validateSearchContext(context);
+            operationListener.validateSearchContext(context, request);
             return context;
         } catch (Exception e) {
             processFailure(context, e);

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) {