Update QueryGroupResourceUsageTrackerService.java

Signed-off-by: Kiran Prakash <[email protected]>
kiranprakash154 · Jul 31, 2024 · 4860e7e · 4860e7e
1 parent 371d3c8
commit 4860e7e
Showing 1 changed file with 18 additions and 38 deletions.
diff --git a/...rc/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/...rc/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java
@@ -20,7 +20,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -29,7 +28,7 @@
 // @ExperimentalApi
 public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTracker, TaskManager.TaskEventListeners {
 
-    public static final List<ResourceType> TRACKED_RESOURCES = List.of(ResourceType.fromName("memory"), ResourceType.fromName("cpu"));
+    public static final List<ResourceType> TRACKED_RESOURCES = List.of(ResourceType.MEMORY, ResourceType.CPU);
     private final TaskManager taskManager;
     private final TaskResourceTrackingService taskResourceTrackingService;
 
@@ -54,18 +53,26 @@ public QueryGroupResourceUsageTrackerService(
      */
     @Override
     public Map<String, QueryGroupLevelResourceUsageView> constructQueryGroupLevelUsageViews() {
-        Map<String, QueryGroupLevelResourceUsageView> queryGroupViews = new HashMap<>();
+        final Map<String, List<Task>> tasksByQueryGroup = getTasksGroupedByQueryGroup();
+        final Map<String, QueryGroupLevelResourceUsageView> queryGroupViews = new HashMap<>();
 
-        Map<String, List<Task>> tasksByQueryGroup = getTasksGroupedByQueryGroup();
-        Map<String, EnumMap<ResourceType, Long>> queryGroupResourceUsage = getResourceUsageOfQueryGroups(tasksByQueryGroup);
+        // Iterate over each QueryGroup entry
+        for (Map.Entry<String, List<Task>> queryGroupEntry : tasksByQueryGroup.entrySet()) {
+            // Compute the QueryGroup usage
+            final EnumMap<ResourceType, Long> queryGroupUsage = new EnumMap<>(ResourceType.class);
+            for (ResourceType resourceType : TRACKED_RESOURCES) {
+                long queryGroupResourceUsage = 0;
+                for (Task task : queryGroupEntry.getValue()) {
+                    queryGroupResourceUsage += QueryGroupHelper.getResourceUsage(resourceType, task);
+                }
+                queryGroupUsage.put(resourceType, queryGroupResourceUsage);
+            }
 
-        for (String queryGroupId : tasksByQueryGroup.keySet()) {
-            QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView(
-                queryGroupId,
-                queryGroupResourceUsage.get(queryGroupId),
-                tasksByQueryGroup.get(queryGroupId)
+            // Add to the QueryGroup View
+            queryGroupViews.put(
+                queryGroupEntry.getKey(),
+                new QueryGroupLevelResourceUsageView(queryGroupEntry.getKey(), queryGroupUsage, queryGroupEntry.getValue())
             );
-            queryGroupViews.put(queryGroupId, queryGroupLevelResourceUsageView);
         }
         return queryGroupViews;
     }
@@ -84,33 +91,6 @@ private Map<String, List<Task>> getTasksGroupedByQueryGroup() {
             .collect(Collectors.groupingBy(QueryGroupTask::getQueryGroupId, Collectors.mapping(task -> (Task) task, Collectors.toList())));
     }
 
-    /**
-     * Calculates the resource usage of each QueryGroup.
-     *
-     * @param tasksByQueryGroup Map of tasks grouped by QueryGroup
-     * @return Map of resource usage for each QueryGroup
-     */
-    private Map<String, EnumMap<ResourceType, Long>> getResourceUsageOfQueryGroups(Map<String, List<Task>> tasksByQueryGroup) {
-        // Prepare a usage map for the QueryGroups
-        Map<String, EnumMap<ResourceType, Long>> resourceUsageOfQueryGroups = tasksByQueryGroup.keySet()
-            .stream()
-            .collect(Collectors.toMap(Function.identity(), k -> new EnumMap<>(ResourceType.class)));
-
-        // Iterate over each QueryGroup entry
-        for (Map.Entry<String, List<Task>> queryGroupEntry : tasksByQueryGroup.entrySet()) {
-            EnumMap<ResourceType, Long> queryGroupUsage = resourceUsageOfQueryGroups.get(queryGroupEntry.getKey());
-            for (ResourceType resourceType : TRACKED_RESOURCES) {
-                long queryGroupResourceUsage = 0;
-                for (Task task : queryGroupEntry.getValue()) {
-                    queryGroupResourceUsage += QueryGroupHelper.getResourceUsage(resourceType, task);
-                }
-                queryGroupUsage.put(resourceType, queryGroupResourceUsage);
-            }
-        }
-
-        return resourceUsageOfQueryGroups;
-    }
-
     /**
      * Handles the completion of a task.
      *