encapsulate querygroupstats in wlmstats

Signed-off-by: Ruirui Zhang <[email protected]>
opensearch-project · Oct 8, 2024 · 9f49a74 · 9f49a74
1 parent abb0428
commit 9f49a74
Show file tree

Hide file tree

Showing 11 changed files with 213 additions and 56 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
@@ -17,17 +17,17 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
 import org.opensearch.wlm.QueryGroupService;
-import org.opensearch.wlm.stats.QueryGroupStats;
+import org.opensearch.wlm.stats.WlmStats;
 
 import java.io.IOException;
 import java.util.List;
 
 /**
- * Transport action for obtaining QueryGroupStats
+ * Transport action for obtaining WlmStats
  *
  * @opensearch.experimental
  */
-public class TransportWlmStatsAction extends TransportNodesAction<WlmStatsRequest, WlmStatsResponse, WlmStatsRequest, QueryGroupStats> {
+public class TransportWlmStatsAction extends TransportNodesAction<WlmStatsRequest, WlmStatsResponse, WlmStatsRequest, WlmStats> {
 
     final QueryGroupService queryGroupService;
 
@@ -48,18 +48,14 @@ public TransportWlmStatsAction(
             WlmStatsRequest::new,
             WlmStatsRequest::new,
             ThreadPool.Names.MANAGEMENT,
-            QueryGroupStats.class
+            WlmStats.class
         );
         this.queryGroupService = queryGroupService;
     }
 
     @Override
-    protected WlmStatsResponse newResponse(
-        WlmStatsRequest request,
-        List<QueryGroupStats> queryGroupStats,
-        List<FailedNodeException> failures
-    ) {
-        return new WlmStatsResponse(clusterService.getClusterName(), queryGroupStats, failures);
+    protected WlmStatsResponse newResponse(WlmStatsRequest request, List<WlmStats> wlmStats, List<FailedNodeException> failures) {
+        return new WlmStatsResponse(clusterService.getClusterName(), wlmStats, failures);
     }
 
     @Override
@@ -68,12 +64,12 @@ protected WlmStatsRequest newNodeRequest(WlmStatsRequest request) {
     }
 
     @Override
-    protected QueryGroupStats newNodeResponse(StreamInput in) throws IOException {
-        return new QueryGroupStats(in);
+    protected WlmStats newNodeResponse(StreamInput in) throws IOException {
+        return new WlmStats(in);
     }
 
     @Override
-    protected QueryGroupStats nodeOperation(WlmStatsRequest wlmStatsRequest) {
+    protected WlmStats nodeOperation(WlmStatsRequest wlmStatsRequest) {
         return queryGroupService.nodeStats(wlmStatsRequest.getQueryGroupIds(), wlmStatsRequest.isBreach());
     }
 }
diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.java
@@ -36,7 +36,7 @@ public WlmStatsRequest(StreamInput in) throws IOException {
      * Get QueryGroup stats from nodes based on the nodes ids specified. If none are passed, stats
      * for all nodes will be returned.
      */
-    public WlmStatsRequest(String[] nodesIds, Set<String> queryGroupIds, boolean breach) {
+    public WlmStatsRequest(String[] nodesIds, Set<String> queryGroupIds, Boolean breach) {
         super(false, nodesIds);
         this.queryGroupIds = queryGroupIds;
         this.breach = breach;
@@ -59,7 +59,7 @@ public Set<String> getQueryGroupIds() {
         return queryGroupIds;
     }
 
-    public boolean isBreach() {
+    public Boolean isBreach() {
         return breach;
     }
 }
diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java
@@ -18,6 +18,7 @@
 import org.opensearch.core.xcontent.ToXContentFragment;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.wlm.stats.QueryGroupStats;
+import org.opensearch.wlm.stats.WlmStats;
 
 import java.io.IOException;
 import java.util.List;
@@ -26,30 +27,31 @@
  * A response for obtaining Workload Management Stats
  */
 @ExperimentalApi
-public class WlmStatsResponse extends BaseNodesResponse<QueryGroupStats> implements ToXContentFragment {
+public class WlmStatsResponse extends BaseNodesResponse<WlmStats> implements ToXContentFragment {
 
     WlmStatsResponse(StreamInput in) throws IOException {
         super(in);
     }
 
-    WlmStatsResponse(ClusterName clusterName, List<QueryGroupStats> nodes, List<FailedNodeException> failures) {
+    WlmStatsResponse(ClusterName clusterName, List<WlmStats> nodes, List<FailedNodeException> failures) {
         super(clusterName, nodes, failures);
     }
 
     @Override
-    protected List<QueryGroupStats> readNodesFrom(StreamInput in) throws IOException {
-        return in.readList(QueryGroupStats::new);
+    protected List<WlmStats> readNodesFrom(StreamInput in) throws IOException {
+        return in.readList(WlmStats::new);
     }
 
     @Override
-    protected void writeNodesTo(StreamOutput out, List<QueryGroupStats> nodes) throws IOException {
+    protected void writeNodesTo(StreamOutput out, List<WlmStats> nodes) throws IOException {
         out.writeList(nodes);
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        for (QueryGroupStats queryGroupStats : getNodes()) {
-            builder.startObject(queryGroupStats.getNode().getId());
+        for (WlmStats wlmStats : getNodes()) {
+            builder.startObject(wlmStats.getNode().getId());
+            QueryGroupStats queryGroupStats = wlmStats.getQueryGroupStats();
             queryGroupStats.toXContent(builder, params);
             builder.endObject();
         }

diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java
@@ -33,6 +33,7 @@
 import org.opensearch.wlm.stats.QueryGroupState;
 import org.opensearch.wlm.stats.QueryGroupStats;
 import org.opensearch.wlm.stats.QueryGroupStats.QueryGroupStatsHolder;
+import org.opensearch.wlm.stats.WlmStats;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -213,7 +214,7 @@ public void incrementFailuresFor(final String queryGroupId) {
     /**
      * @return node level query group stats
      */
-    public QueryGroupStats nodeStats(Set<String> queryGroupIds, Boolean requestedBreached) {
+    public WlmStats nodeStats(Set<String> queryGroupIds, Boolean requestedBreached) {
         final Map<String, QueryGroupStatsHolder> statsHolderMap = new HashMap<>();
         Map<String, QueryGroup> existingGroups = clusterService.state().metadata().queryGroups();
         if (!queryGroupIds.contains("_all")) {
@@ -223,16 +224,20 @@ public QueryGroupStats nodeStats(Set<String> queryGroupIds, Boolean requestedBre
                 }
             }
         }
-        queryGroupsStateAccessor.getQueryGroupStateMap().forEach((queryGroupId, currentState) -> {
-            boolean shouldInclude = (queryGroupIds.size() == 1 && queryGroupIds.contains("_all")) || queryGroupIds.contains(queryGroupId);
-            if (shouldInclude) {
-                boolean breached = resourceLimitBreached(existingGroups.get(queryGroupId), currentState).v1().length() != 0;
-                if (requestedBreached == null || requestedBreached == breached) {
-                    statsHolderMap.put(queryGroupId, QueryGroupStatsHolder.from(currentState));
+        Map<String, QueryGroupState> stateMap = queryGroupsStateAccessor.getQueryGroupStateMap();
+        if (stateMap != null) {
+            stateMap.forEach((queryGroupId, currentState) -> {
+                boolean shouldInclude = queryGroupIds.contains("_all") || queryGroupIds.contains(queryGroupId);
+                if (shouldInclude) {
+                    if (requestedBreached == null
+                        || requestedBreached == (resourceLimitBreached(existingGroups.get(queryGroupId), currentState).v1()
+                            .length() != 0)) {
+                        statsHolderMap.put(queryGroupId, QueryGroupStatsHolder.from(currentState));
+                    }
                 }
-            }
-        });
-        return new QueryGroupStats(transportService.getLocalNode(), statsHolderMap);
+            });
+        }
+        return new WlmStats(transportService.getLocalNode(), new QueryGroupStats(statsHolderMap));
     }
 
     /**

diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java
@@ -8,8 +8,6 @@
 
 package org.opensearch.wlm.stats;
 
-import org.opensearch.action.support.nodes.BaseNodeResponse;
-import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
@@ -39,22 +37,19 @@
  *     ...
  * }
  */
-public class QueryGroupStats extends BaseNodeResponse implements ToXContentObject, Writeable {
+public class QueryGroupStats implements ToXContentObject, Writeable {
     private final Map<String, QueryGroupStatsHolder> stats;
 
-    public QueryGroupStats(DiscoveryNode node, Map<String, QueryGroupStatsHolder> stats) {
-        super(node);
+    public QueryGroupStats(Map<String, QueryGroupStatsHolder> stats) {
         this.stats = stats;
     }
 
     public QueryGroupStats(StreamInput in) throws IOException {
-        super(in);
         stats = in.readMap(StreamInput::readString, QueryGroupStatsHolder::new);
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        super.writeTo(out);
         out.writeMap(stats, StreamOutput::writeString, QueryGroupStatsHolder::writeTo);
     }
 

diff --git a/server/src/main/java/org/opensearch/wlm/stats/WlmStats.java b/server/src/main/java/org/opensearch/wlm/stats/WlmStats.java
@@ -0,0 +1,65 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm.stats;
+
+import org.opensearch.action.support.nodes.BaseNodeResponse;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.xcontent.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+/**
+ * This class contains the stats for Workload Management
+ */
+public class WlmStats extends BaseNodeResponse implements ToXContentObject, Writeable {
+    private final QueryGroupStats queryGroupStats;
+
+    public WlmStats(DiscoveryNode node, QueryGroupStats queryGroupStats) {
+        super(node);
+        this.queryGroupStats = queryGroupStats;
+    }
+
+    public WlmStats(StreamInput in) throws IOException {
+        super(in);
+        queryGroupStats = new QueryGroupStats(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        queryGroupStats.writeTo(out);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        return queryGroupStats.toXContent(builder, params);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        WlmStats that = (WlmStats) o;
+        return Objects.equals(getQueryGroupStats(), that.getQueryGroupStats());
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(queryGroupStats);
+    }
+
+    public QueryGroupStats getQueryGroupStats() {
+        return queryGroupStats;
+    }
+}
diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java
@@ -18,6 +18,7 @@
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.wlm.ResourceType;
 import org.opensearch.wlm.stats.QueryGroupStats;
+import org.opensearch.wlm.stats.WlmStats;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -38,7 +39,6 @@ public class WlmStatsResponseTests extends OpenSearchTestCase {
     );
     Map<String, QueryGroupStats.QueryGroupStatsHolder> statsHolderMap = new HashMap<>();
     QueryGroupStats queryGroupStats = new QueryGroupStats(
-        node,
         Map.of(
             testQueryGroupId,
             new QueryGroupStats.QueryGroupStatsHolder(
@@ -56,11 +56,12 @@ public class WlmStatsResponseTests extends OpenSearchTestCase {
             )
         )
     );
-    List<QueryGroupStats> queryGroupStatsList = List.of(queryGroupStats);
+    WlmStats wlmStats = new WlmStats(node, queryGroupStats);
+    List<WlmStats> wlmStatsList = List.of(wlmStats);
     List<FailedNodeException> failedNodeExceptionList = new ArrayList<>();
 
     public void testSerializationAndDeserialization() throws IOException {
-        WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, queryGroupStatsList, failedNodeExceptionList);
+        WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, wlmStatsList, failedNodeExceptionList);
         BytesStreamOutput out = new BytesStreamOutput();
         queryGroupStatsResponse.writeTo(out);
         StreamInput in = out.bytes().streamInput();
@@ -70,14 +71,15 @@ public void testSerializationAndDeserialization() throws IOException {
     }
 
     public void testToString() {
-        WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, queryGroupStatsList, failedNodeExceptionList);
+        WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, wlmStatsList, failedNodeExceptionList);
         String responseString = queryGroupStatsResponse.toString();
         assertEquals(
             "{\n"
                 + "  \"node-1\" : {\n"
                 + "    \"query_groups\" : {\n"
                 + "      \"safjgagnaeekg-3r3fads\" : {\n"
                 + "        \"completions\" : 0,\n"
+                + "        \"shard_completions\" : 0,\n"
                 + "        \"rejections\" : 0,\n"
                 + "        \"failures\" : 1,\n"
                 + "        \"total_cancellations\" : 0,\n"