revise

Signed-off-by: Ruirui Zhang <[email protected]>
ruai0511 · Oct 4, 2024 · 736dcfa · 736dcfa
1 parent 27d7ea0
commit 736dcfa
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java
@@ -29,7 +29,7 @@
  */
 public class TransportWlmStatsAction extends TransportNodesAction<WlmStatsRequest, WlmStatsResponse, WlmStatsRequest, QueryGroupStats> {
 
-    QueryGroupService queryGroupService;
+    final QueryGroupService queryGroupService;
 
     @Inject
     public TransportWlmStatsAction(

diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsAction.java
@@ -17,7 +17,7 @@
  */
 public class WlmStatsAction extends ActionType<WlmStatsResponse> {
     public static final WlmStatsAction INSTANCE = new WlmStatsAction();
-    public static final String NAME = "cluster:monitor/query_group_stats";
+    public static final String NAME = "cluster:monitor/wlm/stats";
 
     private WlmStatsAction() {
         super(NAME, WlmStatsResponse::new);

diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.wlm;
 
+import org.opensearch.ResourceNotFoundException;
 import org.opensearch.cluster.metadata.QueryGroup;
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.cluster.service.ClusterService;
@@ -63,7 +64,13 @@ public void incrementFailuresFor(final String queryGroupId) {
      */
     public QueryGroupStats nodeStats(Set<String> queryGroupIds, Boolean requestedBreached) {
         final Map<String, QueryGroupStatsHolder> statsHolderMap = new HashMap<>();
-
+        if (!queryGroupIds.contains("_all")) {
+            for (String id : queryGroupIds) {
+                if (!queryGroupStateMap.containsKey(id)) {
+                    throw new ResourceNotFoundException("QueryGroup with id " + id + " does not exist");
+                }
+            }
+        }
         queryGroupStateMap.forEach((queryGroupId, currentState) -> {
             boolean shouldInclude = (queryGroupIds.size() == 1 && queryGroupIds.contains("_all")) || queryGroupIds.contains(queryGroupId);