Backport 2.x [Workload Management] QueryGroup Stats API Logic

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Offline Nodes] Adds offline-tasks library containing various interfaces to be used for Offline Background Tasks. ([#13574](https://github.com/opensearch-project/OpenSearch/pull/13574))
 - Add support for async deletion in S3BlobContainer ([#15621](https://github.com/opensearch-project/OpenSearch/pull/15621))
 - [Workload Management] QueryGroup resource cancellation framework changes ([#15651](https://github.com/opensearch-project/OpenSearch/pull/15651))
+- [Workload Management] Add QueryGroup Stats API Logic ([15777](https://github.com/opensearch-project/OpenSearch/pull/15777))
 - Implement WithFieldName interface in ValuesSourceAggregationBuilder & FieldSortBuilder ([#15916](https://github.com/opensearch-project/OpenSearch/pull/15916))
 - Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967), [#16110](https://github.com/opensearch-project/OpenSearch/pull/16110))
 - Fallback to Remote cluster-state on Term-Version check mismatch - ([#15424](https://github.com/opensearch-project/OpenSearch/pull/15424))

server/src/main/java/org/opensearch/action/ActionModule.java

@@ -123,6 +123,8 @@
 import org.opensearch.action.admin.cluster.storedscripts.TransportPutStoredScriptAction;
 import org.opensearch.action.admin.cluster.tasks.PendingClusterTasksAction;
 import org.opensearch.action.admin.cluster.tasks.TransportPendingClusterTasksAction;
+import org.opensearch.action.admin.cluster.wlm.TransportWlmStatsAction;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsAction;
 import org.opensearch.action.admin.indices.alias.IndicesAliasesAction;
 import org.opensearch.action.admin.indices.alias.IndicesAliasesRequest;
 import org.opensearch.action.admin.indices.alias.TransportIndicesAliasesAction;
@@ -369,6 +371,7 @@
 import org.opensearch.rest.action.admin.cluster.RestRestoreSnapshotAction;
 import org.opensearch.rest.action.admin.cluster.RestSnapshotsStatusAction;
 import org.opensearch.rest.action.admin.cluster.RestVerifyRepositoryAction;
+import org.opensearch.rest.action.admin.cluster.RestWlmStatsAction;
 import org.opensearch.rest.action.admin.cluster.dangling.RestDeleteDanglingIndexAction;
 import org.opensearch.rest.action.admin.cluster.dangling.RestImportDanglingIndexAction;
 import org.opensearch.rest.action.admin.cluster.dangling.RestListDanglingIndicesAction;
@@ -614,6 +617,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         actions.register(NodesInfoAction.INSTANCE, TransportNodesInfoAction.class);
         actions.register(RemoteInfoAction.INSTANCE, TransportRemoteInfoAction.class);
         actions.register(NodesStatsAction.INSTANCE, TransportNodesStatsAction.class);
+        actions.register(WlmStatsAction.INSTANCE, TransportWlmStatsAction.class);
         actions.register(RemoteStoreStatsAction.INSTANCE, TransportRemoteStoreStatsAction.class);
         actions.register(NodesUsageAction.INSTANCE, TransportNodesUsageAction.class);
         actions.register(NodesHotThreadsAction.INSTANCE, TransportNodesHotThreadsAction.class);
@@ -812,6 +816,7 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster) {
         registerHandler.accept(new RestClearVotingConfigExclusionsAction());
         registerHandler.accept(new RestMainAction());
         registerHandler.accept(new RestNodesInfoAction(settingsFilter));
+        registerHandler.accept(new RestWlmStatsAction());
         registerHandler.accept(new RestRemoteClusterInfoAction());
         registerHandler.accept(new RestNodesStatsAction());
         registerHandler.accept(new RestNodesUsageAction());

server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java

@@ -0,0 +1,113 @@
+/*
+ * 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.action.admin.cluster.wlm;
+
+import org.opensearch.action.FailedNodeException;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.nodes.BaseNodeRequest;
+import org.opensearch.action.support.nodes.TransportNodesAction;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportService;
+import org.opensearch.wlm.QueryGroupService;
+import org.opensearch.wlm.stats.WlmStats;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Transport action for obtaining WlmStats
+ *
+ * @opensearch.experimental
+ */
+public class TransportWlmStatsAction extends TransportNodesAction<
+    WlmStatsRequest,
+    WlmStatsResponse,
+    TransportWlmStatsAction.NodeWlmStatsRequest,
+    WlmStats> {
+
+    final QueryGroupService queryGroupService;
+
+    @Inject
+    public TransportWlmStatsAction(
+        ThreadPool threadPool,
+        ClusterService clusterService,
+        TransportService transportService,
+        QueryGroupService queryGroupService,
+        ActionFilters actionFilters
+    ) {
+        super(
+            WlmStatsAction.NAME,
+            threadPool,
+            clusterService,
+            transportService,
+            actionFilters,
+            WlmStatsRequest::new,
+            NodeWlmStatsRequest::new,
+            ThreadPool.Names.MANAGEMENT,
+            WlmStats.class
+        );
+        this.queryGroupService = queryGroupService;
+    }
+
+    @Override
+    protected WlmStatsResponse newResponse(WlmStatsRequest request, List<WlmStats> wlmStats, List<FailedNodeException> failures) {
+        return new WlmStatsResponse(clusterService.getClusterName(), wlmStats, failures);
+    }
+
+    @Override
+    protected NodeWlmStatsRequest newNodeRequest(WlmStatsRequest request) {
+        return new NodeWlmStatsRequest(request);
+    }
+
+    @Override
+    protected WlmStats newNodeResponse(StreamInput in) throws IOException {
+        return new WlmStats(in);
+    }
+
+    @Override
+    protected WlmStats nodeOperation(NodeWlmStatsRequest nodeWlmStatsRequest) {
+        assert transportService.getLocalNode() != null;
+        WlmStatsRequest request = nodeWlmStatsRequest.request;
+        return new WlmStats(transportService.getLocalNode(), queryGroupService.nodeStats(request.getQueryGroupIds(), request.isBreach()));
+    }
+
+    /**
+     * Inner WlmStatsRequest
+     *
+     * @opensearch.experimental
+     */
+    public static class NodeWlmStatsRequest extends BaseNodeRequest {
+
+        protected WlmStatsRequest request;
+
+        public NodeWlmStatsRequest(StreamInput in) throws IOException {
+            super(in);
+            request = new WlmStatsRequest(in);
+        }
+
+        NodeWlmStatsRequest(WlmStatsRequest request) {
+            this.request = request;
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            super.writeTo(out);
+            request.writeTo(out);
+        }
+
+        public DiscoveryNode[] getDiscoveryNodes() {
+            return this.request.concreteNodes();
+        }
+    }
+}

server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsAction.java

@@ -0,0 +1,25 @@
+/*
+ * 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.action.admin.cluster.wlm;
+
+import org.opensearch.action.ActionType;
+
+/**
+ * Transport action for obtaining Workload Management Stats.
+ *
+ * @opensearch.experimental
+ */
+public class WlmStatsAction extends ActionType<WlmStatsResponse> {
+    public static final WlmStatsAction INSTANCE = new WlmStatsAction();
+    public static final String NAME = "cluster:monitor/wlm/stats";
+
+    private WlmStatsAction() {
+        super(NAME, WlmStatsResponse::new);
+    }
+}

server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.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.action.admin.cluster.wlm;
+
+import org.opensearch.action.support.nodes.BaseNodesRequest;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * A request to get Workload Management Stats
+ */
+@ExperimentalApi
+public class WlmStatsRequest extends BaseNodesRequest<WlmStatsRequest> {
+
+    private final Set<String> queryGroupIds;
+    private final Boolean breach;
+
+    public WlmStatsRequest(StreamInput in) throws IOException {
+        super(in);
+        this.queryGroupIds = new HashSet<>(Set.of(in.readStringArray()));
+        this.breach = in.readOptionalBoolean();
+    }
+
+    /**
+     * 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) {
+        super(false, nodesIds);
+        this.queryGroupIds = queryGroupIds;
+        this.breach = breach;
+    }
+
+    public WlmStatsRequest() {
+        super(false, (String[]) null);
+        queryGroupIds = new HashSet<>();
+        this.breach = false;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeStringArray(queryGroupIds.toArray(new String[0]));
+        out.writeOptionalBoolean(breach);
+    }
+
+    public Set<String> getQueryGroupIds() {
+        return queryGroupIds;
+    }
+
+    public Boolean isBreach() {
+        return breach;
+    }
+}

server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java

@@ -0,0 +1,73 @@
+/*
+ * 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.action.admin.cluster.wlm;
+
+import org.opensearch.action.FailedNodeException;
+import org.opensearch.action.support.nodes.BaseNodesResponse;
+import org.opensearch.cluster.ClusterName;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+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;
+
+/**
+ * A response for obtaining Workload Management Stats
+ */
+@ExperimentalApi
+public class WlmStatsResponse extends BaseNodesResponse<WlmStats> implements ToXContentFragment {
+
+    WlmStatsResponse(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    WlmStatsResponse(ClusterName clusterName, List<WlmStats> nodes, List<FailedNodeException> failures) {
+        super(clusterName, nodes, failures);
+    }
+
+    @Override
+    protected List<WlmStats> readNodesFrom(StreamInput in) throws IOException {
+        return in.readList(WlmStats::new);
+    }
+
+    @Override
+    protected void writeNodesTo(StreamOutput out, List<WlmStats> nodes) throws IOException {
+        out.writeList(nodes);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        for (WlmStats wlmStats : getNodes()) {
+            builder.startObject(wlmStats.getNode().getId());
+            QueryGroupStats queryGroupStats = wlmStats.getQueryGroupStats();
+            queryGroupStats.toXContent(builder, params);
+            builder.endObject();
+        }
+        return builder;
+    }
+
+    @Override
+    public String toString() {
+        try {
+            XContentBuilder builder = XContentFactory.jsonBuilder().prettyPrint();
+            builder.startObject();
+            toXContent(builder, EMPTY_PARAMS);
+            builder.endObject();
+            return builder.toString();
+        } catch (IOException e) {
+            return "{ \"error\" : \"" + e.getMessage() + "\"}";
+        }
+    }
+}

server/src/main/java/org/opensearch/action/admin/cluster/wlm/package-info.java

@@ -0,0 +1,10 @@
+/*
+ * 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.
+ */
+
+/** WlmStats transport handlers. */
+package org.opensearch.action.admin.cluster.wlm;

server/src/main/java/org/opensearch/client/ClusterAdminClient.java

@@ -137,6 +137,8 @@
 import org.opensearch.action.admin.cluster.tasks.PendingClusterTasksRequest;
 import org.opensearch.action.admin.cluster.tasks.PendingClusterTasksRequestBuilder;
 import org.opensearch.action.admin.cluster.tasks.PendingClusterTasksResponse;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsRequest;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsResponse;
 import org.opensearch.action.admin.indices.dangling.delete.DeleteDanglingIndexRequest;
 import org.opensearch.action.admin.indices.dangling.import_index.ImportDanglingIndexRequest;
 import org.opensearch.action.admin.indices.dangling.list.ListDanglingIndicesRequest;
@@ -320,6 +322,13 @@ public interface ClusterAdminClient extends OpenSearchClient {
      */
     NodesStatsRequestBuilder prepareNodesStats(String... nodesIds);
 
+    /**
+     * QueryGroup stats of the cluster.
+     * @param request The wlmStatsRequest
+     * @param listener A listener to be notified with a result
+     */
+    void wlmStats(WlmStatsRequest request, ActionListener<WlmStatsResponse> listener);
+
     void remoteStoreStats(RemoteStoreStatsRequest request, ActionListener<RemoteStoreStatsResponse> listener);
 
     RemoteStoreStatsRequestBuilder prepareRemoteStoreStats(String index, String shardId);

server/src/main/java/org/opensearch/client/support/AbstractClient.java

@@ -179,6 +179,9 @@
 import org.opensearch.action.admin.cluster.tasks.PendingClusterTasksRequest;
 import org.opensearch.action.admin.cluster.tasks.PendingClusterTasksRequestBuilder;
 import org.opensearch.action.admin.cluster.tasks.PendingClusterTasksResponse;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsAction;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsRequest;
+import org.opensearch.action.admin.cluster.wlm.WlmStatsResponse;
 import org.opensearch.action.admin.indices.alias.IndicesAliasesAction;
 import org.opensearch.action.admin.indices.alias.IndicesAliasesRequest;
 import org.opensearch.action.admin.indices.alias.IndicesAliasesRequestBuilder;
@@ -892,6 +895,11 @@
             return new NodesStatsRequestBuilder(this, NodesStatsAction.INSTANCE).setNodesIds(nodesIds);
         }
 
+        @Override
+        public void wlmStats(final WlmStatsRequest request, final ActionListener<WlmStatsResponse> listener) {
+            execute(WlmStatsAction.INSTANCE, request, listener);
+        }
+
         @Override
         public void remoteStoreStats(final RemoteStoreStatsRequest request, final ActionListener<RemoteStoreStatsResponse> listener) {
             execute(RemoteStoreStatsAction.INSTANCE, request, listener);

server/src/main/java/org/opensearch/node/Node.java

@@ -353,12 +353,12 @@ public class Node implements Closeable {
     );
 
     /**
-    * controls whether the node is allowed to persist things like metadata to disk
-    * Note that this does not control whether the node stores actual indices (see
-    * {@link #NODE_DATA_SETTING}). However, if this is false, {@link #NODE_DATA_SETTING}
-    * and {@link #NODE_MASTER_SETTING} must also be false.
-    *
-    */
+     * controls whether the node is allowed to persist things like metadata to disk
+     * Note that this does not control whether the node stores actual indices (see
+     * {@link #NODE_DATA_SETTING}). However, if this is false, {@link #NODE_DATA_SETTING}
+     * and {@link #NODE_MASTER_SETTING} must also be false.
+     *
+     */
     public static final Setting<Boolean> NODE_LOCAL_STORAGE_SETTING = Setting.boolSetting(
         "node.local_storage",
         true,