[bug fix] fix incorrect coordinator node search resource usages

Signed-off-by: Chenyang Ji <[email protected]>
opensearch-project · Jul 9, 2024 · d1a6038 · d1a6038
1 parent 0684342
commit d1a6038
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/...hts/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/...hts/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java
@@ -145,6 +145,9 @@ public void onRequestFailure(final SearchPhaseContext context, final SearchReque
     private void constructSearchQueryRecord(final SearchPhaseContext context, final SearchRequestContext searchRequestContext) {
         SearchTask searchTask = context.getTask();
         List<TaskResourceInfo> tasksResourceUsages = searchRequestContext.getPhaseResourceUsage();
+        if (clusterService.getTaskResourceTrackingService() != null) {
+            clusterService.getTaskResourceTrackingService().refreshResourceStats(searchTask);
+        }
         tasksResourceUsages.add(
             new TaskResourceInfo(
                 searchTask.getAction(),

diff --git a/server/src/main/java/org/opensearch/cluster/service/ClusterService.java b/server/src/main/java/org/opensearch/cluster/service/ClusterService.java
@@ -54,6 +54,7 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.index.IndexingPressureService;
 import org.opensearch.node.Node;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.telemetry.metrics.noop.NoopMetricsRegistry;
 import org.opensearch.threadpool.ThreadPool;
 
@@ -92,6 +93,7 @@ public class ClusterService extends AbstractLifecycleComponent {
     private RerouteService rerouteService;
 
     private IndexingPressureService indexingPressureService;
+    private TaskResourceTrackingService taskResourceTrackingService;
 
     public ClusterService(Settings settings, ClusterSettings clusterSettings, ThreadPool threadPool) {
         this(settings, clusterSettings, threadPool, new ClusterManagerMetrics(NoopMetricsRegistry.INSTANCE));
@@ -265,6 +267,24 @@ public IndexingPressureService getIndexingPressureService() {
         return indexingPressureService;
     }
 
+    /**
+     * Getter for {@link TaskResourceTrackingService}, This method exposes task level resource usage for other components to use.
+     *
+     * @return TaskResourceTrackingService
+     */
+    public TaskResourceTrackingService getTaskResourceTrackingService() {
+        return taskResourceTrackingService;
+    }
+
+    /**
+     * Setter for {@link TaskResourceTrackingService}
+     *
+     * @param taskResourceTrackingService taskResourceTrackingService
+     */
+    public void setTaskResourceTrackingService(TaskResourceTrackingService taskResourceTrackingService) {
+        this.taskResourceTrackingService = taskResourceTrackingService;
+    }
+
     public ClusterApplierService getClusterApplierService() {
         return clusterApplierService;
     }

diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
@@ -1109,6 +1109,7 @@ protected Node(
                 clusterService.getClusterSettings(),
                 threadPool
             );
+            clusterService.setTaskResourceTrackingService(taskResourceTrackingService);
 
             final SearchBackpressureSettings searchBackpressureSettings = new SearchBackpressureSettings(
                 settings,

diff --git a/server/src/main/java/org/opensearch/tasks/TaskResourceTrackingService.java b/server/src/main/java/org/opensearch/tasks/TaskResourceTrackingService.java
@@ -16,6 +16,7 @@
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.search.SearchShardTask;
 import org.opensearch.common.SuppressForbidden;
+import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Setting;
@@ -51,6 +52,7 @@
 /**
  * Service that helps track resource usage of tasks running on a node.
  */
+@PublicApi(since = "2.15.0")
 @SuppressForbidden(reason = "ThreadMXBean#getThreadAllocatedBytes")
 public class TaskResourceTrackingService implements RunnableTaskExecutionListener {
 
@@ -357,6 +359,7 @@ public TaskResourceInfo getTaskResourceUsageFromThreadContext() {
     /**
      * Listener that gets invoked when a task execution completes.
      */
+    @PublicApi(since = "2.15.0")
     public interface TaskCompletionListener {
         void onTaskCompleted(Task task);
     }