From d89422bff6206a5654370252c07cd9b2ec2eab31 Mon Sep 17 00:00:00 2001
From: Kaushal Kumar <ravi.kaushal97@gmail.com>
Date: Fri, 30 Aug 2024 11:40:56 -0700
Subject: [PATCH] address comments

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
---
 server/src/main/java/org/opensearch/node/Node.java       | 6 ++++--
 .../src/main/java/org/opensearch/wlm/ResourceType.java   | 7 +++++++
 .../java/org/opensearch/wlm/stats/QueryGroupStats.java   | 9 +++------
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
index e3fb3718e63a2..08cf73b28bee8 100644
--- a/server/src/main/java/org/opensearch/node/Node.java
+++ b/server/src/main/java/org/opensearch/node/Node.java
@@ -1021,8 +1021,10 @@ protected Node(
 
             final QueryGroupService queryGroupService = new QueryGroupService(); // We will need to replace this with actual instance of the
                                                                                  // queryGroupService
-            final QueryGroupRequestOperationListener queryGroupRequestRejectionListener =
-                new QueryGroupRequestOperationListener(queryGroupService, threadPool);
+            final QueryGroupRequestOperationListener queryGroupRequestRejectionListener = new QueryGroupRequestOperationListener(
+                queryGroupService,
+                threadPool
+            );
 
             // register all standard SearchRequestOperationsCompositeListenerFactory to the SearchRequestOperationsCompositeListenerFactory
             final SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory =
diff --git a/server/src/main/java/org/opensearch/wlm/ResourceType.java b/server/src/main/java/org/opensearch/wlm/ResourceType.java
index c3f48f5f793ce..2e8da4f57f36c 100644
--- a/server/src/main/java/org/opensearch/wlm/ResourceType.java
+++ b/server/src/main/java/org/opensearch/wlm/ResourceType.java
@@ -14,6 +14,7 @@
 import org.opensearch.tasks.Task;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.function.Function;
 
 /**
@@ -30,6 +31,8 @@ public enum ResourceType {
     private final Function<Task, Long> getResourceUsage;
     private final boolean statsEnabled;
 
+    private static List<ResourceType> sortedValues = List.of(CPU, MEMORY);
+
     ResourceType(String name, Function<Task, Long> getResourceUsage, boolean statsEnabled) {
         this.name = name;
         this.getResourceUsage = getResourceUsage;
@@ -71,4 +74,8 @@ public long getResourceUsage(Task task) {
     public boolean hasStatsEnabled() {
         return statsEnabled;
     }
+
+    public static List<ResourceType> getSortedValues() {
+        return sortedValues;
+    }
 }
diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java
index 99cbc9348821a..9fc7039cd1852 100644
--- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java
+++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java
@@ -170,12 +170,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(FAILURES, failures);
             builder.field(TOTAL_CANCELLATIONS, totalCancellations);
 
-            List<Map.Entry<ResourceType, ResourceStats>> entryList = new ArrayList<>(resourceStats.entrySet());
-            entryList.sort((k1, k2) -> k1.getKey().compareTo(k2.getKey()));
-
-            for (Map.Entry<ResourceType, ResourceStats> resourceStat : entryList) {
-                ResourceType resourceType = resourceStat.getKey();
-                ResourceStats resourceStats1 = resourceStat.getValue();
+            for (ResourceType resourceType : ResourceType.getSortedValues()) {
+                ResourceStats resourceStats1 = resourceStats.get(resourceType);
+                if (resourceStats1 == null) continue;
                 builder.startObject(resourceType.getName());
                 resourceStats1.toXContent(builder, params);
                 builder.endObject();