From d37094b55534d1f322d7039e22ec4086cae93979 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 16 Jul 2024 12:58:36 -0700 Subject: [PATCH 1/9] Add Update QueryGroup API Logic Signed-off-by: Ruirui Zhang --- .../wlm/TransportUpdateQueryGroupAction.java | 57 +++++ .../plugin/wlm/UpdateQueryGroupAction.java | 36 +++ .../plugin/wlm/UpdateQueryGroupRequest.java | 225 ++++++++++++++++++ .../plugin/wlm/UpdateQueryGroupResponse.java | 73 ++++++ .../plugin/wlm/WorkloadManagementPlugin.java | 7 +- .../wlm/rest/RestUpdateQueryGroupAction.java | 80 +++++++ .../service/QueryGroupPersistenceService.java | 146 +++++++++--- .../plugin/wlm/service/package-info.java | 4 + .../plugin/wlm/QueryGroupTestUtils.java | 6 + .../wlm/UpdateQueryGroupRequestTests.java | 96 ++++++++ .../wlm/UpdateQueryGroupResponseTests.java | 60 +++++ .../QueryGroupPersistenceServiceTests.java | 179 +++++++++++++- .../cluster/metadata/QueryGroup.java | 4 + .../common/settings/ClusterSettings.java | 1 + .../QueryGroupServiceSettings.java | 198 +++++++++++++++ .../search/query_group/package-info.java | 12 + 16 files changed, 1154 insertions(+), 30 deletions(-) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportUpdateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequest.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequestTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponseTests.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/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportUpdateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportUpdateQueryGroupAction.java new file mode 100644 index 0000000000000..0e560718ced62 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportUpdateQueryGroupAction.java @@ -0,0 +1,57 @@ +/* + * 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.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.Persistable; +import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +/** + * Transport action for update QueryGroup + * + * @opensearch.internal + */ +public class TransportUpdateQueryGroupAction extends HandledTransportAction { + + private final ThreadPool threadPool; + private final Persistable queryGroupPersistenceService; + + /** + * Constructor for TransportUpdateQueryGroupAction + * + * @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 TransportUpdateQueryGroupAction( + String actionName, + TransportService transportService, + ActionFilters actionFilters, + ThreadPool threadPool, + Persistable queryGroupPersistenceService + ) { + super(UpdateQueryGroupAction.NAME, transportService, actionFilters, UpdateQueryGroupRequest::new); + this.threadPool = threadPool; + this.queryGroupPersistenceService = queryGroupPersistenceService; + } + + @Override + protected void doExecute(Task task, UpdateQueryGroupRequest request, ActionListener listener) { + threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> queryGroupPersistenceService.update(request, listener)); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupAction.java new file mode 100644 index 0000000000000..5b3600c091944 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupAction.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; + +import org.opensearch.action.ActionType; + +/** + * Transport action to update QueryGroup + * + * @opensearch.api + */ +public class UpdateQueryGroupAction extends ActionType { + + /** + * An instance of UpdateQueryGroupAction + */ + public static final UpdateQueryGroupAction INSTANCE = new UpdateQueryGroupAction(); + + /** + * Name for UpdateQueryGroupAction + */ + public static final String NAME = "cluster:admin/opensearch/wlm/query_group/_update"; + + /** + * Default constructor + */ + private UpdateQueryGroupAction() { + super(NAME, UpdateQueryGroupResponse::new); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequest.java new file mode 100644 index 0000000000000..b33d501d1daf0 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequest.java @@ -0,0 +1,225 @@ +/* + * 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.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; +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.opensearch.search.ResourceType; +import org.joda.time.Instant; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +/** + * A request for update QueryGroup + * + * @opensearch.internal + */ +public class UpdateQueryGroupRequest extends ActionRequest implements Writeable.Reader { + String name; + Map resourceLimits; + private ResiliencyMode resiliencyMode; + long updatedAtInMillis; + + /** + * Default constructor for UpdateQueryGroupRequest + */ + public UpdateQueryGroupRequest() {} + + /** + * Constructor for UpdateQueryGroupRequest + * @param queryGroup - A {@link QueryGroup} object + */ + public UpdateQueryGroupRequest(QueryGroup queryGroup) { + this.name = queryGroup.getName(); + this.resiliencyMode = queryGroup.getResiliencyMode(); + this.resourceLimits = queryGroup.getResourceLimits(); + this.updatedAtInMillis = queryGroup.getUpdatedAtInMillis(); + } + + /** + * Constructor for UpdateQueryGroupRequest + * @param name - QueryGroup name for UpdateQueryGroupRequest + * @param resiliencyMode - QueryGroup mode for UpdateQueryGroupRequest + * @param resourceLimits - QueryGroup resourceLimits for UpdateQueryGroupRequest + * @param updatedAtInMillis - QueryGroup updated time in millis for UpdateQueryGroupRequest + */ + public UpdateQueryGroupRequest( + String name, + ResiliencyMode resiliencyMode, + Map resourceLimits, + long updatedAtInMillis + ) { + this.name = name; + this.resiliencyMode = resiliencyMode; + this.resourceLimits = resourceLimits; + this.updatedAtInMillis = updatedAtInMillis; + } + + /** + * Constructor for UpdateQueryGroupRequest + * @param in - A {@link StreamInput} object + */ + public UpdateQueryGroupRequest(StreamInput in) throws IOException { + super(in); + name = in.readString(); + if (in.readBoolean()) { + resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble); + } else { + resourceLimits = new HashMap<>(); + } + if (in.readBoolean()) { + resiliencyMode = ResiliencyMode.fromName(in.readString()); + } + updatedAtInMillis = in.readLong(); + } + + @Override + public UpdateQueryGroupRequest read(StreamInput in) throws IOException { + return new UpdateQueryGroupRequest(in); + } + + /** + * Generate a UpdateQueryGroupRequest from XContent + * @param parser - A {@link XContentParser} object + * @param name - name of the QueryGroup to be updated + */ + public static UpdateQueryGroupRequest fromXContent(XContentParser parser, String name) 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 = ""; + 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("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(ResourceType.fromName(fieldName), parser.doubleValue()); + } + } + } + } + return new UpdateQueryGroupRequest(name, mode, resourceLimits, Instant.now().getMillis()); + } + + @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; + } + + /** + * resiliencyMode getter + */ + public ResiliencyMode getResiliencyMode() { + return resiliencyMode; + } + + /** + * resiliencyMode setter + * @param resiliencyMode - mode to be set + */ + public void setResiliencyMode(ResiliencyMode resiliencyMode) { + this.resiliencyMode = resiliencyMode; + } + + /** + * updatedAtInMillis getter + */ + public long getUpdatedAtInMillis() { + return updatedAtInMillis; + } + + /** + * updatedAtInMillis setter + * @param updatedAtInMillis - updatedAtInMillis to be set + */ + public void setUpdatedAtInMillis(long updatedAtInMillis) { + this.updatedAtInMillis = updatedAtInMillis; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(name); + if (resourceLimits == null || resourceLimits.isEmpty()) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue); + } + if (resiliencyMode == null) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeString(resiliencyMode.getName()); + } + out.writeLong(updatedAtInMillis); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java new file mode 100644 index 0000000000000..b1fd4b698d017 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java @@ -0,0 +1,73 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm; + +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 update API for QueryGroup + * + * @opensearch.internal + */ +public class UpdateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject { + private final QueryGroup queryGroup; + private final RestStatus restStatus; + + /** + * Constructor for UpdateQueryGroupResponse + * @param queryGroup - The QueryGroup to be updated + */ + public UpdateQueryGroupResponse(final QueryGroup queryGroup, RestStatus restStatus) { + this.queryGroup = queryGroup; + this.restStatus = restStatus; + } + + /** + * Constructor for UpdateQueryGroupResponse + * @param in - A {@link StreamInput} object + */ + public UpdateQueryGroupResponse(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.toXContent(builder, params); + } + + /** + * 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/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index 64f510fa1db67..3b2169d64077a 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 @@ -28,6 +28,8 @@ import org.opensearch.plugin.wlm.rest.RestDeleteQueryGroupAction; import org.opensearch.plugin.wlm.rest.RestGetQueryGroupAction; import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; +import org.opensearch.common.inject.Module; +import org.opensearch.plugin.wlm.rest.RestUpdateQueryGroupAction; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; import org.opensearch.rest.RestController; @@ -52,7 +54,8 @@ public WorkloadManagementPlugin() {} return List.of( new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class), new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class), - new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class) + new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class), + new ActionPlugin.ActionHandler<>(UpdateQueryGroupAction.INSTANCE, TransportUpdateQueryGroupAction.class) ); } @@ -66,7 +69,7 @@ public List getRestHandlers( IndexNameExpressionResolver indexNameExpressionResolver, Supplier nodesInCluster ) { - return List.of(new RestCreateQueryGroupAction(), new RestGetQueryGroupAction(), new RestDeleteQueryGroupAction()); + return List.of(new RestCreateQueryGroupAction(), new RestGetQueryGroupAction(), new RestDeleteQueryGroupAction(), new RestUpdateQueryGroupAction()); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.java new file mode 100644 index 0000000000000..de6162ceda225 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.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.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.UpdateQueryGroupAction; +import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.UpdateQueryGroupResponse; +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 update a QueryGroup + * + * @opensearch.api + */ +public class RestUpdateQueryGroupAction extends BaseRestHandler { + + /** + * Constructor for RestUpdateQueryGroupAction + */ + public RestUpdateQueryGroupAction() {} + + @Override + public String getName() { + return "update_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/{name}"), new Route(PUT, "_query_group/{name}")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + String name = request.param("name"); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest(); + request.applyContentParser((parser) -> parseRestRequest(updateQueryGroupRequest, parser, name)); + return channel -> client.execute(UpdateQueryGroupAction.INSTANCE, updateQueryGroupRequest, updateQueryGroupResponse(channel)); + } + + private void parseRestRequest(UpdateQueryGroupRequest request, XContentParser parser, String name) throws IOException { + final UpdateQueryGroupRequest updateQueryGroupRequest = UpdateQueryGroupRequest.fromXContent(parser, name); + request.setName(updateQueryGroupRequest.getName()); + request.setResourceLimits(updateQueryGroupRequest.getResourceLimits()); + request.setResiliencyMode(updateQueryGroupRequest.getResiliencyMode()); + request.setUpdatedAtInMillis(updateQueryGroupRequest.getUpdatedAtInMillis()); + } + + private RestResponseListener updateQueryGroupResponse(final RestChannel channel) { + return new RestResponseListener<>(channel) { + @Override + public RestResponse buildResponse(final UpdateQueryGroupResponse 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/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index 7561a2f6f99c3..22a2b99efe0cc 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 @@ -18,6 +18,7 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.cluster.service.ClusterManagerTaskThrottler; +import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; import org.opensearch.cluster.service.ClusterManagerTaskThrottler.ThrottlingKey; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.Priority; @@ -29,13 +30,17 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; + import org.opensearch.wlm.ResourceType; +import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.UpdateQueryGroupResponse; import java.util.Collection; import java.util.EnumMap; import java.util.Map; import java.util.Optional; import java.util.stream.Collectors; +import java.util.HashMap; /** * This class defines the functions for QueryGroup persistence @@ -44,6 +49,7 @@ public class QueryGroupPersistenceService { static final String SOURCE = "query-group-persistence-service"; private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group"; private static final String DELETE_QUERY_GROUP_THROTTLING_KEY = "delete-query-group"; + private static final String UPDATE_QUERY_GROUP_THROTTLING_KEY = "update-query-group"; private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); /** * max QueryGroup count setting name @@ -72,6 +78,7 @@ public class QueryGroupPersistenceService { private volatile int maxQueryGroupCount; final ThrottlingKey createQueryGroupThrottlingKey; final ThrottlingKey deleteQueryGroupThrottlingKey; + final ThrottlingKey updateQueryGroupThrottlingKey; /** * Constructor for QueryGroupPersistenceService @@ -89,6 +96,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); setMaxQueryGroupCount(MAX_QUERY_GROUP_COUNT.get(settings)); clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); } @@ -169,39 +177,13 @@ ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final Clu } // check if there's any resource allocation that exceed limit of 1.0 - Map totalUsageMap = calculateTotalUsage(existingQueryGroups, queryGroup); - for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { - 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." - ); - } - } + validateTotalUsage(existingQueryGroups, groupName, queryGroup.getResourceLimits()); return ClusterState.builder(currentClusterState) .metadata(Metadata.builder(currentClusterState.metadata()).put(queryGroup).build()) .build(); } - /** - * 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 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(queryGroup.getName())) { - for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { - map.compute(resourceType, (k, v) -> v + currGroup.getResourceLimits().get(resourceType)); - } - } - } - return map; - } - /** * Get the QueryGroups with the specified name from cluster state * @param name - the QueryGroup name we are getting @@ -264,10 +246,120 @@ ClusterState deleteQueryGroupInClusterState(final String name, final ClusterStat return ClusterState.builder(currentClusterState).metadata(Metadata.builder(metadata).remove(queryGroupToRemove).build()).build(); } + /** + * Modify cluster state to update the QueryGroup + * @param toUpdateGroup {@link QueryGroup} - the QueryGroup that we want to update + * @param listener - ActionListener for UpdateQueryGroupResponse + */ + public void updateInClusterStateMetadata(UpdateQueryGroupRequest toUpdateGroup, ActionListener listener) { + clusterService.submitStateUpdateTask(SOURCE, new ClusterStateUpdateTask(Priority.NORMAL) { + @Override + public ClusterState execute(ClusterState currentState) { + return updateQueryGroupInClusterState(toUpdateGroup, currentState); + } + + @Override + public ThrottlingKey getClusterManagerThrottlingKey() { + return updateQueryGroupThrottlingKey; + } + + @Override + public void onFailure(String source, Exception e) { + logger.warn("Failed to update QueryGroup due to error: {}, for source: {}", e.getMessage(), source); + listener.onFailure(e); + } + + @Override + public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) { + String name = toUpdateGroup.getName(); + Optional findUpdatedGroup = newState.metadata() + .queryGroups() + .values() + .stream() + .filter(group -> group.getName().equals(name)) + .findFirst(); + assert findUpdatedGroup.isPresent(); + QueryGroup updatedGroup = findUpdatedGroup.get(); + UpdateQueryGroupResponse response = new UpdateQueryGroupResponse(updatedGroup, RestStatus.OK); + listener.onResponse(response); + } + }); + } + + /** + * Modify cluster state to update the existing QueryGroup + * @param updateQueryGroupRequest {@link QueryGroup} - the QueryGroup that we want to update + * @param currentState - current cluster state + */ + ClusterState updateQueryGroupInClusterState(UpdateQueryGroupRequest updateQueryGroupRequest, ClusterState currentState) { + final Metadata metadata = currentState.metadata(); + final Map existingGroups = currentState.metadata().queryGroups(); + String name = updateQueryGroupRequest.getName(); + + final QueryGroup existingGroup = existingGroups.values() + .stream() + .filter(group -> group.getName().equals(name)) + .findFirst() + .orElseThrow(() -> new ResourceNotFoundException("No QueryGroup exists with the provided name: " + name)); + + // build the QueryGroup with updated fields + final Map updatedResourceLimits = new HashMap<>(existingGroup.getResourceLimits()); + if (updateQueryGroupRequest.getResourceLimits() != null && !updateQueryGroupRequest.getResourceLimits().isEmpty()) { + validateTotalUsage(existingGroups, name, updateQueryGroupRequest.getResourceLimits()); + updatedResourceLimits.putAll(updateQueryGroupRequest.getResourceLimits()); + } + + final ResiliencyMode mode = Optional.ofNullable(updateQueryGroupRequest.getResiliencyMode()) + .orElse(existingGroup.getResiliencyMode()); + + final QueryGroup updatedGroup = new QueryGroup( + name, + existingGroup.get_id(), + mode, + updatedResourceLimits, + updateQueryGroupRequest.getUpdatedAtInMillis() + ); + return ClusterState.builder(currentState) + .metadata(Metadata.builder(metadata).remove(existingGroup).put(updatedGroup).build()) + .build(); + } + + /** + * This method checks if there's any resource allocation that exceed limit of 1.0 + * @param existingQueryGroups - existing QueryGroups in the system + * @param resourceLimits - the QueryGroup we're creating or updating + */ + private void validateTotalUsage(Map existingQueryGroups, String name, Map resourceLimits) { + final Map totalUsage = new EnumMap<>(ResourceType.class); + totalUsage.putAll(resourceLimits); + for (QueryGroup currGroup : existingQueryGroups.values()) { + if (!currGroup.getName().equals(name)) { + for (ResourceType resourceType : resourceLimits.keySet()) { + totalUsage.compute(resourceType, (k, v) -> v + currGroup.getResourceLimits().getOrDefault(resourceType, 0.0)); + } + } + } + totalUsage.forEach((resourceType, total) -> { + if (total > 1.0) { + 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." + ); + } + }); + } + /** * maxQueryGroupCount getter */ public int getMaxQueryGroupCount() { return maxQueryGroupCount; } + + /** + * clusterService getter + */ + public ClusterService getClusterService() { + return clusterService; + } } 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 5848e9c936623..d21a328b6a0bd 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,10 @@ */ /** +<<<<<<< HEAD * Package for the service classes of WorkloadManagementPlugin +======= + * Package for the service classes for QueryGroup CRUD operations +>>>>>>> 7abcbc98a73 (Add Update QueryGroup API Logic) */ package org.opensearch.plugin.wlm.service; 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 e165645775d5c..f5c10c5ccdd52 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 @@ -20,6 +20,7 @@ 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; @@ -131,6 +132,11 @@ public static Tuple preparePersisten return new Tuple(queryGroupPersistenceService, clusterState); } + public static void assertEqualResourceLimits(Map resourceLimitMapOne, Map resourceLimitMapTwo) { + assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet())); + assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); + } + public static void assertEqualQueryGroups(Collection collectionOne, Collection collectionTwo) { assertEquals(collectionOne.size(), collectionTwo.size()); List listOne = new ArrayList<>(collectionOne); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequestTests.java new file mode 100644 index 0000000000000..6feeb46b1165a --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequestTests.java @@ -0,0 +1,96 @@ +/* + * 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.cluster.metadata.QueryGroup; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.search.ResourceType; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +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_ONE; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.TIMESTAMP_ONE; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualResourceLimits; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; + +public class UpdateQueryGroupRequestTests extends OpenSearchTestCase { + + public void testSerialization() throws IOException { + UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(queryGroupOne); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + UpdateQueryGroupRequest otherRequest = new UpdateQueryGroupRequest(streamInput); + assertEquals(request.getName(), otherRequest.getName()); + assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); + assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); + assertEqualResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); + assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); + } + + public void testSerializationOnlyName() throws IOException { + UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, null, new HashMap<>(), TIMESTAMP_ONE); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + UpdateQueryGroupRequest otherRequest = new UpdateQueryGroupRequest(streamInput); + assertEquals(request.getName(), otherRequest.getName()); + assertEquals(request.getResourceLimits(), otherRequest.getResourceLimits()); + assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); + assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); + } + + public void testSerializationOnlyResourceLimit() throws IOException { + UpdateQueryGroupRequest request = new UpdateQueryGroupRequest( + NAME_ONE, + null, + Map.of(ResourceType.fromName(MEMORY_STRING), 0.4), + TIMESTAMP_ONE + ); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + UpdateQueryGroupRequest otherRequest = new UpdateQueryGroupRequest(streamInput); + assertEquals(request.getName(), otherRequest.getName()); + assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); + assertEqualResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); + assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); + assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); + } + + public void testInvalidResourceLimitList() { + assertThrows( + IllegalArgumentException.class, + () -> new UpdateQueryGroupRequest( + NAME_ONE, + QueryGroup.ResiliencyMode.fromName(MONITOR_STRING), + Map.of(ResourceType.fromName("memory"), 0.3, ResourceType.fromName(MONITOR_STRING), 0.4), + TIMESTAMP_ONE + ) + ); + } + + public void testInvalidEnforcement() { + assertThrows( + IllegalArgumentException.class, + () -> new UpdateQueryGroupRequest( + NAME_ONE, + QueryGroup.ResiliencyMode.fromName("random"), + Map.of(ResourceType.fromName("memory"), 0.3), + TIMESTAMP_ONE + ) + ); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponseTests.java new file mode 100644 index 0000000000000..c21e39160a858 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponseTests.java @@ -0,0 +1,60 @@ +/* + * 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.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.QueryGroupTestUtils.queryGroupOne; +import static org.mockito.Mockito.mock; + +public class UpdateQueryGroupResponseTests extends OpenSearchTestCase { + + public void testSerialization() throws IOException { + UpdateQueryGroupResponse response = new UpdateQueryGroupResponse(queryGroupOne, RestStatus.OK); + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + UpdateQueryGroupResponse otherResponse = new UpdateQueryGroupResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + QueryGroup responseGroup = response.getQueryGroup(); + QueryGroup otherResponseGroup = otherResponse.getQueryGroup(); + List list1 = new ArrayList<>(); + List list2 = new ArrayList<>(); + list1.add(responseGroup); + list2.add(otherResponseGroup); + QueryGroupTestUtils.assertEqualQueryGroups(list1, list2); + } + + public void testToXContentUpdateSingleQueryGroup() throws IOException { + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + UpdateQueryGroupResponse otherResponse = new UpdateQueryGroupResponse(queryGroupOne, RestStatus.OK); + String actual = otherResponse.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/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 5cb3d8fc6d11f..82b4158c5e384 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 @@ -22,11 +22,14 @@ import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.plugin.wlm.QueryGroupTestUtils; +import org.opensearch.plugin.wlm.UpdateQueryGroupResponse; import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.ResourceType; +import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; +import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; import java.util.ArrayList; import java.util.Collection; @@ -35,9 +38,9 @@ import java.util.Map; import java.util.Set; import java.util.stream.Collectors; +import java.util.Optional; 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; @@ -52,6 +55,7 @@ import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterState; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupPersistenceService; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupPersistenceService; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo; @@ -356,4 +360,177 @@ public void testDeleteInClusterStateMetadata() throws Exception { queryGroupPersistenceService.deleteInClusterStateMetadata(request, listener); verify(clusterService).submitStateUpdateTask(eq(SOURCE), any(AckedClusterStateUpdateTask.class)); } + + /** + * Tests updating a QueryGroup with all fields + */ + public void testUpdateQueryGroupAllFields() { + QueryGroup updated = builder().name(NAME_ONE) + ._id(_ID_ONE) + .mode("enforced") + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.15)) + .updatedAt(1690934400000L) + .build(); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + NAME_ONE, + ResiliencyMode.fromName("enforced"), + Map.of(ResourceType.fromName(MEMORY_STRING), 0.15), + 1690934400000L + ); + ClusterState newClusterState = queryGroupPersistenceService().updateQueryGroupInClusterState( + updateQueryGroupRequest, + clusterState() + ); + List updatedQueryGroups = new ArrayList<>(newClusterState.getMetadata().queryGroups().values()); + assertEquals(2, updatedQueryGroups.size()); + List expectedList = new ArrayList<>(); + expectedList.add(queryGroupTwo); + expectedList.add(updated); + assertEqualQueryGroups(expectedList, updatedQueryGroups); + } + + /** + * Tests updating a QueryGroup with only updated resourceLimits + */ + public void testUpdateQueryGroupResourceLimitsOnly() { + QueryGroup updated = builder().name(NAME_ONE) + ._id(_ID_ONE) + .mode(MONITOR_STRING) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.15)) + .updatedAt(1690934400000L) + .build(); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + NAME_ONE, + ResiliencyMode.fromName(MONITOR_STRING), + Map.of(ResourceType.fromName(MEMORY_STRING), 0.15), + 1690934400000L + ); + ClusterState newClusterState = queryGroupPersistenceService().updateQueryGroupInClusterState( + updateQueryGroupRequest, + clusterState() + ); + List updatedQueryGroups = new ArrayList<>(newClusterState.getMetadata().queryGroups().values()); + assertEquals(2, updatedQueryGroups.size()); + Optional findUpdatedGroupOne = newClusterState.metadata() + .queryGroups() + .values() + .stream() + .filter(group -> group.getName().equals(NAME_ONE)) + .findFirst(); + Optional findUpdatedGroupTwo = newClusterState.metadata() + .queryGroups() + .values() + .stream() + .filter(group -> group.getName().equals(NAME_TWO)) + .findFirst(); + assertTrue(findUpdatedGroupOne.isPresent()); + assertTrue(findUpdatedGroupTwo.isPresent()); + List list1 = new ArrayList<>(); + list1.add(updated); + List list2 = new ArrayList<>(); + list2.add(findUpdatedGroupOne.get()); + assertEqualQueryGroups(list1, list2); + } + + /** + * Tests updating a QueryGroup with invalid name + */ + public void testUpdateQueryGroupNonExistedName() { + QueryGroupPersistenceService queryGroupPersistenceService = queryGroupPersistenceService(); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + NAME_NONE_EXISTED, + ResiliencyMode.fromName(MONITOR_STRING), + Map.of(ResourceType.fromName(MEMORY_STRING), 0.15), + 1690934400000L + ); + assertThrows( + RuntimeException.class, + () -> queryGroupPersistenceService.updateQueryGroupInClusterState(updateQueryGroupRequest, clusterState()) + ); + List updatedQueryGroups = new ArrayList<>( + queryGroupPersistenceService.getClusterService().state().metadata().queryGroups().values() + ); + assertEquals(2, updatedQueryGroups.size()); + List expectedList = new ArrayList<>(); + expectedList.add(queryGroupTwo); + expectedList.add(queryGroupOne); + assertEqualQueryGroups(expectedList, updatedQueryGroups); + } + + /** + * Tests UpdateInClusterStateMetadata function + */ + public void testUpdateInClusterStateMetadata() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + queryGroupPersistenceService.updateInClusterStateMetadata(null, listener); + verify(clusterService).submitStateUpdateTask(eq(SOURCE), any()); + } + + /** + * Tests UpdateInClusterStateMetadata function with inner functions + */ + public void testUpdateInClusterStateMetadataInner() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + NAME_TWO, + ResiliencyMode.SOFT, + new HashMap<>(), + 2435465879685L + ); + ArgumentCaptor captor = ArgumentCaptor.forClass(ClusterStateUpdateTask.class); + queryGroupPersistenceService.updateInClusterStateMetadata(updateQueryGroupRequest, listener); + verify(clusterService, times(1)).submitStateUpdateTask(eq(SOURCE), captor.capture()); + ClusterStateUpdateTask capturedTask = captor.getValue(); + assertEquals(queryGroupPersistenceService.updateQueryGroupThrottlingKey, capturedTask.getClusterManagerThrottlingKey()); + + doAnswer(invocation -> { + ClusterStateUpdateTask task = invocation.getArgument(1); + task.clusterStateProcessed(SOURCE, clusterState(), clusterState()); + return null; + }).when(clusterService).submitStateUpdateTask(anyString(), any()); + queryGroupPersistenceService.updateInClusterStateMetadata(updateQueryGroupRequest, listener); + verify(listener).onResponse(any(UpdateQueryGroupResponse.class)); + } + + /** + * Tests UpdateInClusterStateMetadata function with failure + */ + public void testUpdateInClusterStateMetadataFailure() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + NAME_TWO, + ResiliencyMode.SOFT, + new HashMap<>(), + 2435465879685L + ); + 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.updateInClusterStateMetadata(updateQueryGroupRequest, listener); + verify(listener).onFailure(any(RuntimeException.class)); + } } 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 a971aa58940ba..e06e10eabdd88 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -29,7 +29,11 @@ * Class to define the QueryGroup schema * { * "_id": "fafjafjkaf9ag8a9ga9g7ag0aagaga", +<<<<<<< HEAD * "resource_limits": { +======= + * "resourceLimits": { +>>>>>>> 7abcbc98a73 (Add Update QueryGroup API Logic) * "memory": 0.4 * }, * "resiliency_mode": "enforced", 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 8e036209daff0..d3aa68566d978 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -160,6 +160,7 @@ import org.opensearch.search.backpressure.settings.SearchShardTaskSettings; import org.opensearch.search.backpressure.settings.SearchTaskSettings; 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/QueryGroupServiceSettings.java b/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java new file mode 100644 index 0000000000000..425916ce3e924 --- /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 nodeLevelMemoryCancellationThreshold; + private Double nodeLevelMemoryRejectionThreshold; + 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 - 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_CANCELLATION_THRESHOLD.get(settings); + nodeLevelMemoryRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); + maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + + clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelMemoryCancellationThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); + } + + /** + * 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 getNodeLevelMemoryCancellationThreshold() { + return nodeLevelMemoryCancellationThreshold; + } + + /** + * Method to set the 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 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(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + + this.nodeLevelMemoryCancellationThreshold = nodeLevelMemoryCancellationThreshold; + } + + /** + * Method to get the node level rejection threshold + * @return the current node level rejection threshold + */ + public Double getNodeLevelMemoryRejectionThreshold() { + return nodeLevelMemoryRejectionThreshold; + } + + /** + * Method to set the node level rejection threshold + * @param nodeLevelMemoryRejectionThreshold sets the new 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) { + throw new IllegalArgumentException( + NODE_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" + ); + } + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); + + this.nodeLevelMemoryRejectionThreshold = nodeLevelMemoryRejectionThreshold; + } + + private void ensureRejectionThresholdIsLessThanCancellation( + 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 + ); + } + } + + /** + * 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..00b68b0d3306c --- /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. + */ + +/** + * QueryGroup related artifacts + */ +package org.opensearch.search.query_group; From 6e320e9851e0ebdca89a8306322426757912500d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 16 Jul 2024 13:02:14 -0700 Subject: [PATCH 2/9] append to changlog Signed-off-by: Ruirui Zhang --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4df028f12d2fe..c2b97d6926ca4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add took time to request nodes stats ([#15054](https://github.com/opensearch-project/OpenSearch/pull/15054)) - [Workload Management] Add Get QueryGroup API Logic ([14709](https://github.com/opensearch-project/OpenSearch/pull/14709)) - [Workload Management] Add Settings for Workload Management feature ([#15028](https://github.com/opensearch-project/OpenSearch/pull/15028)) +- [Workload Management] Add Update QueryGroup API Logic ([#14775](https://github.com/opensearch-project/OpenSearch/pull/14775)) - [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) - Support filtering on a large list encoded by bitmap ([#14774](https://github.com/opensearch-project/OpenSearch/pull/14774)) - Add slice execution listeners to SearchOperationListener interface ([#15153](https://github.com/opensearch-project/OpenSearch/pull/15153)) From 21c10c6d0806cf3db6ce4bed23c1c39a9c3fefe6 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 16 Jul 2024 13:39:54 -0700 Subject: [PATCH 3/9] add javadoc Signed-off-by: Ruirui Zhang --- .../org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java index b1fd4b698d017..4ddf84a6e3c88 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java @@ -30,7 +30,8 @@ public class UpdateQueryGroupResponse extends ActionResponse implements ToXConte /** * Constructor for UpdateQueryGroupResponse - * @param queryGroup - The QueryGroup to be updated + * @param queryGroup - the QueryGroup to be updated + * @param restStatus - the rest status for the response */ public UpdateQueryGroupResponse(final QueryGroup queryGroup, RestStatus restStatus) { this.queryGroup = queryGroup; @@ -39,7 +40,7 @@ public UpdateQueryGroupResponse(final QueryGroup queryGroup, RestStatus restStat /** * Constructor for UpdateQueryGroupResponse - * @param in - A {@link StreamInput} object + * @param in - a {@link StreamInput} object */ public UpdateQueryGroupResponse(StreamInput in) throws IOException { queryGroup = new QueryGroup(in); From 2a32b173baffccc5a39630ff4ed8d1facac95f7d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Mon, 19 Aug 2024 15:41:18 -0700 Subject: [PATCH 4/9] rebase Signed-off-by: Ruirui Zhang --- .../plugin/wlm/WorkloadManagementPlugin.java | 5 +- .../TransportUpdateQueryGroupAction.java | 22 ++-- .../{ => action}/UpdateQueryGroupAction.java | 4 +- .../{ => action}/UpdateQueryGroupRequest.java | 107 +++--------------- .../UpdateQueryGroupResponse.java | 4 +- .../wlm/rest/RestGetQueryGroupAction.java | 2 +- .../wlm/rest/RestUpdateQueryGroupAction.java | 26 ++--- .../service/QueryGroupPersistenceService.java | 7 +- .../plugin/wlm/service/package-info.java | 4 - .../plugin/wlm/QueryGroupTestUtils.java | 5 +- .../UpdateQueryGroupRequestTests.java | 17 ++- .../UpdateQueryGroupResponseTests.java | 13 ++- .../QueryGroupPersistenceServiceTests.java | 11 +- .../cluster/metadata/QueryGroup.java | 20 +++- .../common/settings/ClusterSettings.java | 1 - 15 files changed, 100 insertions(+), 148 deletions(-) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/TransportUpdateQueryGroupAction.java (65%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/UpdateQueryGroupAction.java (92%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/UpdateQueryGroupRequest.java (55%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => action}/UpdateQueryGroupResponse.java (96%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => action}/UpdateQueryGroupRequestTests.java (89%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => action}/UpdateQueryGroupResponseTests.java (86%) 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 3b2169d64077a..2209e24488691 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 @@ -24,12 +24,13 @@ import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction; import org.opensearch.plugin.wlm.action.TransportDeleteQueryGroupAction; import org.opensearch.plugin.wlm.action.TransportGetQueryGroupAction; +import org.opensearch.plugin.wlm.action.TransportUpdateQueryGroupAction; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupAction; import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction; import org.opensearch.plugin.wlm.rest.RestDeleteQueryGroupAction; import org.opensearch.plugin.wlm.rest.RestGetQueryGroupAction; -import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; -import org.opensearch.common.inject.Module; import org.opensearch.plugin.wlm.rest.RestUpdateQueryGroupAction; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; 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/TransportUpdateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateQueryGroupAction.java similarity index 65% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportUpdateQueryGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateQueryGroupAction.java index 0e560718ced62..a6aa2da8fdc08 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/TransportUpdateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateQueryGroupAction.java @@ -6,27 +6,24 @@ * 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; -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; /** - * Transport action for update QueryGroup + * Transport action to update QueryGroup * - * @opensearch.internal + * @opensearch.experimental */ public class TransportUpdateQueryGroupAction extends HandledTransportAction { - private final ThreadPool threadPool; - private final Persistable queryGroupPersistenceService; + private final QueryGroupPersistenceService queryGroupPersistenceService; /** * Constructor for TransportUpdateQueryGroupAction @@ -34,24 +31,21 @@ public class TransportUpdateQueryGroupAction extends HandledTransportAction queryGroupPersistenceService + QueryGroupPersistenceService queryGroupPersistenceService ) { super(UpdateQueryGroupAction.NAME, transportService, actionFilters, UpdateQueryGroupRequest::new); - this.threadPool = threadPool; this.queryGroupPersistenceService = queryGroupPersistenceService; } @Override protected void doExecute(Task task, UpdateQueryGroupRequest request, ActionListener listener) { - threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> queryGroupPersistenceService.update(request, listener)); + queryGroupPersistenceService.updateInClusterStateMetadata(request, listener); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupAction.java similarity index 92% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupAction.java index 5b3600c091944..ff472f206131c 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupAction.java @@ -6,14 +6,14 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionType; /** * Transport action to update QueryGroup * - * @opensearch.api + * @opensearch.experimental */ public class UpdateQueryGroupAction extends ActionType { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java similarity index 55% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java index b33d501d1daf0..04e0f62720d9a 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.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; @@ -14,7 +14,6 @@ import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; 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.opensearch.search.ResourceType; import org.joda.time.Instant; @@ -26,18 +25,13 @@ /** * A request for update QueryGroup * - * @opensearch.internal + * @opensearch.experimental */ -public class UpdateQueryGroupRequest extends ActionRequest implements Writeable.Reader { - String name; - Map resourceLimits; - private ResiliencyMode resiliencyMode; - long updatedAtInMillis; - - /** - * Default constructor for UpdateQueryGroupRequest - */ - public UpdateQueryGroupRequest() {} +public class UpdateQueryGroupRequest extends ActionRequest { + private final String name; + private final Map resourceLimits; + private final ResiliencyMode resiliencyMode; + private final long updatedAtInMillis; /** * Constructor for UpdateQueryGroupRequest @@ -83,64 +77,29 @@ public UpdateQueryGroupRequest(StreamInput in) throws IOException { } if (in.readBoolean()) { resiliencyMode = ResiliencyMode.fromName(in.readString()); + } else { + resiliencyMode = null; } updatedAtInMillis = in.readLong(); } - @Override - public UpdateQueryGroupRequest read(StreamInput in) throws IOException { - return new UpdateQueryGroupRequest(in); - } - /** * Generate a UpdateQueryGroupRequest from XContent * @param parser - A {@link XContentParser} object * @param name - name of the QueryGroup to be updated */ public static UpdateQueryGroupRequest fromXContent(XContentParser parser, String name) 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 = ""; - 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("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(ResourceType.fromName(fieldName), parser.doubleValue()); - } - } - } - } - return new UpdateQueryGroupRequest(name, mode, resourceLimits, Instant.now().getMillis()); + QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser); + return new UpdateQueryGroupRequest(name, builder.getResiliencyMode(), builder.getResourceLimits(), Instant.now().getMillis()); } @Override public ActionRequestValidationException validate() { + QueryGroup.validateName(name); + if (resourceLimits != null) { + QueryGroup.validateResourceLimits(resourceLimits); + } + assert QueryGroup.isValid(updatedAtInMillis); return null; } @@ -151,14 +110,6 @@ public String getName() { return name; } - /** - * name setter - * @param name - name to be set - */ - public void setName(String name) { - this.name = name; - } - /** * ResourceLimits getter */ @@ -166,14 +117,6 @@ public Map getResourceLimits() { return resourceLimits; } - /** - * ResourceLimits setter - * @param resourceLimits - ResourceLimit to be set - */ - public void setResourceLimits(Map resourceLimits) { - this.resourceLimits = resourceLimits; - } - /** * resiliencyMode getter */ @@ -181,14 +124,6 @@ public ResiliencyMode getResiliencyMode() { return resiliencyMode; } - /** - * resiliencyMode setter - * @param resiliencyMode - mode to be set - */ - public void setResiliencyMode(ResiliencyMode resiliencyMode) { - this.resiliencyMode = resiliencyMode; - } - /** * updatedAtInMillis getter */ @@ -196,14 +131,6 @@ public long getUpdatedAtInMillis() { return updatedAtInMillis; } - /** - * updatedAtInMillis setter - * @param updatedAtInMillis - updatedAtInMillis to be set - */ - public void setUpdatedAtInMillis(long updatedAtInMillis) { - this.updatedAtInMillis = updatedAtInMillis; - } - @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); @@ -212,7 +139,7 @@ public void writeTo(StreamOutput out) throws IOException { out.writeBoolean(false); } else { out.writeBoolean(true); - out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue); + out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeDouble); } if (resiliencyMode == null) { out.writeBoolean(false); diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponse.java similarity index 96% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponse.java index 4ddf84a6e3c88..9071f52ecb5a7 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponse.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; @@ -22,7 +22,7 @@ /** * Response for the update API for QueryGroup * - * @opensearch.internal + * @opensearch.experimental */ public class UpdateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject { private final QueryGroup queryGroup; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java index c250bd2979e98..c87973e113138 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java @@ -27,7 +27,7 @@ import static org.opensearch.rest.RestRequest.Method.GET; /** - * Rest action to get a QueryGroup0 + * Rest action to get a QueryGroup * * @opensearch.experimental */ diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.java index de6162ceda225..55b4bc5a295c4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateQueryGroupAction.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.UpdateQueryGroupAction; -import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; -import org.opensearch.plugin.wlm.UpdateQueryGroupResponse; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupAction; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; @@ -31,7 +31,7 @@ /** * Rest action to update a QueryGroup * - * @opensearch.api + * @opensearch.experimental */ public class RestUpdateQueryGroupAction extends BaseRestHandler { @@ -50,23 +50,15 @@ public String getName() { */ @Override public List routes() { - return List.of(new Route(POST, "_query_group/{name}"), new Route(PUT, "_query_group/{name}")); + return List.of(new Route(POST, "_wlm/query_group/{name}"), new Route(PUT, "_wlm/query_group/{name}")); } @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - String name = request.param("name"); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest(); - request.applyContentParser((parser) -> parseRestRequest(updateQueryGroupRequest, parser, name)); - return channel -> client.execute(UpdateQueryGroupAction.INSTANCE, updateQueryGroupRequest, updateQueryGroupResponse(channel)); - } - - private void parseRestRequest(UpdateQueryGroupRequest request, XContentParser parser, String name) throws IOException { - final UpdateQueryGroupRequest updateQueryGroupRequest = UpdateQueryGroupRequest.fromXContent(parser, name); - request.setName(updateQueryGroupRequest.getName()); - request.setResourceLimits(updateQueryGroupRequest.getResourceLimits()); - request.setResiliencyMode(updateQueryGroupRequest.getResiliencyMode()); - request.setUpdatedAtInMillis(updateQueryGroupRequest.getUpdatedAtInMillis()); + try (XContentParser parser = request.contentParser()) { + UpdateQueryGroupRequest updateQueryGroupRequest = UpdateQueryGroupRequest.fromXContent(parser, request.param("name")); + return channel -> client.execute(UpdateQueryGroupAction.INSTANCE, updateQueryGroupRequest, updateQueryGroupResponse(channel)); + } } private RestResponseListener updateQueryGroupResponse(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 22a2b99efe0cc..75f9f3bebbce9 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 @@ -35,8 +35,13 @@ import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; import org.opensearch.plugin.wlm.UpdateQueryGroupResponse; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; +import org.opensearch.search.ResourceType; + import java.util.Collection; import java.util.EnumMap; +import java.util.HashMap; import java.util.Map; import java.util.Optional; import java.util.stream.Collectors; @@ -244,7 +249,7 @@ ClusterState deleteQueryGroupInClusterState(final String name, final ClusterStat .orElseThrow(() -> new ResourceNotFoundException("No QueryGroup exists with the provided name: " + name)); return ClusterState.builder(currentClusterState).metadata(Metadata.builder(metadata).remove(queryGroupToRemove).build()).build(); - } + } /** * Modify cluster state to update the QueryGroup 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 d21a328b6a0bd..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,10 +7,6 @@ */ /** -<<<<<<< HEAD * Package for the service classes of WorkloadManagementPlugin -======= - * Package for the service classes for QueryGroup CRUD operations ->>>>>>> 7abcbc98a73 (Add Update QueryGroup API Logic) */ package org.opensearch.plugin.wlm.service; 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 f5c10c5ccdd52..19cbc710422d2 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 @@ -132,7 +132,10 @@ public static Tuple preparePersisten return new Tuple(queryGroupPersistenceService, clusterState); } - public static void assertEqualResourceLimits(Map resourceLimitMapOne, Map resourceLimitMapTwo) { + public static void assertEqualResourceLimits( + 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/UpdateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java similarity index 89% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java index 6feeb46b1165a..d804fb7569ba2 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.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; @@ -27,6 +27,9 @@ public class UpdateQueryGroupRequestTests extends OpenSearchTestCase { + /** + * Test case to verify the serialization and deserialization of UpdateQueryGroupRequest. + */ public void testSerialization() throws IOException { UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(queryGroupOne); BytesStreamOutput out = new BytesStreamOutput(); @@ -40,6 +43,9 @@ public void testSerialization() throws IOException { assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); } + /** + * Test case to verify the serialization and deserialization of UpdateQueryGroupRequest with only name field. + */ public void testSerializationOnlyName() throws IOException { UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, null, new HashMap<>(), TIMESTAMP_ONE); BytesStreamOutput out = new BytesStreamOutput(); @@ -52,6 +58,9 @@ public void testSerializationOnlyName() throws IOException { assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); } + /** + * Test case to verify the serialization and deserialization of UpdateQueryGroupRequest with only resourceLimits field. + */ public void testSerializationOnlyResourceLimit() throws IOException { UpdateQueryGroupRequest request = new UpdateQueryGroupRequest( NAME_ONE, @@ -70,6 +79,9 @@ public void testSerializationOnlyResourceLimit() throws IOException { assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); } + /** + * Tests invalid ResourceType. + */ public void testInvalidResourceLimitList() { assertThrows( IllegalArgumentException.class, @@ -82,6 +94,9 @@ public void testInvalidResourceLimitList() { ); } + /** + * Tests invalid resiliencyMode. + */ public void testInvalidEnforcement() { assertThrows( IllegalArgumentException.class, diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java similarity index 86% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java index c21e39160a858..fe3d92763866c 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/UpdateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.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; @@ -26,6 +27,9 @@ public class UpdateQueryGroupResponseTests extends OpenSearchTestCase { + /** + * Test case to verify the serialization and deserialization of UpdateQueryGroupResponse. + */ public void testSerialization() throws IOException { UpdateQueryGroupResponse response = new UpdateQueryGroupResponse(queryGroupOne, RestStatus.OK); BytesStreamOutput out = new BytesStreamOutput(); @@ -42,6 +46,9 @@ public void testSerialization() throws IOException { QueryGroupTestUtils.assertEqualQueryGroups(list1, list2); } + /** + * Test case to verify the toXContent method of UpdateQueryGroupResponse. + */ public void testToXContentUpdateSingleQueryGroup() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); UpdateQueryGroupResponse otherResponse = new UpdateQueryGroupResponse(queryGroupOne, RestStatus.OK); @@ -50,8 +57,8 @@ public void testToXContentUpdateSingleQueryGroup() throws IOException { + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + " \"name\" : \"query_group_one\",\n" + " \"resiliency_mode\" : \"monitor\",\n" - + " \"updatedAt\" : 4513232413,\n" - + " \"resourceLimits\" : {\n" + + " \"updated_at\" : 4513232413,\n" + + " \"resource_limits\" : {\n" + " \"memory\" : 0.3\n" + " }\n" + "}"; 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 82b4158c5e384..ac5dc6cbff05f 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 @@ -16,13 +16,13 @@ import org.opensearch.cluster.ClusterStateUpdateTask; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; 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.UpdateQueryGroupResponse; import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; import org.opensearch.test.OpenSearchTestCase; @@ -30,17 +30,23 @@ import org.opensearch.wlm.ResourceType; import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; +import org.opensearch.search.ResourceType; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; -import java.util.Optional; 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; @@ -55,7 +61,6 @@ import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterState; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupPersistenceService; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupPersistenceService; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo; 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 e06e10eabdd88..9d9c1f50cd5cc 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -29,11 +29,7 @@ * Class to define the QueryGroup schema * { * "_id": "fafjafjkaf9ag8a9ga9g7ag0aagaga", -<<<<<<< HEAD * "resource_limits": { -======= - * "resourceLimits": { ->>>>>>> 7abcbc98a73 (Add Update QueryGroup API Logic) * "memory": 0.4 * }, * "resiliency_mode": "enforced", @@ -83,7 +79,7 @@ public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits) { + public static void validateResourceLimits(Map resourceLimits) { for (Map.Entry resource : resourceLimits.entrySet()) { Double threshold = resource.getValue(); Objects.requireNonNull(resource.getKey(), "resourceName can't be null"); @@ -327,5 +323,17 @@ public Builder resourceLimits(Map resourceLimits) { public QueryGroup build() { return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); } + + public ResiliencyMode getResiliencyMode() { + return resiliencyMode; + } + + public long getUpdatedAt() { + return updatedAt; + } + + public Map getResourceLimits() { + return resourceLimits; + } } } 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 d3aa68566d978..8e036209daff0 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -160,7 +160,6 @@ import org.opensearch.search.backpressure.settings.SearchShardTaskSettings; import org.opensearch.search.backpressure.settings.SearchTaskSettings; 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; From 844282241a1b0c5f69634ac47cf5704028266f5c Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 21 Aug 2024 13:19:54 -0700 Subject: [PATCH 5/9] address comments Signed-off-by: Ruirui Zhang --- .../wlm/action/UpdateQueryGroupRequest.java | 14 +- .../api/update_query_group_context.json | 23 ++ .../rest-api-spec/test/wlm/10_query_group.yml | 42 ++++ .../QueryGroupServiceSettings.java | 198 ------------------ .../search/query_group/package-info.java | 12 -- 5 files changed, 68 insertions(+), 221 deletions(-) create mode 100644 plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/update_query_group_context.json 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/action/UpdateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java index 04e0f62720d9a..d7850fa5cb5a2 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java @@ -75,11 +75,8 @@ public UpdateQueryGroupRequest(StreamInput in) throws IOException { } else { resourceLimits = new HashMap<>(); } - if (in.readBoolean()) { - resiliencyMode = ResiliencyMode.fromName(in.readString()); - } else { - resiliencyMode = null; - } + String updatedResiliencyMode = in.readOptionalString(); + resiliencyMode = updatedResiliencyMode == null ? null : ResiliencyMode.fromName(updatedResiliencyMode); updatedAtInMillis = in.readLong(); } @@ -141,12 +138,7 @@ public void writeTo(StreamOutput out) throws IOException { out.writeBoolean(true); out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeDouble); } - if (resiliencyMode == null) { - out.writeBoolean(false); - } else { - out.writeBoolean(true); - out.writeString(resiliencyMode.getName()); - } + out.writeOptionalString(resiliencyMode == null ? null : resiliencyMode.getName()); out.writeLong(updatedAtInMillis); } } diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/update_query_group_context.json b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/update_query_group_context.json new file mode 100644 index 0000000000000..fbfa2dde292ee --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/update_query_group_context.json @@ -0,0 +1,23 @@ +{ + "update_query_group_context": { + "stability": "experimental", + "url": { + "paths": [ + { + "path": "/_wlm/query_group/{name}", + "methods": ["PUT", "POST"], + "parts": { + "name": { + "type": "string", + "description": "QueryGroup name" + } + } + } + ] + }, + "params":{}, + "body":{ + "description":"The updated QueryGroup schema" + } + } +} diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_query_group.yml b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_query_group.yml index a00314986a5cf..40ec665351094 100644 --- a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_query_group.yml +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_query_group.yml @@ -29,6 +29,48 @@ - match: { query_groups.0.resource_limits.cpu: 0.4 } - match: { query_groups.0.resource_limits.memory: 0.2 } + - do: + update_query_group_context: + name: "analytics" + body: + { + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.42, + "memory": 0.22 + } + } + + - match: { name: "analytics" } + - match: { resiliency_mode: "monitor" } + - match: { resource_limits.cpu: 0.42 } + - match: { resource_limits.memory: 0.22 } + + - do: + catch: /resource_not_found_exception/ + update_query_group_context: + name: "analytics5" + body: + { + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.42, + "memory": 0.22 + } + } + + - do: + catch: /illegal_argument_exception/ + update_query_group_context: + name: "analytics" + body: + { + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 1.1 + } + } + - do: catch: /illegal_argument_exception/ create_query_group_context: 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 deleted file mode 100644 index 425916ce3e924..0000000000000 --- a/server/src/main/java/org/opensearch/search/query_group/QueryGroupServiceSettings.java +++ /dev/null @@ -1,198 +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.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 nodeLevelMemoryCancellationThreshold; - private Double nodeLevelMemoryRejectionThreshold; - 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 - 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_CANCELLATION_THRESHOLD.get(settings); - nodeLevelMemoryRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); - maxQueryGroupCount = MAX_QUERY_GROUP_COUNT.get(settings); - - ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); - - clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelMemoryCancellationThreshold); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); - } - - /** - * 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 getNodeLevelMemoryCancellationThreshold() { - return nodeLevelMemoryCancellationThreshold; - } - - /** - * Method to set the 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 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(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); - - this.nodeLevelMemoryCancellationThreshold = nodeLevelMemoryCancellationThreshold; - } - - /** - * Method to get the node level rejection threshold - * @return the current node level rejection threshold - */ - public Double getNodeLevelMemoryRejectionThreshold() { - return nodeLevelMemoryRejectionThreshold; - } - - /** - * Method to set the node level rejection threshold - * @param nodeLevelMemoryRejectionThreshold sets the new 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) { - throw new IllegalArgumentException( - NODE_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" - ); - } - - ensureRejectionThresholdIsLessThanCancellation(nodeLevelMemoryRejectionThreshold, nodeLevelMemoryCancellationThreshold); - - this.nodeLevelMemoryRejectionThreshold = nodeLevelMemoryRejectionThreshold; - } - - private void ensureRejectionThresholdIsLessThanCancellation( - 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 - ); - } - } - - /** - * 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 3a191bfbbf2102cbd6c3996c87729ba0bee06651 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 28 Aug 2024 14:47:38 -0700 Subject: [PATCH 6/9] address comments Signed-off-by: Ruirui Zhang --- .../plugin/wlm/WorkloadManagementPlugin.java | 7 +- .../wlm/action/UpdateQueryGroupRequest.java | 75 ++----- .../service/QueryGroupPersistenceService.java | 17 +- .../plugin/wlm/QueryGroupTestUtils.java | 29 ++- .../action/CreateQueryGroupRequestTests.java | 2 +- .../action/CreateQueryGroupResponseTests.java | 2 +- .../action/GetQueryGroupResponseTests.java | 4 +- .../action/UpdateQueryGroupRequestTests.java | 25 +-- .../action/UpdateQueryGroupResponseTests.java | 2 +- .../QueryGroupPersistenceServiceTests.java | 69 ++---- .../cluster/metadata/QueryGroup.java | 159 ++++---------- .../opensearch/wlm/ChangeableQueryGroup.java | 202 ++++++++++++++++++ .../metadata/QueryGroupMetadataTests.java | 4 +- .../cluster/metadata/QueryGroupTests.java | 53 ++--- .../wlm/ChangeableQueryGroupTests.java | 65 ++++++ 15 files changed, 417 insertions(+), 298 deletions(-) create mode 100644 server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java create mode 100644 server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.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 2209e24488691..c86490552f2f2 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 @@ -70,7 +70,12 @@ public List getRestHandlers( IndexNameExpressionResolver indexNameExpressionResolver, Supplier nodesInCluster ) { - return List.of(new RestCreateQueryGroupAction(), new RestGetQueryGroupAction(), new RestDeleteQueryGroupAction(), new RestUpdateQueryGroupAction()); + return List.of( + new RestCreateQueryGroupAction(), + new RestGetQueryGroupAction(), + new RestDeleteQueryGroupAction(), + new RestUpdateQueryGroupAction() + ); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java index d7850fa5cb5a2..de13436746b16 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java @@ -11,15 +11,14 @@ 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.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.search.ResourceType; -import org.joda.time.Instant; +import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; import java.io.IOException; -import java.util.HashMap; import java.util.Map; /** @@ -29,38 +28,16 @@ */ public class UpdateQueryGroupRequest extends ActionRequest { private final String name; - private final Map resourceLimits; - private final ResiliencyMode resiliencyMode; - private final long updatedAtInMillis; - - /** - * Constructor for UpdateQueryGroupRequest - * @param queryGroup - A {@link QueryGroup} object - */ - public UpdateQueryGroupRequest(QueryGroup queryGroup) { - this.name = queryGroup.getName(); - this.resiliencyMode = queryGroup.getResiliencyMode(); - this.resourceLimits = queryGroup.getResourceLimits(); - this.updatedAtInMillis = queryGroup.getUpdatedAtInMillis(); - } + private final ChangeableQueryGroup changeableQueryGroup; /** * Constructor for UpdateQueryGroupRequest * @param name - QueryGroup name for UpdateQueryGroupRequest - * @param resiliencyMode - QueryGroup mode for UpdateQueryGroupRequest - * @param resourceLimits - QueryGroup resourceLimits for UpdateQueryGroupRequest - * @param updatedAtInMillis - QueryGroup updated time in millis for UpdateQueryGroupRequest + * @param changeableQueryGroup - ChangeableQueryGroup for UpdateQueryGroupRequest */ - public UpdateQueryGroupRequest( - String name, - ResiliencyMode resiliencyMode, - Map resourceLimits, - long updatedAtInMillis - ) { + public UpdateQueryGroupRequest(String name, ChangeableQueryGroup changeableQueryGroup) { this.name = name; - this.resiliencyMode = resiliencyMode; - this.resourceLimits = resourceLimits; - this.updatedAtInMillis = updatedAtInMillis; + this.changeableQueryGroup = changeableQueryGroup; } /** @@ -68,16 +45,7 @@ public UpdateQueryGroupRequest( * @param in - A {@link StreamInput} object */ public UpdateQueryGroupRequest(StreamInput in) throws IOException { - super(in); - name = in.readString(); - if (in.readBoolean()) { - resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble); - } else { - resourceLimits = new HashMap<>(); - } - String updatedResiliencyMode = in.readOptionalString(); - resiliencyMode = updatedResiliencyMode == null ? null : ResiliencyMode.fromName(updatedResiliencyMode); - updatedAtInMillis = in.readLong(); + this(in.readString(), new ChangeableQueryGroup(in)); } /** @@ -87,16 +55,12 @@ public UpdateQueryGroupRequest(StreamInput in) throws IOException { */ public static UpdateQueryGroupRequest fromXContent(XContentParser parser, String name) throws IOException { QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser); - return new UpdateQueryGroupRequest(name, builder.getResiliencyMode(), builder.getResourceLimits(), Instant.now().getMillis()); + return new UpdateQueryGroupRequest(name, builder.getChangeableQueryGroup()); } @Override public ActionRequestValidationException validate() { QueryGroup.validateName(name); - if (resourceLimits != null) { - QueryGroup.validateResourceLimits(resourceLimits); - } - assert QueryGroup.isValid(updatedAtInMillis); return null; } @@ -108,37 +72,30 @@ public String getName() { } /** - * ResourceLimits getter + * resourceLimits getter */ public Map getResourceLimits() { - return resourceLimits; + return getChangeableQueryGroup().getResourceLimits(); } /** * resiliencyMode getter */ public ResiliencyMode getResiliencyMode() { - return resiliencyMode; + return getChangeableQueryGroup().getResiliencyMode(); } /** - * updatedAtInMillis getter + * changeableQueryGroup getter */ - public long getUpdatedAtInMillis() { - return updatedAtInMillis; + public ChangeableQueryGroup getChangeableQueryGroup() { + return changeableQueryGroup; } @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeString(name); - if (resourceLimits == null || resourceLimits.isEmpty()) { - out.writeBoolean(false); - } else { - out.writeBoolean(true); - out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeDouble); - } - out.writeOptionalString(resiliencyMode == null ? null : resiliencyMode.getName()); - out.writeLong(updatedAtInMillis); + changeableQueryGroup.writeTo(out); } } 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 75f9f3bebbce9..a4fb666c324c8 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 @@ -18,7 +18,6 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.cluster.service.ClusterManagerTaskThrottler; -import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; import org.opensearch.cluster.service.ClusterManagerTaskThrottler.ThrottlingKey; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.Priority; @@ -32,12 +31,12 @@ import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; import org.opensearch.wlm.ResourceType; -import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; -import org.opensearch.plugin.wlm.UpdateQueryGroupResponse; - import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; -import org.opensearch.search.ResourceType; +import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.joda.time.Instant; import java.util.Collection; import java.util.EnumMap; @@ -45,7 +44,6 @@ import java.util.Map; import java.util.Optional; import java.util.stream.Collectors; -import java.util.HashMap; /** * This class defines the functions for QueryGroup persistence @@ -249,7 +247,7 @@ ClusterState deleteQueryGroupInClusterState(final String name, final ClusterStat .orElseThrow(() -> new ResourceNotFoundException("No QueryGroup exists with the provided name: " + name)); return ClusterState.builder(currentClusterState).metadata(Metadata.builder(metadata).remove(queryGroupToRemove).build()).build(); - } + } /** * Modify cluster state to update the QueryGroup @@ -320,9 +318,8 @@ ClusterState updateQueryGroupInClusterState(UpdateQueryGroupRequest updateQueryG final QueryGroup updatedGroup = new QueryGroup( name, existingGroup.get_id(), - mode, - updatedResourceLimits, - updateQueryGroupRequest.getUpdatedAtInMillis() + new ChangeableQueryGroup(mode, updatedResourceLimits), + Instant.now().getMillis() ); return ClusterState.builder(currentState) .metadata(Metadata.builder(metadata).remove(existingGroup).put(updatedGroup).build()) 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 19cbc710422d2..47a1fd9a94a22 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 @@ -20,8 +20,9 @@ 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.wlm.ResourceType; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.ChangeableQueryGroup; import java.util.ArrayList; import java.util.Collection; @@ -32,7 +33,6 @@ import java.util.Set; import static org.opensearch.cluster.metadata.QueryGroup.builder; -import static org.opensearch.wlm.ResourceType.fromName; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -44,21 +44,17 @@ public class QueryGroupTestUtils { 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)) + .changeableQueryGroup(new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 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)) + .changeableQueryGroup(new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.6))) .updatedAt(TIMESTAMP_TWO) .build(); @@ -140,14 +136,27 @@ public static void assertEqualResourceLimits( assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values())); } - public static void assertEqualQueryGroups(Collection collectionOne, Collection collectionTwo) { + public static void assertEqualQueryGroups( + Collection collectionOne, + Collection collectionTwo, + boolean assertUpdateAt + ) { assertEquals(collectionOne.size(), collectionTwo.size()); List listOne = new ArrayList<>(collectionOne); List listTwo = new ArrayList<>(collectionTwo); 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))); + if (assertUpdateAt) { + QueryGroup one = listOne.get(i); + QueryGroup two = listTwo.get(i); + assertEquals(one.getName(), two.getName()); + assertEquals(one.getResourceLimits(), two.getResourceLimits()); + assertEquals(one.getResiliencyMode(), two.getResiliencyMode()); + assertEquals(one.get_id(), two.get_id()); + } else { + assertEquals(listOne.get(i), listTwo.get(i)); + } } } } 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 b0fa96a46df80..dd9de4bf8fb1a 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 @@ -35,6 +35,6 @@ public void testSerialization() throws IOException { List list2 = new ArrayList<>(); list1.add(queryGroupOne); list2.add(otherRequest.getQueryGroup()); - assertEqualQueryGroups(list1, list2); + assertEqualQueryGroups(list1, list2, false); } } 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 ecb9a6b2dc0d2..713c62088b9ca 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.assertEqualQueryGroups(listOne, listTwo); + QueryGroupTestUtils.assertEqualQueryGroups(listOne, listTwo, false); } /** diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java index 774f4b2d8db52..d9192dbb18fe4 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java @@ -41,7 +41,7 @@ public void testSerializationSingleQueryGroup() throws IOException { GetQueryGroupResponse otherResponse = new GetQueryGroupResponse(streamInput); assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); - QueryGroupTestUtils.assertEqualQueryGroups(response.getQueryGroups(), otherResponse.getQueryGroups()); + QueryGroupTestUtils.assertEqualQueryGroups(response.getQueryGroups(), otherResponse.getQueryGroups(), false); } /** @@ -58,7 +58,7 @@ public void testSerializationMultipleQueryGroup() throws IOException { GetQueryGroupResponse otherResponse = new GetQueryGroupResponse(streamInput); assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEquals(2, otherResponse.getQueryGroups().size()); - QueryGroupTestUtils.assertEqualQueryGroups(response.getQueryGroups(), otherResponse.getQueryGroups()); + QueryGroupTestUtils.assertEqualQueryGroups(response.getQueryGroups(), otherResponse.getQueryGroups(), false); } /** diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java index d804fb7569ba2..147759bceebfa 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java @@ -8,20 +8,18 @@ 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.search.ResourceType; import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; import java.io.IOException; import java.util.HashMap; import java.util.Map; -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_ONE; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.TIMESTAMP_ONE; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualResourceLimits; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; @@ -31,7 +29,7 @@ public class UpdateQueryGroupRequestTests extends OpenSearchTestCase { * Test case to verify the serialization and deserialization of UpdateQueryGroupRequest. */ public void testSerialization() throws IOException { - UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(queryGroupOne); + UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, queryGroupOne.getChangeableQueryGroup()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); @@ -40,14 +38,13 @@ public void testSerialization() throws IOException { assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); assertEqualResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); - assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); } /** * Test case to verify the serialization and deserialization of UpdateQueryGroupRequest with only name field. */ public void testSerializationOnlyName() throws IOException { - UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, null, new HashMap<>(), TIMESTAMP_ONE); + UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, new ChangeableQueryGroup(null, new HashMap<>())); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); @@ -55,7 +52,6 @@ public void testSerializationOnlyName() throws IOException { assertEquals(request.getName(), otherRequest.getName()); assertEquals(request.getResourceLimits(), otherRequest.getResourceLimits()); assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); - assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); } /** @@ -64,9 +60,7 @@ public void testSerializationOnlyName() throws IOException { public void testSerializationOnlyResourceLimit() throws IOException { UpdateQueryGroupRequest request = new UpdateQueryGroupRequest( NAME_ONE, - null, - Map.of(ResourceType.fromName(MEMORY_STRING), 0.4), - TIMESTAMP_ONE + new ChangeableQueryGroup(null, Map.of(ResourceType.MEMORY, 0.4)) ); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); @@ -76,7 +70,6 @@ public void testSerializationOnlyResourceLimit() throws IOException { assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); assertEqualResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); - assertEquals(request.getUpdatedAtInMillis(), otherRequest.getUpdatedAtInMillis()); } /** @@ -87,9 +80,7 @@ public void testInvalidResourceLimitList() { IllegalArgumentException.class, () -> new UpdateQueryGroupRequest( NAME_ONE, - QueryGroup.ResiliencyMode.fromName(MONITOR_STRING), - Map.of(ResourceType.fromName("memory"), 0.3, ResourceType.fromName(MONITOR_STRING), 0.4), - TIMESTAMP_ONE + new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.3, ResourceType.fromName("random"), 0.4)) ) ); } @@ -102,9 +93,7 @@ public void testInvalidEnforcement() { IllegalArgumentException.class, () -> new UpdateQueryGroupRequest( NAME_ONE, - QueryGroup.ResiliencyMode.fromName("random"), - Map.of(ResourceType.fromName("memory"), 0.3), - TIMESTAMP_ONE + new ChangeableQueryGroup(ResiliencyMode.fromName("random"), Map.of(ResourceType.fromName("memory"), 0.3)) ) ); } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java index fe3d92763866c..650e0bdc061ce 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java @@ -43,7 +43,7 @@ public void testSerialization() throws IOException { List list2 = new ArrayList<>(); list1.add(responseGroup); list2.add(otherResponseGroup); - QueryGroupTestUtils.assertEqualQueryGroups(list1, list2); + QueryGroupTestUtils.assertEqualQueryGroups(list1, list2, false); } /** 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 ac5dc6cbff05f..b4426335c3169 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 @@ -16,7 +16,6 @@ import org.opensearch.cluster.ClusterStateUpdateTask; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.metadata.QueryGroup; -import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.collect.Tuple; import org.opensearch.common.settings.ClusterSettings; @@ -25,16 +24,13 @@ import org.opensearch.plugin.wlm.QueryGroupTestUtils; import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; -import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.threadpool.ThreadPool; -import org.opensearch.wlm.ResourceType; -import org.opensearch.cluster.metadata.QueryGroup.ResiliencyMode; -import org.opensearch.plugin.wlm.UpdateQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; -import org.opensearch.search.ResourceType; +import org.opensearch.wlm.ResourceType; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; import java.util.ArrayList; import java.util.Collection; @@ -48,8 +44,6 @@ 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; 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; @@ -92,7 +86,7 @@ public void testCreateQueryGroup() { List listTwo = new ArrayList<>(); listOne.add(queryGroupOne); listTwo.add(updatedGroupsMap.get(_ID_ONE)); - assertEqualQueryGroups(listOne, listTwo); + assertEqualQueryGroups(listOne, listTwo, false); } /** @@ -108,7 +102,7 @@ public void testCreateAnotherQueryGroup() { assertEquals(2, updatedGroups.size()); assertTrue(updatedGroups.containsKey(_ID_TWO)); Collection values = updatedGroups.values(); - assertEqualQueryGroups(queryGroupList(), new ArrayList<>(values)); + assertEqualQueryGroups(queryGroupList(), new ArrayList<>(values), false); } /** @@ -120,8 +114,7 @@ public void testCreateQueryGroupDuplicateName() { ClusterState clusterState = setup.v2(); QueryGroup toCreate = builder().name(NAME_ONE) ._id("W5iIqHyhgi4K1qIAAAAIHw==") - .mode(MONITOR_STRING) - .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.3)) + .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.3))) .updatedAt(1690934400000L) .build(); assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); @@ -135,8 +128,7 @@ public void testCreateQueryGroupOverflowAllocation() { 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)) + .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.41))) .updatedAt(1690934400000L) .build(); @@ -152,8 +144,7 @@ public void testCreateQueryGroupOverflowAllocation() { public void testCreateQueryGroupOverflowCount() { QueryGroup toCreate = builder().name(NAME_NONE_EXISTED) ._id("W5iIqHyhgi4K1qIAAAAIHw==") - .mode(MONITOR_STRING) - .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) + .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.5))) .updatedAt(1690934400000L) .build(); Metadata metadata = Metadata.builder().queryGroups(Map.of(_ID_ONE, queryGroupOne, _ID_TWO, queryGroupTwo)).build(); @@ -276,7 +267,7 @@ public void testGetSingleQueryGroup() { List listTwo = new ArrayList<>(); listOne.add(QueryGroupTestUtils.queryGroupOne); listTwo.add(queryGroup); - QueryGroupTestUtils.assertEqualQueryGroups(listOne, listTwo); + QueryGroupTestUtils.assertEqualQueryGroups(listOne, listTwo, false); } /** @@ -290,7 +281,7 @@ public void testGetAllQueryGroups() { Set currentNAME = res.stream().map(QueryGroup::getName).collect(Collectors.toSet()); assertTrue(currentNAME.contains(QueryGroupTestUtils.NAME_ONE)); assertTrue(currentNAME.contains(QueryGroupTestUtils.NAME_TWO)); - QueryGroupTestUtils.assertEqualQueryGroups(QueryGroupTestUtils.queryGroupList(), res); + QueryGroupTestUtils.assertEqualQueryGroups(QueryGroupTestUtils.queryGroupList(), res, false); } /** @@ -325,7 +316,7 @@ public void testDeleteSingleQueryGroup() { assertEquals(1, afterDeletionGroups.size()); List oldQueryGroups = new ArrayList<>(); oldQueryGroups.add(queryGroupOne); - assertEqualQueryGroups(new ArrayList<>(afterDeletionGroups.values()), oldQueryGroups); + assertEqualQueryGroups(new ArrayList<>(afterDeletionGroups.values()), oldQueryGroups, false); } /** @@ -372,16 +363,10 @@ public void testDeleteInClusterStateMetadata() throws Exception { public void testUpdateQueryGroupAllFields() { QueryGroup updated = builder().name(NAME_ONE) ._id(_ID_ONE) - .mode("enforced") - .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.15)) + .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.15))) .updatedAt(1690934400000L) .build(); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( - NAME_ONE, - ResiliencyMode.fromName("enforced"), - Map.of(ResourceType.fromName(MEMORY_STRING), 0.15), - 1690934400000L - ); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest(NAME_ONE, updated.getChangeableQueryGroup()); ClusterState newClusterState = queryGroupPersistenceService().updateQueryGroupInClusterState( updateQueryGroupRequest, clusterState() @@ -391,7 +376,7 @@ public void testUpdateQueryGroupAllFields() { List expectedList = new ArrayList<>(); expectedList.add(queryGroupTwo); expectedList.add(updated); - assertEqualQueryGroups(expectedList, updatedQueryGroups); + assertEqualQueryGroups(expectedList, updatedQueryGroups, true); } /** @@ -400,16 +385,10 @@ public void testUpdateQueryGroupAllFields() { public void testUpdateQueryGroupResourceLimitsOnly() { QueryGroup updated = builder().name(NAME_ONE) ._id(_ID_ONE) - .mode(MONITOR_STRING) - .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.15)) + .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.15))) .updatedAt(1690934400000L) .build(); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( - NAME_ONE, - ResiliencyMode.fromName(MONITOR_STRING), - Map.of(ResourceType.fromName(MEMORY_STRING), 0.15), - 1690934400000L - ); + UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest(NAME_ONE, updated.getChangeableQueryGroup()); ClusterState newClusterState = queryGroupPersistenceService().updateQueryGroupInClusterState( updateQueryGroupRequest, clusterState() @@ -434,7 +413,7 @@ public void testUpdateQueryGroupResourceLimitsOnly() { list1.add(updated); List list2 = new ArrayList<>(); list2.add(findUpdatedGroupOne.get()); - assertEqualQueryGroups(list1, list2); + assertEqualQueryGroups(list1, list2, true); } /** @@ -444,9 +423,7 @@ public void testUpdateQueryGroupNonExistedName() { QueryGroupPersistenceService queryGroupPersistenceService = queryGroupPersistenceService(); UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( NAME_NONE_EXISTED, - ResiliencyMode.fromName(MONITOR_STRING), - Map.of(ResourceType.fromName(MEMORY_STRING), 0.15), - 1690934400000L + new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.15)) ); assertThrows( RuntimeException.class, @@ -459,7 +436,7 @@ public void testUpdateQueryGroupNonExistedName() { List expectedList = new ArrayList<>(); expectedList.add(queryGroupTwo); expectedList.add(queryGroupOne); - assertEqualQueryGroups(expectedList, updatedQueryGroups); + assertEqualQueryGroups(expectedList, updatedQueryGroups, true); } /** @@ -492,9 +469,7 @@ public void testUpdateInClusterStateMetadataInner() { ); UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( NAME_TWO, - ResiliencyMode.SOFT, - new HashMap<>(), - 2435465879685L + new ChangeableQueryGroup(ResiliencyMode.SOFT, new HashMap<>()) ); ArgumentCaptor captor = ArgumentCaptor.forClass(ClusterStateUpdateTask.class); queryGroupPersistenceService.updateInClusterStateMetadata(updateQueryGroupRequest, listener); @@ -525,9 +500,7 @@ public void testUpdateInClusterStateMetadataFailure() { ); UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( NAME_TWO, - ResiliencyMode.SOFT, - new HashMap<>(), - 2435465879685L + new ChangeableQueryGroup(ResiliencyMode.SOFT, new HashMap<>()) ); doAnswer(invocation -> { ClusterStateUpdateTask task = invocation.getArgument(1); 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 9d9c1f50cd5cc..11e8fa695d23d 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -18,10 +18,11 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; import org.joda.time.Instant; import java.io.IOException; -import java.util.HashMap; import java.util.Map; import java.util.Objects; @@ -30,7 +31,8 @@ * { * "_id": "fafjafjkaf9ag8a9ga9g7ag0aagaga", * "resource_limits": { - * "memory": 0.4 + * "memory": 0.4, + * "cpu": 0.2 * }, * "resiliency_mode": "enforced", * "name": "analytics", @@ -42,40 +44,35 @@ public class QueryGroup extends AbstractDiffable implements ToXConte 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 ChangeableQueryGroup changeableQueryGroup; - public QueryGroup(String name, ResiliencyMode resiliencyMode, Map resourceLimits) { - this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis()); + public QueryGroup(String name, ChangeableQueryGroup changeableQueryGroup) { + this(name, UUIDs.randomBase64UUID(), changeableQueryGroup, Instant.now().getMillis()); } - public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits, long updatedAt) { + public QueryGroup(String name, String _id, ChangeableQueryGroup changeableQueryGroup, 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"); + Objects.requireNonNull(changeableQueryGroup.getResourceLimits(), "QueryGroup.resourceLimits can't be null"); + Objects.requireNonNull(changeableQueryGroup.getResiliencyMode(), "QueryGroup.resiliencyMode can't be null"); Objects.requireNonNull(_id, "QueryGroup._id can't be null"); validateName(name); - if (resourceLimits.isEmpty()) { + if (changeableQueryGroup.getResourceLimits().isEmpty()) { throw new IllegalArgumentException("QueryGroup.resourceLimits should at least have 1 resource limit"); } - validateResourceLimits(resourceLimits); if (!isValid(updatedAt)) { throw new IllegalArgumentException("QueryGroup.updatedAtInMillis is not a valid epoch"); } this.name = name; this._id = _id; - this.resiliencyMode = resiliencyMode; - this.resourceLimits = resourceLimits; + this.changeableQueryGroup = changeableQueryGroup; this.updatedAtInMillis = updatedAt; } @@ -90,21 +87,14 @@ public static boolean isValid(long updatedAt) { } public QueryGroup(StreamInput in) throws IOException { - this( - in.readString(), - in.readString(), - ResiliencyMode.fromName(in.readString()), - in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble), - in.readLong() - ); + this(in.readString(), in.readString(), new ChangeableQueryGroup(in), 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::writeDouble); + changeableQueryGroup.writeTo(out); out.writeLong(updatedAtInMillis); } @@ -114,34 +104,15 @@ public static void validateName(String name) { } } - public static 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"); - - 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 { builder.startObject(); 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(RESOURCE_LIMITS_STRING); - for (ResourceType resourceType : ResourceType.values()) { - if (resourceLimits.containsKey(resourceType)) { - builder.field(resourceType.getName(), resourceLimits.get(resourceType)); - } + for (String fieldName : ChangeableQueryGroup.acceptedFieldNames) { + changeableQueryGroup.writeField(builder, fieldName); } - builder.endObject(); - + builder.field(UPDATED_AT_STRING, updatedAtInMillis); builder.endObject(); return builder; } @@ -160,27 +131,30 @@ 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(changeableQueryGroup, that.changeableQueryGroup) && Objects.equals(_id, that._id) && updatedAtInMillis == that.updatedAtInMillis; } @Override public int hashCode() { - return Objects.hash(name, resourceLimits, updatedAtInMillis, _id); + return Objects.hash(name, changeableQueryGroup, updatedAtInMillis, _id); } public String getName() { return name; } + public ChangeableQueryGroup getChangeableQueryGroup() { + return changeableQueryGroup; + } + public ResiliencyMode getResiliencyMode() { - return resiliencyMode; + return getChangeableQueryGroup().getResiliencyMode(); } public Map getResourceLimits() { - return resourceLimits; + return getChangeableQueryGroup().getResourceLimits(); } public String get_id() { @@ -199,37 +173,6 @@ public static Builder builder() { return new Builder(); } - /** - * This enum models the different QueryGroup resiliency modes - * SOFT - means that this query group can consume more than query group resource limits if node is not in duress - * ENFORCED - means that it will never breach the assigned limits and will cancel as soon as the limits are breached - * MONITOR - it will not cause any cancellation but just log the eligible task cancellations - */ - @ExperimentalApi - public enum ResiliencyMode { - SOFT("soft"), - ENFORCED("enforced"), - MONITOR("monitor"); - - private final String name; - - ResiliencyMode(String mode) { - this.name = mode; - } - - public String getName() { - return name; - } - - public static ResiliencyMode fromName(String s) { - for (ResiliencyMode mode : values()) { - if (mode.getName().equalsIgnoreCase(s)) return mode; - - } - throw new IllegalArgumentException("Invalid value for QueryGroupMode: " + s); - } - } - /** * Builder class for {@link QueryGroup} */ @@ -237,9 +180,8 @@ public static ResiliencyMode fromName(String s) { public static class Builder { private String name; private String _id; - private ResiliencyMode resiliencyMode; + private ChangeableQueryGroup changeableQueryGroup; private long updatedAt; - private Map resourceLimits; private Builder() {} @@ -257,8 +199,7 @@ public static Builder fromXContent(XContentParser parser) throws IOException { } String fieldName = ""; - // Map to hold resources - final Map resourceLimits = new HashMap<>(); + ChangeableQueryGroup changeableQueryGroup1 = new ChangeableQueryGroup(); while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { if (token == XContentParser.Token.FIELD_NAME) { fieldName = parser.currentName(); @@ -267,32 +208,21 @@ public static Builder fromXContent(XContentParser parser) throws IOException { 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 (ChangeableQueryGroup.shouldParse(fieldName)) { + changeableQueryGroup1.parseField(parser, fieldName); } 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()); - } + if (!ChangeableQueryGroup.shouldParse(fieldName)) { + throw new IllegalArgumentException(fieldName + " is not a valid object in QueryGroup"); } - + changeableQueryGroup1.parseField(parser, fieldName); } } - return builder.resourceLimits(resourceLimits); + return builder.changeableQueryGroup(changeableQueryGroup1); } public Builder name(String name) { @@ -305,8 +235,8 @@ public Builder _id(String _id) { return this; } - public Builder mode(String mode) { - this.resiliencyMode = ResiliencyMode.fromName(mode); + public Builder changeableQueryGroup(ChangeableQueryGroup changeableQueryGroup) { + this.changeableQueryGroup = changeableQueryGroup; return this; } @@ -315,25 +245,12 @@ public Builder updatedAt(long updatedAt) { return this; } - public Builder resourceLimits(Map resourceLimits) { - this.resourceLimits = resourceLimits; - return this; - } - public QueryGroup build() { - return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); - } - - public ResiliencyMode getResiliencyMode() { - return resiliencyMode; - } - - public long getUpdatedAt() { - return updatedAt; + return new QueryGroup(name, _id, changeableQueryGroup, updatedAt); } - public Map getResourceLimits() { - return resourceLimits; + public ChangeableQueryGroup getChangeableQueryGroup() { + return changeableQueryGroup; } } } diff --git a/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java b/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java new file mode 100644 index 0000000000000..fad653d4fcbf9 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java @@ -0,0 +1,202 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.cluster.AbstractDiffable; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.search.ResourceType; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.function.Function; + +/** + * Class to hold the fields that can be updated in a QueryGroup + */ +@ExperimentalApi +public class ChangeableQueryGroup extends AbstractDiffable { + + public static final String RESILIENCY_MODE_STRING = "resiliency_mode"; + public static final String RESOURCE_LIMITS_STRING = "resource_limits"; + private ResiliencyMode resiliencyMode; + private Map resourceLimits; + + public static final List acceptedFieldNames = List.of(RESILIENCY_MODE_STRING, RESOURCE_LIMITS_STRING); + private final Map> fromXContentMap = Map.of(RESILIENCY_MODE_STRING, (parser) -> { + try { + setResiliencyMode(ResiliencyMode.fromName(parser.text())); + return null; + } catch (IOException e) { + throw new IllegalArgumentException("parsing error encountered for the field " + RESILIENCY_MODE_STRING); + } + }, RESOURCE_LIMITS_STRING, (parser) -> { + try { + String fieldName = ""; + XContentParser.Token token; + final Map resourceLimits = new HashMap<>(); + 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()); + } + } + setResourceLimits(resourceLimits); + return null; + } catch (IOException e) { + throw new IllegalArgumentException("parsing error encountered for the object " + RESOURCE_LIMITS_STRING); + } + }); + + private final Map> toXContentMap = Map.of(RESILIENCY_MODE_STRING, (builder) -> { + try { + builder.field(RESILIENCY_MODE_STRING, resiliencyMode.getName()); + return null; + } catch (IOException e) { + throw new IllegalStateException("writing error encountered for the field " + RESILIENCY_MODE_STRING); + } + }, RESOURCE_LIMITS_STRING, (builder) -> { + try { + builder.startObject(RESOURCE_LIMITS_STRING); + for (ResourceType resourceType : ResourceType.values()) { + if (resourceLimits.containsKey(resourceType)) { + builder.field(resourceType.getName(), resourceLimits.get(resourceType)); + } + } + builder.endObject(); + return null; + } catch (IOException e) { + throw new IllegalStateException("writing error encountered for the field " + RESOURCE_LIMITS_STRING); + } + }); + + public ChangeableQueryGroup() {} + + public ChangeableQueryGroup(ResiliencyMode resiliencyMode, Map resourceLimits) { + validateResourceLimits(resourceLimits); + this.resiliencyMode = resiliencyMode; + this.resourceLimits = resourceLimits; + } + + public ChangeableQueryGroup(StreamInput in) throws IOException { + if (in.readBoolean()) { + resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble); + } else { + resourceLimits = new HashMap<>(); + } + String updatedResiliencyMode = in.readOptionalString(); + resiliencyMode = updatedResiliencyMode == null ? null : ResiliencyMode.fromName(updatedResiliencyMode); + } + + public static boolean shouldParse(String field) { + return acceptedFieldNames.contains(field); + } + + public void parseField(XContentParser parser, String field) { + fromXContentMap.get(field).apply(parser); + } + + public void writeField(XContentBuilder builder, String field) { + toXContentMap.get(field).apply(builder); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + if (resourceLimits == null || resourceLimits.isEmpty()) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeDouble); + } + out.writeOptionalString(resiliencyMode == null ? null : resiliencyMode.getName()); + } + + public static void validateResourceLimits(Map resourceLimits) { + if (resourceLimits == null) { + return; + } + 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"); + + 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 boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + ChangeableQueryGroup that = (ChangeableQueryGroup) o; + return Objects.equals(resiliencyMode, that.resiliencyMode) && Objects.equals(resourceLimits, that.resourceLimits); + } + + @Override + public int hashCode() { + return Objects.hash(resiliencyMode, resourceLimits); + } + + public ResiliencyMode getResiliencyMode() { + return resiliencyMode; + } + + public Map getResourceLimits() { + return resourceLimits; + } + + /** + * This enum models the different QueryGroup resiliency modes + * SOFT - means that this query group can consume more than query group resource limits if node is not in duress + * ENFORCED - means that it will never breach the assigned limits and will cancel as soon as the limits are breached + * MONITOR - it will not cause any cancellation but just log the eligible task cancellations + */ + @ExperimentalApi + public enum ResiliencyMode { + SOFT("soft"), + ENFORCED("enforced"), + MONITOR("monitor"); + + private final String name; + + ResiliencyMode(String mode) { + this.name = mode; + } + + public String getName() { + return name; + } + + public static ResiliencyMode fromName(String s) { + for (ResiliencyMode mode : values()) { + if (mode.getName().equalsIgnoreCase(s)) return mode; + + } + throw new IllegalArgumentException("Invalid value for QueryGroupMode: " + s); + } + } + + public void setResiliencyMode(ResiliencyMode resiliencyMode) { + this.resiliencyMode = resiliencyMode; + } + + public void setResourceLimits(Map resourceLimits) { + validateResourceLimits(resourceLimits); + this.resourceLimits = resourceLimits; + } +} 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 f5e667de73d93..63978a31cad59 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -16,6 +16,7 @@ import org.opensearch.core.xcontent.XContentParser; import org.opensearch.test.AbstractDiffableSerializationTestCase; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.ChangeableQueryGroup; import java.io.IOException; import java.util.Collections; @@ -33,8 +34,7 @@ public void testToXContent() throws IOException { new QueryGroup( "test", "ajakgakg983r92_4242", - QueryGroup.ResiliencyMode.ENFORCED, - Map.of(ResourceType.MEMORY, 0.5), + new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), 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 f4d3e5ceb1784..972f8ed80b6cc 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -16,6 +16,8 @@ import org.opensearch.core.xcontent.XContentParser; import org.opensearch.test.AbstractSerializingTestCase; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; import org.joda.time.Instant; import java.io.IOException; @@ -26,20 +28,16 @@ public class QueryGroupTests extends AbstractSerializingTestCase { - private static final List allowedModes = List.of( - QueryGroup.ResiliencyMode.SOFT, - QueryGroup.ResiliencyMode.ENFORCED, - QueryGroup.ResiliencyMode.MONITOR - ); + private static final List allowedModes = List.of(ResiliencyMode.SOFT, ResiliencyMode.ENFORCED, ResiliencyMode.MONITOR); static QueryGroup createRandomQueryGroup(String _id) { String name = randomAlphaOfLength(10); Map resourceLimit = new HashMap<>(); resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false)); - return new QueryGroup(name, _id, randomMode(), resourceLimit, Instant.now().getMillis()); + return new QueryGroup(name, _id, new ChangeableQueryGroup(randomMode(), resourceLimit), Instant.now().getMillis()); } - private static QueryGroup.ResiliencyMode randomMode() { + private static ResiliencyMode randomMode() { return allowedModes.get(randomIntBetween(0, allowedModes.size() - 1)); } @@ -74,37 +72,50 @@ protected QueryGroup createTestInstance() { public void testNullName() { assertThrows( NullPointerException.class, - () -> new QueryGroup(null, "_id", randomMode(), Collections.emptyMap(), Instant.now().getMillis()) + () -> new QueryGroup(null, "_id", new ChangeableQueryGroup(randomMode(), Collections.emptyMap()), Instant.now().getMillis()) ); } public void testNullId() { assertThrows( NullPointerException.class, - () -> new QueryGroup("Dummy", null, randomMode(), Collections.emptyMap(), Instant.now().getMillis()) + () -> new QueryGroup("Dummy", null, new ChangeableQueryGroup(randomMode(), Collections.emptyMap()), Instant.now().getMillis()) ); } public void testNullResourceLimits() { - assertThrows(NullPointerException.class, () -> new QueryGroup("analytics", "_id", randomMode(), null, Instant.now().getMillis())); + assertThrows( + NullPointerException.class, + () -> new QueryGroup("analytics", "_id", new ChangeableQueryGroup(randomMode(), null), Instant.now().getMillis()) + ); } public void testEmptyResourceLimits() { assertThrows( IllegalArgumentException.class, - () -> new QueryGroup("analytics", "_id", randomMode(), Collections.emptyMap(), Instant.now().getMillis()) + () -> new QueryGroup( + "analytics", + "_id", + new ChangeableQueryGroup(randomMode(), Collections.emptyMap()), + Instant.now().getMillis() + ) ); } public void testIllegalQueryGroupMode() { assertThrows( NullPointerException.class, - () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis()) + () -> new QueryGroup( + "analytics", + "_id", + new ChangeableQueryGroup(null, Map.of(ResourceType.MEMORY, 0.4)), + Instant.now().getMillis() + ) ); } public void testQueryGroupInitiation() { - QueryGroup queryGroup = new QueryGroup("analytics", randomMode(), Map.of(ResourceType.MEMORY, 0.4)); + QueryGroup queryGroup = new QueryGroup("analytics", new ChangeableQueryGroup(randomMode(), Map.of(ResourceType.MEMORY, 0.4))); assertNotNull(queryGroup.getName()); assertNotNull(queryGroup.get_id()); assertNotNull(queryGroup.getResourceLimits()); @@ -117,12 +128,9 @@ public void testQueryGroupInitiation() { 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()) + () -> new QueryGroup("a".repeat(51), "_id", new ChangeableQueryGroup(), Instant.now().getMillis()) ); + assertThrows(NullPointerException.class, () -> new QueryGroup("", "_id", new ChangeableQueryGroup(), Instant.now().getMillis())); } @@ -132,8 +140,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { () -> new QueryGroup( "analytics", "_id", - randomMode(), - Map.of(ResourceType.MEMORY, randomDoubleBetween(1.1, 1.8, false)), + new ChangeableQueryGroup(randomMode(), Map.of(ResourceType.MEMORY, randomDoubleBetween(1.1, 1.8, false))), Instant.now().getMillis() ) ); @@ -143,8 +150,7 @@ public void testValidQueryGroup() { QueryGroup queryGroup = new QueryGroup( "analytics", "_id", - randomMode(), - Map.of(ResourceType.MEMORY, randomDoubleBetween(0.01, 0.8, false)), + new ChangeableQueryGroup(randomMode(), Map.of(ResourceType.MEMORY, randomDoubleBetween(0.01, 0.8, false))), Instant.ofEpochMilli(1717187289).getMillis() ); @@ -163,8 +169,7 @@ public void testToXContent() throws IOException { QueryGroup queryGroup = new QueryGroup( "TestQueryGroup", queryGroupId, - QueryGroup.ResiliencyMode.ENFORCED, - Map.of(ResourceType.CPU, 0.30, ResourceType.MEMORY, 0.40), + new ChangeableQueryGroup(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.30, ResourceType.MEMORY, 0.40)), currentTimeInMillis ); XContentBuilder builder = JsonXContent.contentBuilder(); diff --git a/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java b/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java new file mode 100644 index 0000000000000..77aaa3633ae96 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java @@ -0,0 +1,65 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.search.ResourceType; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +public class ChangeableQueryGroupTests extends OpenSearchTestCase { + + public void testSerializationDeserialization() throws IOException { + Map resourceLimits = new HashMap<>(); + resourceLimits.put(ResourceType.CPU, 0.5); + resourceLimits.put(ResourceType.MEMORY, 0.75); + ChangeableQueryGroup changeableQueryGroup = new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.SOFT, resourceLimits); + BytesStreamOutput out = new BytesStreamOutput(); + changeableQueryGroup.writeTo(out); + StreamInput in = out.bytes().streamInput(); + ChangeableQueryGroup deserializedGroup = new ChangeableQueryGroup(in); + assertEquals(changeableQueryGroup, deserializedGroup); + } + + public void testSerializationDeserializationWithNull() throws IOException { + ChangeableQueryGroup changeableQueryGroup = new ChangeableQueryGroup(); + BytesStreamOutput out = new BytesStreamOutput(); + changeableQueryGroup.writeTo(out); + StreamInput in = out.bytes().streamInput(); + ChangeableQueryGroup deserializedGroup = new ChangeableQueryGroup(in); + assertEquals(0, deserializedGroup.getResourceLimits().size()); + assertEquals(changeableQueryGroup.getResiliencyMode(), deserializedGroup.getResiliencyMode()); + } + + public void testValidateResourceLimits() { + Map invalidLimits = new HashMap<>(); + invalidLimits.put(ResourceType.CPU, 1.5); + Exception exception = assertThrows( + IllegalArgumentException.class, + () -> { ChangeableQueryGroup.validateResourceLimits(invalidLimits); } + ); + String expectedMessage = "resource value should be greater than 0 and less or equal to 1.0"; + String actualMessage = exception.getMessage(); + assertTrue(actualMessage.contains(expectedMessage)); + } + + public void testSetMethodsWithNullAndEmptyValues() { + ChangeableQueryGroup queryGroup = new ChangeableQueryGroup(); + queryGroup.setResiliencyMode(null); + assertNull(queryGroup.getResiliencyMode()); + queryGroup.setResourceLimits(null); + assertNull(queryGroup.getResourceLimits()); + queryGroup.setResourceLimits(new HashMap<>()); + assertEquals(0, queryGroup.getResourceLimits().size()); + } +} From ba63f56a2b9847ebabb79ef095a3edf15c70f1c8 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 28 Aug 2024 15:21:20 -0700 Subject: [PATCH 7/9] fix UT Signed-off-by: Ruirui Zhang --- .../plugin/wlm/action/UpdateQueryGroupRequest.java | 3 +-- .../wlm/service/QueryGroupPersistenceService.java | 4 +--- .../opensearch/plugin/wlm/QueryGroupTestUtils.java | 2 +- .../wlm/action/CreateQueryGroupResponseTests.java | 4 ++-- .../wlm/action/GetQueryGroupResponseTests.java | 12 ++++++------ .../wlm/action/UpdateQueryGroupRequestTests.java | 2 +- .../wlm/action/UpdateQueryGroupResponseTests.java | 4 ++-- .../service/QueryGroupPersistenceServiceTests.java | 2 +- .../org/opensearch/cluster/metadata/QueryGroup.java | 2 +- .../org/opensearch/wlm/ChangeableQueryGroup.java | 3 +-- .../cluster/metadata/QueryGroupMetadataTests.java | 4 ++-- .../opensearch/cluster/metadata/QueryGroupTests.java | 6 +++--- .../opensearch/wlm/ChangeableQueryGroupTests.java | 1 - 13 files changed, 22 insertions(+), 27 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java index de13436746b16..3e16bcda325a8 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java @@ -14,9 +14,9 @@ import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.wlm.ResourceType; import org.opensearch.wlm.ChangeableQueryGroup; import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.ResourceType; import java.io.IOException; import java.util.Map; @@ -94,7 +94,6 @@ public ChangeableQueryGroup getChangeableQueryGroup() { @Override public void writeTo(StreamOutput out) throws IOException { - super.writeTo(out); out.writeString(name); changeableQueryGroup.writeTo(out); } 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 a4fb666c324c8..29cafbc1e6fb8 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 @@ -29,13 +29,11 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; - -import org.opensearch.wlm.ResourceType; import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; -import org.opensearch.wlm.ResourceType; import org.opensearch.wlm.ChangeableQueryGroup; import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.ResourceType; import org.joda.time.Instant; import java.util.Collection; 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 47a1fd9a94a22..2b1b07e357f08 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 @@ -20,9 +20,9 @@ import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; -import org.opensearch.wlm.ResourceType; import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ResourceType; import java.util.ArrayList; import java.util.Collection; 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 713c62088b9ca..3a2ce215d21b5 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 @@ -56,10 +56,10 @@ public void testToXContentCreateQueryGroup() throws IOException { + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + " \"name\" : \"query_group_one\",\n" + " \"resiliency_mode\" : \"monitor\",\n" - + " \"updated_at\" : 4513232413,\n" + " \"resource_limits\" : {\n" + " \"memory\" : 0.3\n" - + " }\n" + + " },\n" + + " \"updated_at\" : 4513232413\n" + "}"; assertEquals(expected, actual); } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java index d9192dbb18fe4..1a2ac282d86a4 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java @@ -93,10 +93,10 @@ public void testToXContentGetSingleQueryGroup() throws IOException { + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + " \"name\" : \"query_group_one\",\n" + " \"resiliency_mode\" : \"monitor\",\n" - + " \"updated_at\" : 4513232413,\n" + " \"resource_limits\" : {\n" + " \"memory\" : 0.3\n" - + " }\n" + + " },\n" + + " \"updated_at\" : 4513232413\n" + " }\n" + " ]\n" + "}"; @@ -119,19 +119,19 @@ public void testToXContentGetMultipleQueryGroup() throws IOException { + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + " \"name\" : \"query_group_one\",\n" + " \"resiliency_mode\" : \"monitor\",\n" - + " \"updated_at\" : 4513232413,\n" + " \"resource_limits\" : {\n" + " \"memory\" : 0.3\n" - + " }\n" + + " },\n" + + " \"updated_at\" : 4513232413\n" + " },\n" + " {\n" + " \"_id\" : \"G5iIqHy4g7eK1qIAAAAIH53=1\",\n" + " \"name\" : \"query_group_two\",\n" + " \"resiliency_mode\" : \"monitor\",\n" - + " \"updated_at\" : 4513232415,\n" + " \"resource_limits\" : {\n" + " \"memory\" : 0.6\n" - + " }\n" + + " },\n" + + " \"updated_at\" : 4513232415\n" + " }\n" + " ]\n" + "}"; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java index 147759bceebfa..2ffa515c12c02 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java @@ -10,10 +10,10 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.search.ResourceType; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.wlm.ChangeableQueryGroup; import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.ResourceType; import java.io.IOException; import java.util.HashMap; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java index 650e0bdc061ce..a7ab4c6a682ef 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupResponseTests.java @@ -57,10 +57,10 @@ public void testToXContentUpdateSingleQueryGroup() throws IOException { + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + " \"name\" : \"query_group_one\",\n" + " \"resiliency_mode\" : \"monitor\",\n" - + " \"updated_at\" : 4513232413,\n" + " \"resource_limits\" : {\n" + " \"memory\" : 0.3\n" - + " }\n" + + " },\n" + + " \"updated_at\" : 4513232413\n" + "}"; assertEquals(expected, actual); } 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 b4426335c3169..aa3427ccbe307 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 @@ -26,11 +26,11 @@ import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; -import org.opensearch.wlm.ResourceType; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.ChangeableQueryGroup; import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.ResourceType; import java.util.ArrayList; import java.util.Collection; 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 11e8fa695d23d..b2755794108ff 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -17,9 +17,9 @@ import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.wlm.ResourceType; import org.opensearch.wlm.ChangeableQueryGroup; import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.ResourceType; import org.joda.time.Instant; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java b/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java index fad653d4fcbf9..f9f499fd84684 100644 --- a/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java +++ b/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java @@ -14,7 +14,6 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.search.ResourceType; import java.io.IOException; import java.util.HashMap; @@ -24,7 +23,7 @@ import java.util.function.Function; /** - * Class to hold the fields that can be updated in a QueryGroup + * Class to hold the fields that can be updated in a QueryGroup. */ @ExperimentalApi public class ChangeableQueryGroup extends AbstractDiffable { 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 63978a31cad59..6e8abdd19a71f 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -15,8 +15,8 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.test.AbstractDiffableSerializationTestCase; -import org.opensearch.wlm.ResourceType; import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.ResourceType; import java.io.IOException; import java.util.Collections; @@ -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\",\"updated_at\":1720047207,\"resource_limits\":{\"memory\":0.5}}}", + "{\"ajakgakg983r92_4242\":{\"_id\":\"ajakgakg983r92_4242\",\"name\":\"test\",\"resiliency_mode\":\"enforced\",\"resource_limits\":{\"memory\":0.5},\"updated_at\":1720047207}}", builder.toString() ); } 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 972f8ed80b6cc..f60276ec9858a 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -15,9 +15,9 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.test.AbstractSerializingTestCase; -import org.opensearch.wlm.ResourceType; import org.opensearch.wlm.ChangeableQueryGroup; import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.ResourceType; import org.joda.time.Instant; import java.io.IOException; @@ -177,9 +177,9 @@ public void testToXContent() throws IOException { assertEquals( "{\"_id\":\"" + queryGroupId - + "\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"updated_at\":" + + "\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"resource_limits\":{\"cpu\":0.3,\"memory\":0.4},\"updated_at\":" + currentTimeInMillis - + ",\"resource_limits\":{\"cpu\":0.3,\"memory\":0.4}}", + + "}", builder.toString() ); } diff --git a/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java b/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java index 77aaa3633ae96..66605d4e91f0b 100644 --- a/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java +++ b/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java @@ -10,7 +10,6 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.search.ResourceType; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; From 05b0a53b148ea9765d849ce7907bd4aba9b99124 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 29 Aug 2024 14:15:41 -0700 Subject: [PATCH 8/9] adress comments Signed-off-by: Ruirui Zhang --- .../opensearch/wlm/ChangeableQueryGroup.java | 90 ++++++++++++------- 1 file changed, 57 insertions(+), 33 deletions(-) diff --git a/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java b/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java index f9f499fd84684..bee3b8fe7827c 100644 --- a/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java +++ b/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java @@ -20,6 +20,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.function.Function; /** @@ -34,15 +35,37 @@ public class ChangeableQueryGroup extends AbstractDiffable private Map resourceLimits; public static final List acceptedFieldNames = List.of(RESILIENCY_MODE_STRING, RESOURCE_LIMITS_STRING); - private final Map> fromXContentMap = Map.of(RESILIENCY_MODE_STRING, (parser) -> { - try { - setResiliencyMode(ResiliencyMode.fromName(parser.text())); - return null; - } catch (IOException e) { - throw new IllegalArgumentException("parsing error encountered for the field " + RESILIENCY_MODE_STRING); + + public ChangeableQueryGroup() {} + + public ChangeableQueryGroup(ResiliencyMode resiliencyMode, Map resourceLimits) { + validateResourceLimits(resourceLimits); + this.resiliencyMode = resiliencyMode; + this.resourceLimits = resourceLimits; + } + + public ChangeableQueryGroup(StreamInput in) throws IOException { + if (in.readBoolean()) { + resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble); + } else { + resourceLimits = new HashMap<>(); } - }, RESOURCE_LIMITS_STRING, (parser) -> { - try { + String updatedResiliencyMode = in.readOptionalString(); + resiliencyMode = updatedResiliencyMode == null ? null : ResiliencyMode.fromName(updatedResiliencyMode); + } + + interface FieldParser { + T parseField(XContentParser parser) throws IOException; + } + + static class ResiliencyModeParser implements FieldParser { + public ResiliencyMode parseField(XContentParser parser) throws IOException { + return ResiliencyMode.fromName(parser.text()); + } + } + + static class ResourceLimitsParser implements FieldParser> { + public Map parseField(XContentParser parser) throws IOException { String fieldName = ""; XContentParser.Token token; final Map resourceLimits = new HashMap<>(); @@ -53,12 +76,20 @@ public class ChangeableQueryGroup extends AbstractDiffable resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); } } - setResourceLimits(resourceLimits); - return null; - } catch (IOException e) { - throw new IllegalArgumentException("parsing error encountered for the object " + RESOURCE_LIMITS_STRING); + return resourceLimits; } - }); + } + + static class FieldParserFactory { + static Optional> fieldParserFor(String fieldName) { + if (fieldName.equals(RESOURCE_LIMITS_STRING)) { + return Optional.of(new ResourceLimitsParser()); + } else if (fieldName.equals(RESILIENCY_MODE_STRING)) { + return Optional.of(new ResiliencyModeParser()); + } + return Optional.empty(); + } + } private final Map> toXContentMap = Map.of(RESILIENCY_MODE_STRING, (builder) -> { try { @@ -82,30 +113,23 @@ public class ChangeableQueryGroup extends AbstractDiffable } }); - public ChangeableQueryGroup() {} - - public ChangeableQueryGroup(ResiliencyMode resiliencyMode, Map resourceLimits) { - validateResourceLimits(resourceLimits); - this.resiliencyMode = resiliencyMode; - this.resourceLimits = resourceLimits; - } - - public ChangeableQueryGroup(StreamInput in) throws IOException { - if (in.readBoolean()) { - resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble); - } else { - resourceLimits = new HashMap<>(); - } - String updatedResiliencyMode = in.readOptionalString(); - resiliencyMode = updatedResiliencyMode == null ? null : ResiliencyMode.fromName(updatedResiliencyMode); - } - public static boolean shouldParse(String field) { - return acceptedFieldNames.contains(field); + return FieldParserFactory.fieldParserFor(field).isPresent(); } public void parseField(XContentParser parser, String field) { - fromXContentMap.get(field).apply(parser); + FieldParserFactory.fieldParserFor(field).ifPresent(fieldParser -> { + try { + Object value = fieldParser.parseField(parser); + if (field.equals(RESILIENCY_MODE_STRING)) { + setResiliencyMode((ResiliencyMode) value); + } else if (field.equals(RESOURCE_LIMITS_STRING)) { + setResourceLimits((Map) value); + } + } catch (IOException e) { + throw new IllegalArgumentException("parsing error encountered for the field " + field); + } + }); } public void writeField(XContentBuilder builder, String field) { From ed850523a6312aed1e6f17c24a3d8349259a89a8 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Sat, 31 Aug 2024 00:56:58 -0700 Subject: [PATCH 9/9] address comments Signed-off-by: Ruirui Zhang --- .../wlm/action/CreateQueryGroupRequest.java | 4 +- .../wlm/action/UpdateQueryGroupRequest.java | 41 +++------- .../service/QueryGroupPersistenceService.java | 35 ++++---- .../plugin/wlm/QueryGroupTestUtils.java | 10 ++- .../wlm/action/QueryGroupActionTestUtils.java | 17 ++++ .../action/UpdateQueryGroupRequestTests.java | 29 +++---- .../QueryGroupPersistenceServiceTests.java | 31 +++---- .../cluster/metadata/QueryGroup.java | 82 +++++++++++-------- ...up.java => MutableQueryGroupFragment.java} | 10 +-- .../metadata/QueryGroupMetadataTests.java | 4 +- .../cluster/metadata/QueryGroupTests.java | 41 ++++++---- ...va => MutableQueryGroupFragmentTests.java} | 30 +++---- 12 files changed, 181 insertions(+), 153 deletions(-) create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupActionTestUtils.java rename server/src/main/java/org/opensearch/wlm/{ChangeableQueryGroup.java => MutableQueryGroupFragment.java} (95%) rename server/src/test/java/org/opensearch/wlm/{ChangeableQueryGroupTests.java => MutableQueryGroupFragmentTests.java} (64%) 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 ff6422be36885..d92283391dd3b 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 @@ -40,7 +40,7 @@ public class CreateQueryGroupRequest extends ActionRequest { * Constructor for CreateQueryGroupRequest * @param queryGroup - A {@link QueryGroup} object */ - public CreateQueryGroupRequest(QueryGroup queryGroup) { + CreateQueryGroupRequest(QueryGroup queryGroup) { this.queryGroup = queryGroup; } @@ -48,7 +48,7 @@ public CreateQueryGroupRequest(QueryGroup queryGroup) { * Constructor for CreateQueryGroupRequest * @param in - A {@link StreamInput} object */ - public CreateQueryGroupRequest(StreamInput in) throws IOException { + CreateQueryGroupRequest(StreamInput in) throws IOException { super(in); queryGroup = new QueryGroup(in); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java index 3e16bcda325a8..048b599f095fd 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java @@ -14,12 +14,9 @@ import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.wlm.ChangeableQueryGroup; -import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; -import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.MutableQueryGroupFragment; import java.io.IOException; -import java.util.Map; /** * A request for update QueryGroup @@ -28,24 +25,24 @@ */ public class UpdateQueryGroupRequest extends ActionRequest { private final String name; - private final ChangeableQueryGroup changeableQueryGroup; + private final MutableQueryGroupFragment mutableQueryGroupFragment; /** * Constructor for UpdateQueryGroupRequest * @param name - QueryGroup name for UpdateQueryGroupRequest - * @param changeableQueryGroup - ChangeableQueryGroup for UpdateQueryGroupRequest + * @param mutableQueryGroupFragment - MutableQueryGroupFragment for UpdateQueryGroupRequest */ - public UpdateQueryGroupRequest(String name, ChangeableQueryGroup changeableQueryGroup) { + UpdateQueryGroupRequest(String name, MutableQueryGroupFragment mutableQueryGroupFragment) { this.name = name; - this.changeableQueryGroup = changeableQueryGroup; + this.mutableQueryGroupFragment = mutableQueryGroupFragment; } /** * Constructor for UpdateQueryGroupRequest * @param in - A {@link StreamInput} object */ - public UpdateQueryGroupRequest(StreamInput in) throws IOException { - this(in.readString(), new ChangeableQueryGroup(in)); + UpdateQueryGroupRequest(StreamInput in) throws IOException { + this(in.readString(), new MutableQueryGroupFragment(in)); } /** @@ -55,7 +52,7 @@ public UpdateQueryGroupRequest(StreamInput in) throws IOException { */ public static UpdateQueryGroupRequest fromXContent(XContentParser parser, String name) throws IOException { QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser); - return new UpdateQueryGroupRequest(name, builder.getChangeableQueryGroup()); + return new UpdateQueryGroupRequest(name, builder.getMutableQueryGroupFragment()); } @Override @@ -72,29 +69,15 @@ public String getName() { } /** - * resourceLimits getter + * mutableQueryGroupFragment getter */ - public Map getResourceLimits() { - return getChangeableQueryGroup().getResourceLimits(); - } - - /** - * resiliencyMode getter - */ - public ResiliencyMode getResiliencyMode() { - return getChangeableQueryGroup().getResiliencyMode(); - } - - /** - * changeableQueryGroup getter - */ - public ChangeableQueryGroup getChangeableQueryGroup() { - return changeableQueryGroup; + public MutableQueryGroupFragment getmMutableQueryGroupFragment() { + return mutableQueryGroupFragment; } @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(name); - changeableQueryGroup.writeTo(out); + mutableQueryGroupFragment.writeTo(out); } } 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 29cafbc1e6fb8..f9332ff3022dc 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 @@ -31,18 +31,17 @@ import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest; import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; -import org.opensearch.wlm.ChangeableQueryGroup; -import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.MutableQueryGroupFragment; import org.opensearch.wlm.ResourceType; -import org.joda.time.Instant; import java.util.Collection; import java.util.EnumMap; -import java.util.HashMap; import java.util.Map; import java.util.Optional; import java.util.stream.Collectors; +import static org.opensearch.cluster.metadata.QueryGroup.updateExistingQueryGroup; + /** * This class defines the functions for QueryGroup persistence */ @@ -296,6 +295,7 @@ ClusterState updateQueryGroupInClusterState(UpdateQueryGroupRequest updateQueryG final Metadata metadata = currentState.metadata(); final Map existingGroups = currentState.metadata().queryGroups(); String name = updateQueryGroupRequest.getName(); + MutableQueryGroupFragment mutableQueryGroupFragment = updateQueryGroupRequest.getmMutableQueryGroupFragment(); final QueryGroup existingGroup = existingGroups.values() .stream() @@ -303,24 +303,14 @@ ClusterState updateQueryGroupInClusterState(UpdateQueryGroupRequest updateQueryG .findFirst() .orElseThrow(() -> new ResourceNotFoundException("No QueryGroup exists with the provided name: " + name)); - // build the QueryGroup with updated fields - final Map updatedResourceLimits = new HashMap<>(existingGroup.getResourceLimits()); - if (updateQueryGroupRequest.getResourceLimits() != null && !updateQueryGroupRequest.getResourceLimits().isEmpty()) { - validateTotalUsage(existingGroups, name, updateQueryGroupRequest.getResourceLimits()); - updatedResourceLimits.putAll(updateQueryGroupRequest.getResourceLimits()); - } - - final ResiliencyMode mode = Optional.ofNullable(updateQueryGroupRequest.getResiliencyMode()) - .orElse(existingGroup.getResiliencyMode()); - - final QueryGroup updatedGroup = new QueryGroup( - name, - existingGroup.get_id(), - new ChangeableQueryGroup(mode, updatedResourceLimits), - Instant.now().getMillis() - ); + validateTotalUsage(existingGroups, name, mutableQueryGroupFragment.getResourceLimits()); return ClusterState.builder(currentState) - .metadata(Metadata.builder(metadata).remove(existingGroup).put(updatedGroup).build()) + .metadata( + Metadata.builder(metadata) + .remove(existingGroup) + .put(updateExistingQueryGroup(existingGroup, mutableQueryGroupFragment)) + .build() + ) .build(); } @@ -330,6 +320,9 @@ ClusterState updateQueryGroupInClusterState(UpdateQueryGroupRequest updateQueryG * @param resourceLimits - the QueryGroup we're creating or updating */ private void validateTotalUsage(Map existingQueryGroups, String name, Map resourceLimits) { + if (resourceLimits == null || resourceLimits.isEmpty()) { + return; + } final Map totalUsage = new EnumMap<>(ResourceType.class); totalUsage.putAll(resourceLimits); for (QueryGroup currGroup : existingQueryGroups.values()) { 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 2b1b07e357f08..c6eb3140e943d 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 @@ -21,7 +21,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; import org.opensearch.threadpool.ThreadPool; -import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.MutableQueryGroupFragment; import org.opensearch.wlm.ResourceType; import java.util.ArrayList; @@ -48,13 +48,17 @@ public class QueryGroupTestUtils { public static final long TIMESTAMP_TWO = 4513232415L; public static final QueryGroup queryGroupOne = builder().name(NAME_ONE) ._id(_ID_ONE) - .changeableQueryGroup(new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.3))) + .mutableQueryGroupFragment( + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.3)) + ) .updatedAt(TIMESTAMP_ONE) .build(); public static final QueryGroup queryGroupTwo = builder().name(NAME_TWO) ._id(_ID_TWO) - .changeableQueryGroup(new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.6))) + .mutableQueryGroupFragment( + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.6)) + ) .updatedAt(TIMESTAMP_TWO) .build(); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupActionTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupActionTestUtils.java new file mode 100644 index 0000000000000..08d128ca7ed59 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupActionTestUtils.java @@ -0,0 +1,17 @@ +/* + * 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.wlm.MutableQueryGroupFragment; + +public class QueryGroupActionTestUtils { + public static UpdateQueryGroupRequest updateQueryGroupRequest(String name, MutableQueryGroupFragment mutableQueryGroupFragment) { + return new UpdateQueryGroupRequest(name, mutableQueryGroupFragment); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java index 2ffa515c12c02..b99f079e81984 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequestTests.java @@ -11,8 +11,8 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.wlm.ChangeableQueryGroup; -import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.MutableQueryGroupFragment; +import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; import java.io.IOException; @@ -20,7 +20,6 @@ import java.util.Map; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_ONE; -import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualResourceLimits; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; public class UpdateQueryGroupRequestTests extends OpenSearchTestCase { @@ -29,29 +28,26 @@ public class UpdateQueryGroupRequestTests extends OpenSearchTestCase { * Test case to verify the serialization and deserialization of UpdateQueryGroupRequest. */ public void testSerialization() throws IOException { - UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, queryGroupOne.getChangeableQueryGroup()); + UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, queryGroupOne.getMutableQueryGroupFragment()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); UpdateQueryGroupRequest otherRequest = new UpdateQueryGroupRequest(streamInput); assertEquals(request.getName(), otherRequest.getName()); - assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); - assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); - assertEqualResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); + assertEquals(request.getmMutableQueryGroupFragment(), otherRequest.getmMutableQueryGroupFragment()); } /** * Test case to verify the serialization and deserialization of UpdateQueryGroupRequest with only name field. */ public void testSerializationOnlyName() throws IOException { - UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, new ChangeableQueryGroup(null, new HashMap<>())); + UpdateQueryGroupRequest request = new UpdateQueryGroupRequest(NAME_ONE, new MutableQueryGroupFragment(null, new HashMap<>())); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); UpdateQueryGroupRequest otherRequest = new UpdateQueryGroupRequest(streamInput); assertEquals(request.getName(), otherRequest.getName()); - assertEquals(request.getResourceLimits(), otherRequest.getResourceLimits()); - assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); + assertEquals(request.getmMutableQueryGroupFragment(), otherRequest.getmMutableQueryGroupFragment()); } /** @@ -60,16 +56,14 @@ public void testSerializationOnlyName() throws IOException { public void testSerializationOnlyResourceLimit() throws IOException { UpdateQueryGroupRequest request = new UpdateQueryGroupRequest( NAME_ONE, - new ChangeableQueryGroup(null, Map.of(ResourceType.MEMORY, 0.4)) + new MutableQueryGroupFragment(null, Map.of(ResourceType.MEMORY, 0.4)) ); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); UpdateQueryGroupRequest otherRequest = new UpdateQueryGroupRequest(streamInput); assertEquals(request.getName(), otherRequest.getName()); - assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size()); - assertEqualResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits()); - assertEquals(request.getResiliencyMode(), otherRequest.getResiliencyMode()); + assertEquals(request.getmMutableQueryGroupFragment(), otherRequest.getmMutableQueryGroupFragment()); } /** @@ -80,7 +74,10 @@ public void testInvalidResourceLimitList() { IllegalArgumentException.class, () -> new UpdateQueryGroupRequest( NAME_ONE, - new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.3, ResourceType.fromName("random"), 0.4)) + new MutableQueryGroupFragment( + ResiliencyMode.MONITOR, + Map.of(ResourceType.MEMORY, 0.3, ResourceType.fromName("random"), 0.4) + ) ) ); } @@ -93,7 +90,7 @@ public void testInvalidEnforcement() { IllegalArgumentException.class, () -> new UpdateQueryGroupRequest( NAME_ONE, - new ChangeableQueryGroup(ResiliencyMode.fromName("random"), Map.of(ResourceType.fromName("memory"), 0.3)) + new MutableQueryGroupFragment(ResiliencyMode.fromName("random"), Map.of(ResourceType.fromName("memory"), 0.3)) ) ); } 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 aa3427ccbe307..08b51fd46cfcf 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 @@ -28,8 +28,8 @@ import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; -import org.opensearch.wlm.ChangeableQueryGroup; -import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.MutableQueryGroupFragment; +import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; import java.util.ArrayList; @@ -58,6 +58,7 @@ 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.plugin.wlm.action.QueryGroupActionTestUtils.updateQueryGroupRequest; 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; @@ -114,7 +115,7 @@ public void testCreateQueryGroupDuplicateName() { ClusterState clusterState = setup.v2(); QueryGroup toCreate = builder().name(NAME_ONE) ._id("W5iIqHyhgi4K1qIAAAAIHw==") - .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.3))) + .mutableQueryGroupFragment(new MutableQueryGroupFragment(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.3))) .updatedAt(1690934400000L) .build(); assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); @@ -128,7 +129,7 @@ public void testCreateQueryGroupOverflowAllocation() { Tuple setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupTwo)); QueryGroup toCreate = builder().name(NAME_ONE) ._id("W5iIqHyhgi4K1qIAAAAIHw==") - .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.41))) + .mutableQueryGroupFragment(new MutableQueryGroupFragment(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.41))) .updatedAt(1690934400000L) .build(); @@ -144,7 +145,7 @@ public void testCreateQueryGroupOverflowAllocation() { public void testCreateQueryGroupOverflowCount() { QueryGroup toCreate = builder().name(NAME_NONE_EXISTED) ._id("W5iIqHyhgi4K1qIAAAAIHw==") - .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.5))) + .mutableQueryGroupFragment(new MutableQueryGroupFragment(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.5))) .updatedAt(1690934400000L) .build(); Metadata metadata = Metadata.builder().queryGroups(Map.of(_ID_ONE, queryGroupOne, _ID_TWO, queryGroupTwo)).build(); @@ -363,10 +364,10 @@ public void testDeleteInClusterStateMetadata() throws Exception { public void testUpdateQueryGroupAllFields() { QueryGroup updated = builder().name(NAME_ONE) ._id(_ID_ONE) - .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.15))) + .mutableQueryGroupFragment(new MutableQueryGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.15))) .updatedAt(1690934400000L) .build(); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest(NAME_ONE, updated.getChangeableQueryGroup()); + UpdateQueryGroupRequest updateQueryGroupRequest = updateQueryGroupRequest(NAME_ONE, updated.getMutableQueryGroupFragment()); ClusterState newClusterState = queryGroupPersistenceService().updateQueryGroupInClusterState( updateQueryGroupRequest, clusterState() @@ -385,10 +386,10 @@ public void testUpdateQueryGroupAllFields() { public void testUpdateQueryGroupResourceLimitsOnly() { QueryGroup updated = builder().name(NAME_ONE) ._id(_ID_ONE) - .changeableQueryGroup(new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.15))) + .mutableQueryGroupFragment(new MutableQueryGroupFragment(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.15))) .updatedAt(1690934400000L) .build(); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest(NAME_ONE, updated.getChangeableQueryGroup()); + UpdateQueryGroupRequest updateQueryGroupRequest = updateQueryGroupRequest(NAME_ONE, updated.getMutableQueryGroupFragment()); ClusterState newClusterState = queryGroupPersistenceService().updateQueryGroupInClusterState( updateQueryGroupRequest, clusterState() @@ -421,9 +422,9 @@ public void testUpdateQueryGroupResourceLimitsOnly() { */ public void testUpdateQueryGroupNonExistedName() { QueryGroupPersistenceService queryGroupPersistenceService = queryGroupPersistenceService(); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + UpdateQueryGroupRequest updateQueryGroupRequest = updateQueryGroupRequest( NAME_NONE_EXISTED, - new ChangeableQueryGroup(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.15)) + new MutableQueryGroupFragment(ResiliencyMode.MONITOR, Map.of(ResourceType.MEMORY, 0.15)) ); assertThrows( RuntimeException.class, @@ -467,9 +468,9 @@ public void testUpdateInClusterStateMetadataInner() { QueryGroupTestUtils.settings(), clusterSettings() ); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + UpdateQueryGroupRequest updateQueryGroupRequest = updateQueryGroupRequest( NAME_TWO, - new ChangeableQueryGroup(ResiliencyMode.SOFT, new HashMap<>()) + new MutableQueryGroupFragment(ResiliencyMode.SOFT, new HashMap<>()) ); ArgumentCaptor captor = ArgumentCaptor.forClass(ClusterStateUpdateTask.class); queryGroupPersistenceService.updateInClusterStateMetadata(updateQueryGroupRequest, listener); @@ -498,9 +499,9 @@ public void testUpdateInClusterStateMetadataFailure() { QueryGroupTestUtils.settings(), clusterSettings() ); - UpdateQueryGroupRequest updateQueryGroupRequest = new UpdateQueryGroupRequest( + UpdateQueryGroupRequest updateQueryGroupRequest = updateQueryGroupRequest( NAME_TWO, - new ChangeableQueryGroup(ResiliencyMode.SOFT, new HashMap<>()) + new MutableQueryGroupFragment(ResiliencyMode.SOFT, new HashMap<>()) ); doAnswer(invocation -> { ClusterStateUpdateTask task = invocation.getArgument(1); 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 b2755794108ff..dcd96dceb4bf1 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -17,14 +17,16 @@ import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.wlm.ChangeableQueryGroup; -import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.MutableQueryGroupFragment; +import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; import org.joda.time.Instant; import java.io.IOException; +import java.util.HashMap; import java.util.Map; import java.util.Objects; +import java.util.Optional; /** * Class to define the QueryGroup schema @@ -50,20 +52,20 @@ public class QueryGroup extends AbstractDiffable implements ToXConte private final String _id; // It is an epoch in millis private final long updatedAtInMillis; - private final ChangeableQueryGroup changeableQueryGroup; + private final MutableQueryGroupFragment mutableQueryGroupFragment; - public QueryGroup(String name, ChangeableQueryGroup changeableQueryGroup) { - this(name, UUIDs.randomBase64UUID(), changeableQueryGroup, Instant.now().getMillis()); + public QueryGroup(String name, MutableQueryGroupFragment mutableQueryGroupFragment) { + this(name, UUIDs.randomBase64UUID(), mutableQueryGroupFragment, Instant.now().getMillis()); } - public QueryGroup(String name, String _id, ChangeableQueryGroup changeableQueryGroup, long updatedAt) { + public QueryGroup(String name, String _id, MutableQueryGroupFragment mutableQueryGroupFragment, long updatedAt) { Objects.requireNonNull(name, "QueryGroup.name can't be null"); - Objects.requireNonNull(changeableQueryGroup.getResourceLimits(), "QueryGroup.resourceLimits can't be null"); - Objects.requireNonNull(changeableQueryGroup.getResiliencyMode(), "QueryGroup.resiliencyMode can't be null"); + Objects.requireNonNull(mutableQueryGroupFragment.getResourceLimits(), "QueryGroup.resourceLimits can't be null"); + Objects.requireNonNull(mutableQueryGroupFragment.getResiliencyMode(), "QueryGroup.resiliencyMode can't be null"); Objects.requireNonNull(_id, "QueryGroup._id can't be null"); validateName(name); - if (changeableQueryGroup.getResourceLimits().isEmpty()) { + if (mutableQueryGroupFragment.getResourceLimits().isEmpty()) { throw new IllegalArgumentException("QueryGroup.resourceLimits should at least have 1 resource limit"); } if (!isValid(updatedAt)) { @@ -72,7 +74,7 @@ public QueryGroup(String name, String _id, ChangeableQueryGroup changeableQueryG this.name = name; this._id = _id; - this.changeableQueryGroup = changeableQueryGroup; + this.mutableQueryGroupFragment = mutableQueryGroupFragment; this.updatedAtInMillis = updatedAt; } @@ -87,14 +89,30 @@ public static boolean isValid(long updatedAt) { } public QueryGroup(StreamInput in) throws IOException { - this(in.readString(), in.readString(), new ChangeableQueryGroup(in), in.readLong()); + this(in.readString(), in.readString(), new MutableQueryGroupFragment(in), in.readLong()); + } + + public static QueryGroup updateExistingQueryGroup(QueryGroup existingGroup, MutableQueryGroupFragment mutableQueryGroupFragment) { + final Map updatedResourceLimits = new HashMap<>(existingGroup.getResourceLimits()); + final Map mutableFragmentResourceLimits = mutableQueryGroupFragment.getResourceLimits(); + if (mutableFragmentResourceLimits != null && !mutableFragmentResourceLimits.isEmpty()) { + updatedResourceLimits.putAll(mutableFragmentResourceLimits); + } + final ResiliencyMode mode = Optional.ofNullable(mutableQueryGroupFragment.getResiliencyMode()) + .orElse(existingGroup.getResiliencyMode()); + return new QueryGroup( + existingGroup.getName(), + existingGroup.get_id(), + new MutableQueryGroupFragment(mode, updatedResourceLimits), + Instant.now().getMillis() + ); } @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(name); out.writeString(_id); - changeableQueryGroup.writeTo(out); + mutableQueryGroupFragment.writeTo(out); out.writeLong(updatedAtInMillis); } @@ -109,8 +127,8 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa builder.startObject(); builder.field(_ID_STRING, _id); builder.field(NAME_STRING, name); - for (String fieldName : ChangeableQueryGroup.acceptedFieldNames) { - changeableQueryGroup.writeField(builder, fieldName); + for (String fieldName : MutableQueryGroupFragment.acceptedFieldNames) { + mutableQueryGroupFragment.writeField(builder, fieldName); } builder.field(UPDATED_AT_STRING, updatedAtInMillis); builder.endObject(); @@ -131,30 +149,30 @@ 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(changeableQueryGroup, that.changeableQueryGroup) + && Objects.equals(mutableQueryGroupFragment, that.mutableQueryGroupFragment) && Objects.equals(_id, that._id) && updatedAtInMillis == that.updatedAtInMillis; } @Override public int hashCode() { - return Objects.hash(name, changeableQueryGroup, updatedAtInMillis, _id); + return Objects.hash(name, mutableQueryGroupFragment, updatedAtInMillis, _id); } public String getName() { return name; } - public ChangeableQueryGroup getChangeableQueryGroup() { - return changeableQueryGroup; + public MutableQueryGroupFragment getMutableQueryGroupFragment() { + return mutableQueryGroupFragment; } public ResiliencyMode getResiliencyMode() { - return getChangeableQueryGroup().getResiliencyMode(); + return getMutableQueryGroupFragment().getResiliencyMode(); } public Map getResourceLimits() { - return getChangeableQueryGroup().getResourceLimits(); + return getMutableQueryGroupFragment().getResourceLimits(); } public String get_id() { @@ -180,7 +198,7 @@ public static Builder builder() { public static class Builder { private String name; private String _id; - private ChangeableQueryGroup changeableQueryGroup; + private MutableQueryGroupFragment mutableQueryGroupFragment; private long updatedAt; private Builder() {} @@ -199,7 +217,7 @@ public static Builder fromXContent(XContentParser parser) throws IOException { } String fieldName = ""; - ChangeableQueryGroup changeableQueryGroup1 = new ChangeableQueryGroup(); + MutableQueryGroupFragment mutableQueryGroupFragment1 = new MutableQueryGroupFragment(); while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { if (token == XContentParser.Token.FIELD_NAME) { fieldName = parser.currentName(); @@ -208,21 +226,21 @@ public static Builder fromXContent(XContentParser parser) throws IOException { builder._id(parser.text()); } else if (fieldName.equals(NAME_STRING)) { builder.name(parser.text()); - } else if (ChangeableQueryGroup.shouldParse(fieldName)) { - changeableQueryGroup1.parseField(parser, fieldName); + } else if (MutableQueryGroupFragment.shouldParse(fieldName)) { + mutableQueryGroupFragment1.parseField(parser, fieldName); } 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 (!ChangeableQueryGroup.shouldParse(fieldName)) { + if (!MutableQueryGroupFragment.shouldParse(fieldName)) { throw new IllegalArgumentException(fieldName + " is not a valid object in QueryGroup"); } - changeableQueryGroup1.parseField(parser, fieldName); + mutableQueryGroupFragment1.parseField(parser, fieldName); } } - return builder.changeableQueryGroup(changeableQueryGroup1); + return builder.mutableQueryGroupFragment(mutableQueryGroupFragment1); } public Builder name(String name) { @@ -235,8 +253,8 @@ public Builder _id(String _id) { return this; } - public Builder changeableQueryGroup(ChangeableQueryGroup changeableQueryGroup) { - this.changeableQueryGroup = changeableQueryGroup; + public Builder mutableQueryGroupFragment(MutableQueryGroupFragment mutableQueryGroupFragment) { + this.mutableQueryGroupFragment = mutableQueryGroupFragment; return this; } @@ -246,11 +264,11 @@ public Builder updatedAt(long updatedAt) { } public QueryGroup build() { - return new QueryGroup(name, _id, changeableQueryGroup, updatedAt); + return new QueryGroup(name, _id, mutableQueryGroupFragment, updatedAt); } - public ChangeableQueryGroup getChangeableQueryGroup() { - return changeableQueryGroup; + public MutableQueryGroupFragment getMutableQueryGroupFragment() { + return mutableQueryGroupFragment; } } } diff --git a/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java b/server/src/main/java/org/opensearch/wlm/MutableQueryGroupFragment.java similarity index 95% rename from server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java rename to server/src/main/java/org/opensearch/wlm/MutableQueryGroupFragment.java index bee3b8fe7827c..8ea240132fea2 100644 --- a/server/src/main/java/org/opensearch/wlm/ChangeableQueryGroup.java +++ b/server/src/main/java/org/opensearch/wlm/MutableQueryGroupFragment.java @@ -27,7 +27,7 @@ * Class to hold the fields that can be updated in a QueryGroup. */ @ExperimentalApi -public class ChangeableQueryGroup extends AbstractDiffable { +public class MutableQueryGroupFragment extends AbstractDiffable { public static final String RESILIENCY_MODE_STRING = "resiliency_mode"; public static final String RESOURCE_LIMITS_STRING = "resource_limits"; @@ -36,15 +36,15 @@ public class ChangeableQueryGroup extends AbstractDiffable public static final List acceptedFieldNames = List.of(RESILIENCY_MODE_STRING, RESOURCE_LIMITS_STRING); - public ChangeableQueryGroup() {} + public MutableQueryGroupFragment() {} - public ChangeableQueryGroup(ResiliencyMode resiliencyMode, Map resourceLimits) { + public MutableQueryGroupFragment(ResiliencyMode resiliencyMode, Map resourceLimits) { validateResourceLimits(resourceLimits); this.resiliencyMode = resiliencyMode; this.resourceLimits = resourceLimits; } - public ChangeableQueryGroup(StreamInput in) throws IOException { + public MutableQueryGroupFragment(StreamInput in) throws IOException { if (in.readBoolean()) { resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble); } else { @@ -166,7 +166,7 @@ public static void validateResourceLimits(Map resourceLimi public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - ChangeableQueryGroup that = (ChangeableQueryGroup) o; + MutableQueryGroupFragment that = (MutableQueryGroupFragment) o; return Objects.equals(resiliencyMode, that.resiliencyMode) && Objects.equals(resourceLimits, that.resourceLimits); } 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 6e8abdd19a71f..3f8d231ffb91e 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -15,7 +15,7 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.test.AbstractDiffableSerializationTestCase; -import org.opensearch.wlm.ChangeableQueryGroup; +import org.opensearch.wlm.MutableQueryGroupFragment; import org.opensearch.wlm.ResourceType; import java.io.IOException; @@ -34,7 +34,7 @@ public void testToXContent() throws IOException { new QueryGroup( "test", "ajakgakg983r92_4242", - new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), 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 f60276ec9858a..ce1b1270fc94e 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -15,8 +15,8 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.test.AbstractSerializingTestCase; -import org.opensearch.wlm.ChangeableQueryGroup; -import org.opensearch.wlm.ChangeableQueryGroup.ResiliencyMode; +import org.opensearch.wlm.MutableQueryGroupFragment; +import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; import org.joda.time.Instant; @@ -34,7 +34,7 @@ static QueryGroup createRandomQueryGroup(String _id) { String name = randomAlphaOfLength(10); Map resourceLimit = new HashMap<>(); resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false)); - return new QueryGroup(name, _id, new ChangeableQueryGroup(randomMode(), resourceLimit), Instant.now().getMillis()); + return new QueryGroup(name, _id, new MutableQueryGroupFragment(randomMode(), resourceLimit), Instant.now().getMillis()); } private static ResiliencyMode randomMode() { @@ -72,21 +72,31 @@ protected QueryGroup createTestInstance() { public void testNullName() { assertThrows( NullPointerException.class, - () -> new QueryGroup(null, "_id", new ChangeableQueryGroup(randomMode(), Collections.emptyMap()), Instant.now().getMillis()) + () -> new QueryGroup( + null, + "_id", + new MutableQueryGroupFragment(randomMode(), Collections.emptyMap()), + Instant.now().getMillis() + ) ); } public void testNullId() { assertThrows( NullPointerException.class, - () -> new QueryGroup("Dummy", null, new ChangeableQueryGroup(randomMode(), Collections.emptyMap()), Instant.now().getMillis()) + () -> new QueryGroup( + "Dummy", + null, + new MutableQueryGroupFragment(randomMode(), Collections.emptyMap()), + Instant.now().getMillis() + ) ); } public void testNullResourceLimits() { assertThrows( NullPointerException.class, - () -> new QueryGroup("analytics", "_id", new ChangeableQueryGroup(randomMode(), null), Instant.now().getMillis()) + () -> new QueryGroup("analytics", "_id", new MutableQueryGroupFragment(randomMode(), null), Instant.now().getMillis()) ); } @@ -96,7 +106,7 @@ public void testEmptyResourceLimits() { () -> new QueryGroup( "analytics", "_id", - new ChangeableQueryGroup(randomMode(), Collections.emptyMap()), + new MutableQueryGroupFragment(randomMode(), Collections.emptyMap()), Instant.now().getMillis() ) ); @@ -108,14 +118,14 @@ public void testIllegalQueryGroupMode() { () -> new QueryGroup( "analytics", "_id", - new ChangeableQueryGroup(null, Map.of(ResourceType.MEMORY, 0.4)), + new MutableQueryGroupFragment(null, Map.of(ResourceType.MEMORY, 0.4)), Instant.now().getMillis() ) ); } public void testQueryGroupInitiation() { - QueryGroup queryGroup = new QueryGroup("analytics", new ChangeableQueryGroup(randomMode(), Map.of(ResourceType.MEMORY, 0.4))); + QueryGroup queryGroup = new QueryGroup("analytics", new MutableQueryGroupFragment(randomMode(), Map.of(ResourceType.MEMORY, 0.4))); assertNotNull(queryGroup.getName()); assertNotNull(queryGroup.get_id()); assertNotNull(queryGroup.getResourceLimits()); @@ -128,9 +138,12 @@ public void testQueryGroupInitiation() { public void testIllegalQueryGroupName() { assertThrows( NullPointerException.class, - () -> new QueryGroup("a".repeat(51), "_id", new ChangeableQueryGroup(), Instant.now().getMillis()) + () -> new QueryGroup("a".repeat(51), "_id", new MutableQueryGroupFragment(), Instant.now().getMillis()) + ); + assertThrows( + NullPointerException.class, + () -> new QueryGroup("", "_id", new MutableQueryGroupFragment(), Instant.now().getMillis()) ); - assertThrows(NullPointerException.class, () -> new QueryGroup("", "_id", new ChangeableQueryGroup(), Instant.now().getMillis())); } @@ -140,7 +153,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { () -> new QueryGroup( "analytics", "_id", - new ChangeableQueryGroup(randomMode(), Map.of(ResourceType.MEMORY, randomDoubleBetween(1.1, 1.8, false))), + new MutableQueryGroupFragment(randomMode(), Map.of(ResourceType.MEMORY, randomDoubleBetween(1.1, 1.8, false))), Instant.now().getMillis() ) ); @@ -150,7 +163,7 @@ public void testValidQueryGroup() { QueryGroup queryGroup = new QueryGroup( "analytics", "_id", - new ChangeableQueryGroup(randomMode(), Map.of(ResourceType.MEMORY, randomDoubleBetween(0.01, 0.8, false))), + new MutableQueryGroupFragment(randomMode(), Map.of(ResourceType.MEMORY, randomDoubleBetween(0.01, 0.8, false))), Instant.ofEpochMilli(1717187289).getMillis() ); @@ -169,7 +182,7 @@ public void testToXContent() throws IOException { QueryGroup queryGroup = new QueryGroup( "TestQueryGroup", queryGroupId, - new ChangeableQueryGroup(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.30, ResourceType.MEMORY, 0.40)), + new MutableQueryGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.30, ResourceType.MEMORY, 0.40)), currentTimeInMillis ); XContentBuilder builder = JsonXContent.contentBuilder(); diff --git a/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java b/server/src/test/java/org/opensearch/wlm/MutableQueryGroupFragmentTests.java similarity index 64% rename from server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java rename to server/src/test/java/org/opensearch/wlm/MutableQueryGroupFragmentTests.java index 66605d4e91f0b..cfe53ddbd2c14 100644 --- a/server/src/test/java/org/opensearch/wlm/ChangeableQueryGroupTests.java +++ b/server/src/test/java/org/opensearch/wlm/MutableQueryGroupFragmentTests.java @@ -16,44 +16,46 @@ import java.util.HashMap; import java.util.Map; -public class ChangeableQueryGroupTests extends OpenSearchTestCase { +public class MutableQueryGroupFragmentTests extends OpenSearchTestCase { public void testSerializationDeserialization() throws IOException { Map resourceLimits = new HashMap<>(); resourceLimits.put(ResourceType.CPU, 0.5); resourceLimits.put(ResourceType.MEMORY, 0.75); - ChangeableQueryGroup changeableQueryGroup = new ChangeableQueryGroup(ChangeableQueryGroup.ResiliencyMode.SOFT, resourceLimits); + MutableQueryGroupFragment mutableQueryGroupFragment = new MutableQueryGroupFragment( + MutableQueryGroupFragment.ResiliencyMode.SOFT, + resourceLimits + ); BytesStreamOutput out = new BytesStreamOutput(); - changeableQueryGroup.writeTo(out); + mutableQueryGroupFragment.writeTo(out); StreamInput in = out.bytes().streamInput(); - ChangeableQueryGroup deserializedGroup = new ChangeableQueryGroup(in); - assertEquals(changeableQueryGroup, deserializedGroup); + MutableQueryGroupFragment deserializedGroup = new MutableQueryGroupFragment(in); + assertEquals(mutableQueryGroupFragment, deserializedGroup); } public void testSerializationDeserializationWithNull() throws IOException { - ChangeableQueryGroup changeableQueryGroup = new ChangeableQueryGroup(); + MutableQueryGroupFragment mutableQueryGroupFragment = new MutableQueryGroupFragment(); BytesStreamOutput out = new BytesStreamOutput(); - changeableQueryGroup.writeTo(out); + mutableQueryGroupFragment.writeTo(out); StreamInput in = out.bytes().streamInput(); - ChangeableQueryGroup deserializedGroup = new ChangeableQueryGroup(in); + MutableQueryGroupFragment deserializedGroup = new MutableQueryGroupFragment(in); assertEquals(0, deserializedGroup.getResourceLimits().size()); - assertEquals(changeableQueryGroup.getResiliencyMode(), deserializedGroup.getResiliencyMode()); + assertEquals(mutableQueryGroupFragment.getResiliencyMode(), deserializedGroup.getResiliencyMode()); } public void testValidateResourceLimits() { Map invalidLimits = new HashMap<>(); invalidLimits.put(ResourceType.CPU, 1.5); - Exception exception = assertThrows( - IllegalArgumentException.class, - () -> { ChangeableQueryGroup.validateResourceLimits(invalidLimits); } - ); + Exception exception = assertThrows(IllegalArgumentException.class, () -> { + MutableQueryGroupFragment.validateResourceLimits(invalidLimits); + }); String expectedMessage = "resource value should be greater than 0 and less or equal to 1.0"; String actualMessage = exception.getMessage(); assertTrue(actualMessage.contains(expectedMessage)); } public void testSetMethodsWithNullAndEmptyValues() { - ChangeableQueryGroup queryGroup = new ChangeableQueryGroup(); + MutableQueryGroupFragment queryGroup = new MutableQueryGroupFragment(); queryGroup.setResiliencyMode(null); assertNull(queryGroup.getResiliencyMode()); queryGroup.setResourceLimits(null);