[POC] Accurate query-level resource usages instrumentation on Cluster Manager

libs/core/src/main/java/org/opensearch/core/tasks/resourcetracker/TaskResourceUsage.java

@@ -90,7 +90,8 @@ public static TaskResourceUsage fromXContent(XContentParser parser) {
 
     @Override
     public String toString() {
-        return Strings.toString(MediaTypeRegistry.JSON, this, true, true);
+        // TODO revert after debugging
+        return Strings.toString(MediaTypeRegistry.JSON, this, false, true);
     }
 
     // Implements equals and hashcode for testing

modules/analysis-common/src/main/java/org/opensearch/analysis/common/CommonAnalysisModulePlugin.java

@@ -153,6 +153,7 @@
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.script.ScriptContext;
 import org.opensearch.script.ScriptService;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.watcher.ResourceWatcherService;
 
@@ -187,7 +188,8 @@ public Collection<Object> createComponents(
         NodeEnvironment nodeEnvironment,
         NamedWriteableRegistry namedWriteableRegistry,
         IndexNameExpressionResolver expressionResolver,
-        Supplier<RepositoriesService> repositoriesServiceSupplier
+        Supplier<RepositoriesService> repositoriesServiceSupplier,
+        TaskResourceTrackingService taskResourceTrackingService
     ) {
         this.scriptService.set(scriptService);
         return Collections.emptyList();

modules/lang-painless/src/main/java/org/opensearch/painless/PainlessModulePlugin.java

@@ -66,6 +66,7 @@
 import org.opensearch.script.ScriptEngine;
 import org.opensearch.script.ScriptService;
 import org.opensearch.search.aggregations.pipeline.MovingFunctionScript;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.watcher.ResourceWatcherService;
 
@@ -140,7 +141,8 @@ public Collection<Object> createComponents(
         NodeEnvironment nodeEnvironment,
         NamedWriteableRegistry namedWriteableRegistry,
         IndexNameExpressionResolver expressionResolver,
-        Supplier<RepositoriesService> repositoriesServiceSupplier
+        Supplier<RepositoriesService> repositoriesServiceSupplier,
+        TaskResourceTrackingService taskResourceTrackingService
     ) {
         // this is a hack to bind the painless script engine in guice (all components are added to guice), so that
         // the painless context api. this is a temporary measure until transport actions do no require guice

modules/reindex/src/main/java/org/opensearch/index/reindex/ReindexModulePlugin.java

@@ -58,6 +58,7 @@
 import org.opensearch.rest.RestHandler;
 import org.opensearch.script.ScriptService;
 import org.opensearch.tasks.Task;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.watcher.ResourceWatcherService;
 
@@ -122,7 +123,8 @@ public Collection<Object> createComponents(
         NodeEnvironment nodeEnvironment,
         NamedWriteableRegistry namedWriteableRegistry,
         IndexNameExpressionResolver expressionResolver,
-        Supplier<RepositoriesService> repositoriesServiceSupplier
+        Supplier<RepositoriesService> repositoriesServiceSupplier,
+        TaskResourceTrackingService taskResourceTrackingService
     ) {
         return Collections.singletonList(new ReindexSslConfig(environment.settings(), environment, resourceWatcherService));
     }

modules/systemd/src/main/java/org/opensearch/systemd/SystemdModulePlugin.java

@@ -47,6 +47,7 @@
 import org.opensearch.plugins.Plugin;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.script.ScriptService;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.threadpool.Scheduler;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.watcher.ResourceWatcherService;
@@ -100,7 +101,8 @@ public Collection<Object> createComponents(
         final NodeEnvironment nodeEnvironment,
         final NamedWriteableRegistry namedWriteableRegistry,
         final IndexNameExpressionResolver expressionResolver,
-        final Supplier<RepositoriesService> repositoriesServiceSupplier
+        final Supplier<RepositoriesService> repositoriesServiceSupplier,
+        TaskResourceTrackingService taskResourceTrackingService
     ) {
         if (enabled == false) {
             extender.set(null);

plugins/query-insights/src/main/java/org/opensearch/plugin/insights/QueryInsightsPlugin.java

@@ -26,9 +26,12 @@
 import org.opensearch.env.Environment;
 import org.opensearch.env.NodeEnvironment;
 import org.opensearch.plugin.insights.core.listener.QueryInsightsListener;
+import org.opensearch.plugin.insights.core.listener.ResourceTrackingListener;
 import org.opensearch.plugin.insights.core.service.QueryInsightsService;
+import org.opensearch.plugin.insights.rules.action.top_queries.SearchMetadataAction;
 import org.opensearch.plugin.insights.rules.action.top_queries.TopQueriesAction;
 import org.opensearch.plugin.insights.rules.resthandler.top_queries.RestTopQueriesAction;
+import org.opensearch.plugin.insights.rules.transport.top_queries.TransportSearchMetadataAction;
 import org.opensearch.plugin.insights.rules.transport.top_queries.TransportTopQueriesAction;
 import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
 import org.opensearch.plugins.ActionPlugin;
@@ -37,6 +40,7 @@
 import org.opensearch.rest.RestController;
 import org.opensearch.rest.RestHandler;
 import org.opensearch.script.ScriptService;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.threadpool.ExecutorBuilder;
 import org.opensearch.threadpool.ScalingExecutorBuilder;
 import org.opensearch.threadpool.ThreadPool;
@@ -67,11 +71,16 @@ public Collection<Object> createComponents(
         final NodeEnvironment nodeEnvironment,
         final NamedWriteableRegistry namedWriteableRegistry,
         final IndexNameExpressionResolver indexNameExpressionResolver,
-        final Supplier<RepositoriesService> repositoriesServiceSupplier
+        final Supplier<RepositoriesService> repositoriesServiceSupplier,
+        final TaskResourceTrackingService taskResourceTrackingService
     ) {
         // create top n queries service
-        final QueryInsightsService queryInsightsService = new QueryInsightsService(threadPool);
-        return List.of(queryInsightsService, new QueryInsightsListener(clusterService, queryInsightsService));
+        final QueryInsightsService queryInsightsService = new QueryInsightsService(threadPool, client, clusterService);
+        return List.of(
+            queryInsightsService,
+            new QueryInsightsListener(clusterService, queryInsightsService),
+            new ResourceTrackingListener(queryInsightsService, taskResourceTrackingService)
+        );
     }
 
     @Override
@@ -96,12 +105,15 @@ public List<RestHandler> getRestHandlers(
         final IndexNameExpressionResolver indexNameExpressionResolver,
         final Supplier<DiscoveryNodes> nodesInCluster
     ) {
-        return List.of(new RestTopQueriesAction());
+        return List.of(new RestTopQueriesAction(nodesInCluster));
     }
 
     @Override
     public List<ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
-        return List.of(new ActionPlugin.ActionHandler<>(TopQueriesAction.INSTANCE, TransportTopQueriesAction.class));
+        return List.of(
+            new ActionPlugin.ActionHandler<>(TopQueriesAction.INSTANCE, TransportTopQueriesAction.class),
+            new ActionPlugin.ActionHandler<>(SearchMetadataAction.INSTANCE, TransportSearchMetadataAction.class)
+        );
     }
 
     @Override

plugins/query-insights/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java

@@ -123,6 +123,10 @@ public void onRequestStart(SearchRequestContext searchRequestContext) {}
 
     @Override
     public void onRequestEnd(final SearchPhaseContext context, final SearchRequestContext searchRequestContext) {
+        long taskGroupId = context.getTask().getParentTaskId().getId();
+        if (taskGroupId == -1) {
+            taskGroupId = context.getTask().getId();
+        }
         final SearchRequest request = context.getRequest();
         try {
             Map<MetricType, Number> measurements = new HashMap<>();
@@ -138,7 +142,8 @@ public void onRequestEnd(final SearchPhaseContext context, final SearchRequestCo
             attributes.put(Attribute.TOTAL_SHARDS, context.getNumShards());
             attributes.put(Attribute.INDICES, request.indices());
             attributes.put(Attribute.PHASE_LATENCY_MAP, searchRequestContext.phaseTookMap());
-            SearchQueryRecord record = new SearchQueryRecord(request.getOrCreateAbsoluteStartMillis(), measurements, attributes);
+            attributes.put(Attribute.NODE_ID, this.queryInsightsService.clusterService.localNode().getId());
+            SearchQueryRecord record = new SearchQueryRecord(request.getOrCreateAbsoluteStartMillis(), taskGroupId, measurements, attributes);
             queryInsightsService.addRecord(record);
         } catch (Exception e) {
             log.error(String.format(Locale.ROOT, "fail to ingest query insight data, error: %s", e));

plugins/query-insights/src/main/java/org/opensearch/plugin/insights/core/listener/ResourceTrackingListener.java

@@ -0,0 +1,62 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.insights.core.listener;
+
+import org.opensearch.core.tasks.resourcetracker.ResourceStats;
+import org.opensearch.plugin.insights.core.service.QueryInsightsService;
+import org.opensearch.plugin.insights.rules.model.SearchTaskMetadata;
+import org.opensearch.tasks.Task;
+import org.opensearch.tasks.TaskResourceTrackingService;
+import org.opensearch.tasks.TaskResourceTrackingService.TaskCompletionListener;
+import org.opensearch.tasks.TaskResourceTrackingService.TaskStartListener;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class ResourceTrackingListener implements TaskCompletionListener, TaskStartListener {
+
+    private final TaskResourceTrackingService taskResourceTrackingService;
+    private final QueryInsightsService queryInsightsService;
+
+    public ResourceTrackingListener (
+        QueryInsightsService queryInsightsService,
+        TaskResourceTrackingService taskResourceTrackingService
+    ) {
+        this.queryInsightsService = queryInsightsService;
+        this.taskResourceTrackingService = taskResourceTrackingService;
+        this.taskResourceTrackingService.addTaskCompletionListener(this);
+        this.taskResourceTrackingService.addTaskStartListener(this);
+    }
+    @Override
+    public void onTaskCompleted(Task task) {
+        long taskGroupId = task.getParentTaskId().getId();
+        if (taskGroupId == -1) {
+            taskGroupId = task.getId();
+        }
+        SearchTaskMetadata info = new SearchTaskMetadata(
+            task.getAction(), task.getId(), taskGroupId, task.getTotalResourceStats()
+        );
+
+        int pendingTaskCount = this.queryInsightsService.taskStatusMap.get(taskGroupId).decrementAndGet();
+        if (pendingTaskCount == 0) {
+            this.queryInsightsService.taskStatusMap.remove(taskGroupId);
+        }
+        queryInsightsService.taskRecordsQueue.add(info);
+        System.out.println(String.format("id = %s, parent = %s, resource = %s, action = %s, total CPU and MEM: %s, %s", task.getId(), task.getParentTaskId(), task.getResourceStats(), task.getAction(),task.getTotalResourceUtilization(ResourceStats.CPU),task.getTotalResourceUtilization(ResourceStats.MEMORY) ));
+    }
+
+    @Override
+    public void onTaskStarts(Task task) {
+        long taskGroupId = task.getParentTaskId().getId();
+        if (taskGroupId == -1) {
+            taskGroupId = task.getId();
+        }
+        this.queryInsightsService.taskStatusMap.putIfAbsent(taskGroupId, new AtomicInteger(0));
+        this.queryInsightsService.taskStatusMap.get(taskGroupId).incrementAndGet();
+    }
+}