From b27f37032d901f3955086ef927947a5dc726b67d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 10 Oct 2024 00:28:11 -0700 Subject: [PATCH] [Workload Management] QueryGroup Stats API Logic (#15777) * changelog Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * add tests Signed-off-by: Ruirui Zhang * modify uri Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * modify based on comments Signed-off-by: Ruirui Zhang * changelog Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * add tests Signed-off-by: Ruirui Zhang * modify uri Signed-off-by: Ruirui Zhang * modify based on comments Signed-off-by: Ruirui Zhang * modify based on comments Signed-off-by: Ruirui Zhang * revise Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * changelog Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * add tests Signed-off-by: Ruirui Zhang * modify uri Signed-off-by: Ruirui Zhang * modify based on comments Signed-off-by: Ruirui Zhang * modify based on comments Signed-off-by: Ruirui Zhang * git pull Signed-off-by: Ruirui Zhang * rebase Signed-off-by: Ruirui Zhang * encapsulate querygroupstats in wlmstats Signed-off-by: Ruirui Zhang * fix UT Signed-off-by: Ruirui Zhang --- CHANGELOG.md | 1 + .../org/opensearch/action/ActionModule.java | 5 + .../cluster/wlm/TransportWlmStatsAction.java | 79 ++++++++++++++ .../admin/cluster/wlm/WlmStatsAction.java | 25 +++++ .../admin/cluster/wlm/WlmStatsRequest.java | 65 +++++++++++ .../admin/cluster/wlm/WlmStatsResponse.java | 73 +++++++++++++ .../admin/cluster/wlm/package-info.java | 10 ++ .../opensearch/client/ClusterAdminClient.java | 9 ++ .../client/support/AbstractClient.java | 8 ++ .../main/java/org/opensearch/node/Node.java | 12 +- .../admin/cluster/RestWlmStatsAction.java | 58 ++++++++++ .../org/opensearch/wlm/QueryGroupService.java | 47 ++++++-- .../opensearch/wlm/stats/QueryGroupStats.java | 4 + .../org/opensearch/wlm/stats/WlmStats.java | 65 +++++++++++ .../cluster/wlm/WlmStatsResponseTests.java | 103 ++++++++++++++++++ .../nodes/TransportWlmStatsActionTests.java | 80 ++++++++++++++ ...eryGroupRequestOperationListenerTests.java | 14 ++- .../wlm/stats/QueryGroupStatsTests.java | 14 +++ .../opensearch/wlm/stats/WlmStatsTests.java | 75 +++++++++++++ 19 files changed, 731 insertions(+), 16 deletions(-) create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/wlm/package-info.java create mode 100644 server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java create mode 100644 server/src/main/java/org/opensearch/wlm/stats/WlmStats.java create mode 100644 server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java create mode 100644 server/src/test/java/org/opensearch/action/support/nodes/TransportWlmStatsActionTests.java create mode 100644 server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 0067d69837eed..5d295c9a199f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - MultiTermQueries in keyword fields now default to `indexed` approach and gated behind cluster setting ([#15637](https://github.com/opensearch-project/OpenSearch/pull/15637)) - [Workload Management] QueryGroup resource cancellation framework changes ([#15651](https://github.com/opensearch-project/OpenSearch/pull/15651)) - [Workload Management] Add orchestrator for wlm resiliency (QueryGroupService) ([#15925](https://github.com/opensearch-project/OpenSearch/pull/15925)) +- [Workload Management] Add QueryGroup Stats API Logic ([15777](https://github.com/opensearch-project/OpenSearch/pull/15777)) - Fallback to Remote cluster-state on Term-Version check mismatch - ([#15424](https://github.com/opensearch-project/OpenSearch/pull/15424)) - 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)) diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java index 12c7626e385a4..1e6eae87af53a 100644 --- a/server/src/main/java/org/opensearch/action/ActionModule.java +++ b/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; @@ -375,6 +377,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; @@ -622,6 +625,7 @@ public 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); @@ -828,6 +832,7 @@ public void initRestHandlers(Supplier 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()); 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 new file mode 100644 index 0000000000000..9c2fb3f1689ec --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java @@ -0,0 +1,79 @@ +/* + * 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.TransportNodesAction; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.common.io.stream.StreamInput; +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 { + + 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, + WlmStatsRequest::new, + ThreadPool.Names.MANAGEMENT, + WlmStats.class + ); + this.queryGroupService = queryGroupService; + } + + @Override + protected WlmStatsResponse newResponse(WlmStatsRequest request, List wlmStats, List failures) { + return new WlmStatsResponse(clusterService.getClusterName(), wlmStats, failures); + } + + @Override + protected WlmStatsRequest newNodeRequest(WlmStatsRequest request) { + return request; + } + + @Override + protected WlmStats newNodeResponse(StreamInput in) throws IOException { + return new WlmStats(in); + } + + @Override + protected WlmStats nodeOperation(WlmStatsRequest wlmStatsRequest) { + assert transportService.getLocalNode() != null; + return new WlmStats( + transportService.getLocalNode(), + queryGroupService.nodeStats(wlmStatsRequest.getQueryGroupIds(), wlmStatsRequest.isBreach()) + ); + } +} 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 new file mode 100644 index 0000000000000..2dfb10fc5dc90 --- /dev/null +++ b/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 { + public static final WlmStatsAction INSTANCE = new WlmStatsAction(); + public static final String NAME = "cluster:monitor/wlm/stats"; + + private WlmStatsAction() { + super(NAME, WlmStatsResponse::new); + } +} 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 new file mode 100644 index 0000000000000..bf4f79faff478 --- /dev/null +++ b/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 { + + private final Set 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 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 getQueryGroupIds() { + return queryGroupIds; + } + + 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 new file mode 100644 index 0000000000000..2ce1b09a61fc6 --- /dev/null +++ b/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 implements ToXContentFragment { + + WlmStatsResponse(StreamInput in) throws IOException { + super(in); + } + + WlmStatsResponse(ClusterName clusterName, List nodes, List failures) { + super(clusterName, nodes, failures); + } + + @Override + protected List readNodesFrom(StreamInput in) throws IOException { + return in.readList(WlmStats::new); + } + + @Override + protected void writeNodesTo(StreamOutput out, List 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() + "\"}"; + } + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/package-info.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/package-info.java new file mode 100644 index 0000000000000..13724b335e7c6 --- /dev/null +++ b/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; diff --git a/server/src/main/java/org/opensearch/client/ClusterAdminClient.java b/server/src/main/java/org/opensearch/client/ClusterAdminClient.java index 05f09c1a6e661..5ce4d442fbe0b 100644 --- a/server/src/main/java/org/opensearch/client/ClusterAdminClient.java +++ b/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 listener); + void remoteStoreStats(RemoteStoreStatsRequest request, ActionListener listener); RemoteStoreStatsRequestBuilder prepareRemoteStoreStats(String index, String shardId); diff --git a/server/src/main/java/org/opensearch/client/support/AbstractClient.java b/server/src/main/java/org/opensearch/client/support/AbstractClient.java index 509cd732357d6..f4683ab516cef 100644 --- a/server/src/main/java/org/opensearch/client/support/AbstractClient.java +++ b/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; @@ -918,6 +921,11 @@ public NodesStatsRequestBuilder prepareNodesStats(String... nodesIds) { return new NodesStatsRequestBuilder(this, NodesStatsAction.INSTANCE).setNodesIds(nodesIds); } + @Override + public void wlmStats(final WlmStatsRequest request, final ActionListener listener) { + execute(WlmStatsAction.INSTANCE, request, listener); + } + @Override public void remoteStoreStats(final RemoteStoreStatsRequest request, final ActionListener listener) { execute(RemoteStoreStatsAction.INSTANCE, request, listener); diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index 5c965b06a4b69..584d95b9ff6b5 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -354,12 +354,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 NODE_LOCAL_STORAGE_SETTING = Setting.boolSetting( "node.local_storage", true, diff --git a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java new file mode 100644 index 0000000000000..51bd313e74df0 --- /dev/null +++ b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestWlmStatsAction.java @@ -0,0 +1,58 @@ +/* + * 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.rest.action.admin.cluster; + +import org.opensearch.action.admin.cluster.wlm.WlmStatsRequest; +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.common.Strings; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestActions; + +import java.io.IOException; +import java.util.List; +import java.util.Set; + +import static java.util.Arrays.asList; +import static java.util.Collections.unmodifiableList; +import static org.opensearch.rest.RestRequest.Method.GET; + +/** + * Transport action to get Workload Management stats + * + * @opensearch.experimental + */ +public class RestWlmStatsAction extends BaseRestHandler { + + @Override + public List routes() { + return unmodifiableList( + asList( + new Route(GET, "_wlm/stats"), + new Route(GET, "_wlm/{nodeId}/stats"), + new Route(GET, "_wlm/stats/{queryGroupId}"), + new Route(GET, "_wlm/{nodeId}/stats/{queryGroupId}") + ) + ); + } + + @Override + public String getName() { + return "wlm_stats_action"; + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + String[] nodesIds = Strings.splitStringByCommaToArray(request.param("nodeId")); + Set queryGroupIds = Strings.tokenizeByCommaToSet(request.param("queryGroupId", "_all")); + Boolean breach = request.hasParam("breach") ? Boolean.parseBoolean(request.param("boolean")) : null; + WlmStatsRequest wlmStatsRequest = new WlmStatsRequest(nodesIds, queryGroupIds, breach); + return channel -> client.admin().cluster().wlmStats(wlmStatsRequest, new RestActions.NodesResponseRestListener<>(channel)); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java index cda5916db26f3..30d39686931be 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java @@ -10,6 +10,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceNotFoundException; import org.opensearch.action.search.SearchShardTask; import org.opensearch.cluster.ClusterChangedEvent; import org.opensearch.cluster.ClusterStateListener; @@ -42,6 +43,7 @@ /** * As of now this is a stub and main implementation PR will be raised soon.Coming PR will collate these changes with core QueryGroupService changes + * @opensearch.experimental */ public class QueryGroupService extends AbstractLifecycleComponent implements @@ -49,7 +51,6 @@ public class QueryGroupService extends AbstractLifecycleComponent TaskResourceTrackingService.TaskCompletionListener { private static final Logger logger = LogManager.getLogger(QueryGroupService.class); - private final QueryGroupTaskCancellationService taskCancellationService; private volatile Scheduler.Cancellable scheduledFuture; private final ThreadPool threadPool; @@ -205,16 +206,48 @@ public void incrementFailuresFor(final String queryGroupId) { /** * @return node level query group stats */ - public QueryGroupStats nodeStats() { + public QueryGroupStats nodeStats(Set queryGroupIds, Boolean requestedBreached) { final Map statsHolderMap = new HashMap<>(); - for (Map.Entry queryGroupsState : queryGroupsStateAccessor.getQueryGroupStateMap().entrySet()) { - final String queryGroupId = queryGroupsState.getKey(); - final QueryGroupState currentState = queryGroupsState.getValue(); + Map existingStateMap = queryGroupsStateAccessor.getQueryGroupStateMap(); + if (!queryGroupIds.contains("_all")) { + for (String id : queryGroupIds) { + if (!existingStateMap.containsKey(id)) { + throw new ResourceNotFoundException("QueryGroup with id " + id + " does not exist"); + } + } + } + if (existingStateMap != null) { + existingStateMap.forEach((queryGroupId, currentState) -> { + boolean shouldInclude = queryGroupIds.contains("_all") || queryGroupIds.contains(queryGroupId); + if (shouldInclude) { + if (requestedBreached == null || requestedBreached == resourceLimitBreached(queryGroupId, currentState)) { + statsHolderMap.put(queryGroupId, QueryGroupStatsHolder.from(currentState)); + } + } + }); + } + return new QueryGroupStats(statsHolderMap); + } - statsHolderMap.put(queryGroupId, QueryGroupStatsHolder.from(currentState)); + /** + * @return if the QueryGroup breaches any resource limit based on the LastRecordedUsage + */ + public boolean resourceLimitBreached(String id, QueryGroupState currentState) { + QueryGroup queryGroup = clusterService.state().metadata().queryGroups().get(id); + if (queryGroup == null) { + throw new ResourceNotFoundException("QueryGroup with id " + id + " does not exist"); } - return new QueryGroupStats(statsHolderMap); + for (ResourceType resourceType : TRACKED_RESOURCES) { + if (queryGroup.getResourceLimits().containsKey(resourceType)) { + final double threshold = getNormalisedRejectionThreshold(queryGroup.getResourceLimits().get(resourceType), resourceType); + final double lastRecordedUsage = currentState.getResourceState().get(resourceType).getLastRecordedUsage(); + if (threshold < lastRecordedUsage) { + return true; + } + } + } + return false; } /** diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java index 9d74201de252b..12e0f0ff96f65 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java @@ -82,6 +82,10 @@ public int hashCode() { return Objects.hash(stats); } + public Map getStats() { + return stats; + } + /** * This is a stats holder object which will hold the data for a query group at a point in time * the instance will only be created on demand through stats api diff --git a/server/src/main/java/org/opensearch/wlm/stats/WlmStats.java b/server/src/main/java/org/opensearch/wlm/stats/WlmStats.java new file mode 100644 index 0000000000000..3313021bfae52 --- /dev/null +++ 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 new file mode 100644 index 0000000000000..86fbc341691d2 --- /dev/null +++ b/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java @@ -0,0 +1,103 @@ +/* + * 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.Version; +import org.opensearch.action.FailedNodeException; +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +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; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +public class WlmStatsResponseTests extends OpenSearchTestCase { + ClusterName clusterName = new ClusterName("test-cluster"); + String testQueryGroupId = "safjgagnaeekg-3r3fads"; + DiscoveryNode node = new DiscoveryNode( + "node-1", + buildNewFakeTransportAddress(), + new HashMap<>(), + Set.of(DiscoveryNodeRole.DATA_ROLE), + Version.CURRENT + ); + Map statsHolderMap = new HashMap<>(); + QueryGroupStats queryGroupStats = new QueryGroupStats( + Map.of( + testQueryGroupId, + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 1, + 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ) + ) + ); + WlmStats wlmStats = new WlmStats(node, queryGroupStats); + List wlmStatsList = List.of(wlmStats); + List failedNodeExceptionList = new ArrayList<>(); + + public void testSerializationAndDeserialization() throws IOException { + WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, wlmStatsList, failedNodeExceptionList); + BytesStreamOutput out = new BytesStreamOutput(); + queryGroupStatsResponse.writeTo(out); + StreamInput in = out.bytes().streamInput(); + WlmStatsResponse deserializedResponse = new WlmStatsResponse(in); + assertEquals(queryGroupStatsResponse.getClusterName(), deserializedResponse.getClusterName()); + assertEquals(queryGroupStatsResponse.getNodes().size(), deserializedResponse.getNodes().size()); + } + + public void testToString() { + 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" + + " \"cpu\" : {\n" + + " \"current_usage\" : 0.0,\n" + + " \"cancellations\" : 0,\n" + + " \"rejections\" : 0\n" + + " },\n" + + " \"memory\" : {\n" + + " \"current_usage\" : 0.0,\n" + + " \"cancellations\" : 0,\n" + + " \"rejections\" : 0\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + "}", + responseString + ); + } +} diff --git a/server/src/test/java/org/opensearch/action/support/nodes/TransportWlmStatsActionTests.java b/server/src/test/java/org/opensearch/action/support/nodes/TransportWlmStatsActionTests.java new file mode 100644 index 0000000000000..49d2cc4d23e62 --- /dev/null +++ b/server/src/test/java/org/opensearch/action/support/nodes/TransportWlmStatsActionTests.java @@ -0,0 +1,80 @@ +/* + * 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.support.nodes; + +import org.opensearch.action.admin.cluster.wlm.TransportWlmStatsAction; +import org.opensearch.action.admin.cluster.wlm.WlmStatsRequest; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.PlainActionFuture; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.test.transport.CapturingTransport; +import org.opensearch.wlm.QueryGroupService; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.mockito.Mockito.mock; + +public class TransportWlmStatsActionTests extends TransportNodesActionTests { + + /** + * We don't want to send discovery nodes list to each request that is sent across from the coordinator node. + * This behavior is asserted in this test. + */ + public void testWlmStatsActionWithRetentionOfDiscoveryNodesList() { + WlmStatsRequest request = new WlmStatsRequest(); + Map> combinedSentRequest = performWlmStatsAction(request); + + assertNotNull(combinedSentRequest); + combinedSentRequest.forEach((node, capturedRequestList) -> { + assertNotNull(capturedRequestList); + capturedRequestList.forEach(sentRequest -> { assertNull(sentRequest.concreteNodes()); }); + }); + } + + private Map> performWlmStatsAction(WlmStatsRequest request) { + TransportNodesAction action = new TransportWlmStatsAction( + THREAD_POOL, + clusterService, + transportService, + mock(QueryGroupService.class), + new ActionFilters(Collections.emptySet()) + ); + PlainActionFuture listener = new PlainActionFuture<>(); + action.new AsyncAction(null, request, listener).start(); + Map> capturedRequests = transport.getCapturedRequestsByTargetNodeAndClear(); + Map> combinedSentRequest = new HashMap<>(); + + capturedRequests.forEach((node, capturedRequestList) -> { + List sentRequestList = new ArrayList<>(); + + capturedRequestList.forEach(preSentRequest -> { + BytesStreamOutput out = new BytesStreamOutput(); + try { + WlmStatsRequest wlmStatsRequestFromCoordinator = (WlmStatsRequest) preSentRequest.request; + wlmStatsRequestFromCoordinator.writeTo(out); + StreamInput in = out.bytes().streamInput(); + WlmStatsRequest wlmStatsRequest = new WlmStatsRequest(in); + sentRequestList.add(wlmStatsRequest); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + + combinedSentRequest.put(node, sentRequestList); + }); + + return combinedSentRequest; + } +} diff --git a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java index 1127b50399d24..4d61e89b281f7 100644 --- a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java @@ -29,6 +29,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; @@ -56,6 +57,10 @@ public void setUp() throws Exception { queryGroupStateMap = new HashMap<>(); testQueryGroupId = "safjgagnakg-3r3fads"; testThreadPool = new TestThreadPool("RejectionTestThreadPool"); + ClusterState mockClusterState = mock(ClusterState.class); + when(mockClusterService.state()).thenReturn(mockClusterState); + Metadata mockMetaData = mock(Metadata.class); + when(mockClusterState.metadata()).thenReturn(mockMetaData); queryGroupService = mock(QueryGroupService.class); sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); } @@ -155,7 +160,9 @@ public void testMultiThreadedValidQueryGroupRequestFailures() { } }); - QueryGroupStats actualStats = queryGroupService.nodeStats(); + HashSet set = new HashSet<>(); + set.add("_all"); + QueryGroupStats actualStats = queryGroupService.nodeStats(set, null); QueryGroupStats expectedStats = new QueryGroupStats( Map.of( @@ -254,11 +261,12 @@ private void assertSuccess( Collections.emptySet(), Collections.emptySet() ); - sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); sut.onRequestFailure(null, null); - QueryGroupStats actualStats = queryGroupService.nodeStats(); + HashSet set = new HashSet<>(); + set.add("_all"); + QueryGroupStats actualStats = queryGroupService.nodeStats(set, null); assertEquals(expectedStats, actualStats); } diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java index ac6d19580dacb..bcab0c4d124f5 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java @@ -8,17 +8,24 @@ package org.opensearch.wlm.stats; +import org.opensearch.Version; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.test.AbstractWireSerializingTestCase; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.VersionUtils; import org.opensearch.wlm.ResourceType; import java.io.IOException; import java.util.HashMap; import java.util.Map; +import static java.util.Collections.emptyMap; + public class QueryGroupStatsTests extends AbstractWireSerializingTestCase { public void testToXContent() throws IOException { @@ -72,6 +79,13 @@ protected QueryGroupStats createTestInstance() { ) ) ); + DiscoveryNode discoveryNode = new DiscoveryNode( + "node", + OpenSearchTestCase.buildNewFakeTransportAddress(), + emptyMap(), + DiscoveryNodeRole.BUILT_IN_ROLES, + VersionUtils.randomCompatibleVersion(random(), Version.CURRENT) + ); return new QueryGroupStats(stats); } } diff --git a/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java new file mode 100644 index 0000000000000..75b8f570b7975 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java @@ -0,0 +1,75 @@ +/* + * 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.Version; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.test.AbstractWireSerializingTestCase; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.VersionUtils; +import org.opensearch.wlm.ResourceType; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +import static java.util.Collections.emptyMap; +import static org.mockito.Mockito.mock; + +public class WlmStatsTests extends AbstractWireSerializingTestCase { + + public void testToXContent() throws IOException { + final Map stats = new HashMap<>(); + final String queryGroupId = "afakjklaj304041-afaka"; + stats.put( + queryGroupId, + new QueryGroupStats.QueryGroupStatsHolder( + 123456789, + 13, + 2, + 0, + 1213718, + Map.of(ResourceType.CPU, new QueryGroupStats.ResourceStats(0.3, 13, 2)) + ) + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + QueryGroupStats queryGroupStats = new QueryGroupStats(stats); + WlmStats wlmStats = new WlmStats(mock(DiscoveryNode.class), queryGroupStats); + builder.startObject(); + wlmStats.toXContent(builder, ToXContent.EMPTY_PARAMS); + builder.endObject(); + assertEquals( + "{\"query_groups\":{\"afakjklaj304041-afaka\":{\"completions\":123456789,\"shard_completions\":1213718,\"rejections\":13,\"failures\":2,\"total_cancellations\":0,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", + builder.toString() + ); + } + + @Override + protected Writeable.Reader instanceReader() { + return WlmStats::new; + } + + @Override + protected WlmStats createTestInstance() { + DiscoveryNode discoveryNode = new DiscoveryNode( + "node", + OpenSearchTestCase.buildNewFakeTransportAddress(), + emptyMap(), + DiscoveryNodeRole.BUILT_IN_ROLES, + VersionUtils.randomCompatibleVersion(random(), Version.CURRENT) + ); + QueryGroupStatsTests queryGroupStatsTests = new QueryGroupStatsTests(); + return new WlmStats(discoveryNode, queryGroupStatsTests.createTestInstance()); + } +}