include SearchRequest in SearchContext and rename to SearchRequestOpe…

…rationsCompositeListenerFactory Signed-off-by: Chenyang Ji <cyji@amazon.com>
ansjcy · Jan 11, 2024 · 96c8388 · 96c8388
1 parent bb64cd2
commit 96c8388
Show file tree

Hide file tree

Showing 17 changed files with 257 additions and 100 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -193,7 +193,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add deleted doc count in _cat/shards ([#11678](https://github.com/opensearch-project/OpenSearch/pull/11678))
 - Capture information for additional query types and aggregation types ([#11582](https://github.com/opensearch-project/OpenSearch/pull/11582))
 - Use slice_size == shard_size heuristic in terms aggs for concurrent segment search and properly calculate the doc_count_error ([#11732](https://github.com/opensearch-project/OpenSearch/pull/11732))
-- Added Support for dynamically adding SearchRequestOperationsListeners ([#11526](https://github.com/opensearch-project/OpenSearch/pull/11526))
+- Added Support for dynamically adding SearchRequestOperationsListeners with SearchRequestOperationsCompositeListenerFactory ([#11526](https://github.com/opensearch-project/OpenSearch/pull/11526))
 
 ### Deprecated
 

diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
@@ -8,13 +8,11 @@
 
 package org.opensearch.action.search;
 
-import org.apache.logging.log4j.LogManager;
 import org.apache.lucene.search.TotalHits;
 import org.opensearch.common.annotation.InternalApi;
 
 import java.util.EnumMap;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 
@@ -31,25 +29,14 @@ class SearchRequestContext {
     private TotalHits totalHits;
     private final EnumMap<ShardStatsFieldNames, Integer> shardStats;
 
-    private final boolean phaseTookEnabled;
+    private final SearchRequest searchRequest;
 
-    /**
-     * This constructor is for testing only
-     */
-    SearchRequestContext() {
-        this(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), false);
-    }
-
-    SearchRequestContext(SearchRequestOperationsListener searchRequestOperationsListener, boolean phaseTookEnabled) {
+    SearchRequestContext(final SearchRequestOperationsListener searchRequestOperationsListener, final SearchRequest searchRequest) {
         this.searchRequestOperationsListener = searchRequestOperationsListener;
         this.absoluteStartNanos = System.nanoTime();
         this.phaseTookMap = new HashMap<>();
         this.shardStats = new EnumMap<>(ShardStatsFieldNames.class);
-        this.phaseTookEnabled = phaseTookEnabled;
-    }
-
-    SearchRequestContext(SearchRequestOperationsListener searchRequestOperationsListener) {
-        this(searchRequestOperationsListener, false);
+        this.searchRequest = searchRequest;
     }
 
     SearchRequestOperationsListener getSearchRequestOperationsListener() {
@@ -65,7 +52,7 @@ Map<String, Long> phaseTookMap() {
     }
 
     SearchResponse.PhaseTook getPhaseTook() {
-        if (phaseTookEnabled) {
+        if (searchRequest != null && searchRequest.isPhaseTook() != null && searchRequest.isPhaseTook()) {
             return new SearchResponse.PhaseTook(phaseTookMap);
         } else {
             return null;

diff --git a/...rch/SearchRequestOperationsListeners.java → ...stOperationsCompositeListenerFactory.java b/...rch/SearchRequestOperationsListeners.java → ...stOperationsCompositeListenerFactory.java
@@ -17,25 +17,25 @@
 import java.util.stream.Stream;
 
 /**
- * SearchRequestOperationsListeners contains listeners registered to search requests,
+ * SearchRequestOperationsCompositeListenerFactory contains listeners registered to search requests,
  * and is responsible for creating the {@link SearchRequestOperationsListener.CompositeListener}
  * with the all listeners enabled at cluster-level and request-level.
  *
  *
  * @opensearch.internal
  */
-public class SearchRequestOperationsListeners {
+public final class SearchRequestOperationsCompositeListenerFactory {
     private final List<SearchRequestOperationsListener> searchRequestListenersList;
 
     /**
-     * Create the SearchRequestOperationsListeners and add multiple {@link SearchRequestOperationsListener}
+     * Create the SearchRequestOperationsCompositeListenerFactory and add multiple {@link SearchRequestOperationsListener}
      * to the searchRequestListenersList.
      * Those enabled listeners will be executed during each search request.
      *
      * @param listeners Multiple SearchRequestOperationsListener object to add.
      * @throws IllegalArgumentException if any input listener is null.
      */
-    public SearchRequestOperationsListeners(SearchRequestOperationsListener... listeners) {
+    public SearchRequestOperationsCompositeListenerFactory(final SearchRequestOperationsListener... listeners) {
         searchRequestListenersList = new ArrayList<>();
         for (SearchRequestOperationsListener listener : listeners) {
             if (listener == null) {
@@ -49,7 +49,6 @@ public SearchRequestOperationsListeners(SearchRequestOperationsListener... liste
      * Get searchRequestListenersList,
      *
      * @return List of SearchRequestOperationsListener
-     * @throws IllegalArgumentException if the input listener is null or already exists in the list.
      */
     public List<SearchRequestOperationsListener> getListeners() {
         return searchRequestListenersList;
@@ -65,9 +64,9 @@ public List<SearchRequestOperationsListener> getListeners() {
      * @return SearchRequestOperationsListener.CompositeListener
      */
     public SearchRequestOperationsListener.CompositeListener buildCompositeListener(
-        SearchRequest searchRequest,
-        Logger logger,
-        SearchRequestOperationsListener... perRequestListeners
+        final SearchRequest searchRequest,
+        final Logger logger,
+        final SearchRequestOperationsListener... perRequestListeners
     ) {
         final List<SearchRequestOperationsListener> searchListenersList = Stream.concat(
             searchRequestListenersList.stream(),

diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java b/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java
@@ -27,7 +27,7 @@ protected SearchRequestOperationsListener() {
         this.enabled = true;
     }
 
-    protected SearchRequestOperationsListener(boolean enabled) {
+    protected SearchRequestOperationsListener(final boolean enabled) {
         this.enabled = enabled;
     }
 
@@ -49,7 +49,7 @@ boolean isEnabled() {
         return enabled;
     }
 
-    protected void setEnabled(boolean enabled) {
+    protected void setEnabled(final boolean enabled) {
         this.enabled = enabled;
     }
 

diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -172,7 +172,7 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
     private final NamedWriteableRegistry namedWriteableRegistry;
     private final CircuitBreaker circuitBreaker;
     private final SearchPipelineService searchPipelineService;
-    private final SearchRequestOperationsListeners searchRequestOperationsListeners;
+    private final SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory;
 
     private volatile boolean searchQueryMetricsEnabled;
 
@@ -195,7 +195,7 @@ public TransportSearchAction(
         NamedWriteableRegistry namedWriteableRegistry,
         SearchPipelineService searchPipelineService,
         MetricsRegistry metricsRegistry,
-        SearchRequestOperationsListeners searchRequestOperationsListeners
+        SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory
     ) {
         super(SearchAction.NAME, transportService, actionFilters, (Writeable.Reader<SearchRequest>) SearchRequest::new);
         this.client = client;
@@ -212,7 +212,7 @@ public TransportSearchAction(
         this.searchPipelineService = searchPipelineService;
         this.metricsRegistry = metricsRegistry;
         this.searchQueryMetricsEnabled = clusterService.getClusterSettings().get(SEARCH_QUERY_METRICS_ENABLED_SETTING);
-        this.searchRequestOperationsListeners = searchRequestOperationsListeners;
+        this.searchRequestOperationsCompositeListenerFactory = searchRequestOperationsCompositeListenerFactory;
         clusterService.getClusterSettings()
             .addSettingsUpdateConsumer(SEARCH_QUERY_METRICS_ENABLED_SETTING, this::setSearchQueryMetricsEnabled);
     }
@@ -428,15 +428,12 @@ private void executeRequest(
             relativeStartNanos,
             System::nanoTime
         );
-        final boolean phaseTookEnabled;
         if (originalSearchRequest.isPhaseTook() == null) {
-            phaseTookEnabled = clusterService.getClusterSettings().get(SEARCH_PHASE_TOOK_ENABLED);
-        } else {
-            phaseTookEnabled = originalSearchRequest.isPhaseTook();
+            originalSearchRequest.setPhaseTook(clusterService.getClusterSettings().get(SEARCH_PHASE_TOOK_ENABLED));
         }
-        SearchRequestOperationsListener.CompositeListener requestOperationsListeners = searchRequestOperationsListeners
+        SearchRequestOperationsListener.CompositeListener requestOperationsListeners = searchRequestOperationsCompositeListenerFactory
             .buildCompositeListener(originalSearchRequest, logger);
-        SearchRequestContext searchRequestContext = new SearchRequestContext(requestOperationsListeners, phaseTookEnabled);
+        SearchRequestContext searchRequestContext = new SearchRequestContext(requestOperationsListeners, originalSearchRequest);
         searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
 
         PipelinedRequest searchRequest;

diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
@@ -46,8 +46,8 @@
 import org.opensearch.action.admin.cluster.snapshots.status.TransportNodesSnapshotsStatus;
 import org.opensearch.action.search.SearchExecutionStatsCollector;
 import org.opensearch.action.search.SearchPhaseController;
+import org.opensearch.action.search.SearchRequestOperationsCompositeListenerFactory;
 import org.opensearch.action.search.SearchRequestOperationsListener;
-import org.opensearch.action.search.SearchRequestOperationsListeners;
 import org.opensearch.action.search.SearchRequestSlowLog;
 import org.opensearch.action.search.SearchRequestStats;
 import org.opensearch.action.search.SearchTransportService;
@@ -881,15 +881,16 @@ protected Node(
                 )
                 .collect(Collectors.toList());
 
-            // register all standard SearchRequestOperationsListeners to the SearchRequestOperationsListeners
-            final SearchRequestOperationsListeners searchRequestOperationsListeners = new SearchRequestOperationsListeners(
-                Stream.concat(
-                    Stream.of(searchRequestStats, searchRequestSlowLog),
-                    pluginComponents.stream()
-                        .filter(p -> p instanceof SearchRequestOperationsListener)
-                        .map(p -> (SearchRequestOperationsListener) p)
-                ).toArray(SearchRequestOperationsListener[]::new)
-            );
+            // register all standard SearchRequestOperationsCompositeListenerFactory to the SearchRequestOperationsCompositeListenerFactory
+            final SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory =
+                new SearchRequestOperationsCompositeListenerFactory(
+                    Stream.concat(
+                        Stream.of(searchRequestStats, searchRequestSlowLog),
+                        pluginComponents.stream()
+                            .filter(p -> p instanceof SearchRequestOperationsListener)
+                            .map(p -> (SearchRequestOperationsListener) p)
+                    ).toArray(SearchRequestOperationsListener[]::new)
+                );
 
             ActionModule actionModule = new ActionModule(
                 settings,
@@ -1287,7 +1288,7 @@ protected Node(
                 b.bind(RemoteClusterStateService.class).toProvider(() -> remoteClusterStateService);
                 b.bind(PersistedStateRegistry.class).toInstance(persistedStateRegistry);
                 b.bind(SegmentReplicationStatsTracker.class).toInstance(segmentReplicationStatsTracker);
-                b.bind(SearchRequestOperationsListeners.class).toInstance(searchRequestOperationsListeners);
+                b.bind(SearchRequestOperationsCompositeListenerFactory.class).toInstance(searchRequestOperationsCompositeListenerFactory);
             });
             injector = modules.createInjector();
 

diff --git a/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.action.search;
 
+import org.apache.logging.log4j.LogManager;
 import org.opensearch.action.OriginalIndices;
 import org.opensearch.action.support.IndicesOptions;
 import org.opensearch.cluster.ClusterState;
@@ -177,7 +178,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
             results,
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext()
+            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), request)
         ) {
             @Override
             protected SearchPhase getNextPhase(final SearchPhaseResults<SearchPhaseResult> results, SearchPhaseContext context) {
@@ -715,7 +716,10 @@ private SearchDfsQueryThenFetchAsyncAction createSearchDfsQueryThenFetchAsyncAct
             null,
             task,
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger))
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger),
+                searchRequest
+            )
         );
     }
 
@@ -765,7 +769,10 @@ private SearchQueryThenFetchAsyncAction createSearchQueryThenFetchAsyncAction(
             null,
             task,
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger))
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger),
+                searchRequest
+            )
         ) {
             @Override
             ShardSearchFailure[] buildShardFailures() {

diff --git a/server/src/test/java/org/opensearch/action/search/CanMatchPreFilterSearchPhaseTests.java b/server/src/test/java/org/opensearch/action/search/CanMatchPreFilterSearchPhaseTests.java
@@ -31,6 +31,7 @@
 
 package org.opensearch.action.search;
 
+import org.apache.logging.log4j.LogManager;
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.Version;
 import org.opensearch.action.OriginalIndices;
@@ -137,7 +138,10 @@ public void run() throws IOException {
                 }
             },
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext()
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                searchRequest
+            )
         );
 
         canMatchPhase.start();
@@ -229,7 +233,10 @@ public void run() throws IOException {
                 }
             },
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext()
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                searchRequest
+            )
         );
 
         canMatchPhase.start();
@@ -320,7 +327,10 @@ public void sendCanMatch(
                 new ArraySearchPhaseResults<>(iter.size()),
                 randomIntBetween(1, 32),
                 SearchResponse.Clusters.EMPTY,
-                new SearchRequestContext()
+                new SearchRequestContext(
+                    new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                    searchRequest
+                )
             ) {
 
                 @Override
@@ -348,7 +358,10 @@ protected void executePhaseOnShard(
                 }
             },
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext()
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                searchRequest
+            )
         );
 
         canMatchPhase.start();
@@ -433,7 +446,10 @@ public void run() {
                     }
                 },
                 SearchResponse.Clusters.EMPTY,
-                new SearchRequestContext()
+                new SearchRequestContext(
+                    new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                    searchRequest
+                )
             );
 
             canMatchPhase.start();
@@ -533,7 +549,10 @@ public void run() {
                     }
                 },
                 SearchResponse.Clusters.EMPTY,
-                new SearchRequestContext()
+                new SearchRequestContext(
+                    new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                    searchRequest
+                )
             );
 
             canMatchPhase.start();