From b88ae354117b7ed97e2fd224b4e4b7c8c83fed97 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Mon, 8 Jul 2024 17:51:30 -0700 Subject: [PATCH 01/20] add logic for Create QueryGroup API Signed-off-by: Ruirui Zhang --- CHANGELOG.md | 1 + plugins/workload-management/build.gradle | 18 ++ .../wlm/action/CreateQueryGroupAction.java | 36 +++ .../wlm/action/CreateQueryGroupRequest.java | 252 ++++++++++++++++++ .../wlm/action/CreateQueryGroupResponse.java | 74 +++++ .../TransportCreateQueryGroupAction.java | 74 +++++ .../wlm/action/WorkloadManagementPlugin.java | 62 +++++ .../WorkloadManagementPluginModule.java | 32 +++ .../plugin/wlm/action/package-info.java | 12 + .../rest/RestCreateQueryGroupAction.java | 80 ++++++ .../plugin/wlm/action/rest/package-info.java | 12 + .../wlm/action/service/Persistable.java | 25 ++ .../service/QueryGroupPersistenceService.java | 228 ++++++++++++++++ .../wlm/action/service/package-info.java | 12 + .../action/CreateQueryGroupRequestTests.java | 34 +++ .../action/CreateQueryGroupResponseTests.java | 38 +++ .../wlm/action/QueryGroupTestUtils.java | 132 +++++++++ .../QueryGroupPersistenceServiceTests.java | 104 ++++++++ .../opensearch/cluster/metadata/Metadata.java | 6 + .../cluster/metadata/QueryGroup.java | 19 +- .../common/settings/ClusterSettings.java | 1 + .../QueryGroupServiceSettings.java | 198 ++++++++++++++ .../search/query_group/package-info.java | 12 + 23 files changed, 1455 insertions(+), 7 deletions(-) create mode 100644 plugins/workload-management/build.gradle create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPlugin.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPluginModule.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/package-info.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/package-info.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java create mode 100644 server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java create mode 100644 server/src/main/java/org/opensearch/search/query_group/package-info.java diff --git a/CHANGELOG.md b/CHANGELOG.md index f44949bf38511..148e823f6ce95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add setting to ignore throttling nodes for allocation of unassigned primaries in remote restore ([#14991](https://github.com/opensearch-project/OpenSearch/pull/14991)) - [Streaming Indexing] Enhance RestClient with a new streaming API support ([#14437](https://github.com/opensearch-project/OpenSearch/pull/14437)) - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618)) +- [Workload Management] Add Create QueryGroup API Logic ([#14680](https://github.com/opensearch-project/OpenSearch/pull/14680))- [Workload Management] Add Create QueryGroup API Logic ([#14680](https://github.com/opensearch-project/OpenSearch/pull/14680)) - Add ThreadContextPermission for markAsSystemContext and allow core to perform the method ([#15016](https://github.com/opensearch-project/OpenSearch/pull/15016)) - Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039)) - [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) diff --git a/plugins/workload-management/build.gradle b/plugins/workload-management/build.gradle new file mode 100644 index 0000000000000..89e13c079795e --- /dev/null +++ b/plugins/workload-management/build.gradle @@ -0,0 +1,18 @@ +/* + * 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. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +opensearchplugin { + description 'OpenSearch Workload Management Plugin.' + classname 'org.opensearch.plugin.wlm.action.WorkloadManagementPlugin' +} + +dependencies { +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java new file mode 100644 index 0000000000000..003b8b60c4ff9 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java @@ -0,0 +1,36 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.action.ActionType; + +/** + * Transport action to create QueryGroup + * + * @opensearch.api + */ +public class CreateQueryGroupAction extends ActionType { + + /** + * An instance of CreateQueryGroupAction + */ + public static final CreateQueryGroupAction INSTANCE = new CreateQueryGroupAction(); + + /** + * Name for CreateQueryGroupAction + */ + public static final String NAME = "cluster:admin/opensearch/query_group/wlm/_create"; + + /** + * Default constructor + */ + private CreateQueryGroupAction() { + super(NAME, CreateQueryGroupResponse::new); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java new file mode 100644 index 0000000000000..2ce70ed8ac931 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -0,0 +1,252 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; + +import org.opensearch.common.UUIDs; +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.XContentParser; +import org.joda.time.Instant; +import org.opensearch.search.ResourceType; + +import java.io.IOException; +import java.util.Map; +import java.util.HashMap; + +/** + * A request for create QueryGroup + * User input schema: + * { + * "name": "analytics", + * "resiliency_mode": "enforced", + * "resourceLimits": { + * "cpu" : 0.4, + * "memory" : 0.2 + * } + * } + * + * @opensearch.internal + */ +public class CreateQueryGroupRequest extends ActionRequest implements Writeable.Reader { + private String name; + private String _id; + private ResiliencyMode resiliencyMode; + private Map resourceLimits; + private long updatedAtInMillis; + + /** + * Default constructor for CreateQueryGroupRequest + */ + public CreateQueryGroupRequest() {} + + /** + * Constructor for CreateQueryGroupRequest + * @param queryGroup - A {@link QueryGroup} object + */ + public CreateQueryGroupRequest(QueryGroup queryGroup) { + this.name = queryGroup.getName(); + this._id = queryGroup.get_id(); + Map resourceTypesMap = queryGroup.getResourceLimits(); + Map resourceLimits_ = new HashMap<>(); + for (Map.Entry resource : resourceTypesMap.entrySet()) { + resourceLimits_.put(resource.getKey().getName(), resource.getValue()); + } + this.resourceLimits = resourceLimits_; + this.resiliencyMode = queryGroup.getResiliencyMode(); + this.updatedAtInMillis = queryGroup.getUpdatedAtInMillis(); + } + + /** + * Constructor for CreateQueryGroupRequest + * @param name - QueryGroup name for CreateQueryGroupRequest + * @param _id - QueryGroup _id for CreateQueryGroupRequest + * @param mode - QueryGroup mode for CreateQueryGroupRequest + * @param resourceLimits - QueryGroup resourceLimits for CreateQueryGroupRequest + * @param updatedAtInMillis - QueryGroup updated time in millis for CreateQueryGroupRequest + */ + public CreateQueryGroupRequest( + String name, + String _id, + ResiliencyMode mode, + Map resourceLimits, + long updatedAtInMillis + ) { + this.name = name; + this._id = _id; + this.resourceLimits = resourceLimits; + this.resiliencyMode = mode; + this.updatedAtInMillis = updatedAtInMillis; + } + + /** + * Constructor for CreateQueryGroupRequest + * @param in - A {@link StreamInput} object + */ + public CreateQueryGroupRequest(StreamInput in) throws IOException { + super(in); + name = in.readString(); + _id = in.readString(); + resiliencyMode = ResiliencyMode.fromName(in.readString()); + resourceLimits = in.readMap(); + updatedAtInMillis = in.readLong(); + } + + @Override + public CreateQueryGroupRequest read(StreamInput in) throws IOException { + return new CreateQueryGroupRequest(in); + } + + /** + * Generate a CreateQueryGroupRequest from XContent + * @param parser - A {@link XContentParser} object + */ + public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws IOException { + + while (parser.currentToken() != XContentParser.Token.START_OBJECT) { + parser.nextToken(); + } + + if (parser.currentToken() != XContentParser.Token.START_OBJECT) { + throw new IllegalArgumentException("expected start object but got a " + parser.currentToken()); + } + + XContentParser.Token token; + String fieldName = ""; + String name = null; + String _id = UUIDs.randomBase64UUID(); + ResiliencyMode mode = null; + long updatedAtInMillis = Instant.now().getMillis(); + + // Map to hold resources + final Map resourceLimits = new HashMap<>(); + while ((token = parser.nextToken()) != null) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else if (token.isValue()) { + if (fieldName.equals("name")) { + name = parser.text(); + } else if (fieldName.equals("resiliency_mode")) { + mode = ResiliencyMode.fromName(parser.text()); + } else { + throw new IllegalArgumentException("unrecognised [field=" + fieldName + " in QueryGroup"); + } + } else if (token == XContentParser.Token.START_OBJECT) { + if (!fieldName.equals("resourceLimits")) { + throw new IllegalArgumentException( + "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token + ); + } + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else { + resourceLimits.put(fieldName, parser.doubleValue()); + } + } + } + } + return new CreateQueryGroupRequest(name, _id, mode, resourceLimits, updatedAtInMillis); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + /** + * name getter + */ + public String getName() { + return name; + } + + /** + * name setter + * @param name - name to be set + */ + public void setName(String name) { + this.name = name; + } + + /** + * resourceLimits getter + */ + public Map getResourceLimits() { + return resourceLimits; + } + + /** + * resourceLimits setter + * @param resourceLimits - resourceLimit to be set + */ + public void setResourceLimits(Map resourceLimits) { + this.resourceLimits = resourceLimits; + } + + /** + * mode getter + */ + public ResiliencyMode getResiliencyMode() { + return resiliencyMode; + } + + /** + * mode setter + * @param resiliencyMode - mode to be set + */ + public void setResiliencyMode(ResiliencyMode resiliencyMode) { + this.resiliencyMode = resiliencyMode; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(name); + out.writeString(_id); + out.writeString(resiliencyMode.getName()); + out.writeMap(resourceLimits); + out.writeLong(updatedAtInMillis); + } + + /** + * _id getter + */ + public String get_id() { + return _id; + } + + /** + * UUID setter + * @param _id - _id to be set + */ + public void set_id(String _id) { + this._id = _id; + } + + /** + * updatedAtInMillis getter + */ + public long getUpdatedAtInMillis() { + return updatedAtInMillis; + } + + /** + * updatedAtInMillis setter + * @param updatedAtInMillis - updatedAtInMillis to be set + */ + public void setUpdatedAtInMillis(long updatedAtInMillis) { + this.updatedAtInMillis = updatedAtInMillis; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java new file mode 100644 index 0000000000000..4a446698495ff --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java @@ -0,0 +1,74 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; + +import java.io.IOException; + +/** + * Response for the create API for QueryGroup + * + * @opensearch.internal + */ +public class CreateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject { + private final QueryGroup queryGroup; + private RestStatus restStatus; + + /** + * Constructor for CreateQueryGroupResponse + * @param queryGroup - The QueryGroup to be included in the response + * @param restStatus - The restStatus for the response + */ + public CreateQueryGroupResponse(final QueryGroup queryGroup, RestStatus restStatus) { + this.queryGroup = queryGroup; + this.restStatus = restStatus; + } + + /** + * Constructor for CreateQueryGroupResponse + * @param in - A {@link StreamInput} object + */ + public CreateQueryGroupResponse(StreamInput in) throws IOException { + queryGroup = new QueryGroup(in); + restStatus = RestStatus.readFrom(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + queryGroup.writeTo(out); + RestStatus.writeTo(out, restStatus); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return QueryGroup.writeToXContent(queryGroup, builder); + } + + /** + * queryGroup getter + */ + public QueryGroup getQueryGroup() { + return queryGroup; + } + + /** + * restStatus getter + */ + public RestStatus getRestStatus() { + return restStatus; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java new file mode 100644 index 0000000000000..245cd7e04c1fd --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java @@ -0,0 +1,74 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.action.service.Persistable; +import org.opensearch.search.ResourceType; +import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.util.HashMap; +import java.util.Map; + +import static org.opensearch.cluster.metadata.QueryGroup.builder; + +/** + * Transport action to create QueryGroup + * + * @opensearch.internal + */ +public class TransportCreateQueryGroupAction extends HandledTransportAction { + + private final ThreadPool threadPool; + private final Persistable queryGroupPersistenceService; + + /** + * Constructor for TransportCreateQueryGroupAction + * + * @param actionName - action name + * @param transportService - a {@link TransportService} object + * @param actionFilters - a {@link ActionFilters} object + * @param threadPool - a {@link ThreadPool} object + * @param queryGroupPersistenceService - a {@link Persistable} object + */ + @Inject + public TransportCreateQueryGroupAction( + String actionName, + TransportService transportService, + ActionFilters actionFilters, + ThreadPool threadPool, + Persistable queryGroupPersistenceService + ) { + super(CreateQueryGroupAction.NAME, transportService, actionFilters, CreateQueryGroupRequest::new); + this.threadPool = threadPool; + this.queryGroupPersistenceService = queryGroupPersistenceService; + } + + @Override + protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener listener) { + Map resourceTypesMap = new HashMap<>(); + Map resourceLimitsStringMap = request.getResourceLimits(); + for (Map.Entry resource : resourceLimitsStringMap.entrySet()) { + resourceTypesMap.put(ResourceType.fromName(resource.getKey()), resource.getValue()); + } + QueryGroup queryGroup = builder().name(request.getName()) + ._id(request.get_id()) + .mode(request.getResiliencyMode().getName()) + .resourceLimits(resourceTypesMap) + .updatedAt(request.getUpdatedAtInMillis()) + .build(); + threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> queryGroupPersistenceService.persist(queryGroup, listener)); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPlugin.java new file mode 100644 index 0000000000000..c568d3e6bc2ef --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPlugin.java @@ -0,0 +1,62 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.action.ActionRequest; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.common.inject.Module; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.IndexScopedSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.settings.SettingsFilter; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.plugin.wlm.action.rest.RestCreateQueryGroupAction; +import org.opensearch.plugins.ActionPlugin; +import org.opensearch.plugins.Plugin; +import org.opensearch.rest.RestController; +import org.opensearch.rest.RestHandler; + +import java.util.Collection; +import java.util.List; +import java.util.function.Supplier; + +/** + * Plugin class for WorkloadManagement + */ +public class WorkloadManagementPlugin extends Plugin implements ActionPlugin { + + /** + * Default constructor + */ + public WorkloadManagementPlugin() {} + + @Override + public List> getActions() { + return List.of(new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class)); + } + + @Override + public List getRestHandlers( + Settings settings, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster + ) { + return List.of(new RestCreateQueryGroupAction()); + } + + @Override + public Collection createGuiceModules() { + return List.of(new WorkloadManagementPluginModule()); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPluginModule.java new file mode 100644 index 0000000000000..65f92a59a576b --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPluginModule.java @@ -0,0 +1,32 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.inject.AbstractModule; +import org.opensearch.common.inject.TypeLiteral; +import org.opensearch.plugin.wlm.action.service.Persistable; +import org.opensearch.plugin.wlm.action.service.QueryGroupPersistenceService; + +/** + * Guice Module to manage WorkloadManagement related objects + */ +public class WorkloadManagementPluginModule extends AbstractModule { + + /** + * Constructor for WorkloadManagementPluginModule + */ + public WorkloadManagementPluginModule() {} + + @Override + protected void configure() { + bind(new TypeLiteral>() { + }).to(QueryGroupPersistenceService.class).asEagerSingleton(); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java new file mode 100644 index 0000000000000..8f7d2647546f5 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * Base Package of CRUD API of QueryGroup + */ +package org.opensearch.plugin.wlm.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java new file mode 100644 index 0000000000000..669e928f0391b --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.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.plugin.wlm.action.rest; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.action.CreateQueryGroupRequest; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.RestResponse; +import org.opensearch.rest.action.RestResponseListener; + +import java.io.IOException; +import java.util.List; + +import static org.opensearch.rest.RestRequest.Method.POST; +import static org.opensearch.rest.RestRequest.Method.PUT; + +/** + * Rest action to create a QueryGroup + * + * @opensearch.api + */ +public class RestCreateQueryGroupAction extends BaseRestHandler { + + /** + * Constructor for RestCreateQueryGroupAction + */ + public RestCreateQueryGroupAction() {} + + @Override + public String getName() { + return "create_query_group"; + } + + /** + * The list of {@link Route}s that this RestHandler is responsible for handling. + */ + @Override + public List routes() { + return List.of(new Route(POST, "_query_group/"), new Route(PUT, "_query_group/")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + CreateQueryGroupRequest createQueryGroupRequest = new CreateQueryGroupRequest(); + request.applyContentParser((parser) -> parseRestRequest(createQueryGroupRequest, parser)); + return channel -> client.execute(CreateQueryGroupAction.INSTANCE, createQueryGroupRequest, createQueryGroupResponse(channel)); + } + + private void parseRestRequest(CreateQueryGroupRequest request, XContentParser parser) throws IOException { + final CreateQueryGroupRequest createQueryGroupRequest = CreateQueryGroupRequest.fromXContent(parser); + request.setName(createQueryGroupRequest.getName()); + request.set_id(createQueryGroupRequest.get_id()); + request.setResiliencyMode(createQueryGroupRequest.getResiliencyMode()); + request.setResourceLimits(createQueryGroupRequest.getResourceLimits()); + request.setUpdatedAtInMillis(createQueryGroupRequest.getUpdatedAtInMillis()); + } + + private RestResponseListener createQueryGroupResponse(final RestChannel channel) { + return new RestResponseListener<>(channel) { + @Override + public RestResponse buildResponse(final CreateQueryGroupResponse response) throws Exception { + return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS)); + } + }; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/package-info.java new file mode 100644 index 0000000000000..783826608c517 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/package-info.java @@ -0,0 +1,12 @@ +/* + * 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 for the rest classes for QueryGroup CRUD operations + */ +package org.opensearch.plugin.wlm.action.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java new file mode 100644 index 0000000000000..4515893ad5458 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.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.plugin.wlm.action.service; + +import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; + +/** + * This interface defines the key APIs for implementing QueruGroup persistence + */ +public interface Persistable { + + /** + * persists the QueryGroup in a durable storage + * @param queryGroup + * @param listener + */ + void persist(T queryGroup, ActionListener listener); +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java new file mode 100644 index 0000000000000..fb1862d045b7b --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java @@ -0,0 +1,228 @@ +/* + * 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.plugin.wlm.action.service; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.ClusterStateUpdateTask; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterManagerTaskThrottler.ThrottlingKey; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Priority; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.search.ResourceType; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.DoubleAdder; + +import static org.opensearch.search.query_group.QueryGroupServiceSettings.MAX_QUERY_GROUP_COUNT; +import static org.opensearch.search.query_group.QueryGroupServiceSettings.QUERY_GROUP_COUNT_SETTING_NAME; + +/** + * This class defines the functions for QueryGroup persistence + */ +public class QueryGroupPersistenceService implements Persistable { + private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); + private final ClusterService clusterService; + private static final String SOURCE = "query-group-persistence-service"; + private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group"; + private static final String UPDATE_QUERY_GROUP_THROTTLING_KEY = "update-query-group"; + private static final String DELETE_QUERY_GROUP_THROTTLING_KEY = "delete-query-group"; + private final AtomicInteger inflightCreateQueryGroupRequestCount; + private final Map inflightResourceLimitValues; + private volatile int maxQueryGroupCount; + final ThrottlingKey createQueryGroupThrottlingKey; + final ThrottlingKey updateQueryGroupThrottlingKey; + final ThrottlingKey deleteQueryGroupThrottlingKey; + + /** + * Constructor for QueryGroupPersistenceService + * + * @param clusterService {@link ClusterService} - The cluster service to be used by QueryGroupPersistenceService + * @param settings {@link Settings} - The settings to be used by QueryGroupPersistenceService + * @param clusterSettings {@link ClusterSettings} - The cluster settings to be used by QueryGroupPersistenceService + */ + @Inject + public QueryGroupPersistenceService( + final ClusterService clusterService, + final Settings settings, + final ClusterSettings clusterSettings + ) { + this.clusterService = clusterService; + this.createQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(CREATE_QUERY_GROUP_THROTTLING_KEY, true); + this.deleteQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(DELETE_QUERY_GROUP_THROTTLING_KEY, true); + this.updateQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(UPDATE_QUERY_GROUP_THROTTLING_KEY, true); + maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); + clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); + inflightCreateQueryGroupRequestCount = new AtomicInteger(); + inflightResourceLimitValues = new HashMap<>(); + } + + /** + * Set maxQueryGroupCount to be newMaxQueryGroupCount + * @param newMaxQueryGroupCount - the max number of QueryGroup allowed + */ + public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { + if (newMaxQueryGroupCount < 0) { + throw new IllegalArgumentException("node.query_group.max_count can't be negative"); + } + this.maxQueryGroupCount = newMaxQueryGroupCount; + } + + @Override + public void persist(QueryGroup queryGroup, ActionListener listener) { + persistInClusterStateMetadata(queryGroup, listener); + } + + /** + * Update cluster state to include the new QueryGroup + * @param queryGroup {@link QueryGroup} - the QueryGroup we're currently creating + */ + void persistInClusterStateMetadata(QueryGroup queryGroup, ActionListener listener) { + clusterService.submitStateUpdateTask(SOURCE, new ClusterStateUpdateTask(Priority.URGENT) { + @Override + public ClusterState execute(ClusterState currentState) throws Exception { + return saveQueryGroupInClusterState(queryGroup, currentState); + } + + @Override + public ThrottlingKey getClusterManagerThrottlingKey() { + return createQueryGroupThrottlingKey; + } + + @Override + public void onFailure(String source, Exception e) { + restoreInflightValues(queryGroup.getResourceLimits()); + inflightCreateQueryGroupRequestCount.decrementAndGet(); + logger.warn("failed to save QueryGroup object due to error: {}, for source: {}", e.getMessage(), source); + listener.onFailure(e); + } + + @Override + public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) { + restoreInflightValues(queryGroup.getResourceLimits()); + inflightCreateQueryGroupRequestCount.decrementAndGet(); + CreateQueryGroupResponse response = new CreateQueryGroupResponse(queryGroup, RestStatus.OK); + listener.onResponse(response); + } + }); + } + + /** + * Get the allocation value for resourceName for the QueryGroup + * @param resourceName - the resourceName we want to get the usage for + * @param resourceLimits - the resource limit from which to get the allocation value for resourceName + */ + private double getResourceLimitValue(String resourceName, final Map resourceLimits) { + for (ResourceType resourceType: resourceLimits.keySet()) { + if (resourceType.getName().equals(resourceName)) { + return (double) resourceLimits.get(resourceType); + } + } + return 0.0; + } + + /** + * This method will be executed before we submit the new cluster state + * @param queryGroup - the QueryGroup we're currently creating + * @param currentClusterState - the cluster state before the update + */ + ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) { + final Metadata metadata = currentClusterState.metadata(); + String groupName = queryGroup.getName(); + final Map previousGroups = metadata.queryGroups(); + + // check if there's any resource allocation that exceed limit of 1.0 + String resourceNameWithThresholdExceeded = ""; + for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { + String resourceName = resourceType.getName(); + double existingUsage = calculateExistingUsage(resourceName, previousGroups, groupName); + double newGroupUsage = getResourceLimitValue(resourceName, queryGroup.getResourceLimits()); + inflightResourceLimitValues.computeIfAbsent(resourceName, k -> new DoubleAdder()).add(newGroupUsage); + double totalUsage = existingUsage + inflightResourceLimitValues.get(resourceName).doubleValue(); + if (totalUsage > 1) { + resourceNameWithThresholdExceeded = resourceName; + } + } + // check if group count exceed max + boolean groupCountExceeded = inflightCreateQueryGroupRequestCount.incrementAndGet() + previousGroups.size() > maxQueryGroupCount; + + if (previousGroups.containsKey(groupName)) { + logger.warn("QueryGroup with name {} already exists. Not creating a new one.", groupName); + throw new RuntimeException("QueryGroup with name " + groupName + " already exists. Not creating a new one."); + } + if (!resourceNameWithThresholdExceeded.isEmpty()) { + logger.error("Total resource allocation for {} will go above the max limit of 1.0", resourceNameWithThresholdExceeded); + throw new RuntimeException( + "Total resource allocation for " + resourceNameWithThresholdExceeded + " will go above the max limit of 1.0" + ); + } + if (groupCountExceeded) { + logger.error("{} value exceeded its assigned limit of {}", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); + throw new RuntimeException("Can't create more than " + maxQueryGroupCount + " QueryGroups in the system"); + } + Metadata newData = Metadata.builder(metadata).put(queryGroup).build(); + + return ClusterState.builder(currentClusterState).metadata(newData).build(); + } + + /** + * This method restores the inflight values to be before the QueryGroup is processed + * @param resourceLimits - the resourceLimits we're currently restoring + */ + void restoreInflightValues(Map resourceLimits) { + if (resourceLimits == null || resourceLimits.isEmpty()) { + return; + } + for (ResourceType resourceType : resourceLimits.keySet()) { + String currResourceName = resourceType.getName(); + inflightResourceLimitValues.get(currResourceName).add(-getResourceLimitValue(currResourceName, resourceLimits)); + } + } + + /** + * This method calculates the existing total usage of the resource (except the group that we're updating here) + * @param resourceName - the resource name we're calculating + * @param groupsMap - existing QueryGroups + * @param groupName - the QueryGroup name we're updating + */ + private double calculateExistingUsage(String resourceName, Map groupsMap, String groupName) { + double existingUsage = 0; + for (String currGroupId : groupsMap.keySet()) { + QueryGroup currGroup = groupsMap.get(currGroupId); + if (!currGroup.getName().equals(groupName)) { + existingUsage += getResourceLimitValue(resourceName, currGroup.getResourceLimits()); + } + } + return existingUsage; + } + + /** + * inflightCreateQueryGroupRequestCount getter + */ + public AtomicInteger getInflightCreateQueryGroupRequestCount() { + return inflightCreateQueryGroupRequestCount; + } + + /** + * inflightResourceLimitValues getter + */ + public Map getInflightResourceLimitValues() { + return inflightResourceLimitValues; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/package-info.java new file mode 100644 index 0000000000000..e70ac3afb81b5 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/package-info.java @@ -0,0 +1,12 @@ +/* + * 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 for the service classes for QueryGroup CRUD operations + */ +package org.opensearch.plugin.wlm.action.service; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java new file mode 100644 index 0000000000000..5cd0443fb1387 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java @@ -0,0 +1,34 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; + +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.compareResourceLimits; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupOne; + +public class CreateQueryGroupRequestTests extends OpenSearchTestCase { + + public void testSerialization() throws IOException { + CreateQueryGroupRequest request = new CreateQueryGroupRequest(queryGroupOne); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + CreateQueryGroupRequest otherRequest = new CreateQueryGroupRequest(streamInput); + assertEquals(request.getName(), otherRequest.getName()); + assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); + assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); + compareResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); + assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java new file mode 100644 index 0000000000000..3074f7480e94c --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java @@ -0,0 +1,38 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +public class CreateQueryGroupResponseTests extends OpenSearchTestCase { + + public void testSerialization() throws IOException { + CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + CreateQueryGroupResponse otherResponse = new CreateQueryGroupResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + QueryGroup responseGroup = response.getQueryGroup(); + QueryGroup otherResponseGroup = otherResponse.getQueryGroup(); + List listOne = new ArrayList<>(); + List listTwo = new ArrayList<>(); + listOne.add(responseGroup); + listTwo.add(otherResponseGroup); + QueryGroupTestUtils.compareQueryGroups(listOne, listTwo); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java new file mode 100644 index 0000000000000..fce686d257af4 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java @@ -0,0 +1,132 @@ +/* + * 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.plugin.wlm.action; + +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.plugin.wlm.action.service.QueryGroupPersistenceService; +import org.opensearch.search.ResourceType; +import org.opensearch.threadpool.ThreadPool; + +import java.util.*; +import java.util.concurrent.atomic.DoubleAdder; + +import static org.opensearch.cluster.metadata.QueryGroup.builder; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.opensearch.search.ResourceType.fromName; + +public class QueryGroupTestUtils { + public static final String NAME_ONE = "query_group_one"; + public static final String NAME_TWO = "query_group_two"; + public static final String _ID_ONE = "AgfUO5Ja9yfsYlONlYi3TQ=="; + public static final String _ID_TWO = "G5iIqHy4g7eK1qIAAAAIH53=1"; + public static final String NAME_NONE_EXISTED = "query_group_none_existed"; + public static final String MEMORY_STRING = "memory"; + public static final String MONITOR_STRING = "monitor"; + public static final long TIMESTAMP_ONE = 4513232413L; + public static final long TIMESTAMP_TWO = 4513232415L; + public static final QueryGroup queryGroupOne = builder().name(NAME_ONE) + ._id(_ID_ONE) + .mode(MONITOR_STRING) + .resourceLimits(Map.of(fromName(MEMORY_STRING), 0.3)) + .updatedAt(TIMESTAMP_ONE) + .build(); + + public static final QueryGroup queryGroupTwo = builder().name(NAME_TWO) + ._id(_ID_TWO) + .mode(MONITOR_STRING) + .resourceLimits(Map.of(fromName(MEMORY_STRING), 0.6)) + .updatedAt(TIMESTAMP_TWO) + .build(); + + public static final Map queryGroupMap = Map.of(NAME_ONE, queryGroupOne, NAME_TWO, queryGroupTwo); + + public static List queryGroupList() { + List list = new ArrayList<>(); + list.add(queryGroupOne); + list.add(queryGroupTwo); + return list; + } + + public static ClusterState clusterState() { + final Metadata metadata = Metadata.builder().queryGroups(Map.of(NAME_ONE, queryGroupOne, NAME_TWO, queryGroupTwo)).build(); + return ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); + } + + public static Settings settings() { + return Settings.builder().build(); + } + + public static ClusterSettings clusterSettings() { + return new ClusterSettings(settings(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + } + + public static QueryGroupPersistenceService queryGroupPersistenceService() { + ClusterService clusterService = new ClusterService(settings(), clusterSettings(), mock(ThreadPool.class)); + return new QueryGroupPersistenceService(clusterService, settings(), clusterSettings()); + } + + public static List prepareSandboxPersistenceService(Map queryGroups) { + Metadata metadata = Metadata.builder().queryGroups(queryGroups).build(); + Settings settings = Settings.builder().build(); + ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + ClusterService clusterService = new ClusterService(settings, clusterSettings, mock(ThreadPool.class)); + ClusterState clusterState = ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings + ); + return List.of(queryGroupPersistenceService, clusterState); + } + + public static void compareResourceTypes(Map resourceLimitMapOne, Map resourceLimitMapTwo) { + assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet())); + assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); +// for (Map.Entry entryOne : resourceLimitMapOne.entrySet()) { +// String resourceName = entryOne.getKey().getName(); +// Optional> entryTwo = resourceLimitMapTwo.entrySet().stream() +// .filter(e -> e.getKey().getName().equals(resourceName)) +// .findFirst(); +// assertTrue(entryTwo.isPresent()); +// assertEquals(entryOne.getValue(), entryTwo.get().getValue()); +// } + } + + public static void compareResourceLimits(Map resourceLimitMapOne, Map resourceLimitMapTwo) { + assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet())); + assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); + } + + public static void compareQueryGroups(List listOne, List listTwo) { + assertEquals(listOne.size(), listTwo.size()); + listOne.sort(Comparator.comparing(QueryGroup::getName)); + listTwo.sort(Comparator.comparing(QueryGroup::getName)); + for (int i = 0; i < listOne.size(); i++) { + assertTrue(listOne.get(i).equals(listTwo.get(i))); + } + } + + public static void assertInflightValuesAreZero(QueryGroupPersistenceService queryGroupPersistenceService) { + assertEquals(0, queryGroupPersistenceService.getInflightCreateQueryGroupRequestCount().get()); + Map inflightResourceMap = queryGroupPersistenceService.getInflightResourceLimitValues(); + if (inflightResourceMap != null) { + for (String resourceName : inflightResourceMap.keySet()) { + assertEquals(0, inflightResourceMap.get(resourceName).intValue()); + } + } + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java new file mode 100644 index 0000000000000..c37da96b7ebad --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java @@ -0,0 +1,104 @@ +/* + * 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.plugin.wlm.action.service; + +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.search.ResourceType; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; + +import java.util.*; + +import static org.opensearch.cluster.metadata.QueryGroup.builder; +import static org.mockito.Mockito.mock; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.*; +import static org.opensearch.search.query_group.QueryGroupServiceSettings.QUERY_GROUP_COUNT_SETTING_NAME; + +public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase { + + public void testCreateQueryGroup() { + List setup = prepareSandboxPersistenceService(new HashMap<>()); + QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); + ClusterState clusterState = (ClusterState) setup.get(1); + ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupOne, clusterState); + Map updatedGroupsMap = newClusterState.getMetadata().queryGroups(); + assertEquals(1, updatedGroupsMap.size()); + assertTrue(updatedGroupsMap.containsKey(_ID_ONE)); + List listOne = new ArrayList<>(); + List listTwo = new ArrayList<>(); + listOne.add(queryGroupOne); + listTwo.add(updatedGroupsMap.get(_ID_ONE)); + compareQueryGroups(listOne, listTwo); + assertInflightValuesAreZero(queryGroupPersistenceService()); + } + + public void testCreateAnotherQueryGroup() { + List setup = prepareSandboxPersistenceService(Map.of(NAME_ONE, queryGroupOne)); + QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); + ClusterState clusterState = (ClusterState) setup.get(1); + ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupTwo, clusterState); + Map updatedGroups = newClusterState.getMetadata().queryGroups(); + assertEquals(2, updatedGroups.size()); + Collection values = updatedGroups.values(); + compareQueryGroups(queryGroupList(), new ArrayList<>(values)); + assertInflightValuesAreZero(queryGroupPersistenceService()); + } + + public void testCreateQueryGroupDuplicateName() { + List setup = prepareSandboxPersistenceService(Map.of(NAME_ONE, queryGroupOne)); + QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); + ClusterState clusterState = (ClusterState) setup.get(1); + QueryGroup toCreate = builder().name(NAME_ONE) + ._id("W5iIqHyhgi4K1qIAAAAIHw==") + .mode(MONITOR_STRING) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.3)) + .updatedAt(1690934400000L) + .build(); + assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); + } + + public void testCreateQueryGroupOverflowAllocation() { + List setup = prepareSandboxPersistenceService(Map.of(NAME_TWO, queryGroupOne)); + QueryGroup toCreate = builder().name(NAME_TWO) + ._id("W5iIqHyhgi4K1qIAAAAIHw==") + .mode(MONITOR_STRING) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) + .updatedAt(1690934400000L) + .build(); + QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); + ClusterState clusterState = (ClusterState) setup.get(1); + assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); + } + + public void testCreateQueryGroupOverflowCount() { + QueryGroup toCreate = builder().name(NAME_NONE_EXISTED) + ._id("W5iIqHyhgi4K1qIAAAAIHw==") + .mode(MONITOR_STRING) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) + .updatedAt(1690934400000L) + .build(); + Metadata metadata = Metadata.builder().queryGroups(Map.of(NAME_ONE, queryGroupOne, NAME_TWO, queryGroupTwo)).build(); + Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build(); + ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + ClusterService clusterService = new ClusterService(settings, clusterSettings, mock(ThreadPool.class)); + ClusterState clusterState = ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); + QueryGroupPersistenceService queryGroupPersistenceService1 = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings + ); + assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); + } +} diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java index 09bef2ddf9ee6..4da6c68b40733 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -853,6 +853,12 @@ public Map views() { return Optional.ofNullable((ViewMetadata) this.custom(ViewMetadata.TYPE)).map(ViewMetadata::views).orElse(Collections.emptyMap()); } + public Map queryGroups() { + return Optional.ofNullable((QueryGroupMetadata) this.custom(QueryGroupMetadata.TYPE)) + .map(QueryGroupMetadata::queryGroups) + .orElse(Collections.emptyMap()); + } + public DecommissionAttributeMetadata decommissionAttributeMetadata() { return custom(DecommissionAttributeMetadata.TYPE); } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index beaab198073df..8677e28d6b257 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -30,7 +30,7 @@ * { * "_id": "fafjafjkaf9ag8a9ga9g7ag0aagaga", * "resourceLimits": { - * "jvm": 0.4 + * "memory": 0.4 * }, * "resiliency_mode": "enforced", * "name": "analytics", @@ -112,21 +112,26 @@ private void validateResourceLimits(Map resourceLimits) { Objects.requireNonNull(resource.getKey(), "resourceName can't be null"); Objects.requireNonNull(threshold, "resource limit threshold for" + resource.getKey().getName() + " : can't be null"); - if (Double.compare(threshold, 1.0) > 0) { - throw new IllegalArgumentException("resource value should be less than 1.0"); + if (Double.compare(threshold, 0.0) <= 0 || Double.compare(threshold, 1.0) > 0) { + throw new IllegalArgumentException("resource value should be greater than 0 and less or equal to 1.0"); } } } @Override public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { + return writeToXContent(this, builder); + } + + public static XContentBuilder writeToXContent(QueryGroup queryGroup, final XContentBuilder builder) throws IOException { builder.startObject(); - builder.field("_id", _id); - builder.field("name", name); - builder.field("resiliency_mode", resiliencyMode.getName()); - builder.field("updatedAt", updatedAtInMillis); + builder.field("_id", queryGroup.get_id()); + builder.field("name", queryGroup.getName()); + builder.field("resiliency_mode", queryGroup.getResiliencyMode().getName()); + builder.field("updatedAt", queryGroup.getUpdatedAtInMillis()); // write resource limits builder.startObject("resourceLimits"); + Map resourceLimits = queryGroup.getResourceLimits(); for (ResourceType resourceType : ResourceType.values()) { if (resourceLimits.containsKey(resourceType)) { builder.field(resourceType.getName(), resourceLimits.get(resourceType)); diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index a73e5d44b7e02..ff341ae039b7f 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -158,6 +158,7 @@ import org.opensearch.search.backpressure.settings.SearchBackpressureSettings; import org.opensearch.search.backpressure.settings.SearchShardTaskSettings; import org.opensearch.search.backpressure.settings.SearchTaskSettings; +import org.opensearch.search.query_group.QueryGroupServiceSettings; import org.opensearch.search.fetch.subphase.highlight.FastVectorHighlighter; import org.opensearch.snapshots.InternalSnapshotsInfoService; import org.opensearch.snapshots.SnapshotsService; diff --git a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java new file mode 100644 index 0000000000000..bfde3e1daf426 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java @@ -0,0 +1,198 @@ +/* + * 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.search.query_group; + +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; + +/** + * Main class to declare the QueryGroup feature related settings + */ +public class QueryGroupServiceSettings { + private static final Long DEFAULT_RUN_INTERVAL_MILLIS = 1000l; + private static final Double DEFAULT_NODE_LEVEL_REJECTION_THRESHOLD = 0.8; + private static final Double DEFAULT_NODE_LEVEL_CANCELLATION_THRESHOLD = 0.9; + /** + * default max queryGroup count on any node at any given point in time + */ + public static final int DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE = 100; + + public static final String QUERY_GROUP_COUNT_SETTING_NAME = "node.query_group.max_count"; + public static final double NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE = 0.95; + public static final double NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE = 0.90; + + private TimeValue runIntervalMillis; + private Double nodeLevelJvmCancellationThreshold; + private Double nodeLevelJvmRejectionThreshold; + private volatile int maxQueryGroupCount; + /** + * max QueryGroup count setting + */ + public static final Setting MAX_QUERY_GROUP_COUNT = Setting.intSetting( + QUERY_GROUP_COUNT_SETTING_NAME, + DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE, + 0, + (newVal) -> { + if (newVal > 100 || newVal < 1) throw new IllegalArgumentException( + QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]" + ); + }, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** + * Setting name for default QueryGroup count + */ + public static final String SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME = "query_group.service.run_interval_millis"; + /** + * Setting to control the run interval of QSB service + */ + private static final Setting QUERY_GROUP_RUN_INTERVAL_SETTING = Setting.longSetting( + SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME, + DEFAULT_RUN_INTERVAL_MILLIS, + 1, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Setting name for node level rejection threshold for QSB + */ + public static final String NODE_REJECTION_THRESHOLD_SETTING_NAME = "query_group.node.rejection_threshold"; + /** + * Setting to control the rejection threshold + */ + public static final Setting NODE_LEVEL_REJECTION_THRESHOLD = Setting.doubleSetting( + NODE_REJECTION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_REJECTION_THRESHOLD, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** + * Setting name for node level cancellation threshold + */ + public static final String NODE_CANCELLATION_THRESHOLD_SETTING_NAME = "query_group.node.cancellation_threshold"; + /** + * Setting name for node level cancellation threshold + */ + public static final Setting NODE_LEVEL_CANCELLATION_THRESHOLD = Setting.doubleSetting( + NODE_CANCELLATION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_CANCELLATION_THRESHOLD, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * QueryGroup service settings constructor + * @param settings + * @param clusterSettings + */ + public QueryGroupServiceSettings(Settings settings, ClusterSettings clusterSettings) { + runIntervalMillis = new TimeValue(QUERY_GROUP_RUN_INTERVAL_SETTING.get(settings)); + nodeLevelJvmCancellationThreshold = NODE_LEVEL_CANCELLATION_THRESHOLD.get(settings); + nodeLevelJvmRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); + maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + + clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelJvmCancellationThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelJvmRejectionThreshold); + } + + /** + * Method to get runInterval for QSB + * @return runInterval in milliseconds for QSB Service + */ + public TimeValue getRunIntervalMillis() { + return runIntervalMillis; + } + + /** + * Method to set the new QueryGroup count + * @param newMaxQueryGroupCount is the new maxQueryGroupCount per node + */ + public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { + if (newMaxQueryGroupCount < 0) { + throw new IllegalArgumentException("node.node.query_group.max_count can't be negative"); + } + this.maxQueryGroupCount = newMaxQueryGroupCount; + } + + /** + * Method to get the node level cancellation threshold + * @return current node level cancellation threshold + */ + public Double getNodeLevelJvmCancellationThreshold() { + return nodeLevelJvmCancellationThreshold; + } + + /** + * Method to set the node level cancellation threshold + * @param nodeLevelJvmCancellationThreshold sets the new node level cancellation threshold + * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold + */ + public void setNodeLevelJvmCancellationThreshold(Double nodeLevelJvmCancellationThreshold) { + if (Double.compare(nodeLevelJvmCancellationThreshold, NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { + throw new IllegalArgumentException( + NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" + ); + } + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + + this.nodeLevelJvmCancellationThreshold = nodeLevelJvmCancellationThreshold; + } + + /** + * Method to get the node level rejection threshold + * @return the current node level rejection threshold + */ + public Double getNodeLevelJvmRejectionThreshold() { + return nodeLevelJvmRejectionThreshold; + } + + /** + * Method to set the node level rejection threshold + * @param nodeLevelJvmRejectionThreshold sets the new rejection threshold + * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold + */ + public void setNodeLevelJvmRejectionThreshold(Double nodeLevelJvmRejectionThreshold) { + if (Double.compare(nodeLevelJvmRejectionThreshold, NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE) > 0) { + throw new IllegalArgumentException( + NODE_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" + ); + } + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + + this.nodeLevelJvmRejectionThreshold = nodeLevelJvmRejectionThreshold; + } + + private void ensureRejectionThresholdIsLessThanCancellation( + Double nodeLevelJvmRejectionThreshold, + Double nodeLevelJvmCancellationThreshold + ) { + if (Double.compare(nodeLevelJvmCancellationThreshold, nodeLevelJvmRejectionThreshold) < 0) { + throw new IllegalArgumentException( + NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be less than " + NODE_REJECTION_THRESHOLD_SETTING_NAME + ); + } + } + + /** + * Method to get the current QueryGroup count + * @return the current max QueryGroup count + */ + public int getMaxQueryGroupCount() { + return maxQueryGroupCount; + } +} diff --git a/server/src/main/java/org/opensearch/search/query_group/package-info.java b/server/src/main/java/org/opensearch/search/query_group/package-info.java new file mode 100644 index 0000000000000..e7b8443799e83 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/query_group/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * Query Sandboxing related artifacts + */ +package org.opensearch.search.query_group; From 3322ab8cd80be26a370c602940979a4c6dedae2f Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 9 Jul 2024 13:04:38 -0700 Subject: [PATCH 02/20] remove wildcard imports Signed-off-by: Ruirui Zhang --- .../wlm/action/CreateQueryGroupRequest.java | 5 ++- .../service/QueryGroupPersistenceService.java | 2 +- .../wlm/action/QueryGroupTestUtils.java | 22 ++++-------- .../QueryGroupPersistenceServiceTests.java | 36 +++++++++++++------ .../cluster/metadata/QueryGroup.java | 2 +- .../common/settings/ClusterSettings.java | 2 +- .../search/query_group/package-info.java | 2 +- 7 files changed, 38 insertions(+), 33 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java index 2ce70ed8ac931..6057c78b1b0ea 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -12,18 +12,17 @@ import org.opensearch.action.ActionRequestValidationException; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; - import org.opensearch.common.UUIDs; 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.XContentParser; -import org.joda.time.Instant; import org.opensearch.search.ResourceType; +import org.joda.time.Instant; import java.io.IOException; -import java.util.Map; import java.util.HashMap; +import java.util.Map; /** * A request for create QueryGroup diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java index fb1862d045b7b..895938ae9f7da 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java @@ -129,7 +129,7 @@ public void clusterStateProcessed(String source, ClusterState oldState, ClusterS * @param resourceLimits - the resource limit from which to get the allocation value for resourceName */ private double getResourceLimitValue(String resourceName, final Map resourceLimits) { - for (ResourceType resourceType: resourceLimits.keySet()) { + for (ResourceType resourceType : resourceLimits.keySet()) { if (resourceType.getName().equals(resourceName)) { return (double) resourceLimits.get(resourceType); } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java index fce686d257af4..1a67aea92efa2 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java @@ -19,14 +19,17 @@ import org.opensearch.search.ResourceType; import org.opensearch.threadpool.ThreadPool; -import java.util.*; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.DoubleAdder; import static org.opensearch.cluster.metadata.QueryGroup.builder; +import static org.opensearch.search.ResourceType.fromName; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; -import static org.opensearch.search.ResourceType.fromName; public class QueryGroupTestUtils { public static final String NAME_ONE = "query_group_one"; @@ -79,7 +82,7 @@ public static QueryGroupPersistenceService queryGroupPersistenceService() { return new QueryGroupPersistenceService(clusterService, settings(), clusterSettings()); } - public static List prepareSandboxPersistenceService(Map queryGroups) { + public static List preparePersistenceServiceSetup(Map queryGroups) { Metadata metadata = Metadata.builder().queryGroups(queryGroups).build(); Settings settings = Settings.builder().build(); ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); @@ -93,19 +96,6 @@ public static List prepareSandboxPersistenceService(Map resourceLimitMapOne, Map resourceLimitMapTwo) { - assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet())); - assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); -// for (Map.Entry entryOne : resourceLimitMapOne.entrySet()) { -// String resourceName = entryOne.getKey().getName(); -// Optional> entryTwo = resourceLimitMapTwo.entrySet().stream() -// .filter(e -> e.getKey().getName().equals(resourceName)) -// .findFirst(); -// assertTrue(entryTwo.isPresent()); -// assertEquals(entryOne.getValue(), entryTwo.get().getValue()); -// } - } - public static void compareResourceLimits(Map resourceLimitMapOne, Map resourceLimitMapTwo) { assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet())); assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java index c37da96b7ebad..0034c04df90ed 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java @@ -19,17 +19,33 @@ import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; -import java.util.*; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import static org.opensearch.cluster.metadata.QueryGroup.builder; -import static org.mockito.Mockito.mock; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.*; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.MEMORY_STRING; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.MONITOR_STRING; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_NONE_EXISTED; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_ONE; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_TWO; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.assertInflightValuesAreZero; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.compareQueryGroups; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.preparePersistenceServiceSetup; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupList; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupOne; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupPersistenceService; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupTwo; import static org.opensearch.search.query_group.QueryGroupServiceSettings.QUERY_GROUP_COUNT_SETTING_NAME; +import static org.mockito.Mockito.mock; public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase { public void testCreateQueryGroup() { - List setup = prepareSandboxPersistenceService(new HashMap<>()); + List setup = preparePersistenceServiceSetup(new HashMap<>()); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); ClusterState clusterState = (ClusterState) setup.get(1); ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupOne, clusterState); @@ -45,7 +61,7 @@ public void testCreateQueryGroup() { } public void testCreateAnotherQueryGroup() { - List setup = prepareSandboxPersistenceService(Map.of(NAME_ONE, queryGroupOne)); + List setup = preparePersistenceServiceSetup(Map.of(NAME_ONE, queryGroupOne)); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); ClusterState clusterState = (ClusterState) setup.get(1); ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupTwo, clusterState); @@ -57,24 +73,24 @@ public void testCreateAnotherQueryGroup() { } public void testCreateQueryGroupDuplicateName() { - List setup = prepareSandboxPersistenceService(Map.of(NAME_ONE, queryGroupOne)); + List setup = preparePersistenceServiceSetup(Map.of(NAME_ONE, queryGroupOne)); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); ClusterState clusterState = (ClusterState) setup.get(1); QueryGroup toCreate = builder().name(NAME_ONE) ._id("W5iIqHyhgi4K1qIAAAAIHw==") .mode(MONITOR_STRING) - .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.3)) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.3)) .updatedAt(1690934400000L) .build(); assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } public void testCreateQueryGroupOverflowAllocation() { - List setup = prepareSandboxPersistenceService(Map.of(NAME_TWO, queryGroupOne)); + List setup = preparePersistenceServiceSetup(Map.of(NAME_TWO, queryGroupOne)); QueryGroup toCreate = builder().name(NAME_TWO) ._id("W5iIqHyhgi4K1qIAAAAIHw==") .mode(MONITOR_STRING) - .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) .updatedAt(1690934400000L) .build(); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); @@ -86,7 +102,7 @@ public void testCreateQueryGroupOverflowCount() { QueryGroup toCreate = builder().name(NAME_NONE_EXISTED) ._id("W5iIqHyhgi4K1qIAAAAIHw==") .mode(MONITOR_STRING) - .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) .updatedAt(1690934400000L) .build(); Metadata metadata = Metadata.builder().queryGroups(Map.of(NAME_ONE, queryGroupOne, NAME_TWO, queryGroupTwo)).build(); diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index 8677e28d6b257..9df5cf73e57ed 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -206,6 +206,7 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; QueryGroup that = (QueryGroup) o; return Objects.equals(name, that.name) + && Objects.equals(resiliencyMode, that.resiliencyMode) && Objects.equals(resourceLimits, that.resourceLimits) && Objects.equals(_id, that._id) && updatedAtInMillis == that.updatedAtInMillis; @@ -273,7 +274,6 @@ public static ResiliencyMode fromName(String s) { } throw new IllegalArgumentException("Invalid value for QueryGroupMode: " + s); } - } /** diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index ff341ae039b7f..8f213ccd67021 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -158,8 +158,8 @@ import org.opensearch.search.backpressure.settings.SearchBackpressureSettings; import org.opensearch.search.backpressure.settings.SearchShardTaskSettings; import org.opensearch.search.backpressure.settings.SearchTaskSettings; -import org.opensearch.search.query_group.QueryGroupServiceSettings; import org.opensearch.search.fetch.subphase.highlight.FastVectorHighlighter; +import org.opensearch.search.query_group.QueryGroupServiceSettings; import org.opensearch.snapshots.InternalSnapshotsInfoService; import org.opensearch.snapshots.SnapshotsService; import org.opensearch.tasks.TaskCancellationMonitoringSettings; diff --git a/server/src/main/java/org/opensearch/search/query_group/package-info.java b/server/src/main/java/org/opensearch/search/query_group/package-info.java index e7b8443799e83..00b68b0d3306c 100644 --- a/server/src/main/java/org/opensearch/search/query_group/package-info.java +++ b/server/src/main/java/org/opensearch/search/query_group/package-info.java @@ -7,6 +7,6 @@ */ /** - * Query Sandboxing related artifacts + * QueryGroup related artifacts */ package org.opensearch.search.query_group; From 8e0930a40acd8110864e1e2e04aea99368855aeb Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 9 Jul 2024 14:33:31 -0700 Subject: [PATCH 03/20] change jvm to memeory Signed-off-by: Ruirui Zhang --- .../QueryGroupServiceSettings.java | 48 +++++++++---------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java index bfde3e1daf426..7f6e4955cf22f 100644 --- a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java +++ b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java @@ -30,8 +30,8 @@ public class QueryGroupServiceSettings { public static final double NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE = 0.90; private TimeValue runIntervalMillis; - private Double nodeLevelJvmCancellationThreshold; - private Double nodeLevelJvmRejectionThreshold; + private Double nodeLevelMemoryCancellationThreshold; + private Double nodeLevelMemoryRejectionThreshold; private volatile int maxQueryGroupCount; /** * max QueryGroup count setting @@ -97,15 +97,15 @@ public class QueryGroupServiceSettings { */ public QueryGroupServiceSettings(Settings settings, ClusterSettings clusterSettings) { runIntervalMillis = new TimeValue(QUERY_GROUP_RUN_INTERVAL_SETTING.get(settings)); - nodeLevelJvmCancellationThreshold = NODE_LEVEL_CANCELLATION_THRESHOLD.get(settings); - nodeLevelJvmRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); + nodeLevelMemoryCancellationThreshold = NODE_LEVEL_CANCELLATION_THRESHOLD.get(settings); + nodeLevelMemoryRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); - ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelJvmCancellationThreshold); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelJvmRejectionThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelMemoryCancellationThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); } /** @@ -131,57 +131,57 @@ public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { * Method to get the node level cancellation threshold * @return current node level cancellation threshold */ - public Double getNodeLevelJvmCancellationThreshold() { - return nodeLevelJvmCancellationThreshold; + public Double getNodeLevelMemoryCancellationThreshold() { + return nodeLevelMemoryCancellationThreshold; } /** * Method to set the node level cancellation threshold - * @param nodeLevelJvmCancellationThreshold sets the new node level cancellation threshold + * @param nodeLevelMemoryCancellationThreshold sets the new node level cancellation threshold * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold */ - public void setNodeLevelJvmCancellationThreshold(Double nodeLevelJvmCancellationThreshold) { - if (Double.compare(nodeLevelJvmCancellationThreshold, NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { + public void setNodeLevelMemoryCancellationThreshold(Double nodeLevelMemoryCancellationThreshold) { + if (Double.compare(nodeLevelMemoryCancellationThreshold, NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { throw new IllegalArgumentException( NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" ); } - ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); - this.nodeLevelJvmCancellationThreshold = nodeLevelJvmCancellationThreshold; + this.nodeLevelMemoryCancellationThreshold = nodeLevelMemoryCancellationThreshold; } /** * Method to get the node level rejection threshold * @return the current node level rejection threshold */ - public Double getNodeLevelJvmRejectionThreshold() { - return nodeLevelJvmRejectionThreshold; + public Double getNodeLevelMemoryRejectionThreshold() { + return nodeLevelMemoryRejectionThreshold; } /** * Method to set the node level rejection threshold - * @param nodeLevelJvmRejectionThreshold sets the new rejection threshold + * @param nodeLevelMemoryRejectionThreshold sets the new rejection threshold * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold */ - public void setNodeLevelJvmRejectionThreshold(Double nodeLevelJvmRejectionThreshold) { - if (Double.compare(nodeLevelJvmRejectionThreshold, NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE) > 0) { + public void setNodeLevelMemoryRejectionThreshold(Double nodeLevelMemoryRejectionThreshold) { + if (Double.compare(nodeLevelMemoryRejectionThreshold, NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE) > 0) { throw new IllegalArgumentException( NODE_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" ); } - ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); - this.nodeLevelJvmRejectionThreshold = nodeLevelJvmRejectionThreshold; + this.nodeLevelMemoryRejectionThreshold = nodeLevelMemoryRejectionThreshold; } private void ensureRejectionThresholdIsLessThanCancellation( - Double nodeLevelJvmRejectionThreshold, - Double nodeLevelJvmCancellationThreshold + Double nodeLevelMemoryRejectionThreshold, + Double nodeLevelMemoryCancellationThreshold ) { - if (Double.compare(nodeLevelJvmCancellationThreshold, nodeLevelJvmRejectionThreshold) < 0) { + if (Double.compare(nodeLevelMemoryCancellationThreshold, nodeLevelMemoryRejectionThreshold) < 0) { throw new IllegalArgumentException( NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be less than " + NODE_REJECTION_THRESHOLD_SETTING_NAME ); From 4047d0d3f82dab984a1a0ceb48d4f99d7c31b675 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 9 Jul 2024 16:37:46 -0700 Subject: [PATCH 04/20] modify querygroup Signed-off-by: Ruirui Zhang --- .../plugin/wlm/action/CreateQueryGroupResponse.java | 2 +- .../plugin/wlm/action/QueryGroupTestUtils.java | 1 - .../org/opensearch/cluster/metadata/QueryGroup.java | 13 ++++--------- 3 files changed, 5 insertions(+), 11 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java index 4a446698495ff..d4a2cfad5db3d 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java @@ -55,7 +55,7 @@ public void writeTo(StreamOutput out) throws IOException { @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - return QueryGroup.writeToXContent(queryGroup, builder); + return queryGroup.toXContent(builder, params); } /** diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java index 1a67aea92efa2..4ee84e6a419ae 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java @@ -16,7 +16,6 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.plugin.wlm.action.service.QueryGroupPersistenceService; -import org.opensearch.search.ResourceType; import org.opensearch.threadpool.ThreadPool; import java.util.ArrayList; diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index 9df5cf73e57ed..96bb098cbb6e6 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -120,18 +120,13 @@ private void validateResourceLimits(Map resourceLimits) { @Override public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { - return writeToXContent(this, builder); - } - - public static XContentBuilder writeToXContent(QueryGroup queryGroup, final XContentBuilder builder) throws IOException { builder.startObject(); - builder.field("_id", queryGroup.get_id()); - builder.field("name", queryGroup.getName()); - builder.field("resiliency_mode", queryGroup.getResiliencyMode().getName()); - builder.field("updatedAt", queryGroup.getUpdatedAtInMillis()); + builder.field("_id", _id); + builder.field("name", name); + builder.field("resiliency_mode", resiliencyMode.getName()); + builder.field("updatedAt", updatedAtInMillis); // write resource limits builder.startObject("resourceLimits"); - Map resourceLimits = queryGroup.getResourceLimits(); for (ResourceType resourceType : ResourceType.values()) { if (resourceLimits.containsKey(resourceType)) { builder.field(resourceType.getName(), resourceLimits.get(resourceType)); From ed6566c966bc56c3420ee23ab19a7b2ba8fb3419 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 10 Jul 2024 15:53:17 -0700 Subject: [PATCH 05/20] fix javadoc and add more tests Signed-off-by: Ruirui Zhang --- .../wlm/action/CreateQueryGroupAction.java | 2 +- .../wlm/action/service/Persistable.java | 4 ++-- .../action/CreateQueryGroupResponseTests.java | 24 ++++++++++++++++++- .../wlm/action/QueryGroupTestUtils.java | 4 +--- .../QueryGroupServiceSettings.java | 4 ++-- 5 files changed, 29 insertions(+), 9 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java index 003b8b60c4ff9..c5b3ca05643fe 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java @@ -25,7 +25,7 @@ public class CreateQueryGroupAction extends ActionType /** * Name for CreateQueryGroupAction */ - public static final String NAME = "cluster:admin/opensearch/query_group/wlm/_create"; + public static final String NAME = "cluster:admin/opensearch/wlm/query_group/_create"; /** * Default constructor diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java index 4515893ad5458..c136ef0219734 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java @@ -18,8 +18,8 @@ public interface Persistable { /** * persists the QueryGroup in a durable storage - * @param queryGroup - * @param listener + * @param queryGroup - queryGroup to be persisted + * @param listener - ActionListener for CreateQueryGroupResponse */ void persist(T queryGroup, ActionListener listener); } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java index 3074f7480e94c..0efd989d6bc6f 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java @@ -10,18 +10,24 @@ import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.ArrayList; import java.util.List; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupOne; +import static org.mockito.Mockito.mock; + public class CreateQueryGroupResponseTests extends OpenSearchTestCase { public void testSerialization() throws IOException { - CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); + CreateQueryGroupResponse response = new CreateQueryGroupResponse(queryGroupOne, RestStatus.OK); BytesStreamOutput out = new BytesStreamOutput(); response.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); @@ -35,4 +41,20 @@ public void testSerialization() throws IOException { listTwo.add(otherResponseGroup); QueryGroupTestUtils.compareQueryGroups(listOne, listTwo); } + + public void testToXContentCreateQueryGroup() throws IOException { + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + CreateQueryGroupResponse response = new CreateQueryGroupResponse(queryGroupOne, RestStatus.OK); + String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + + " \"name\" : \"query_group_one\",\n" + + " \"resiliency_mode\" : \"monitor\",\n" + + " \"updatedAt\" : 4513232413,\n" + + " \"resourceLimits\" : {\n" + + " \"memory\" : 0.3\n" + + " }\n" + + "}"; + assertEquals(expected, actual); + } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java index 4ee84e6a419ae..2a82ebcfd163d 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java @@ -54,8 +54,6 @@ public class QueryGroupTestUtils { .updatedAt(TIMESTAMP_TWO) .build(); - public static final Map queryGroupMap = Map.of(NAME_ONE, queryGroupOne, NAME_TWO, queryGroupTwo); - public static List queryGroupList() { List list = new ArrayList<>(); list.add(queryGroupOne); @@ -64,7 +62,7 @@ public static List queryGroupList() { } public static ClusterState clusterState() { - final Metadata metadata = Metadata.builder().queryGroups(Map.of(NAME_ONE, queryGroupOne, NAME_TWO, queryGroupTwo)).build(); + final Metadata metadata = Metadata.builder().queryGroups(Map.of(_ID_ONE, queryGroupOne, _ID_TWO, queryGroupTwo)).build(); return ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); } diff --git a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java index 7f6e4955cf22f..425916ce3e924 100644 --- a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java +++ b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java @@ -92,8 +92,8 @@ public class QueryGroupServiceSettings { /** * QueryGroup service settings constructor - * @param settings - * @param clusterSettings + * @param settings - QueryGroup service settings + * @param clusterSettings - QueryGroup cluster settings */ public QueryGroupServiceSettings(Settings settings, ClusterSettings clusterSettings) { runIntervalMillis = new TimeValue(QUERY_GROUP_RUN_INTERVAL_SETTING.get(settings)); From 123620bb1a12e304434213113c4a13e6ac0fed7d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 12 Jul 2024 16:03:23 -0700 Subject: [PATCH 06/20] add more tests Signed-off-by: Ruirui Zhang --- .../service/QueryGroupPersistenceService.java | 7 ++++ .../wlm/action/QueryGroupTestUtils.java | 32 +++++++++++++++++-- .../QueryGroupPersistenceServiceTests.java | 23 ++++++++++--- 3 files changed, 55 insertions(+), 7 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java index 895938ae9f7da..2c855b86356f0 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java @@ -225,4 +225,11 @@ public AtomicInteger getInflightCreateQueryGroupRequestCount() { public Map getInflightResourceLimitValues() { return inflightResourceLimitValues; } + + /** + * clusterService getter + */ + public ClusterService getClusterService() { + return clusterService; + } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java index 2a82ebcfd163d..5a070b0127db3 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java @@ -12,6 +12,8 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterApplierService; +import org.opensearch.cluster.service.ClusterManagerService; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; @@ -75,16 +77,40 @@ public static ClusterSettings clusterSettings() { } public static QueryGroupPersistenceService queryGroupPersistenceService() { - ClusterService clusterService = new ClusterService(settings(), clusterSettings(), mock(ThreadPool.class)); + ClusterApplierService clusterApplierService = new ClusterApplierService( + "name", + settings(), + clusterSettings(), + mock(ThreadPool.class) + ); + clusterApplierService.setInitialState(clusterState()); + ClusterService clusterService = new ClusterService( + settings(), + clusterSettings(), + mock(ClusterManagerService.class), + clusterApplierService + ); return new QueryGroupPersistenceService(clusterService, settings(), clusterSettings()); } public static List preparePersistenceServiceSetup(Map queryGroups) { Metadata metadata = Metadata.builder().queryGroups(queryGroups).build(); Settings settings = Settings.builder().build(); - ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); - ClusterService clusterService = new ClusterService(settings, clusterSettings, mock(ThreadPool.class)); ClusterState clusterState = ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); + ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + ClusterApplierService clusterApplierService = new ClusterApplierService( + "name", + settings(), + clusterSettings(), + mock(ThreadPool.class) + ); + clusterApplierService.setInitialState(clusterState); + ClusterService clusterService = new ClusterService( + settings(), + clusterSettings(), + mock(ClusterManagerService.class), + clusterApplierService + ); QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( clusterService, settings, diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java index 0034c04df90ed..727066a043645 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java @@ -15,6 +15,8 @@ import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; @@ -32,6 +34,7 @@ import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_ONE; import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_TWO; import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils._ID_TWO; import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.assertInflightValuesAreZero; import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.compareQueryGroups; import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.preparePersistenceServiceSetup; @@ -61,19 +64,20 @@ public void testCreateQueryGroup() { } public void testCreateAnotherQueryGroup() { - List setup = preparePersistenceServiceSetup(Map.of(NAME_ONE, queryGroupOne)); + List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); ClusterState clusterState = (ClusterState) setup.get(1); ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupTwo, clusterState); Map updatedGroups = newClusterState.getMetadata().queryGroups(); assertEquals(2, updatedGroups.size()); + assertTrue(updatedGroups.containsKey(_ID_TWO)); Collection values = updatedGroups.values(); compareQueryGroups(queryGroupList(), new ArrayList<>(values)); assertInflightValuesAreZero(queryGroupPersistenceService()); } public void testCreateQueryGroupDuplicateName() { - List setup = preparePersistenceServiceSetup(Map.of(NAME_ONE, queryGroupOne)); + List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); ClusterState clusterState = (ClusterState) setup.get(1); QueryGroup toCreate = builder().name(NAME_ONE) @@ -86,7 +90,7 @@ public void testCreateQueryGroupDuplicateName() { } public void testCreateQueryGroupOverflowAllocation() { - List setup = preparePersistenceServiceSetup(Map.of(NAME_TWO, queryGroupOne)); + List setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupOne)); QueryGroup toCreate = builder().name(NAME_TWO) ._id("W5iIqHyhgi4K1qIAAAAIHw==") .mode(MONITOR_STRING) @@ -105,7 +109,7 @@ public void testCreateQueryGroupOverflowCount() { .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) .updatedAt(1690934400000L) .build(); - Metadata metadata = Metadata.builder().queryGroups(Map.of(NAME_ONE, queryGroupOne, NAME_TWO, queryGroupTwo)).build(); + Metadata metadata = Metadata.builder().queryGroups(Map.of(_ID_ONE, queryGroupOne, _ID_TWO, queryGroupTwo)).build(); Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build(); ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); ClusterService clusterService = new ClusterService(settings, clusterSettings, mock(ThreadPool.class)); @@ -117,4 +121,15 @@ public void testCreateQueryGroupOverflowCount() { ); assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } + + @SuppressWarnings("unchecked") + public void testPersist() { + List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); + QueryGroupPersistenceService queryGroupPersistenceService = (QueryGroupPersistenceService) setup.get(0); + ActionListener mockListener = mock(ActionListener.class); + queryGroupPersistenceService.persist(queryGroupTwo, mockListener); + Map newQueryGroups = queryGroupPersistenceService.getClusterService().state().metadata().queryGroups(); + assertEquals(2, newQueryGroups.size()); + assertTrue(newQueryGroups.containsKey(_ID_TWO)); + } } From c40235184575bdbf470ccf0a95c0ab8924b3398d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Mon, 15 Jul 2024 13:44:16 -0700 Subject: [PATCH 07/20] address comments Signed-off-by: Ruirui Zhang --- plugins/workload-management/build.gradle | 2 +- .../{action => }/CreateQueryGroupAction.java | 2 +- .../{action => }/CreateQueryGroupRequest.java | 29 +++++++---------- .../CreateQueryGroupResponse.java | 2 +- .../TransportCreateQueryGroupAction.java | 15 ++------- .../WorkloadManagementPlugin.java | 4 +-- .../WorkloadManagementPluginModule.java | 6 ++-- .../plugin/wlm/{action => }/package-info.java | 4 +-- .../rest/RestCreateQueryGroupAction.java | 10 +++--- .../wlm/{action => }/rest/package-info.java | 2 +- .../wlm/{action => }/service/Persistable.java | 4 +-- .../service/QueryGroupPersistenceService.java | 4 +-- .../{action => }/service/package-info.java | 2 +- .../CreateQueryGroupRequestTests.java | 9 ++---- .../CreateQueryGroupResponseTests.java | 7 ++-- .../wlm/{action => }/QueryGroupTestUtils.java | 7 ++-- .../QueryGroupPersistenceServiceTests.java | 32 +++++++++---------- 17 files changed, 61 insertions(+), 80 deletions(-) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/CreateQueryGroupAction.java (95%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/CreateQueryGroupRequest.java (86%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/CreateQueryGroupResponse.java (98%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/TransportCreateQueryGroupAction.java (80%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/WorkloadManagementPlugin.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/WorkloadManagementPluginModule.java (81%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/package-info.java (71%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/rest/RestCreateQueryGroupAction.java (89%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/rest/package-info.java (85%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/service/Persistable.java (84%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/service/QueryGroupPersistenceService.java (98%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => }/service/package-info.java (84%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{action => }/CreateQueryGroupRequestTests.java (78%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{action => }/CreateQueryGroupResponseTests.java (92%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{action => }/QueryGroupTestUtils.java (95%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{action => }/service/QueryGroupPersistenceServiceTests.java (83%) diff --git a/plugins/workload-management/build.gradle b/plugins/workload-management/build.gradle index 89e13c079795e..e91ef90005c74 100644 --- a/plugins/workload-management/build.gradle +++ b/plugins/workload-management/build.gradle @@ -11,7 +11,7 @@ opensearchplugin { description 'OpenSearch Workload Management Plugin.' - classname 'org.opensearch.plugin.wlm.action.WorkloadManagementPlugin' + classname 'org.opensearch.plugin.wlm.WorkloadManagementPlugin' } dependencies { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java similarity index 95% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java index c5b3ca05643fe..1a990db56520a 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java similarity index 86% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java index 6057c78b1b0ea..d36b70158f5d1 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; @@ -42,7 +42,7 @@ public class CreateQueryGroupRequest extends ActionRequest implements Writeable. private String name; private String _id; private ResiliencyMode resiliencyMode; - private Map resourceLimits; + private Map resourceLimits; private long updatedAtInMillis; /** @@ -57,12 +57,7 @@ public CreateQueryGroupRequest() {} public CreateQueryGroupRequest(QueryGroup queryGroup) { this.name = queryGroup.getName(); this._id = queryGroup.get_id(); - Map resourceTypesMap = queryGroup.getResourceLimits(); - Map resourceLimits_ = new HashMap<>(); - for (Map.Entry resource : resourceTypesMap.entrySet()) { - resourceLimits_.put(resource.getKey().getName(), resource.getValue()); - } - this.resourceLimits = resourceLimits_; + this.resourceLimits = queryGroup.getResourceLimits(); this.resiliencyMode = queryGroup.getResiliencyMode(); this.updatedAtInMillis = queryGroup.getUpdatedAtInMillis(); } @@ -79,7 +74,7 @@ public CreateQueryGroupRequest( String name, String _id, ResiliencyMode mode, - Map resourceLimits, + Map resourceLimits, long updatedAtInMillis ) { this.name = name; @@ -98,7 +93,7 @@ public CreateQueryGroupRequest(StreamInput in) throws IOException { name = in.readString(); _id = in.readString(); resiliencyMode = ResiliencyMode.fromName(in.readString()); - resourceLimits = in.readMap(); + resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readGenericValue); updatedAtInMillis = in.readLong(); } @@ -124,12 +119,10 @@ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws XContentParser.Token token; String fieldName = ""; String name = null; - String _id = UUIDs.randomBase64UUID(); ResiliencyMode mode = null; - long updatedAtInMillis = Instant.now().getMillis(); // Map to hold resources - final Map resourceLimits = new HashMap<>(); + final Map resourceLimits = new HashMap<>(); while ((token = parser.nextToken()) != null) { if (token == XContentParser.Token.FIELD_NAME) { fieldName = parser.currentName(); @@ -151,12 +144,12 @@ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws if (token == XContentParser.Token.FIELD_NAME) { fieldName = parser.currentName(); } else { - resourceLimits.put(fieldName, parser.doubleValue()); + resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); } } } } - return new CreateQueryGroupRequest(name, _id, mode, resourceLimits, updatedAtInMillis); + return new CreateQueryGroupRequest(name, UUIDs.randomBase64UUID(), mode, resourceLimits, Instant.now().getMillis()); } @Override @@ -182,7 +175,7 @@ public void setName(String name) { /** * resourceLimits getter */ - public Map getResourceLimits() { + public Map getResourceLimits() { return resourceLimits; } @@ -190,7 +183,7 @@ public Map getResourceLimits() { * resourceLimits setter * @param resourceLimits - resourceLimit to be set */ - public void setResourceLimits(Map resourceLimits) { + public void setResourceLimits(Map resourceLimits) { this.resourceLimits = resourceLimits; } @@ -215,7 +208,7 @@ public void writeTo(StreamOutput out) throws IOException { out.writeString(name); out.writeString(_id); out.writeString(resiliencyMode.getName()); - out.writeMap(resourceLimits); + out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue); out.writeLong(updatedAtInMillis); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java similarity index 98% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java index d4a2cfad5db3d..2a19486d52136 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java similarity index 80% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java index 245cd7e04c1fd..7f3fdb823e5b2 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java @@ -6,22 +6,18 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.HandledTransportAction; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.action.service.Persistable; -import org.opensearch.search.ResourceType; +import org.opensearch.plugin.wlm.service.Persistable; import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; -import java.util.HashMap; -import java.util.Map; - import static org.opensearch.cluster.metadata.QueryGroup.builder; /** @@ -58,15 +54,10 @@ public TransportCreateQueryGroupAction( @Override protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener listener) { - Map resourceTypesMap = new HashMap<>(); - Map resourceLimitsStringMap = request.getResourceLimits(); - for (Map.Entry resource : resourceLimitsStringMap.entrySet()) { - resourceTypesMap.put(ResourceType.fromName(resource.getKey()), resource.getValue()); - } QueryGroup queryGroup = builder().name(request.getName()) ._id(request.get_id()) .mode(request.getResiliencyMode().getName()) - .resourceLimits(resourceTypesMap) + .resourceLimits(request.getResourceLimits()) .updatedAt(request.getUpdatedAtInMillis()) .build(); threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> queryGroupPersistenceService.persist(queryGroup, listener)); diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPlugin.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index c568d3e6bc2ef..065c2f3229d15 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.action.ActionRequest; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; @@ -17,7 +17,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; import org.opensearch.core.action.ActionResponse; -import org.opensearch.plugin.wlm.action.rest.RestCreateQueryGroupAction; +import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; import org.opensearch.rest.RestController; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java similarity index 81% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPluginModule.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java index 65f92a59a576b..58208af293f14 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/WorkloadManagementPluginModule.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java @@ -6,13 +6,13 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.inject.AbstractModule; import org.opensearch.common.inject.TypeLiteral; -import org.opensearch.plugin.wlm.action.service.Persistable; -import org.opensearch.plugin.wlm.action.service.QueryGroupPersistenceService; +import org.opensearch.plugin.wlm.service.Persistable; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; /** * Guice Module to manage WorkloadManagement related objects diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java similarity index 71% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java index 8f7d2647546f5..d85cf8d417a5f 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java @@ -7,6 +7,6 @@ */ /** - * Base Package of CRUD API of QueryGroup + * Base package for CRUD API of QueryGroup */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java similarity index 89% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java index 669e928f0391b..32b06b4f98cdc 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java @@ -6,15 +6,15 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action.rest; +package org.opensearch.plugin.wlm.rest; import org.opensearch.client.node.NodeClient; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; -import org.opensearch.plugin.wlm.action.CreateQueryGroupRequest; -import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.plugin.wlm.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.CreateQueryGroupRequest; +import org.opensearch.plugin.wlm.CreateQueryGroupResponse; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; @@ -50,7 +50,7 @@ public String getName() { */ @Override public List routes() { - return List.of(new Route(POST, "_query_group/"), new Route(PUT, "_query_group/")); + return List.of(new Route(POST, "_wlm/_query_group/"), new Route(PUT, "_wlm/_query_group/")); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java similarity index 85% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java index 783826608c517..6d1853f04ce58 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/rest/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java @@ -9,4 +9,4 @@ /** * Package for the rest classes for QueryGroup CRUD operations */ -package org.opensearch.plugin.wlm.action.rest; +package org.opensearch.plugin.wlm.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/Persistable.java similarity index 84% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/Persistable.java index c136ef0219734..9683a5b0ecc42 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/Persistable.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/Persistable.java @@ -6,10 +6,10 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action.service; +package org.opensearch.plugin.wlm.service; import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.plugin.wlm.CreateQueryGroupResponse; /** * This interface defines the key APIs for implementing QueruGroup persistence diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java similarity index 98% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 2c855b86356f0..7a442d3025a99 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action.service; +package org.opensearch.plugin.wlm.service; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -22,7 +22,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; -import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.plugin.wlm.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; import java.util.HashMap; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java similarity index 84% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java index e70ac3afb81b5..ecfc94568888a 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/service/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java @@ -9,4 +9,4 @@ /** * Package for the service classes for QueryGroup CRUD operations */ -package org.opensearch.plugin.wlm.action.service; +package org.opensearch.plugin.wlm.service; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupRequestTests.java similarity index 78% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupRequestTests.java index 5cd0443fb1387..ccc208f3c86d1 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; @@ -14,13 +14,10 @@ import java.io.IOException; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.compareResourceLimits; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupOne; - public class CreateQueryGroupRequestTests extends OpenSearchTestCase { public void testSerialization() throws IOException { - CreateQueryGroupRequest request = new CreateQueryGroupRequest(queryGroupOne); + CreateQueryGroupRequest request = new CreateQueryGroupRequest(QueryGroupTestUtils.queryGroupOne); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); @@ -28,7 +25,7 @@ public void testSerialization() throws IOException { assertEquals(request.getName(), otherRequest.getName()); assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); - compareResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); + QueryGroupTestUtils.compareResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupResponseTests.java similarity index 92% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupResponseTests.java index 0efd989d6bc6f..038b8bfad87a4 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.io.stream.BytesStreamOutput; @@ -21,13 +21,12 @@ import java.util.ArrayList; import java.util.List; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupOne; import static org.mockito.Mockito.mock; public class CreateQueryGroupResponseTests extends OpenSearchTestCase { public void testSerialization() throws IOException { - CreateQueryGroupResponse response = new CreateQueryGroupResponse(queryGroupOne, RestStatus.OK); + CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); BytesStreamOutput out = new BytesStreamOutput(); response.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); @@ -44,7 +43,7 @@ public void testSerialization() throws IOException { public void testToXContentCreateQueryGroup() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - CreateQueryGroupResponse response = new CreateQueryGroupResponse(queryGroupOne, RestStatus.OK); + CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java similarity index 95% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java index 5a070b0127db3..202969d3e0c39 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm; import org.opensearch.cluster.ClusterName; import org.opensearch.cluster.ClusterState; @@ -17,7 +17,8 @@ import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; -import org.opensearch.plugin.wlm.action.service.QueryGroupPersistenceService; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; +import org.opensearch.search.ResourceType; import org.opensearch.threadpool.ThreadPool; import java.util.ArrayList; @@ -119,7 +120,7 @@ public static List preparePersistenceServiceSetup(Map resourceLimitMapOne, Map resourceLimitMapTwo) { + public static void compareResourceLimits(Map resourceLimitMapOne, Map resourceLimitMapTwo) { assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet())); assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java similarity index 83% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 727066a043645..1bd2801e7a91c 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action.service; +package org.opensearch.plugin.wlm.service; import org.opensearch.cluster.ClusterName; import org.opensearch.cluster.ClusterState; @@ -16,7 +16,7 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.plugin.wlm.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; @@ -28,20 +28,20 @@ import java.util.Map; import static org.opensearch.cluster.metadata.QueryGroup.builder; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.MEMORY_STRING; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.MONITOR_STRING; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_NONE_EXISTED; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_ONE; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.NAME_TWO; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils._ID_ONE; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils._ID_TWO; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.assertInflightValuesAreZero; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.compareQueryGroups; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.preparePersistenceServiceSetup; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupList; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupOne; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupPersistenceService; -import static org.opensearch.plugin.wlm.action.QueryGroupTestUtils.queryGroupTwo; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MEMORY_STRING; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MONITOR_STRING; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_NONE_EXISTED; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_ONE; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_TWO; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertInflightValuesAreZero; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.compareQueryGroups; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupPersistenceService; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo; import static org.opensearch.search.query_group.QueryGroupServiceSettings.QUERY_GROUP_COUNT_SETTING_NAME; import static org.mockito.Mockito.mock; From e9e8d1b4499a8763cc717485a53d80cff4b468b4 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 16 Jul 2024 13:37:02 -0700 Subject: [PATCH 08/20] fix the persist logic Signed-off-by: Ruirui Zhang --- .../service/QueryGroupPersistenceService.java | 8 +++++++- .../QueryGroupPersistenceServiceTests.java | 18 ++---------------- 2 files changed, 9 insertions(+), 17 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 7a442d3025a99..032e40612ea9b 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -27,6 +27,7 @@ import java.util.HashMap; import java.util.Map; +import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.DoubleAdder; @@ -162,7 +163,12 @@ ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final Clu // check if group count exceed max boolean groupCountExceeded = inflightCreateQueryGroupRequestCount.incrementAndGet() + previousGroups.size() > maxQueryGroupCount; - if (previousGroups.containsKey(groupName)) { + Optional findExistingGroup = previousGroups.values() + .stream() + .filter(group -> group.getName().equals(groupName)) + .findFirst(); + + if (findExistingGroup.isPresent()) { logger.warn("QueryGroup with name {} already exists. Not creating a new one.", groupName); throw new RuntimeException("QueryGroup with name " + groupName + " already exists. Not creating a new one."); } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 1bd2801e7a91c..b0c344670580e 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -15,8 +15,6 @@ import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; -import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; @@ -32,7 +30,6 @@ import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MONITOR_STRING; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_NONE_EXISTED; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_ONE; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_TWO; import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertInflightValuesAreZero; @@ -90,8 +87,8 @@ public void testCreateQueryGroupDuplicateName() { } public void testCreateQueryGroupOverflowAllocation() { - List setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupOne)); - QueryGroup toCreate = builder().name(NAME_TWO) + List setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupTwo)); + QueryGroup toCreate = builder().name(NAME_ONE) ._id("W5iIqHyhgi4K1qIAAAAIHw==") .mode(MONITOR_STRING) .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) @@ -121,15 +118,4 @@ public void testCreateQueryGroupOverflowCount() { ); assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } - - @SuppressWarnings("unchecked") - public void testPersist() { - List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); - QueryGroupPersistenceService queryGroupPersistenceService = (QueryGroupPersistenceService) setup.get(0); - ActionListener mockListener = mock(ActionListener.class); - queryGroupPersistenceService.persist(queryGroupTwo, mockListener); - Map newQueryGroups = queryGroupPersistenceService.getClusterService().state().metadata().queryGroups(); - assertEquals(2, newQueryGroups.size()); - assertTrue(newQueryGroups.containsKey(_ID_TWO)); - } } From a5d14b6a8dd27b1da54c5a32354a1bbf13d7d264 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 18 Jul 2024 12:43:12 -0700 Subject: [PATCH 09/20] remove inflight checks as they are not necessary Signed-off-by: Ruirui Zhang --- .../plugin/wlm/CreateQueryGroupAction.java | 2 +- .../plugin/wlm/CreateQueryGroupRequest.java | 11 +-- .../plugin/wlm/CreateQueryGroupResponse.java | 2 +- .../wlm/TransportCreateQueryGroupAction.java | 2 +- .../wlm/rest/RestCreateQueryGroupAction.java | 2 +- .../service/QueryGroupPersistenceService.java | 97 +++++-------------- .../plugin/wlm/QueryGroupTestUtils.java | 11 --- .../QueryGroupPersistenceServiceTests.java | 4 - 8 files changed, 34 insertions(+), 97 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java index 1a990db56520a..bfc931ba1a2f9 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java @@ -13,7 +13,7 @@ /** * Transport action to create QueryGroup * - * @opensearch.api + * @opensearch.experimental */ public class CreateQueryGroupAction extends ActionType { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java index d36b70158f5d1..c8728c20bf000 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java @@ -16,6 +16,7 @@ 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.XContentParseException; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.search.ResourceType; import org.joda.time.Instant; @@ -36,7 +37,7 @@ * } * } * - * @opensearch.internal + * @opensearch.experimental */ public class CreateQueryGroupRequest extends ActionRequest implements Writeable.Reader { private String name; @@ -113,7 +114,7 @@ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws } if (parser.currentToken() != XContentParser.Token.START_OBJECT) { - throw new IllegalArgumentException("expected start object but got a " + parser.currentToken()); + throw new XContentParseException("expected start object but got a " + parser.currentToken()); } XContentParser.Token token; @@ -132,13 +133,11 @@ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws } else if (fieldName.equals("resiliency_mode")) { mode = ResiliencyMode.fromName(parser.text()); } else { - throw new IllegalArgumentException("unrecognised [field=" + fieldName + " in QueryGroup"); + throw new XContentParseException("unrecognised [field=" + fieldName + " in QueryGroup"); } } else if (token == XContentParser.Token.START_OBJECT) { if (!fieldName.equals("resourceLimits")) { - throw new IllegalArgumentException( - "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token - ); + throw new XContentParseException("Invalid field passed. QueryGroup does not support " + fieldName + "."); } while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { if (token == XContentParser.Token.FIELD_NAME) { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java index 2a19486d52136..30e052b5a4d31 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java @@ -22,7 +22,7 @@ /** * Response for the create API for QueryGroup * - * @opensearch.internal + * @opensearch.experimental */ public class CreateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject { private final QueryGroup queryGroup; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java index 7f3fdb823e5b2..670fb870e44ce 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java @@ -23,7 +23,7 @@ /** * Transport action to create QueryGroup * - * @opensearch.internal + * @opensearch.experimental */ public class TransportCreateQueryGroupAction extends HandledTransportAction { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java index 32b06b4f98cdc..0f488be15bd11 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java @@ -31,7 +31,7 @@ /** * Rest action to create a QueryGroup * - * @opensearch.api + * @opensearch.experimental */ public class RestCreateQueryGroupAction extends BaseRestHandler { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 032e40612ea9b..cb833f02db569 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -25,11 +25,8 @@ import org.opensearch.plugin.wlm.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; -import java.util.HashMap; import java.util.Map; import java.util.Optional; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.DoubleAdder; import static org.opensearch.search.query_group.QueryGroupServiceSettings.MAX_QUERY_GROUP_COUNT; import static org.opensearch.search.query_group.QueryGroupServiceSettings.QUERY_GROUP_COUNT_SETTING_NAME; @@ -44,8 +41,6 @@ public class QueryGroupPersistenceService implements Persistable { private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group"; private static final String UPDATE_QUERY_GROUP_THROTTLING_KEY = "update-query-group"; private static final String DELETE_QUERY_GROUP_THROTTLING_KEY = "delete-query-group"; - private final AtomicInteger inflightCreateQueryGroupRequestCount; - private final Map inflightResourceLimitValues; private volatile int maxQueryGroupCount; final ThrottlingKey createQueryGroupThrottlingKey; final ThrottlingKey updateQueryGroupThrottlingKey; @@ -70,8 +65,6 @@ public QueryGroupPersistenceService( this.updateQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(UPDATE_QUERY_GROUP_THROTTLING_KEY, true); maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); - inflightCreateQueryGroupRequestCount = new AtomicInteger(); - inflightResourceLimitValues = new HashMap<>(); } /** @@ -95,7 +88,7 @@ public void persist(QueryGroup queryGroup, ActionListener listener) { - clusterService.submitStateUpdateTask(SOURCE, new ClusterStateUpdateTask(Priority.URGENT) { + clusterService.submitStateUpdateTask(SOURCE, new ClusterStateUpdateTask(Priority.NORMAL) { @Override public ClusterState execute(ClusterState currentState) throws Exception { return saveQueryGroupInClusterState(queryGroup, currentState); @@ -108,36 +101,18 @@ public ThrottlingKey getClusterManagerThrottlingKey() { @Override public void onFailure(String source, Exception e) { - restoreInflightValues(queryGroup.getResourceLimits()); - inflightCreateQueryGroupRequestCount.decrementAndGet(); logger.warn("failed to save QueryGroup object due to error: {}, for source: {}", e.getMessage(), source); listener.onFailure(e); } @Override public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) { - restoreInflightValues(queryGroup.getResourceLimits()); - inflightCreateQueryGroupRequestCount.decrementAndGet(); CreateQueryGroupResponse response = new CreateQueryGroupResponse(queryGroup, RestStatus.OK); listener.onResponse(response); } }); } - /** - * Get the allocation value for resourceName for the QueryGroup - * @param resourceName - the resourceName we want to get the usage for - * @param resourceLimits - the resource limit from which to get the allocation value for resourceName - */ - private double getResourceLimitValue(String resourceName, final Map resourceLimits) { - for (ResourceType resourceType : resourceLimits.keySet()) { - if (resourceType.getName().equals(resourceName)) { - return (double) resourceLimits.get(resourceType); - } - } - return 0.0; - } - /** * This method will be executed before we submit the new cluster state * @param queryGroup - the QueryGroup we're currently creating @@ -148,57 +123,49 @@ ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final Clu String groupName = queryGroup.getName(); final Map previousGroups = metadata.queryGroups(); - // check if there's any resource allocation that exceed limit of 1.0 - String resourceNameWithThresholdExceeded = ""; - for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { - String resourceName = resourceType.getName(); - double existingUsage = calculateExistingUsage(resourceName, previousGroups, groupName); - double newGroupUsage = getResourceLimitValue(resourceName, queryGroup.getResourceLimits()); - inflightResourceLimitValues.computeIfAbsent(resourceName, k -> new DoubleAdder()).add(newGroupUsage); - double totalUsage = existingUsage + inflightResourceLimitValues.get(resourceName).doubleValue(); - if (totalUsage > 1) { - resourceNameWithThresholdExceeded = resourceName; - } + // check if maxQueryGroupCount will breach + if (previousGroups.size() + 1 > maxQueryGroupCount) { + logger.error("{} value exceeded its assigned limit of {}", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); + throw new RuntimeException("Can't create more than " + maxQueryGroupCount + " QueryGroups in the system"); } - // check if group count exceed max - boolean groupCountExceeded = inflightCreateQueryGroupRequestCount.incrementAndGet() + previousGroups.size() > maxQueryGroupCount; + // check for duplicate name Optional findExistingGroup = previousGroups.values() .stream() .filter(group -> group.getName().equals(groupName)) .findFirst(); - if (findExistingGroup.isPresent()) { logger.warn("QueryGroup with name {} already exists. Not creating a new one.", groupName); throw new RuntimeException("QueryGroup with name " + groupName + " already exists. Not creating a new one."); } - if (!resourceNameWithThresholdExceeded.isEmpty()) { - logger.error("Total resource allocation for {} will go above the max limit of 1.0", resourceNameWithThresholdExceeded); - throw new RuntimeException( - "Total resource allocation for " + resourceNameWithThresholdExceeded + " will go above the max limit of 1.0" - ); - } - if (groupCountExceeded) { - logger.error("{} value exceeded its assigned limit of {}", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); - throw new RuntimeException("Can't create more than " + maxQueryGroupCount + " QueryGroups in the system"); + + // check if there's any resource allocation that exceed limit of 1.0 + for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { + String resourceName = resourceType.getName(); + double existingUsage = calculateExistingUsage(resourceName, previousGroups, groupName); + double newGroupUsage = getResourceLimitValue(resourceName, queryGroup.getResourceLimits()); + if (existingUsage + newGroupUsage > 1) { + logger.error("Total resource allocation for {} will go above the max limit of 1.0", resourceName); + throw new RuntimeException("Total resource allocation for " + resourceName + " will go above the max limit of 1.0"); + } } - Metadata newData = Metadata.builder(metadata).put(queryGroup).build(); + Metadata newData = Metadata.builder(metadata).put(queryGroup).build(); return ClusterState.builder(currentClusterState).metadata(newData).build(); } /** - * This method restores the inflight values to be before the QueryGroup is processed - * @param resourceLimits - the resourceLimits we're currently restoring + * Get the allocation value for resourceName from the resourceLimits of a QueryGroup + * @param resourceName - the resourceName we want to get the usage for + * @param resourceLimits - the resource limit from which to get the allocation value for resourceName */ - void restoreInflightValues(Map resourceLimits) { - if (resourceLimits == null || resourceLimits.isEmpty()) { - return; - } + private double getResourceLimitValue(String resourceName, final Map resourceLimits) { for (ResourceType resourceType : resourceLimits.keySet()) { - String currResourceName = resourceType.getName(); - inflightResourceLimitValues.get(currResourceName).add(-getResourceLimitValue(currResourceName, resourceLimits)); + if (resourceType.getName().equals(resourceName)) { + return (double) resourceLimits.get(resourceType); + } } + return 0.0; } /** @@ -218,20 +185,6 @@ private double calculateExistingUsage(String resourceName, Map getInflightResourceLimitValues() { - return inflightResourceLimitValues; - } - /** * clusterService getter */ diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java index 202969d3e0c39..5936927db1ca4 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java @@ -25,7 +25,6 @@ import java.util.Comparator; import java.util.List; import java.util.Map; -import java.util.concurrent.atomic.DoubleAdder; import static org.opensearch.cluster.metadata.QueryGroup.builder; import static org.opensearch.search.ResourceType.fromName; @@ -133,14 +132,4 @@ public static void compareQueryGroups(List listOne, List assertTrue(listOne.get(i).equals(listTwo.get(i))); } } - - public static void assertInflightValuesAreZero(QueryGroupPersistenceService queryGroupPersistenceService) { - assertEquals(0, queryGroupPersistenceService.getInflightCreateQueryGroupRequestCount().get()); - Map inflightResourceMap = queryGroupPersistenceService.getInflightResourceLimitValues(); - if (inflightResourceMap != null) { - for (String resourceName : inflightResourceMap.keySet()) { - assertEquals(0, inflightResourceMap.get(resourceName).intValue()); - } - } - } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index b0c344670580e..7ae6250181f4d 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -32,12 +32,10 @@ import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_ONE; import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertInflightValuesAreZero; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.compareQueryGroups; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupPersistenceService; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo; import static org.opensearch.search.query_group.QueryGroupServiceSettings.QUERY_GROUP_COUNT_SETTING_NAME; import static org.mockito.Mockito.mock; @@ -57,7 +55,6 @@ public void testCreateQueryGroup() { listOne.add(queryGroupOne); listTwo.add(updatedGroupsMap.get(_ID_ONE)); compareQueryGroups(listOne, listTwo); - assertInflightValuesAreZero(queryGroupPersistenceService()); } public void testCreateAnotherQueryGroup() { @@ -70,7 +67,6 @@ public void testCreateAnotherQueryGroup() { assertTrue(updatedGroups.containsKey(_ID_TWO)); Collection values = updatedGroups.values(); compareQueryGroups(queryGroupList(), new ArrayList<>(values)); - assertInflightValuesAreZero(queryGroupPersistenceService()); } public void testCreateQueryGroupDuplicateName() { From 0c6a609ac59eb16a03ade58335e2da7fbf6c1304 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 19 Jul 2024 16:20:07 -0700 Subject: [PATCH 10/20] remove persistable interface Signed-off-by: Ruirui Zhang --- .../wlm/TransportCreateQueryGroupAction.java | 11 +- .../wlm/WorkloadManagementPluginModule.java | 9 +- .../wlm/rest/RestCreateQueryGroupAction.java | 2 +- .../plugin/wlm/service/Persistable.java | 25 --- .../service/QueryGroupPersistenceService.java | 17 +- .../QueryGroupServiceSettings.java | 174 ++++++++++++++---- 6 files changed, 148 insertions(+), 90 deletions(-) delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/Persistable.java diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java index 670fb870e44ce..aa3d724c20c2d 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java @@ -13,7 +13,7 @@ import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.service.Persistable; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; @@ -28,7 +28,7 @@ public class TransportCreateQueryGroupAction extends HandledTransportAction { private final ThreadPool threadPool; - private final Persistable queryGroupPersistenceService; + private final QueryGroupPersistenceService queryGroupPersistenceService; /** * Constructor for TransportCreateQueryGroupAction @@ -37,7 +37,7 @@ public class TransportCreateQueryGroupAction extends HandledTransportAction queryGroupPersistenceService + QueryGroupPersistenceService queryGroupPersistenceService ) { super(CreateQueryGroupAction.NAME, transportService, actionFilters, CreateQueryGroupRequest::new); this.threadPool = threadPool; @@ -60,6 +60,7 @@ protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListe .resourceLimits(request.getResourceLimits()) .updatedAt(request.getUpdatedAtInMillis()) .build(); - threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> queryGroupPersistenceService.persist(queryGroup, listener)); + threadPool.executor(ThreadPool.Names.GENERIC) + .execute(() -> queryGroupPersistenceService.persistInClusterStateMetadata(queryGroup, listener)); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java index 58208af293f14..87109d06259bb 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java @@ -8,11 +8,7 @@ package org.opensearch.plugin.wlm; -import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.inject.AbstractModule; -import org.opensearch.common.inject.TypeLiteral; -import org.opensearch.plugin.wlm.service.Persistable; -import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; /** * Guice Module to manage WorkloadManagement related objects @@ -25,8 +21,5 @@ public class WorkloadManagementPluginModule extends AbstractModule { public WorkloadManagementPluginModule() {} @Override - protected void configure() { - bind(new TypeLiteral>() { - }).to(QueryGroupPersistenceService.class).asEagerSingleton(); - } + protected void configure() {} } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java index 0f488be15bd11..45d190fdd7a78 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java @@ -50,7 +50,7 @@ public String getName() { */ @Override public List routes() { - return List.of(new Route(POST, "_wlm/_query_group/"), new Route(PUT, "_wlm/_query_group/")); + return List.of(new Route(POST, "_wlm/query_group/"), new Route(PUT, "_wlm/query_group/")); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/Persistable.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/Persistable.java deleted file mode 100644 index 9683a5b0ecc42..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/Persistable.java +++ /dev/null @@ -1,25 +0,0 @@ -/* - * 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.plugin.wlm.service; - -import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.CreateQueryGroupResponse; - -/** - * This interface defines the key APIs for implementing QueruGroup persistence - */ -public interface Persistable { - - /** - * persists the QueryGroup in a durable storage - * @param queryGroup - queryGroup to be persisted - * @param listener - ActionListener for CreateQueryGroupResponse - */ - void persist(T queryGroup, ActionListener listener); -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index cb833f02db569..e903b947ad929 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -34,7 +34,7 @@ /** * This class defines the functions for QueryGroup persistence */ -public class QueryGroupPersistenceService implements Persistable { +public class QueryGroupPersistenceService { private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); private final ClusterService clusterService; private static final String SOURCE = "query-group-persistence-service"; @@ -78,16 +78,12 @@ public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { this.maxQueryGroupCount = newMaxQueryGroupCount; } - @Override - public void persist(QueryGroup queryGroup, ActionListener listener) { - persistInClusterStateMetadata(queryGroup, listener); - } - /** * Update cluster state to include the new QueryGroup * @param queryGroup {@link QueryGroup} - the QueryGroup we're currently creating + * @param listener - ActionListener for CreateQueryGroupResponse */ - void persistInClusterStateMetadata(QueryGroup queryGroup, ActionListener listener) { + public void persistInClusterStateMetadata(QueryGroup queryGroup, ActionListener listener) { clusterService.submitStateUpdateTask(SOURCE, new ClusterStateUpdateTask(Priority.NORMAL) { @Override public ClusterState execute(ClusterState currentState) throws Exception { @@ -184,11 +180,4 @@ private double calculateExistingUsage(String resourceName, Map NODE_LEVEL_REJECTION_THRESHOLD = Setting.doubleSetting( - NODE_REJECTION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_REJECTION_THRESHOLD, + public static final Setting NODE_LEVEL_MEMORY_REJECTION_THRESHOLD = Setting.doubleSetting( + NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_MEMORY_REJECTION_THRESHOLD, Setting.Property.Dynamic, Setting.Property.NodeScope ); /** - * Setting name for node level cancellation threshold + * Setting name for node level cpu rejection threshold for QSB */ - public static final String NODE_CANCELLATION_THRESHOLD_SETTING_NAME = "query_group.node.cancellation_threshold"; + public static final String NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME = "query_group.node.cpu_rejection_threshold"; /** - * Setting name for node level cancellation threshold + * Setting to control the cpu rejection threshold */ - public static final Setting NODE_LEVEL_CANCELLATION_THRESHOLD = Setting.doubleSetting( - NODE_CANCELLATION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_CANCELLATION_THRESHOLD, + public static final Setting NODE_LEVEL_CPU_REJECTION_THRESHOLD = Setting.doubleSetting( + NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_CPU_REJECTION_THRESHOLD, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** + * Setting name for node level memory cancellation threshold + */ + public static final String NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME = "query_group.node.memory_cancellation_threshold"; + /** + * Setting name for node level memory cancellation threshold + */ + public static final Setting NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD = Setting.doubleSetting( + NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** + * Setting name for node level cpu cancellation threshold + */ + public static final String NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME = "query_group.node.cpu_cancellation_threshold"; + /** + * Setting name for node level cpu cancellation threshold + */ + public static final Setting NODE_LEVEL_CPU_CANCELLATION_THRESHOLD = Setting.doubleSetting( + NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_CPU_CANCELLATION_THRESHOLD, Setting.Property.Dynamic, Setting.Property.NodeScope ); @@ -97,15 +129,20 @@ public class QueryGroupServiceSettings { */ public QueryGroupServiceSettings(Settings settings, ClusterSettings clusterSettings) { runIntervalMillis = new TimeValue(QUERY_GROUP_RUN_INTERVAL_SETTING.get(settings)); - nodeLevelMemoryCancellationThreshold = NODE_LEVEL_CANCELLATION_THRESHOLD.get(settings); - nodeLevelMemoryRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); + nodeLevelMemoryCancellationThreshold = NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD.get(settings); + nodeLevelMemoryRejectionThreshold = NODE_LEVEL_MEMORY_REJECTION_THRESHOLD.get(settings); + nodeLevelCpuCancellationThreshold = NODE_LEVEL_CPU_CANCELLATION_THRESHOLD.get(settings); + nodeLevelCpuRejectionThreshold = NODE_LEVEL_CPU_REJECTION_THRESHOLD.get(settings); maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); - ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + ensureMemoryRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + ensureCpuRejectionThresholdIsLessThanCancellation(nodeLevelCpuRejectionThreshold, nodeLevelCpuCancellationThreshold); clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelMemoryCancellationThreshold); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD, this::setNodeLevelMemoryCancellationThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_MEMORY_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_CANCELLATION_THRESHOLD, this::setNodeLevelCpuCancellationThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_REJECTION_THRESHOLD, this::setNodeLevelCpuRejectionThreshold); } /** @@ -128,62 +165,125 @@ public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { } /** - * Method to get the node level cancellation threshold - * @return current node level cancellation threshold + * Method to get the node level memory cancellation threshold + * @return current node level memory cancellation threshold */ public Double getNodeLevelMemoryCancellationThreshold() { return nodeLevelMemoryCancellationThreshold; } /** - * Method to set the node level cancellation threshold - * @param nodeLevelMemoryCancellationThreshold sets the new node level cancellation threshold + * Method to set the node level memory cancellation threshold + * @param nodeLevelMemoryCancellationThreshold sets the new node level memory cancellation threshold * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold */ public void setNodeLevelMemoryCancellationThreshold(Double nodeLevelMemoryCancellationThreshold) { - if (Double.compare(nodeLevelMemoryCancellationThreshold, NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { + if (Double.compare(nodeLevelMemoryCancellationThreshold, NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { throw new IllegalArgumentException( - NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" + NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" ); } - ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + ensureMemoryRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); this.nodeLevelMemoryCancellationThreshold = nodeLevelMemoryCancellationThreshold; } /** - * Method to get the node level rejection threshold - * @return the current node level rejection threshold + * Method to get the node level cpu cancellation threshold + * @return current node level cpu cancellation threshold + */ + public Double getNodeLevelCpuCancellationThreshold() { + return nodeLevelCpuCancellationThreshold; + } + + /** + * Method to set the node level cpu cancellation threshold + * @param nodeLevelCpuCancellationThreshold sets the new node level cpu cancellation threshold + * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold + */ + public void setNodeLevelCpuCancellationThreshold(Double nodeLevelCpuCancellationThreshold) { + if (Double.compare(nodeLevelCpuCancellationThreshold, NODE_LEVEL_CPU_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { + throw new IllegalArgumentException( + NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" + ); + } + + ensureCpuRejectionThresholdIsLessThanCancellation(nodeLevelCpuRejectionThreshold, nodeLevelCpuCancellationThreshold); + + this.nodeLevelCpuCancellationThreshold = nodeLevelCpuCancellationThreshold; + } + + /** + * Method to get the memory node level rejection threshold + * @return the current memory node level rejection threshold */ public Double getNodeLevelMemoryRejectionThreshold() { return nodeLevelMemoryRejectionThreshold; } /** - * Method to set the node level rejection threshold - * @param nodeLevelMemoryRejectionThreshold sets the new rejection threshold + * Method to set the node level memory rejection threshold + * @param nodeLevelMemoryRejectionThreshold sets the new memory rejection threshold * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold */ public void setNodeLevelMemoryRejectionThreshold(Double nodeLevelMemoryRejectionThreshold) { - if (Double.compare(nodeLevelMemoryRejectionThreshold, NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE) > 0) { + if (Double.compare(nodeLevelMemoryRejectionThreshold, NODE_LEVEL_MEMORY_REJECTION_THRESHOLD_MAX_VALUE) > 0) { throw new IllegalArgumentException( - NODE_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" + NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" ); } - ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + ensureMemoryRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); this.nodeLevelMemoryRejectionThreshold = nodeLevelMemoryRejectionThreshold; } - private void ensureRejectionThresholdIsLessThanCancellation( + /** + * Method to get the cpu node level rejection threshold + * @return the current cpu node level rejection threshold + */ + public Double getNodeLevelCpuRejectionThreshold() { + return nodeLevelCpuRejectionThreshold; + } + + /** + * Method to set the node level cpu rejection threshold + * @param nodeLevelCpuRejectionThreshold sets the new cpu rejection threshold + * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold + */ + public void setNodeLevelCpuRejectionThreshold(Double nodeLevelCpuRejectionThreshold) { + if (Double.compare(nodeLevelCpuRejectionThreshold, NODE_LEVEL_CPU_REJECTION_THRESHOLD_MAX_VALUE) > 0) { + throw new IllegalArgumentException( + NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" + ); + } + + ensureCpuRejectionThresholdIsLessThanCancellation(nodeLevelCpuRejectionThreshold, nodeLevelCpuCancellationThreshold); + + this.nodeLevelCpuRejectionThreshold = nodeLevelCpuRejectionThreshold; + } + + private void ensureMemoryRejectionThresholdIsLessThanCancellation( Double nodeLevelMemoryRejectionThreshold, Double nodeLevelMemoryCancellationThreshold ) { if (Double.compare(nodeLevelMemoryCancellationThreshold, nodeLevelMemoryRejectionThreshold) < 0) { throw new IllegalArgumentException( - NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be less than " + NODE_REJECTION_THRESHOLD_SETTING_NAME + NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME + + " value should not be less than " + + NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME + ); + } + } + + private void ensureCpuRejectionThresholdIsLessThanCancellation( + Double nodeLevelCpuRejectionThreshold, + Double nodeLevelCpuCancellationThreshold + ) { + if (Double.compare(nodeLevelCpuCancellationThreshold, nodeLevelCpuRejectionThreshold) < 0) { + throw new IllegalArgumentException( + NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be less than " + NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME ); } } From 08efbbdcc0b39f7b742819b5ff664b8fe5d675b6 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 23 Jul 2024 12:14:38 -0700 Subject: [PATCH 11/20] modify QueryGroupServiceSettings Signed-off-by: Ruirui Zhang --- .../QueryGroupServiceSettings.java | 67 ++++++++++++------- 1 file changed, 43 insertions(+), 24 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java index 3a067cef53f37..00489414b5262 100644 --- a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java +++ b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java @@ -135,8 +135,18 @@ public QueryGroupServiceSettings(Settings settings, ClusterSettings clusterSetti nodeLevelCpuRejectionThreshold = NODE_LEVEL_CPU_REJECTION_THRESHOLD.get(settings); maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); - ensureMemoryRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); - ensureCpuRejectionThresholdIsLessThanCancellation(nodeLevelCpuRejectionThreshold, nodeLevelCpuCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation( + nodeLevelMemoryRejectionThreshold, + nodeLevelMemoryCancellationThreshold, + NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, + NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME + ); + ensureRejectionThresholdIsLessThanCancellation( + nodeLevelCpuRejectionThreshold, + nodeLevelCpuCancellationThreshold, + NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, + NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME + ); clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD, this::setNodeLevelMemoryCancellationThreshold); @@ -184,7 +194,12 @@ public void setNodeLevelMemoryCancellationThreshold(Double nodeLevelMemoryCancel ); } - ensureMemoryRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation( + nodeLevelMemoryRejectionThreshold, + nodeLevelMemoryCancellationThreshold, + NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, + NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME + ); this.nodeLevelMemoryCancellationThreshold = nodeLevelMemoryCancellationThreshold; } @@ -209,7 +224,12 @@ public void setNodeLevelCpuCancellationThreshold(Double nodeLevelCpuCancellation ); } - ensureCpuRejectionThresholdIsLessThanCancellation(nodeLevelCpuRejectionThreshold, nodeLevelCpuCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation( + nodeLevelCpuRejectionThreshold, + nodeLevelCpuCancellationThreshold, + NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, + NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME + ); this.nodeLevelCpuCancellationThreshold = nodeLevelCpuCancellationThreshold; } @@ -234,7 +254,12 @@ public void setNodeLevelMemoryRejectionThreshold(Double nodeLevelMemoryRejection ); } - ensureMemoryRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation( + nodeLevelMemoryRejectionThreshold, + nodeLevelMemoryCancellationThreshold, + NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, + NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME + ); this.nodeLevelMemoryRejectionThreshold = nodeLevelMemoryRejectionThreshold; } @@ -259,31 +284,25 @@ public void setNodeLevelCpuRejectionThreshold(Double nodeLevelCpuRejectionThresh ); } - ensureCpuRejectionThresholdIsLessThanCancellation(nodeLevelCpuRejectionThreshold, nodeLevelCpuCancellationThreshold); + ensureRejectionThresholdIsLessThanCancellation( + nodeLevelCpuRejectionThreshold, + nodeLevelCpuCancellationThreshold, + NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, + NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME + ); this.nodeLevelCpuRejectionThreshold = nodeLevelCpuRejectionThreshold; } - private void ensureMemoryRejectionThresholdIsLessThanCancellation( - Double nodeLevelMemoryRejectionThreshold, - Double nodeLevelMemoryCancellationThreshold + private void ensureRejectionThresholdIsLessThanCancellation( + Double nodeLevelRejectionThreshold, + Double nodeLevelCancellationThreshold, + String rejectionThresholdSettingName, + String cancellationThresholdSettingName ) { - if (Double.compare(nodeLevelMemoryCancellationThreshold, nodeLevelMemoryRejectionThreshold) < 0) { + if (Double.compare(nodeLevelCancellationThreshold, nodeLevelRejectionThreshold) < 0) { throw new IllegalArgumentException( - NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME - + " value should not be less than " - + NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME - ); - } - } - - private void ensureCpuRejectionThresholdIsLessThanCancellation( - Double nodeLevelCpuRejectionThreshold, - Double nodeLevelCpuCancellationThreshold - ) { - if (Double.compare(nodeLevelCpuCancellationThreshold, nodeLevelCpuRejectionThreshold) < 0) { - throw new IllegalArgumentException( - NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be less than " + NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME + cancellationThresholdSettingName + " value should not be less than " + rejectionThresholdSettingName ); } } From bfaf500264e61a9d8f197acc7f44ebcbbddb3977 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 23 Jul 2024 13:45:24 -0700 Subject: [PATCH 12/20] add in an action package in the plugin Signed-off-by: Ruirui Zhang --- .../plugin/wlm/WorkloadManagementPlugin.java | 2 ++ .../wlm/{ => action}/CreateQueryGroupAction.java | 2 +- .../wlm/{ => action}/CreateQueryGroupRequest.java | 2 +- .../wlm/{ => action}/CreateQueryGroupResponse.java | 2 +- .../TransportCreateQueryGroupAction.java | 2 +- .../opensearch/plugin/wlm/action/package-info.java | 12 ++++++++++++ .../plugin/wlm/rest/RestCreateQueryGroupAction.java | 6 +++--- .../wlm/service/QueryGroupPersistenceService.java | 2 +- .../{ => action}/CreateQueryGroupRequestTests.java | 3 ++- .../{ => action}/CreateQueryGroupResponseTests.java | 3 ++- 10 files changed, 26 insertions(+), 10 deletions(-) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/CreateQueryGroupAction.java (95%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/CreateQueryGroupRequest.java (99%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/CreateQueryGroupResponse.java (98%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/TransportCreateQueryGroupAction.java (98%) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => action}/CreateQueryGroupRequestTests.java (93%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => action}/CreateQueryGroupResponseTests.java (96%) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index 065c2f3229d15..de2dfc23f1c8f 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -17,6 +17,8 @@ import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; import org.opensearch.core.action.ActionResponse; +import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction; import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java similarity index 95% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java index bfc931ba1a2f9..14cb8cfcd125a 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java similarity index 99% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java index c8728c20bf000..3ba0e31ac1ead 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java similarity index 98% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java index 30e052b5a4d31..0c3de28a5bd51 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/CreateQueryGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java similarity index 98% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java index aa3d724c20c2d..a028effc8b5d1 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.HandledTransportAction; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java new file mode 100644 index 0000000000000..adcc5706ace3b --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java @@ -0,0 +1,12 @@ +/* + * 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 for the action classes for QueryGroup CRUD operations + */ +package org.opensearch.plugin.wlm.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java index 45d190fdd7a78..5950853e7263e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java @@ -12,9 +12,9 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.plugin.wlm.CreateQueryGroupAction; -import org.opensearch.plugin.wlm.CreateQueryGroupRequest; -import org.opensearch.plugin.wlm.CreateQueryGroupResponse; +import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.action.CreateQueryGroupRequest; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index e903b947ad929..7065ba66313ef 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -22,7 +22,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; -import org.opensearch.plugin.wlm.CreateQueryGroupResponse; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; import java.util.Map; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java similarity index 93% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java index ccc208f3c86d1..4711927d28cec 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java @@ -6,10 +6,11 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.action; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.plugin.wlm.QueryGroupTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java similarity index 96% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java index 038b8bfad87a4..d54753567d5ce 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/CreateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.io.stream.BytesStreamOutput; @@ -15,6 +15,7 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.plugin.wlm.QueryGroupTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; From 8f13d41871b3b7338aea316d7056c0150b084a9b Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 24 Jul 2024 17:34:48 -0700 Subject: [PATCH 13/20] modify based on commments Signed-off-by: Ruirui Zhang --- .../plugin/wlm/WorkloadManagementPlugin.java | 7 - .../wlm/WorkloadManagementPluginModule.java | 25 --- .../wlm/action/CreateQueryGroupRequest.java | 187 ++---------------- .../TransportCreateQueryGroupAction.java | 11 +- .../plugin/wlm/action/package-info.java | 2 +- .../opensearch/plugin/wlm/package-info.java | 2 +- .../wlm/rest/RestCreateQueryGroupAction.java | 18 +- .../plugin/wlm/rest/package-info.java | 2 +- .../service/QueryGroupPersistenceService.java | 4 +- .../plugin/wlm/service/package-info.java | 2 +- .../action/CreateQueryGroupRequestTests.java | 19 +- .../cluster/metadata/QueryGroup.java | 72 ++++--- .../cluster/metadata/QueryGroupTests.java | 10 +- 13 files changed, 85 insertions(+), 276 deletions(-) delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index de2dfc23f1c8f..59a740cc06049 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -11,7 +11,6 @@ import org.opensearch.action.ActionRequest; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.node.DiscoveryNodes; -import org.opensearch.common.inject.Module; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.IndexScopedSettings; import org.opensearch.common.settings.Settings; @@ -25,7 +24,6 @@ import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; -import java.util.Collection; import java.util.List; import java.util.function.Supplier; @@ -56,9 +54,4 @@ public List getRestHandlers( ) { return List.of(new RestCreateQueryGroupAction()); } - - @Override - public Collection createGuiceModules() { - return List.of(new WorkloadManagementPluginModule()); - } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java deleted file mode 100644 index 87109d06259bb..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java +++ /dev/null @@ -1,25 +0,0 @@ -/* - * 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.plugin.wlm; - -import org.opensearch.common.inject.AbstractModule; - -/** - * Guice Module to manage WorkloadManagement related objects - */ -public class WorkloadManagementPluginModule extends AbstractModule { - - /** - * Constructor for WorkloadManagementPluginModule - */ - public WorkloadManagementPluginModule() {} - - @Override - protected void configure() {} -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java index 3ba0e31ac1ead..4e04ad1ce222c 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -11,19 +11,11 @@ import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.cluster.metadata.QueryGroup; -import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; -import org.opensearch.common.UUIDs; 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.XContentParseException; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.search.ResourceType; -import org.joda.time.Instant; import java.io.IOException; -import java.util.HashMap; -import java.util.Map; /** * A request for create QueryGroup @@ -31,7 +23,7 @@ * { * "name": "analytics", * "resiliency_mode": "enforced", - * "resourceLimits": { + * "resource_limits": { * "cpu" : 0.4, * "memory" : 0.2 * } @@ -39,50 +31,15 @@ * * @opensearch.experimental */ -public class CreateQueryGroupRequest extends ActionRequest implements Writeable.Reader { - private String name; - private String _id; - private ResiliencyMode resiliencyMode; - private Map resourceLimits; - private long updatedAtInMillis; - - /** - * Default constructor for CreateQueryGroupRequest - */ - public CreateQueryGroupRequest() {} +public class CreateQueryGroupRequest extends ActionRequest { + final QueryGroup queryGroup; /** * Constructor for CreateQueryGroupRequest * @param queryGroup - A {@link QueryGroup} object */ public CreateQueryGroupRequest(QueryGroup queryGroup) { - this.name = queryGroup.getName(); - this._id = queryGroup.get_id(); - this.resourceLimits = queryGroup.getResourceLimits(); - this.resiliencyMode = queryGroup.getResiliencyMode(); - this.updatedAtInMillis = queryGroup.getUpdatedAtInMillis(); - } - - /** - * Constructor for CreateQueryGroupRequest - * @param name - QueryGroup name for CreateQueryGroupRequest - * @param _id - QueryGroup _id for CreateQueryGroupRequest - * @param mode - QueryGroup mode for CreateQueryGroupRequest - * @param resourceLimits - QueryGroup resourceLimits for CreateQueryGroupRequest - * @param updatedAtInMillis - QueryGroup updated time in millis for CreateQueryGroupRequest - */ - public CreateQueryGroupRequest( - String name, - String _id, - ResiliencyMode mode, - Map resourceLimits, - long updatedAtInMillis - ) { - this.name = name; - this._id = _id; - this.resourceLimits = resourceLimits; - this.resiliencyMode = mode; - this.updatedAtInMillis = updatedAtInMillis; + this.queryGroup = queryGroup; } /** @@ -91,16 +48,7 @@ public CreateQueryGroupRequest( */ public CreateQueryGroupRequest(StreamInput in) throws IOException { super(in); - name = in.readString(); - _id = in.readString(); - resiliencyMode = ResiliencyMode.fromName(in.readString()); - resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readGenericValue); - updatedAtInMillis = in.readLong(); - } - - @Override - public CreateQueryGroupRequest read(StreamInput in) throws IOException { - return new CreateQueryGroupRequest(in); + queryGroup = new QueryGroup(in); } /** @@ -108,47 +56,10 @@ public CreateQueryGroupRequest read(StreamInput in) throws IOException { * @param parser - A {@link XContentParser} object */ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws IOException { - - while (parser.currentToken() != XContentParser.Token.START_OBJECT) { - parser.nextToken(); - } - - if (parser.currentToken() != XContentParser.Token.START_OBJECT) { - throw new XContentParseException("expected start object but got a " + parser.currentToken()); - } - - XContentParser.Token token; - String fieldName = ""; - String name = null; - ResiliencyMode mode = null; - - // Map to hold resources - final Map resourceLimits = new HashMap<>(); - while ((token = parser.nextToken()) != null) { - if (token == XContentParser.Token.FIELD_NAME) { - fieldName = parser.currentName(); - } else if (token.isValue()) { - if (fieldName.equals("name")) { - name = parser.text(); - } else if (fieldName.equals("resiliency_mode")) { - mode = ResiliencyMode.fromName(parser.text()); - } else { - throw new XContentParseException("unrecognised [field=" + fieldName + " in QueryGroup"); - } - } else if (token == XContentParser.Token.START_OBJECT) { - if (!fieldName.equals("resourceLimits")) { - throw new XContentParseException("Invalid field passed. QueryGroup does not support " + fieldName + "."); - } - while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { - if (token == XContentParser.Token.FIELD_NAME) { - fieldName = parser.currentName(); - } else { - resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); - } - } - } - } - return new CreateQueryGroupRequest(name, UUIDs.randomBase64UUID(), mode, resourceLimits, Instant.now().getMillis()); + QueryGroup queryGroup1 = QueryGroup.fromXContent(parser); + // creating this queryGroup to ensure forceful creation of _id and updatedAt + QueryGroup queryGroup = new QueryGroup(queryGroup1.getName(), queryGroup1.getResiliencyMode(), queryGroup1.getResourceLimits()); + return new CreateQueryGroupRequest(queryGroup); } @Override @@ -156,88 +67,16 @@ public ActionRequestValidationException validate() { return null; } - /** - * name getter - */ - public String getName() { - return name; - } - - /** - * name setter - * @param name - name to be set - */ - public void setName(String name) { - this.name = name; - } - - /** - * resourceLimits getter - */ - public Map getResourceLimits() { - return resourceLimits; - } - - /** - * resourceLimits setter - * @param resourceLimits - resourceLimit to be set - */ - public void setResourceLimits(Map resourceLimits) { - this.resourceLimits = resourceLimits; - } - - /** - * mode getter - */ - public ResiliencyMode getResiliencyMode() { - return resiliencyMode; - } - - /** - * mode setter - * @param resiliencyMode - mode to be set - */ - public void setResiliencyMode(ResiliencyMode resiliencyMode) { - this.resiliencyMode = resiliencyMode; - } - @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); - out.writeString(name); - out.writeString(_id); - out.writeString(resiliencyMode.getName()); - out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue); - out.writeLong(updatedAtInMillis); - } - - /** - * _id getter - */ - public String get_id() { - return _id; - } - - /** - * UUID setter - * @param _id - _id to be set - */ - public void set_id(String _id) { - this._id = _id; - } - - /** - * updatedAtInMillis getter - */ - public long getUpdatedAtInMillis() { - return updatedAtInMillis; + QueryGroup.writeToOutput(out, queryGroup); } /** - * updatedAtInMillis setter - * @param updatedAtInMillis - updatedAtInMillis to be set + * QueryGroup getter */ - public void setUpdatedAtInMillis(long updatedAtInMillis) { - this.updatedAtInMillis = updatedAtInMillis; + public QueryGroup getQueryGroup() { + return queryGroup; } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java index a028effc8b5d1..f3e913ad019e3 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java @@ -10,7 +10,6 @@ import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.HandledTransportAction; -import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; @@ -18,8 +17,6 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; -import static org.opensearch.cluster.metadata.QueryGroup.builder; - /** * Transport action to create QueryGroup * @@ -54,13 +51,7 @@ public TransportCreateQueryGroupAction( @Override protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener listener) { - QueryGroup queryGroup = builder().name(request.getName()) - ._id(request.get_id()) - .mode(request.getResiliencyMode().getName()) - .resourceLimits(request.getResourceLimits()) - .updatedAt(request.getUpdatedAtInMillis()) - .build(); threadPool.executor(ThreadPool.Names.GENERIC) - .execute(() -> queryGroupPersistenceService.persistInClusterStateMetadata(queryGroup, listener)); + .execute(() -> queryGroupPersistenceService.persistInClusterStateMetadata(request.getQueryGroup(), listener)); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java index adcc5706ace3b..9921500df8a81 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java @@ -7,6 +7,6 @@ */ /** - * Package for the action classes for QueryGroup CRUD operations + * Package for the action classes of WorkloadManagementPlugin */ package org.opensearch.plugin.wlm.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java index d85cf8d417a5f..84c99967b226b 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java @@ -7,6 +7,6 @@ */ /** - * Base package for CRUD API of QueryGroup + * Base package for WorkloadManagementPlugin */ package org.opensearch.plugin.wlm; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java index 5950853e7263e..ae328ba6025cc 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java @@ -55,18 +55,12 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - CreateQueryGroupRequest createQueryGroupRequest = new CreateQueryGroupRequest(); - request.applyContentParser((parser) -> parseRestRequest(createQueryGroupRequest, parser)); - return channel -> client.execute(CreateQueryGroupAction.INSTANCE, createQueryGroupRequest, createQueryGroupResponse(channel)); - } - - private void parseRestRequest(CreateQueryGroupRequest request, XContentParser parser) throws IOException { - final CreateQueryGroupRequest createQueryGroupRequest = CreateQueryGroupRequest.fromXContent(parser); - request.setName(createQueryGroupRequest.getName()); - request.set_id(createQueryGroupRequest.get_id()); - request.setResiliencyMode(createQueryGroupRequest.getResiliencyMode()); - request.setResourceLimits(createQueryGroupRequest.getResourceLimits()); - request.setUpdatedAtInMillis(createQueryGroupRequest.getUpdatedAtInMillis()); + XContentParser parser = request.contentParser(); + return channel -> client.execute( + CreateQueryGroupAction.INSTANCE, + CreateQueryGroupRequest.fromXContent(parser), + createQueryGroupResponse(channel) + ); } private RestResponseListener createQueryGroupResponse(final RestChannel channel) { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java index 6d1853f04ce58..7d7cb9028fdb8 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java @@ -7,6 +7,6 @@ */ /** - * Package for the rest classes for QueryGroup CRUD operations + * Package for the rest classes of WorkloadManagementPlugin */ package org.opensearch.plugin.wlm.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 7065ba66313ef..6684199d57dff 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -155,10 +155,10 @@ ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final Clu * @param resourceName - the resourceName we want to get the usage for * @param resourceLimits - the resource limit from which to get the allocation value for resourceName */ - private double getResourceLimitValue(String resourceName, final Map resourceLimits) { + private double getResourceLimitValue(String resourceName, final Map resourceLimits) { for (ResourceType resourceType : resourceLimits.keySet()) { if (resourceType.getName().equals(resourceName)) { - return (double) resourceLimits.get(resourceType); + return resourceLimits.get(resourceType); } } return 0.0; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java index ecfc94568888a..5848e9c936623 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java @@ -7,6 +7,6 @@ */ /** - * Package for the service classes for QueryGroup CRUD operations + * Package for the service classes of WorkloadManagementPlugin */ package org.opensearch.plugin.wlm.service; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java index 4711927d28cec..a7ec3a49c6eb3 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java @@ -8,25 +8,30 @@ package org.opensearch.plugin.wlm.action; +import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.plugin.wlm.QueryGroupTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.compareQueryGroups; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; public class CreateQueryGroupRequestTests extends OpenSearchTestCase { public void testSerialization() throws IOException { - CreateQueryGroupRequest request = new CreateQueryGroupRequest(QueryGroupTestUtils.queryGroupOne); + CreateQueryGroupRequest request = new CreateQueryGroupRequest(queryGroupOne); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); CreateQueryGroupRequest otherRequest = new CreateQueryGroupRequest(streamInput); - assertEquals(request.getName(), otherRequest.getName()); - assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); - assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); - QueryGroupTestUtils.compareResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); - assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); + List list1 = new ArrayList<>(); + List list2 = new ArrayList<>(); + list1.add(queryGroupOne); + list2.add(otherRequest.getQueryGroup()); + compareQueryGroups(list1, list2); } } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index 96bb098cbb6e6..e71df169de5db 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -29,30 +29,35 @@ * Class to define the QueryGroup schema * { * "_id": "fafjafjkaf9ag8a9ga9g7ag0aagaga", - * "resourceLimits": { + * "resource_limits": { * "memory": 0.4 * }, * "resiliency_mode": "enforced", * "name": "analytics", - * "updatedAt": 4513232415 + * "updated_at": 4513232415 * } */ @ExperimentalApi public class QueryGroup extends AbstractDiffable implements ToXContentObject { + public static final String _ID_STRING = "_id"; + public static final String NAME_STRING = "name"; + public static final String RESILIENCY_MODE_STRING = "resiliency_mode"; + public static final String UPDATED_AT_STRING = "updated_at"; + public static final String RESOURCE_LIMITS_STRING = "resource_limits"; private static final int MAX_CHARS_ALLOWED_IN_NAME = 50; private final String name; private final String _id; private final ResiliencyMode resiliencyMode; // It is an epoch in millis private final long updatedAtInMillis; - private final Map resourceLimits; + private final Map resourceLimits; - public QueryGroup(String name, ResiliencyMode resiliencyMode, Map resourceLimits) { + public QueryGroup(String name, ResiliencyMode resiliencyMode, Map resourceLimits) { this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis()); } - public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits, long updatedAt) { + public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits, long updatedAt) { Objects.requireNonNull(name, "QueryGroup.name can't be null"); Objects.requireNonNull(resourceLimits, "QueryGroup.resourceLimits can't be null"); Objects.requireNonNull(resiliencyMode, "QueryGroup.resiliencyMode can't be null"); @@ -92,23 +97,27 @@ public QueryGroup(StreamInput in) throws IOException { in.readString(), in.readString(), ResiliencyMode.fromName(in.readString()), - in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readGenericValue), + in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble), in.readLong() ); } @Override public void writeTo(StreamOutput out) throws IOException { - out.writeString(name); - out.writeString(_id); - out.writeString(resiliencyMode.getName()); - out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue); - out.writeLong(updatedAtInMillis); + writeToOutput(out, this); } - private void validateResourceLimits(Map resourceLimits) { - for (Map.Entry resource : resourceLimits.entrySet()) { - Double threshold = (Double) resource.getValue(); + public static void writeToOutput(StreamOutput out, QueryGroup queryGroup) throws IOException { + out.writeString(queryGroup.getName()); + out.writeString(queryGroup.get_id()); + out.writeString(queryGroup.getResiliencyMode().getName()); + out.writeMap(queryGroup.getResourceLimits(), ResourceType::writeTo, StreamOutput::writeDouble); + out.writeLong(queryGroup.getUpdatedAtInMillis()); + } + + private void validateResourceLimits(Map resourceLimits) { + for (Map.Entry resource : resourceLimits.entrySet()) { + Double threshold = resource.getValue(); Objects.requireNonNull(resource.getKey(), "resourceName can't be null"); Objects.requireNonNull(threshold, "resource limit threshold for" + resource.getKey().getName() + " : can't be null"); @@ -121,12 +130,12 @@ private void validateResourceLimits(Map resourceLimits) { @Override public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { builder.startObject(); - builder.field("_id", _id); - builder.field("name", name); - builder.field("resiliency_mode", resiliencyMode.getName()); - builder.field("updatedAt", updatedAtInMillis); + builder.field(_ID_STRING, _id); + builder.field(NAME_STRING, name); + builder.field(RESILIENCY_MODE_STRING, resiliencyMode.getName()); + builder.field(UPDATED_AT_STRING, updatedAtInMillis); // write resource limits - builder.startObject("resourceLimits"); + builder.startObject(RESOURCE_LIMITS_STRING); for (ResourceType resourceType : ResourceType.values()) { if (resourceLimits.containsKey(resourceType)) { builder.field(resourceType.getName(), resourceLimits.get(resourceType)); @@ -153,25 +162,25 @@ public static QueryGroup fromXContent(final XContentParser parser) throws IOExce String fieldName = ""; // Map to hold resources - final Map resourceLimits = new HashMap<>(); + final Map resourceLimits = new HashMap<>(); while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { if (token == XContentParser.Token.FIELD_NAME) { fieldName = parser.currentName(); } else if (token.isValue()) { - if (fieldName.equals("_id")) { + if (fieldName.equals(_ID_STRING)) { builder._id(parser.text()); - } else if (fieldName.equals("name")) { + } else if (fieldName.equals(NAME_STRING)) { builder.name(parser.text()); - } else if (fieldName.equals("resiliency_mode")) { + } else if (fieldName.equals(RESILIENCY_MODE_STRING)) { builder.mode(parser.text()); - } else if (fieldName.equals("updatedAt")) { + } else if (fieldName.equals(UPDATED_AT_STRING)) { builder.updatedAt(parser.longValue()); } else { throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup"); } } else if (token == XContentParser.Token.START_OBJECT) { - if (!fieldName.equals("resourceLimits")) { + if (!fieldName.equals(RESOURCE_LIMITS_STRING)) { throw new IllegalArgumentException( "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token ); @@ -220,7 +229,7 @@ public ResiliencyMode getResiliencyMode() { return resiliencyMode; } - public Map getResourceLimits() { + public Map getResourceLimits() { return resourceLimits; } @@ -280,7 +289,7 @@ public static class Builder { private String _id; private ResiliencyMode resiliencyMode; private long updatedAt; - private Map resourceLimits; + private Map resourceLimits; private Builder() {} @@ -304,14 +313,17 @@ public Builder updatedAt(long updatedAt) { return this; } - public Builder resourceLimits(Map resourceLimits) { + public Builder resourceLimits(Map resourceLimits) { this.resourceLimits = resourceLimits; return this; } public QueryGroup build() { - return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); + if (_id == null && updatedAt == 0L) { + return new QueryGroup(name, resiliencyMode, resourceLimits); + } else { + return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); + } } - } } diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java index c564f0778e6f0..db2ff239271b0 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -34,7 +34,7 @@ public class QueryGroupTests extends AbstractSerializingTestCase { static QueryGroup createRandomQueryGroup(String _id) { String name = randomAlphaOfLength(10); - Map resourceLimit = new HashMap<>(); + Map resourceLimit = new HashMap<>(); resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false)); return new QueryGroup(name, _id, randomMode(), resourceLimit, Instant.now().getMillis()); } @@ -99,7 +99,7 @@ public void testEmptyResourceLimits() { public void testIllegalQueryGroupMode() { assertThrows( NullPointerException.class, - () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, (Object) 0.4), Instant.now().getMillis()) + () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis()) ); } @@ -110,7 +110,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { "analytics", "_id", randomMode(), - Map.of(ResourceType.MEMORY, (Object) randomDoubleBetween(1.1, 1.8, false)), + Map.of(ResourceType.MEMORY, randomDoubleBetween(1.1, 1.8, false)), Instant.now().getMillis() ) ); @@ -149,9 +149,9 @@ public void testToXContent() throws IOException { assertEquals( "{\"_id\":\"" + queryGroupId - + "\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"updatedAt\":" + + "\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"updated_at\":" + currentTimeInMillis - + ",\"resourceLimits\":{\"cpu\":0.3,\"memory\":0.4}}", + + ",\"resource_limits\":{\"cpu\":0.3,\"memory\":0.4}}", builder.toString() ); } From 9b8c1cedbe683473bd9db06fa466075cc03d5e6b Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 25 Jul 2024 16:20:28 -0700 Subject: [PATCH 14/20] address comments on QueryGroupPersistenceService Signed-off-by: Ruirui Zhang --- .../plugin/wlm/WorkloadManagementPlugin.java | 7 + .../service/QueryGroupPersistenceService.java | 99 +++--- .../plugin/wlm/QueryGroupTestUtils.java | 15 +- .../action/CreateQueryGroupResponseTests.java | 4 +- .../QueryGroupPersistenceServiceTests.java | 7 +- .../common/settings/ClusterSettings.java | 1 - .../QueryGroupServiceSettings.java | 317 ------------------ .../search/query_group/package-info.java | 12 - 8 files changed, 74 insertions(+), 388 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java delete mode 100644 server/src/main/java/org/opensearch/search/query_group/package-info.java diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index 59a740cc06049..80807f0d5bc37 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -13,12 +13,14 @@ import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.IndexScopedSettings; +import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; import org.opensearch.core.action.ActionResponse; import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction; import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; import org.opensearch.rest.RestController; @@ -54,4 +56,9 @@ public List getRestHandlers( ) { return List.of(new RestCreateQueryGroupAction()); } + + @Override + public List> getSettings() { + return List.of(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT); + } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 6684199d57dff..7273a0048cae9 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -19,39 +19,47 @@ import org.opensearch.common.Priority; import org.opensearch.common.inject.Inject; import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; +import java.util.HashMap; import java.util.Map; import java.util.Optional; -import static org.opensearch.search.query_group.QueryGroupServiceSettings.MAX_QUERY_GROUP_COUNT; -import static org.opensearch.search.query_group.QueryGroupServiceSettings.QUERY_GROUP_COUNT_SETTING_NAME; - /** * This class defines the functions for QueryGroup persistence */ public class QueryGroupPersistenceService { - private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); - private final ClusterService clusterService; + public static final String QUERY_GROUP_COUNT_SETTING_NAME = "node.query_group.max_count"; private static final String SOURCE = "query-group-persistence-service"; private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group"; - private static final String UPDATE_QUERY_GROUP_THROTTLING_KEY = "update-query-group"; - private static final String DELETE_QUERY_GROUP_THROTTLING_KEY = "delete-query-group"; + public static final int DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE = 100; + public static final Setting MAX_QUERY_GROUP_COUNT = Setting.intSetting( + QUERY_GROUP_COUNT_SETTING_NAME, + DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE, + 0, + (newVal) -> { + if (newVal > 100 || newVal < 1) throw new IllegalArgumentException( + QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]" + ); + }, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); + private final ClusterService clusterService; private volatile int maxQueryGroupCount; final ThrottlingKey createQueryGroupThrottlingKey; - final ThrottlingKey updateQueryGroupThrottlingKey; - final ThrottlingKey deleteQueryGroupThrottlingKey; /** * Constructor for QueryGroupPersistenceService * * @param clusterService {@link ClusterService} - The cluster service to be used by QueryGroupPersistenceService * @param settings {@link Settings} - The settings to be used by QueryGroupPersistenceService - * @param clusterSettings {@link ClusterSettings} - The cluster settings to be used by QueryGroupPersistenceService */ @Inject public QueryGroupPersistenceService( @@ -61,9 +69,7 @@ public QueryGroupPersistenceService( ) { this.clusterService = clusterService; this.createQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(CREATE_QUERY_GROUP_THROTTLING_KEY, true); - this.deleteQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(DELETE_QUERY_GROUP_THROTTLING_KEY, true); - this.updateQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(UPDATE_QUERY_GROUP_THROTTLING_KEY, true); - maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); + setMaxQueryGroupCount(MAX_QUERY_GROUP_COUNT.get(settings)); clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); } @@ -72,8 +78,8 @@ public QueryGroupPersistenceService( * @param newMaxQueryGroupCount - the max number of QueryGroup allowed */ public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { - if (newMaxQueryGroupCount < 0) { - throw new IllegalArgumentException("node.query_group.max_count can't be negative"); + if (newMaxQueryGroupCount > 100 || newMaxQueryGroupCount < 1) { + throw new IllegalArgumentException(QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]"); } this.maxQueryGroupCount = newMaxQueryGroupCount; } @@ -115,18 +121,17 @@ public void clusterStateProcessed(String source, ClusterState oldState, ClusterS * @param currentClusterState - the cluster state before the update */ ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) { - final Metadata metadata = currentClusterState.metadata(); + final Map existingQueryGroups = currentClusterState.metadata().queryGroups(); String groupName = queryGroup.getName(); - final Map previousGroups = metadata.queryGroups(); // check if maxQueryGroupCount will breach - if (previousGroups.size() + 1 > maxQueryGroupCount) { - logger.error("{} value exceeded its assigned limit of {}", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); + if (existingQueryGroups.size() + 1 > maxQueryGroupCount) { + logger.warn("{} value exceeded its assigned limit of {}", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); throw new RuntimeException("Can't create more than " + maxQueryGroupCount + " QueryGroups in the system"); } // check for duplicate name - Optional findExistingGroup = previousGroups.values() + Optional findExistingGroup = existingQueryGroups.values() .stream() .filter(group -> group.getName().equals(groupName)) .findFirst(); @@ -136,48 +141,40 @@ ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final Clu } // check if there's any resource allocation that exceed limit of 1.0 + Map existingUsageMap = calculateExistingUsage(existingQueryGroups, groupName); for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { - String resourceName = resourceType.getName(); - double existingUsage = calculateExistingUsage(resourceName, previousGroups, groupName); - double newGroupUsage = getResourceLimitValue(resourceName, queryGroup.getResourceLimits()); - if (existingUsage + newGroupUsage > 1) { - logger.error("Total resource allocation for {} will go above the max limit of 1.0", resourceName); - throw new RuntimeException("Total resource allocation for " + resourceName + " will go above the max limit of 1.0"); + double totalUsage = existingUsageMap.getOrDefault(resourceType, 0.0) + queryGroup.getResourceLimits().get(resourceType); + if (totalUsage > 1) { + logger.warn("Total resource allocation for {} will go above the max limit of 1.0", resourceType.getName()); + throw new RuntimeException( + "Total resource allocation for " + resourceType.getName() + " will go above the max limit of 1.0" + ); } } - Metadata newData = Metadata.builder(metadata).put(queryGroup).build(); - return ClusterState.builder(currentClusterState).metadata(newData).build(); - } - - /** - * Get the allocation value for resourceName from the resourceLimits of a QueryGroup - * @param resourceName - the resourceName we want to get the usage for - * @param resourceLimits - the resource limit from which to get the allocation value for resourceName - */ - private double getResourceLimitValue(String resourceName, final Map resourceLimits) { - for (ResourceType resourceType : resourceLimits.keySet()) { - if (resourceType.getName().equals(resourceName)) { - return resourceLimits.get(resourceType); - } - } - return 0.0; + return ClusterState.builder(currentClusterState) + .metadata( + Metadata.builder(currentClusterState.metadata()) + .put(queryGroup) + .build() + ) + .build(); } /** - * This method calculates the existing total usage of the resource (except the group that we're updating here) - * @param resourceName - the resource name we're calculating - * @param groupsMap - existing QueryGroups + * This method calculates the existing total usage of the all the resource limits (except the group that we're updating here) + * @param existingQueryGroups - existing QueryGroups * @param groupName - the QueryGroup name we're updating */ - private double calculateExistingUsage(String resourceName, Map groupsMap, String groupName) { - double existingUsage = 0; - for (String currGroupId : groupsMap.keySet()) { - QueryGroup currGroup = groupsMap.get(currGroupId); + private Map calculateExistingUsage(Map existingQueryGroups, String groupName) { + Map map = new HashMap<>(); + for (QueryGroup currGroup : existingQueryGroups.values()) { if (!currGroup.getName().equals(groupName)) { - existingUsage += getResourceLimitValue(resourceName, currGroup.getResourceLimits()); + for (ResourceType resourceType : currGroup.getResourceLimits().keySet()) { + map.put(resourceType, map.getOrDefault(resourceType, 0.0) + currGroup.getResourceLimits().get(resourceType)); + } } } - return existingUsage; + return map; } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java index 5936927db1ca4..777eb14570376 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java @@ -16,6 +16,7 @@ import org.opensearch.cluster.service.ClusterManagerService; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; import org.opensearch.search.ResourceType; @@ -23,12 +24,15 @@ import java.util.ArrayList; import java.util.Comparator; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import static org.opensearch.cluster.metadata.QueryGroup.builder; import static org.opensearch.search.ResourceType.fromName; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; @@ -68,12 +72,19 @@ public static ClusterState clusterState() { return ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); } + public static Set> clusterSettingsSet() { + Set> set = new HashSet<>(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + set.add(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT); + assertFalse(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS.contains(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT)); + return set; + } + public static Settings settings() { return Settings.builder().build(); } public static ClusterSettings clusterSettings() { - return new ClusterSettings(settings(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + return new ClusterSettings(settings(), clusterSettingsSet()); } public static QueryGroupPersistenceService queryGroupPersistenceService() { @@ -97,7 +108,7 @@ public static List preparePersistenceServiceSetup(Map MAX_QUERY_GROUP_COUNT = Setting.intSetting( - QUERY_GROUP_COUNT_SETTING_NAME, - DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE, - 0, - (newVal) -> { - if (newVal > 100 || newVal < 1) throw new IllegalArgumentException( - QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]" - ); - }, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - /** - * Setting name for default QueryGroup count - */ - public static final String SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME = "query_group.service.run_interval_millis"; - /** - * Setting to control the run interval of QSB service - */ - private static final Setting QUERY_GROUP_RUN_INTERVAL_SETTING = Setting.longSetting( - SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME, - DEFAULT_RUN_INTERVAL_MILLIS, - 1, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - - /** - * Setting name for node level memory rejection threshold for QSB - */ - public static final String NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME = "query_group.node.memory_rejection_threshold"; - /** - * Setting to control the memory rejection threshold - */ - public static final Setting NODE_LEVEL_MEMORY_REJECTION_THRESHOLD = Setting.doubleSetting( - NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_MEMORY_REJECTION_THRESHOLD, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - /** - * Setting name for node level cpu rejection threshold for QSB - */ - public static final String NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME = "query_group.node.cpu_rejection_threshold"; - /** - * Setting to control the cpu rejection threshold - */ - public static final Setting NODE_LEVEL_CPU_REJECTION_THRESHOLD = Setting.doubleSetting( - NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_CPU_REJECTION_THRESHOLD, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - /** - * Setting name for node level memory cancellation threshold - */ - public static final String NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME = "query_group.node.memory_cancellation_threshold"; - /** - * Setting name for node level memory cancellation threshold - */ - public static final Setting NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD = Setting.doubleSetting( - NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - /** - * Setting name for node level cpu cancellation threshold - */ - public static final String NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME = "query_group.node.cpu_cancellation_threshold"; - /** - * Setting name for node level cpu cancellation threshold - */ - public static final Setting NODE_LEVEL_CPU_CANCELLATION_THRESHOLD = Setting.doubleSetting( - NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_CPU_CANCELLATION_THRESHOLD, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - - /** - * QueryGroup service settings constructor - * @param settings - QueryGroup service settings - * @param clusterSettings - QueryGroup cluster settings - */ - public QueryGroupServiceSettings(Settings settings, ClusterSettings clusterSettings) { - runIntervalMillis = new TimeValue(QUERY_GROUP_RUN_INTERVAL_SETTING.get(settings)); - nodeLevelMemoryCancellationThreshold = NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD.get(settings); - nodeLevelMemoryRejectionThreshold = NODE_LEVEL_MEMORY_REJECTION_THRESHOLD.get(settings); - nodeLevelCpuCancellationThreshold = NODE_LEVEL_CPU_CANCELLATION_THRESHOLD.get(settings); - nodeLevelCpuRejectionThreshold = NODE_LEVEL_CPU_REJECTION_THRESHOLD.get(settings); - maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); - - ensureRejectionThresholdIsLessThanCancellation( - nodeLevelMemoryRejectionThreshold, - nodeLevelMemoryCancellationThreshold, - NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, - NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME - ); - ensureRejectionThresholdIsLessThanCancellation( - nodeLevelCpuRejectionThreshold, - nodeLevelCpuCancellationThreshold, - NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, - NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME - ); - - clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD, this::setNodeLevelMemoryCancellationThreshold); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_MEMORY_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_CANCELLATION_THRESHOLD, this::setNodeLevelCpuCancellationThreshold); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_REJECTION_THRESHOLD, this::setNodeLevelCpuRejectionThreshold); - } - - /** - * Method to get runInterval for QSB - * @return runInterval in milliseconds for QSB Service - */ - public TimeValue getRunIntervalMillis() { - return runIntervalMillis; - } - - /** - * Method to set the new QueryGroup count - * @param newMaxQueryGroupCount is the new maxQueryGroupCount per node - */ - public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { - if (newMaxQueryGroupCount < 0) { - throw new IllegalArgumentException("node.node.query_group.max_count can't be negative"); - } - this.maxQueryGroupCount = newMaxQueryGroupCount; - } - - /** - * Method to get the node level memory cancellation threshold - * @return current node level memory cancellation threshold - */ - public Double getNodeLevelMemoryCancellationThreshold() { - return nodeLevelMemoryCancellationThreshold; - } - - /** - * Method to set the node level memory cancellation threshold - * @param nodeLevelMemoryCancellationThreshold sets the new node level memory cancellation threshold - * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold - */ - public void setNodeLevelMemoryCancellationThreshold(Double nodeLevelMemoryCancellationThreshold) { - if (Double.compare(nodeLevelMemoryCancellationThreshold, NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { - throw new IllegalArgumentException( - NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" - ); - } - - ensureRejectionThresholdIsLessThanCancellation( - nodeLevelMemoryRejectionThreshold, - nodeLevelMemoryCancellationThreshold, - NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, - NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME - ); - - this.nodeLevelMemoryCancellationThreshold = nodeLevelMemoryCancellationThreshold; - } - - /** - * Method to get the node level cpu cancellation threshold - * @return current node level cpu cancellation threshold - */ - public Double getNodeLevelCpuCancellationThreshold() { - return nodeLevelCpuCancellationThreshold; - } - - /** - * Method to set the node level cpu cancellation threshold - * @param nodeLevelCpuCancellationThreshold sets the new node level cpu cancellation threshold - * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold - */ - public void setNodeLevelCpuCancellationThreshold(Double nodeLevelCpuCancellationThreshold) { - if (Double.compare(nodeLevelCpuCancellationThreshold, NODE_LEVEL_CPU_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { - throw new IllegalArgumentException( - NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" - ); - } - - ensureRejectionThresholdIsLessThanCancellation( - nodeLevelCpuRejectionThreshold, - nodeLevelCpuCancellationThreshold, - NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, - NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME - ); - - this.nodeLevelCpuCancellationThreshold = nodeLevelCpuCancellationThreshold; - } - - /** - * Method to get the memory node level rejection threshold - * @return the current memory node level rejection threshold - */ - public Double getNodeLevelMemoryRejectionThreshold() { - return nodeLevelMemoryRejectionThreshold; - } - - /** - * Method to set the node level memory rejection threshold - * @param nodeLevelMemoryRejectionThreshold sets the new memory rejection threshold - * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold - */ - public void setNodeLevelMemoryRejectionThreshold(Double nodeLevelMemoryRejectionThreshold) { - if (Double.compare(nodeLevelMemoryRejectionThreshold, NODE_LEVEL_MEMORY_REJECTION_THRESHOLD_MAX_VALUE) > 0) { - throw new IllegalArgumentException( - NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" - ); - } - - ensureRejectionThresholdIsLessThanCancellation( - nodeLevelMemoryRejectionThreshold, - nodeLevelMemoryCancellationThreshold, - NODE_MEMORY_REJECTION_THRESHOLD_SETTING_NAME, - NODE_MEMORY_CANCELLATION_THRESHOLD_SETTING_NAME - ); - - this.nodeLevelMemoryRejectionThreshold = nodeLevelMemoryRejectionThreshold; - } - - /** - * Method to get the cpu node level rejection threshold - * @return the current cpu node level rejection threshold - */ - public Double getNodeLevelCpuRejectionThreshold() { - return nodeLevelCpuRejectionThreshold; - } - - /** - * Method to set the node level cpu rejection threshold - * @param nodeLevelCpuRejectionThreshold sets the new cpu rejection threshold - * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold - */ - public void setNodeLevelCpuRejectionThreshold(Double nodeLevelCpuRejectionThreshold) { - if (Double.compare(nodeLevelCpuRejectionThreshold, NODE_LEVEL_CPU_REJECTION_THRESHOLD_MAX_VALUE) > 0) { - throw new IllegalArgumentException( - NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" - ); - } - - ensureRejectionThresholdIsLessThanCancellation( - nodeLevelCpuRejectionThreshold, - nodeLevelCpuCancellationThreshold, - NODE_CPU_REJECTION_THRESHOLD_SETTING_NAME, - NODE_CPU_CANCELLATION_THRESHOLD_SETTING_NAME - ); - - this.nodeLevelCpuRejectionThreshold = nodeLevelCpuRejectionThreshold; - } - - private void ensureRejectionThresholdIsLessThanCancellation( - Double nodeLevelRejectionThreshold, - Double nodeLevelCancellationThreshold, - String rejectionThresholdSettingName, - String cancellationThresholdSettingName - ) { - if (Double.compare(nodeLevelCancellationThreshold, nodeLevelRejectionThreshold) < 0) { - throw new IllegalArgumentException( - cancellationThresholdSettingName + " value should not be less than " + rejectionThresholdSettingName - ); - } - } - - /** - * Method to get the current QueryGroup count - * @return the current max QueryGroup count - */ - public int getMaxQueryGroupCount() { - return maxQueryGroupCount; - } -} diff --git a/server/src/main/java/org/opensearch/search/query_group/package-info.java b/server/src/main/java/org/opensearch/search/query_group/package-info.java deleted file mode 100644 index 00b68b0d3306c..0000000000000 --- a/server/src/main/java/org/opensearch/search/query_group/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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. - */ - -/** - * QueryGroup related artifacts - */ -package org.opensearch.search.query_group; From 322b2662b5f605b264474d195a1dd171b3b27282 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 25 Jul 2024 16:32:26 -0700 Subject: [PATCH 15/20] address comments on persistence service Signed-off-by: Ruirui Zhang --- .../wlm/action/CreateQueryGroupRequest.java | 2 +- .../wlm/action/CreateQueryGroupResponse.java | 2 +- .../wlm/rest/RestCreateQueryGroupAction.java | 10 +-- .../service/QueryGroupPersistenceService.java | 90 +++++++++++-------- 4 files changed, 58 insertions(+), 46 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java index 4e04ad1ce222c..35e3446fcedfe 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -32,7 +32,7 @@ * @opensearch.experimental */ public class CreateQueryGroupRequest extends ActionRequest { - final QueryGroup queryGroup; + private final QueryGroup queryGroup; /** * Constructor for CreateQueryGroupRequest diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java index 0c3de28a5bd51..9a2a8178c0a29 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java @@ -26,7 +26,7 @@ */ public class CreateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject { private final QueryGroup queryGroup; - private RestStatus restStatus; + private final RestStatus restStatus; /** * Constructor for CreateQueryGroupResponse diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java index ae328ba6025cc..b0e0af4f9d17f 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java @@ -55,12 +55,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - XContentParser parser = request.contentParser(); - return channel -> client.execute( - CreateQueryGroupAction.INSTANCE, - CreateQueryGroupRequest.fromXContent(parser), - createQueryGroupResponse(channel) - ); + try (XContentParser parser = request.contentParser()) { + CreateQueryGroupRequest createQueryGroupRequest = CreateQueryGroupRequest.fromXContent(parser); + return channel -> client.execute(CreateQueryGroupAction.INSTANCE, createQueryGroupRequest, createQueryGroupResponse(channel)); + } } private RestResponseListener createQueryGroupResponse(final RestChannel channel) { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 7273a0048cae9..f087d98ad2bfc 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -26,7 +26,7 @@ import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; -import java.util.HashMap; +import java.util.EnumMap; import java.util.Map; import java.util.Optional; @@ -34,32 +34,42 @@ * This class defines the functions for QueryGroup persistence */ public class QueryGroupPersistenceService { - public static final String QUERY_GROUP_COUNT_SETTING_NAME = "node.query_group.max_count"; private static final String SOURCE = "query-group-persistence-service"; private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group"; - public static final int DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE = 100; + private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); + /** + * max QueryGroup count setting name + */ + public static final String QUERY_GROUP_COUNT_SETTING_NAME = "node.query_group.max_count"; + /** + * default max queryGroup count on any node at any given point in time + */ + private static final int DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE = 100; + /** + * min queryGroup count on any node at any given point in time + */ + private static final int MIN_QUERY_GROUP_COUNT_VALUE = 1; + /** + * max QueryGroup count setting + */ public static final Setting MAX_QUERY_GROUP_COUNT = Setting.intSetting( QUERY_GROUP_COUNT_SETTING_NAME, DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE, 0, - (newVal) -> { - if (newVal > 100 || newVal < 1) throw new IllegalArgumentException( - QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]" - ); - }, + QueryGroupPersistenceService::validateMaxQueryGroupCount, Setting.Property.Dynamic, Setting.Property.NodeScope ); - private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); private final ClusterService clusterService; private volatile int maxQueryGroupCount; - final ThrottlingKey createQueryGroupThrottlingKey; + private final ThrottlingKey createQueryGroupThrottlingKey; /** * Constructor for QueryGroupPersistenceService * * @param clusterService {@link ClusterService} - The cluster service to be used by QueryGroupPersistenceService * @param settings {@link Settings} - The settings to be used by QueryGroupPersistenceService + * @param clusterSettings {@link ClusterSettings} - The cluster settings to be used by QueryGroupPersistenceService */ @Inject public QueryGroupPersistenceService( @@ -78,12 +88,20 @@ public QueryGroupPersistenceService( * @param newMaxQueryGroupCount - the max number of QueryGroup allowed */ public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { - if (newMaxQueryGroupCount > 100 || newMaxQueryGroupCount < 1) { - throw new IllegalArgumentException(QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]"); - } + validateMaxQueryGroupCount(newMaxQueryGroupCount); this.maxQueryGroupCount = newMaxQueryGroupCount; } + /** + * Validator for maxQueryGroupCount + * @param maxQueryGroupCount - the maxQueryGroupCount number to be verified + */ + private static void validateMaxQueryGroupCount(int maxQueryGroupCount) { + if (maxQueryGroupCount > DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE || maxQueryGroupCount < MIN_QUERY_GROUP_COUNT_VALUE) { + throw new IllegalArgumentException(QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]."); + } + } + /** * Update cluster state to include the new QueryGroup * @param queryGroup {@link QueryGroup} - the QueryGroup we're currently creating @@ -103,7 +121,7 @@ public ThrottlingKey getClusterManagerThrottlingKey() { @Override public void onFailure(String source, Exception e) { - logger.warn("failed to save QueryGroup object due to error: {}, for source: {}", e.getMessage(), source); + logger.warn("failed to save QueryGroup object due to error: {}, for source: {}.", e.getMessage(), source); listener.onFailure(e); } @@ -120,14 +138,14 @@ public void clusterStateProcessed(String source, ClusterState oldState, ClusterS * @param queryGroup - the QueryGroup we're currently creating * @param currentClusterState - the cluster state before the update */ - ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) { + public ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) { final Map existingQueryGroups = currentClusterState.metadata().queryGroups(); String groupName = queryGroup.getName(); // check if maxQueryGroupCount will breach - if (existingQueryGroups.size() + 1 > maxQueryGroupCount) { - logger.warn("{} value exceeded its assigned limit of {}", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); - throw new RuntimeException("Can't create more than " + maxQueryGroupCount + " QueryGroups in the system"); + if (existingQueryGroups.size() == maxQueryGroupCount) { + logger.warn("{} value exceeded its assigned limit of {}.", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); + throw new IllegalStateException("Can't create more than " + maxQueryGroupCount + " QueryGroups in the system."); } // check for duplicate name @@ -137,41 +155,37 @@ ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final Clu .findFirst(); if (findExistingGroup.isPresent()) { logger.warn("QueryGroup with name {} already exists. Not creating a new one.", groupName); - throw new RuntimeException("QueryGroup with name " + groupName + " already exists. Not creating a new one."); + throw new IllegalArgumentException("QueryGroup with name " + groupName + " already exists. Not creating a new one."); } // check if there's any resource allocation that exceed limit of 1.0 - Map existingUsageMap = calculateExistingUsage(existingQueryGroups, groupName); + Map totalUsageMap = calculateTotalUsage(existingQueryGroups, queryGroup); for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { - double totalUsage = existingUsageMap.getOrDefault(resourceType, 0.0) + queryGroup.getResourceLimits().get(resourceType); - if (totalUsage > 1) { - logger.warn("Total resource allocation for {} will go above the max limit of 1.0", resourceType.getName()); - throw new RuntimeException( - "Total resource allocation for " + resourceType.getName() + " will go above the max limit of 1.0" + if (totalUsageMap.get(resourceType) > 1) { + logger.warn("Total resource allocation for {} will go above the max limit of 1.0.", resourceType.getName()); + throw new IllegalArgumentException( + "Total resource allocation for " + resourceType.getName() + " will go above the max limit of 1.0." ); } } return ClusterState.builder(currentClusterState) - .metadata( - Metadata.builder(currentClusterState.metadata()) - .put(queryGroup) - .build() - ) + .metadata(Metadata.builder(currentClusterState.metadata()).put(queryGroup).build()) .build(); } /** - * This method calculates the existing total usage of the all the resource limits (except the group that we're updating here) - * @param existingQueryGroups - existing QueryGroups - * @param groupName - the QueryGroup name we're updating + * This method calculates the existing total usage of the all the resource limits + * @param existingQueryGroups - existing QueryGroups in the system + * @param queryGroup - the QueryGroup we're creating or updating */ - private Map calculateExistingUsage(Map existingQueryGroups, String groupName) { - Map map = new HashMap<>(); + private Map calculateTotalUsage(Map existingQueryGroups, QueryGroup queryGroup) { + final Map map = new EnumMap<>(ResourceType.class); + map.putAll(queryGroup.getResourceLimits()); for (QueryGroup currGroup : existingQueryGroups.values()) { - if (!currGroup.getName().equals(groupName)) { - for (ResourceType resourceType : currGroup.getResourceLimits().keySet()) { - map.put(resourceType, map.getOrDefault(resourceType, 0.0) + currGroup.getResourceLimits().get(resourceType)); + if (!currGroup.getName().equals(queryGroup.getName())) { + for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { + map.compute(resourceType, (k, v) -> v + currGroup.getResourceLimits().get(resourceType)); } } } From e8bf2b8dc48644f18441d30b2fa044c1360750b6 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Mon, 29 Jul 2024 16:36:41 -0700 Subject: [PATCH 16/20] address comments Signed-off-by: Ruirui Zhang --- .../wlm/action/CreateQueryGroupRequest.java | 8 +- .../service/QueryGroupPersistenceService.java | 2 +- .../action/CreateQueryGroupRequestTests.java | 3 + .../action/CreateQueryGroupResponseTests.java | 6 + .../QueryGroupPersistenceServiceTests.java | 18 +++ .../cluster/metadata/QueryGroup.java | 109 +++++++++--------- 6 files changed, 86 insertions(+), 60 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java index 35e3446fcedfe..b967d9744bd77 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -11,9 +11,11 @@ import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.UUIDs; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.XContentParser; +import org.joda.time.Instant; import java.io.IOException; @@ -56,10 +58,8 @@ public CreateQueryGroupRequest(StreamInput in) throws IOException { * @param parser - A {@link XContentParser} object */ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws IOException { - QueryGroup queryGroup1 = QueryGroup.fromXContent(parser); - // creating this queryGroup to ensure forceful creation of _id and updatedAt - QueryGroup queryGroup = new QueryGroup(queryGroup1.getName(), queryGroup1.getResiliencyMode(), queryGroup1.getResourceLimits()); - return new CreateQueryGroupRequest(queryGroup); + QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser); + return new CreateQueryGroupRequest(builder._id(UUIDs.randomBase64UUID()).updatedAt(Instant.now().getMillis()).build()); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index f087d98ad2bfc..226cc3639fafb 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -138,7 +138,7 @@ public void clusterStateProcessed(String source, ClusterState oldState, ClusterS * @param queryGroup - the QueryGroup we're currently creating * @param currentClusterState - the cluster state before the update */ - public ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) { + ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) { final Map existingQueryGroups = currentClusterState.metadata().queryGroups(); String groupName = queryGroup.getName(); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java index a7ec3a49c6eb3..6940fca791595 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java @@ -22,6 +22,9 @@ public class CreateQueryGroupRequestTests extends OpenSearchTestCase { + /** + * Test case to verify the serialization and deserialization of CreateQueryGroupRequest. + */ public void testSerialization() throws IOException { CreateQueryGroupRequest request = new CreateQueryGroupRequest(queryGroupOne); BytesStreamOutput out = new BytesStreamOutput(); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java index 20aab264adffe..d9319217b786f 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java @@ -26,6 +26,9 @@ public class CreateQueryGroupResponseTests extends OpenSearchTestCase { + /** + * Test case to verify the serialization and deserialization of CreateQueryGroupResponse. + */ public void testSerialization() throws IOException { CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); BytesStreamOutput out = new BytesStreamOutput(); @@ -42,6 +45,9 @@ public void testSerialization() throws IOException { QueryGroupTestUtils.compareQueryGroups(listOne, listTwo); } + /** + * Test case to verify the toXContent method of CreateQueryGroupResponse. + */ public void testToXContentCreateQueryGroup() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 6c389d47cfa7c..7d932a6249f46 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -43,6 +43,9 @@ public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase { + /** + * Test case to validate the creation logic of a single QueryGroup + */ public void testCreateQueryGroup() { List setup = preparePersistenceServiceSetup(new HashMap<>()); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); @@ -58,6 +61,10 @@ public void testCreateQueryGroup() { compareQueryGroups(listOne, listTwo); } + /** + * Test case to validate the logic for adding a new QueryGroup to a cluster state that already contains + * an existing QueryGroup + */ public void testCreateAnotherQueryGroup() { List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); @@ -70,6 +77,9 @@ public void testCreateAnotherQueryGroup() { compareQueryGroups(queryGroupList(), new ArrayList<>(values)); } + /** + * Test case to ensure the error is thrown when we try to create another QueryGroup with duplicate name + */ public void testCreateQueryGroupDuplicateName() { List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); @@ -83,6 +93,10 @@ public void testCreateQueryGroupDuplicateName() { assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } + /** + * Test case to ensure the error is thrown when we try to create another QueryGroup that will make + * the total resource limits go above 1 + */ public void testCreateQueryGroupOverflowAllocation() { List setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupTwo)); QueryGroup toCreate = builder().name(NAME_ONE) @@ -96,6 +110,10 @@ public void testCreateQueryGroupOverflowAllocation() { assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } + /** + * Test case to ensure the error is thrown when we already have the max allowed number of QueryGroups, but + * we try to create another one + */ public void testCreateQueryGroupOverflowCount() { QueryGroup toCreate = builder().name(NAME_NONE_EXISTED) ._id("W5iIqHyhgi4K1qIAAAAIHw==") diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index e71df169de5db..d4b8d9abbc9a6 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -148,56 +148,7 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa } public static QueryGroup fromXContent(final XContentParser parser) throws IOException { - if (parser.currentToken() == null) { // fresh parser? move to the first token - parser.nextToken(); - } - - Builder builder = builder(); - - XContentParser.Token token = parser.currentToken(); - - if (token != XContentParser.Token.START_OBJECT) { - throw new IllegalArgumentException("Expected START_OBJECT token but found [" + parser.currentName() + "]"); - } - - String fieldName = ""; - // Map to hold resources - final Map resourceLimits = new HashMap<>(); - while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { - if (token == XContentParser.Token.FIELD_NAME) { - fieldName = parser.currentName(); - } else if (token.isValue()) { - if (fieldName.equals(_ID_STRING)) { - builder._id(parser.text()); - } else if (fieldName.equals(NAME_STRING)) { - builder.name(parser.text()); - } else if (fieldName.equals(RESILIENCY_MODE_STRING)) { - builder.mode(parser.text()); - } else if (fieldName.equals(UPDATED_AT_STRING)) { - builder.updatedAt(parser.longValue()); - } else { - throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup"); - } - } else if (token == XContentParser.Token.START_OBJECT) { - - if (!fieldName.equals(RESOURCE_LIMITS_STRING)) { - throw new IllegalArgumentException( - "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token - ); - } - - while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { - if (token == XContentParser.Token.FIELD_NAME) { - fieldName = parser.currentName(); - } else { - resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); - } - } - - } - } - builder.resourceLimits(resourceLimits); - return builder.build(); + return Builder.fromXContent(parser).build(); } public static Diff readDiff(final StreamInput in) throws IOException { @@ -293,6 +244,58 @@ public static class Builder { private Builder() {} + public static Builder fromXContent(XContentParser parser) throws IOException { + if (parser.currentToken() == null) { // fresh parser? move to the first token + parser.nextToken(); + } + + Builder builder = builder(); + + XContentParser.Token token = parser.currentToken(); + + if (token != XContentParser.Token.START_OBJECT) { + throw new IllegalArgumentException("Expected START_OBJECT token but found [" + parser.currentName() + "]"); + } + + String fieldName = ""; + // Map to hold resources + final Map resourceLimits = new HashMap<>(); + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else if (token.isValue()) { + if (fieldName.equals(_ID_STRING)) { + builder._id(parser.text()); + } else if (fieldName.equals(NAME_STRING)) { + builder.name(parser.text()); + } else if (fieldName.equals(RESILIENCY_MODE_STRING)) { + builder.mode(parser.text()); + } else if (fieldName.equals(UPDATED_AT_STRING)) { + builder.updatedAt(parser.longValue()); + } else { + throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup"); + } + } else if (token == XContentParser.Token.START_OBJECT) { + + if (!fieldName.equals(RESOURCE_LIMITS_STRING)) { + throw new IllegalArgumentException( + "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token + ); + } + + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else { + resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); + } + } + + } + } + return builder.resourceLimits(resourceLimits); + } + public Builder name(String name) { this.name = name; return this; @@ -319,11 +322,7 @@ public Builder resourceLimits(Map resourceLimits) { } public QueryGroup build() { - if (_id == null && updatedAt == 0L) { - return new QueryGroup(name, resiliencyMode, resourceLimits); - } else { - return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); - } + return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); } } } From 1b09c34d4fec1a18f12b796604c29bb868ecc763 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 31 Jul 2024 15:25:21 -0700 Subject: [PATCH 17/20] fix unit test Signed-off-by: Ruirui Zhang --- .../service/QueryGroupPersistenceServiceTests.java | 14 +++++++++++++- .../cluster/metadata/QueryGroupMetadataTests.java | 2 +- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 7d932a6249f46..4acf09a6f4c86 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -112,7 +112,7 @@ public void testCreateQueryGroupOverflowAllocation() { /** * Test case to ensure the error is thrown when we already have the max allowed number of QueryGroups, but - * we try to create another one + * we want to create another one */ public void testCreateQueryGroupOverflowCount() { QueryGroup toCreate = builder().name(NAME_NONE_EXISTED) @@ -133,4 +133,16 @@ public void testCreateQueryGroupOverflowCount() { ); assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } + + public void testInvalidMaxQueryGroupCount() { + Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build(); + ClusterSettings clusterSettings = new ClusterSettings(settings, clusterSettingsSet()); + ClusterService clusterService = new ClusterService(settings, clusterSettings, mock(ThreadPool.class)); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings + ); + assertThrows(IllegalArgumentException.class, () -> queryGroupPersistenceService.setMaxQueryGroupCount(-1)); + } } diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java index d70a9ce5e10cd..06734b8e0bac2 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -44,7 +44,7 @@ public void testToXContent() throws IOException { queryGroupMetadata.toXContent(builder, null); builder.endObject(); assertEquals( - "{\"ajakgakg983r92_4242\":{\"_id\":\"ajakgakg983r92_4242\",\"name\":\"test\",\"resiliency_mode\":\"enforced\",\"updatedAt\":1720047207,\"resourceLimits\":{\"memory\":0.5}}}", + "{\"ajakgakg983r92_4242\":{\"_id\":\"ajakgakg983r92_4242\",\"name\":\"test\",\"resiliency_mode\":\"enforced\",\"updated_at\":1720047207,\"resource_limits\":{\"memory\":0.5}}}", builder.toString() ); } From 494a025c33a06b1dcc9c7400f68b1207211eb0f6 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 6 Aug 2024 13:57:14 -0700 Subject: [PATCH 18/20] address comments Signed-off-by: Ruirui Zhang --- .../wlm/action/CreateQueryGroupRequest.java | 2 +- .../plugin/wlm/QueryGroupTestUtils.java | 13 +++----- .../action/CreateQueryGroupRequestTests.java | 4 +-- .../action/CreateQueryGroupResponseTests.java | 2 +- .../QueryGroupPersistenceServiceTests.java | 32 ++++++++++--------- .../cluster/metadata/QueryGroup.java | 14 +++----- 6 files changed, 30 insertions(+), 37 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java index b967d9744bd77..ff6422be36885 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -70,7 +70,7 @@ public ActionRequestValidationException validate() { @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); - QueryGroup.writeToOutput(out, queryGroup); + queryGroup.writeTo(out); } /** diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java index 777eb14570376..fc324853d9b34 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java @@ -15,11 +15,11 @@ import org.opensearch.cluster.service.ClusterApplierService; import org.opensearch.cluster.service.ClusterManagerService; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.collect.Tuple; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; -import org.opensearch.search.ResourceType; import org.opensearch.threadpool.ThreadPool; import java.util.ArrayList; @@ -104,7 +104,7 @@ public static QueryGroupPersistenceService queryGroupPersistenceService() { return new QueryGroupPersistenceService(clusterService, settings(), clusterSettings()); } - public static List preparePersistenceServiceSetup(Map queryGroups) { + public static Tuple preparePersistenceServiceSetup(Map queryGroups) { Metadata metadata = Metadata.builder().queryGroups(queryGroups).build(); Settings settings = Settings.builder().build(); ClusterState clusterState = ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); @@ -127,15 +127,10 @@ public static List preparePersistenceServiceSetup(Map(queryGroupPersistenceService, clusterState); } - public static void compareResourceLimits(Map resourceLimitMapOne, Map resourceLimitMapTwo) { - assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet())); - assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); - } - - public static void compareQueryGroups(List listOne, List listTwo) { + public static void assertEqualQueryGroups(List listOne, List listTwo) { assertEquals(listOne.size(), listTwo.size()); listOne.sort(Comparator.comparing(QueryGroup::getName)); listTwo.sort(Comparator.comparing(QueryGroup::getName)); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java index 6940fca791595..b0fa96a46df80 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java @@ -17,7 +17,7 @@ import java.util.ArrayList; import java.util.List; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.compareQueryGroups; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; public class CreateQueryGroupRequestTests extends OpenSearchTestCase { @@ -35,6 +35,6 @@ public void testSerialization() throws IOException { List list2 = new ArrayList<>(); list1.add(queryGroupOne); list2.add(otherRequest.getQueryGroup()); - compareQueryGroups(list1, list2); + assertEqualQueryGroups(list1, list2); } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java index d9319217b786f..038f015713c5b 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java @@ -42,7 +42,7 @@ public void testSerialization() throws IOException { List listTwo = new ArrayList<>(); listOne.add(responseGroup); listTwo.add(otherResponseGroup); - QueryGroupTestUtils.compareQueryGroups(listOne, listTwo); + QueryGroupTestUtils.assertEqualQueryGroups(listOne, listTwo); } /** diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 4acf09a6f4c86..1eb2b6e01f775 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -13,6 +13,7 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.collect.Tuple; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.search.ResourceType; @@ -32,8 +33,8 @@ import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_ONE; import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettingsSet; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.compareQueryGroups; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; @@ -47,9 +48,9 @@ public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase { * Test case to validate the creation logic of a single QueryGroup */ public void testCreateQueryGroup() { - List setup = preparePersistenceServiceSetup(new HashMap<>()); - QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); - ClusterState clusterState = (ClusterState) setup.get(1); + Tuple setup = preparePersistenceServiceSetup(new HashMap<>()); + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupOne, clusterState); Map updatedGroupsMap = newClusterState.getMetadata().queryGroups(); assertEquals(1, updatedGroupsMap.size()); @@ -58,7 +59,7 @@ public void testCreateQueryGroup() { List listTwo = new ArrayList<>(); listOne.add(queryGroupOne); listTwo.add(updatedGroupsMap.get(_ID_ONE)); - compareQueryGroups(listOne, listTwo); + assertEqualQueryGroups(listOne, listTwo); } /** @@ -66,24 +67,24 @@ public void testCreateQueryGroup() { * an existing QueryGroup */ public void testCreateAnotherQueryGroup() { - List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); - QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); - ClusterState clusterState = (ClusterState) setup.get(1); + Tuple setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupTwo, clusterState); Map updatedGroups = newClusterState.getMetadata().queryGroups(); assertEquals(2, updatedGroups.size()); assertTrue(updatedGroups.containsKey(_ID_TWO)); Collection values = updatedGroups.values(); - compareQueryGroups(queryGroupList(), new ArrayList<>(values)); + assertEqualQueryGroups(queryGroupList(), new ArrayList<>(values)); } /** * Test case to ensure the error is thrown when we try to create another QueryGroup with duplicate name */ public void testCreateQueryGroupDuplicateName() { - List setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); - QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); - ClusterState clusterState = (ClusterState) setup.get(1); + Tuple setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); QueryGroup toCreate = builder().name(NAME_ONE) ._id("W5iIqHyhgi4K1qIAAAAIHw==") .mode(MONITOR_STRING) @@ -98,15 +99,16 @@ public void testCreateQueryGroupDuplicateName() { * the total resource limits go above 1 */ public void testCreateQueryGroupOverflowAllocation() { - List setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupTwo)); + Tuple setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupTwo)); QueryGroup toCreate = builder().name(NAME_ONE) ._id("W5iIqHyhgi4K1qIAAAAIHw==") .mode(MONITOR_STRING) .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.41)) .updatedAt(1690934400000L) .build(); - QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0); - ClusterState clusterState = (ClusterState) setup.get(1); + + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index d4b8d9abbc9a6..a9ed3fc86b2f2 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -104,15 +104,11 @@ public QueryGroup(StreamInput in) throws IOException { @Override public void writeTo(StreamOutput out) throws IOException { - writeToOutput(out, this); - } - - public static void writeToOutput(StreamOutput out, QueryGroup queryGroup) throws IOException { - out.writeString(queryGroup.getName()); - out.writeString(queryGroup.get_id()); - out.writeString(queryGroup.getResiliencyMode().getName()); - out.writeMap(queryGroup.getResourceLimits(), ResourceType::writeTo, StreamOutput::writeDouble); - out.writeLong(queryGroup.getUpdatedAtInMillis()); + out.writeString(name); + out.writeString(_id); + out.writeString(resiliencyMode.getName()); + out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeDouble); + out.writeLong(updatedAtInMillis); } private void validateResourceLimits(Map resourceLimits) { From c4bc6f5ae3feca3fabcc628d5c2229537300426c Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 7 Aug 2024 15:55:43 -0700 Subject: [PATCH 19/20] add IT Signed-off-by: Ruirui Zhang --- plugins/workload-management/build.gradle | 3 + ...rkloadManagementClientYamlTestSuiteIT.java | 52 +++++++++ .../api/create_query_group_context.json | 18 ++++ .../test/wlm/10_create_query_group.yml | 102 ++++++++++++++++++ .../cluster/metadata/QueryGroup.java | 4 +- 5 files changed, 177 insertions(+), 2 deletions(-) create mode 100644 plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java create mode 100644 plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json create mode 100644 plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml diff --git a/plugins/workload-management/build.gradle b/plugins/workload-management/build.gradle index e91ef90005c74..cb14d22ef149f 100644 --- a/plugins/workload-management/build.gradle +++ b/plugins/workload-management/build.gradle @@ -9,6 +9,9 @@ * GitHub history for details. */ +apply plugin: 'opensearch.yaml-rest-test' +apply plugin: 'opensearch.internal-cluster-test' + opensearchplugin { description 'OpenSearch Workload Management Plugin.' classname 'org.opensearch.plugin.wlm.WorkloadManagementPlugin' diff --git a/plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java b/plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java new file mode 100644 index 0000000000000..9ec4a36ff6a5b --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java @@ -0,0 +1,52 @@ +/* + * 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. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.plugin.wlm; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.opensearch.test.rest.yaml.ClientYamlTestCandidate; +import org.opensearch.test.rest.yaml.OpenSearchClientYamlSuiteTestCase; + +/** Runs yaml rest tests */ +public class WorkloadManagementClientYamlTestSuiteIT extends OpenSearchClientYamlSuiteTestCase { + + public WorkloadManagementClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) { + super(testCandidate); + } + + @ParametersFactory + public static Iterable parameters() throws Exception { + return OpenSearchClientYamlSuiteTestCase.createParameters(); + } +} diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json new file mode 100644 index 0000000000000..bb4620c01f2d6 --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json @@ -0,0 +1,18 @@ +{ + "create_query_group_context": { + "stability": "experimental", + "url": { + "paths": [ + { + "path": "/_wlm/query_group", + "methods": ["PUT", "POST"], + "parts": {} + } + ] + }, + "params":{}, + "body":{ + "description":"The QueryGroup schema" + } + } +} diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml new file mode 100644 index 0000000000000..9e0e7390cee0a --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml @@ -0,0 +1,102 @@ +"test create a QueryGroup API successfully": + - skip: + version: " - 2.16.99" + reason: "QueryGroup WorkloadManagement feature was added in 2.17" + + - do: + create_query_group_context: + body: + { + "name": "analytics", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.4, + "memory": 0.2 + } + } + + - match: { name: "analytics" } + - match: { resiliency_mode: "monitor" } + - match: { resource_limits.cpu: 0.4 } + - match: { resource_limits.memory: 0.2 } + +--- +"test bad QueryGroup API arguments": + - skip: + version: " - 2.16.99" + reason: "QueryGroup WorkloadManagement feature was added in 2.17" + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "analytics", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.4, + "memory": 0.2 + } + } + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "analytics2", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.61, + "memory": 0.2 + } + } + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "analytics2", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": -0.1, + "memory": 0.2 + } + } + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.1, + "memory": 0.2 + } + } + +--- +"test create another QueryGroup API successfully": + - skip: + version: " - 2.16.99" + reason: "QueryGroup WorkloadManagement feature was added in 2.17" + + - do: + create_query_group_context: + body: + { + "name": "analytics2", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.35, + "memory": 0.25 + } + } + + - match: { name: "analytics2" } + - match: { resiliency_mode: "monitor" } + - match: { resource_limits.cpu: 0.35 } + - match: { resource_limits.memory: 0.25 } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index a9ed3fc86b2f2..6ab11b1d6f150 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -63,8 +63,8 @@ public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map MAX_CHARS_ALLOWED_IN_NAME) { - throw new IllegalArgumentException("QueryGroup.name shouldn't be more than 50 chars long"); + if (name.length() > MAX_CHARS_ALLOWED_IN_NAME || name.isEmpty()) { + throw new IllegalArgumentException("QueryGroup.name shouldn't be empty or more than 50 chars long"); } if (resourceLimits.isEmpty()) { From 8e123f7b504843bc98a4038de497c45c92631fde Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 9 Aug 2024 14:00:13 -0700 Subject: [PATCH 20/20] add coverage Signed-off-by: Ruirui Zhang --- .../TransportCreateQueryGroupAction.java | 2 +- .../service/QueryGroupPersistenceService.java | 11 ++- .../QueryGroupPersistenceServiceTests.java | 99 ++++++++++++++++++- .../test/wlm/10_create_query_group.yml | 14 +-- .../cluster/metadata/QueryGroupTests.java | 23 +++++ 5 files changed, 132 insertions(+), 17 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java index f3e913ad019e3..01aa8cfb5e610 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java @@ -51,7 +51,7 @@ public TransportCreateQueryGroupAction( @Override protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener listener) { - threadPool.executor(ThreadPool.Names.GENERIC) + threadPool.executor(ThreadPool.Names.SAME) .execute(() -> queryGroupPersistenceService.persistInClusterStateMetadata(request.getQueryGroup(), listener)); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 226cc3639fafb..b2164df561bf9 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -34,7 +34,7 @@ * This class defines the functions for QueryGroup persistence */ public class QueryGroupPersistenceService { - private static final String SOURCE = "query-group-persistence-service"; + static final String SOURCE = "query-group-persistence-service"; private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group"; private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); /** @@ -62,7 +62,7 @@ public class QueryGroupPersistenceService { ); private final ClusterService clusterService; private volatile int maxQueryGroupCount; - private final ThrottlingKey createQueryGroupThrottlingKey; + final ThrottlingKey createQueryGroupThrottlingKey; /** * Constructor for QueryGroupPersistenceService @@ -191,4 +191,11 @@ private Map calculateTotalUsage(Map ex } return map; } + + /** + * maxQueryGroupCount getter + */ + public int getMaxQueryGroupCount() { + return maxQueryGroupCount; + } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 1eb2b6e01f775..533c98b44685d 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -10,12 +10,16 @@ import org.opensearch.cluster.ClusterName; import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.ClusterStateUpdateTask; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.collect.Tuple; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.QueryGroupTestUtils; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; @@ -26,6 +30,8 @@ import java.util.List; import java.util.Map; +import org.mockito.ArgumentCaptor; + import static org.opensearch.cluster.metadata.QueryGroup.builder; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MEMORY_STRING; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MONITOR_STRING; @@ -34,18 +40,26 @@ import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettings; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettingsSet; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo; import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.QUERY_GROUP_COUNT_SETTING_NAME; +import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.SOURCE; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase { /** - * Test case to validate the creation logic of a single QueryGroup + * Test case to validate the creation logic of a QueryGroup */ public void testCreateQueryGroup() { Tuple setup = preparePersistenceServiceSetup(new HashMap<>()); @@ -136,6 +150,9 @@ public void testCreateQueryGroupOverflowCount() { assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); } + /** + * Tests the invalid value of {@code node.query_group.max_count} + */ public void testInvalidMaxQueryGroupCount() { Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build(); ClusterSettings clusterSettings = new ClusterSettings(settings, clusterSettingsSet()); @@ -147,4 +164,84 @@ public void testInvalidMaxQueryGroupCount() { ); assertThrows(IllegalArgumentException.class, () -> queryGroupPersistenceService.setMaxQueryGroupCount(-1)); } + + /** + * Tests the valid value of {@code node.query_group.max_count} + */ + public void testValidMaxSandboxCountSetting() { + Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 100).build(); + ClusterService clusterService = new ClusterService(settings, clusterSettings(), mock(ThreadPool.class)); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings() + ); + queryGroupPersistenceService.setMaxQueryGroupCount(50); + assertEquals(50, queryGroupPersistenceService.getMaxQueryGroupCount()); + } + + /** + * Tests PersistInClusterStateMetadata function + */ + public void testPersistInClusterStateMetadata() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(clusterService).submitStateUpdateTask(eq(SOURCE), any()); + } + + /** + * Tests PersistInClusterStateMetadata function with inner functions + */ + public void testPersistInClusterStateMetadataInner() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + ArgumentCaptor captor = ArgumentCaptor.forClass(ClusterStateUpdateTask.class); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(clusterService, times(1)).submitStateUpdateTask(eq(SOURCE), captor.capture()); + ClusterStateUpdateTask capturedTask = captor.getValue(); + assertEquals(queryGroupPersistenceService.createQueryGroupThrottlingKey, capturedTask.getClusterManagerThrottlingKey()); + + doAnswer(invocation -> { + ClusterStateUpdateTask task = invocation.getArgument(1); + task.clusterStateProcessed(SOURCE, mock(ClusterState.class), mock(ClusterState.class)); + return null; + }).when(clusterService).submitStateUpdateTask(anyString(), any()); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(listener).onResponse(any(CreateQueryGroupResponse.class)); + } + + /** + * Tests PersistInClusterStateMetadata function with failure + */ + public void testPersistInClusterStateMetadataFailure() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + doAnswer(invocation -> { + ClusterStateUpdateTask task = invocation.getArgument(1); + Exception exception = new RuntimeException("Test Exception"); + task.onFailure(SOURCE, exception); + return null; + }).when(clusterService).submitStateUpdateTask(anyString(), any()); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(listener).onFailure(any(RuntimeException.class)); + } } diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml index 9e0e7390cee0a..ae82a8146e9cd 100644 --- a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml @@ -1,4 +1,4 @@ -"test create a QueryGroup API successfully": +"test create QueryGroup API": - skip: version: " - 2.16.99" reason: "QueryGroup WorkloadManagement feature was added in 2.17" @@ -20,12 +20,6 @@ - match: { resource_limits.cpu: 0.4 } - match: { resource_limits.memory: 0.2 } ---- -"test bad QueryGroup API arguments": - - skip: - version: " - 2.16.99" - reason: "QueryGroup WorkloadManagement feature was added in 2.17" - - do: catch: /illegal_argument_exception/ create_query_group_context: @@ -78,12 +72,6 @@ } } ---- -"test create another QueryGroup API successfully": - - skip: - version: " - 2.16.99" - reason: "QueryGroup WorkloadManagement feature was added in 2.17" - - do: create_query_group_context: body: diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java index db2ff239271b0..884b364fb26b8 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -103,6 +103,29 @@ public void testIllegalQueryGroupMode() { ); } + public void testQueryGroupInitiation() { + QueryGroup queryGroup = new QueryGroup("analytics", randomMode(), Map.of(ResourceType.MEMORY, 0.4)); + assertNotNull(queryGroup.getName()); + assertNotNull(queryGroup.get_id()); + assertNotNull(queryGroup.getResourceLimits()); + assertFalse(queryGroup.getResourceLimits().isEmpty()); + assertEquals(1, queryGroup.getResourceLimits().size()); + assertTrue(allowedModes.contains(queryGroup.getResiliencyMode())); + assertTrue(queryGroup.getUpdatedAtInMillis() != 0); + } + + public void testIllegalQueryGroupName() { + assertThrows( + NullPointerException.class, + () -> new QueryGroup("a".repeat(51), "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis()) + ); + assertThrows( + NullPointerException.class, + () -> new QueryGroup("", "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis()) + ); + + } + public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { assertThrows( IllegalArgumentException.class,