From 974ee3e753c8606aefc83df04732c93ad3980ebd Mon Sep 17 00:00:00 2001
From: Siddhant Deshmukh <deshsid@amazon.com>
Date: Mon, 22 Jul 2024 18:39:36 -0700
Subject: [PATCH] Make sure listener runs before plugin listeners

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
---
 .../SearchTaskRequestOperationsListener.java     |  4 ++--
 .../action/search/TransportSearchAction.java     |  5 +----
 .../src/main/java/org/opensearch/node/Node.java  | 16 +++++++++-------
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java b/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java
index 1a8ca2c5c24df..4434d71793b23 100644
--- a/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java
+++ b/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java
@@ -16,10 +16,10 @@
  * upon request completion.
  *
  */
-final class SearchTaskRequestOperationsListener extends SearchRequestOperationsListener {
+public final class SearchTaskRequestOperationsListener extends SearchRequestOperationsListener {
     private final TaskResourceTrackingService taskResourceTrackingService;
 
-    SearchTaskRequestOperationsListener(TaskResourceTrackingService taskResourceTrackingService) {
+    public SearchTaskRequestOperationsListener(TaskResourceTrackingService taskResourceTrackingService) {
         this.taskResourceTrackingService = taskResourceTrackingService;
     }
 
diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
index d97d50fc3b7fe..7d3237d43cd5c 100644
--- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
+++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -174,7 +174,6 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
     private final SearchPipelineService searchPipelineService;
     private final SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory;
     private final Tracer tracer;
-    private final SearchTaskRequestOperationsListener searchTaskRequestOperationsListener;
 
     private final MetricsRegistry metricsRegistry;
 
@@ -216,7 +215,6 @@ public TransportSearchAction(
         this.searchRequestOperationsCompositeListenerFactory = searchRequestOperationsCompositeListenerFactory;
         this.tracer = tracer;
         this.taskResourceTrackingService = taskResourceTrackingService;
-        this.searchTaskRequestOperationsListener = new SearchTaskRequestOperationsListener(taskResourceTrackingService);
     }
 
     private Map<String, AliasFilter> buildPerIndexAliasFilter(
@@ -435,8 +433,7 @@ private void executeRequest(
             requestOperationsListeners = searchRequestOperationsCompositeListenerFactory.buildCompositeListener(
                 originalSearchRequest,
                 logger,
-                TraceableSearchRequestOperationsListener.create(tracer, requestSpan),
-                searchTaskRequestOperationsListener
+                TraceableSearchRequestOperationsListener.create(tracer, requestSpan)
             );
             SearchRequestContext searchRequestContext = new SearchRequestContext(
                 requestOperationsListeners,
diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
index d91b2a45a48c6..04faae00a6638 100644
--- a/server/src/main/java/org/opensearch/node/Node.java
+++ b/server/src/main/java/org/opensearch/node/Node.java
@@ -52,6 +52,7 @@
 import org.opensearch.action.search.SearchRequestOperationsListener;
 import org.opensearch.action.search.SearchRequestSlowLog;
 import org.opensearch.action.search.SearchRequestStats;
+import org.opensearch.action.search.SearchTaskRequestOperationsListener;
 import org.opensearch.action.search.SearchTransportService;
 import org.opensearch.action.support.TransportAction;
 import org.opensearch.action.update.UpdateHelper;
@@ -855,8 +856,15 @@ protected Node(
                 threadPool
             );
 
+            final TaskResourceTrackingService taskResourceTrackingService = new TaskResourceTrackingService(
+                settings,
+                clusterService.getClusterSettings(),
+                threadPool
+            );
+
             final SearchRequestStats searchRequestStats = new SearchRequestStats(clusterService.getClusterSettings());
             final SearchRequestSlowLog searchRequestSlowLog = new SearchRequestSlowLog(clusterService);
+            final SearchTaskRequestOperationsListener searchTaskRequestOperationsListener = new SearchTaskRequestOperationsListener(taskResourceTrackingService);
 
             remoteStoreStatsTrackerFactory = new RemoteStoreStatsTrackerFactory(clusterService, settings);
             CacheModule cacheModule = new CacheModule(pluginsService.filterPlugins(CachePlugin.class), settings);
@@ -988,7 +996,7 @@ protected Node(
             final SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory =
                 new SearchRequestOperationsCompositeListenerFactory(
                     Stream.concat(
-                        Stream.of(searchRequestStats, searchRequestSlowLog),
+                        Stream.of(searchRequestStats, searchRequestSlowLog, searchTaskRequestOperationsListener),
                         pluginComponents.stream()
                             .filter(p -> p instanceof SearchRequestOperationsListener)
                             .map(p -> (SearchRequestOperationsListener) p)
@@ -1117,12 +1125,6 @@ protected Node(
             // development. Then we can deprecate Getter and Setter for IndexingPressureService in ClusterService (#478).
             clusterService.setIndexingPressureService(indexingPressureService);
 
-            final TaskResourceTrackingService taskResourceTrackingService = new TaskResourceTrackingService(
-                settings,
-                clusterService.getClusterSettings(),
-                threadPool
-            );
-
             final SearchBackpressureSettings searchBackpressureSettings = new SearchBackpressureSettings(
                 settings,
                 clusterService.getClusterSettings()