diff --git a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 6dec836aa02ca..749c5c4b8e2ca 100644 --- a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/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 TRACKED_RESOURCES = List.of(ResourceType.fromName("memory"), ResourceType.fromName("cpu")); + public static final List 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 constructQueryGroupLevelUsageViews() { - Map queryGroupViews = new HashMap<>(); + final Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); + final Map queryGroupViews = new HashMap<>(); - Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); - Map> queryGroupResourceUsage = getResourceUsageOfQueryGroups(tasksByQueryGroup); + // Iterate over each QueryGroup entry + for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { + // Compute the QueryGroup usage + final EnumMap 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> 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> getResourceUsageOfQueryGroups(Map> tasksByQueryGroup) { - // Prepare a usage map for the QueryGroups - Map> resourceUsageOfQueryGroups = tasksByQueryGroup.keySet() - .stream() - .collect(Collectors.toMap(Function.identity(), k -> new EnumMap<>(ResourceType.class))); - - // Iterate over each QueryGroup entry - for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { - EnumMap 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. *