address comments

Signed-off-by: Kaushal Kumar <[email protected]>
kaushalmahi12 · Aug 30, 2024 · d89422b · d89422b
1 parent fd8f3f5
commit d89422b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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();