rebase

Signed-off-by: Ruirui Zhang <[email protected]>
ruai0511 · Oct 4, 2024 · 362e7e7 · 362e7e7
1 parent 429569b
commit 362e7e7
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
@@ -1135,8 +1135,13 @@ protected Node(
                 tracer
             );
 
-            final QueryGroupService queryGroupService = new QueryGroupService(transportService.getLocalNode(), clusterService); // We will need to replace
-                                                                                                                                // this with actual instance of the
+            final QueryGroupService queryGroupService = new QueryGroupService(transportService.getLocalNode(), clusterService); // We will
+                                                                                                                                // need to
+                                                                                                                                // replace
+                                                                                                                                // this with
+                                                                                                                                // actual
+                                                                                                                                // instance
+                                                                                                                                // of the
                                                                                                                                 // queryGroupService
             queryGroupServiceSetOnce.set(queryGroupService);
 

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);