address comments

Signed-off-by: Ruirui Zhang <[email protected]>
opensearch-project · Jul 15, 2024 · a846e58 · a846e58
1 parent 7badf0a
commit a846e58
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 30 deletions.
diff --git a/...ad-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/...ad-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java
@@ -42,7 +42,7 @@ public class CreateQueryGroupRequest extends ActionRequest implements Writeable.
     private String name;
     private String _id;
     private ResiliencyMode resiliencyMode;
-    private Map<String, Object> resourceLimits;
+    private Map<ResourceType, Object> resourceLimits;
     private long updatedAtInMillis;
 
     /**
@@ -57,12 +57,7 @@ public CreateQueryGroupRequest() {}
     public CreateQueryGroupRequest(QueryGroup queryGroup) {
         this.name = queryGroup.getName();
         this._id = queryGroup.get_id();
-        Map<ResourceType, Object> resourceTypesMap = queryGroup.getResourceLimits();
-        Map<String, Object> resourceLimits_ = new HashMap<>();
-        for (Map.Entry<ResourceType, Object> resource : resourceTypesMap.entrySet()) {
-            resourceLimits_.put(resource.getKey().getName(), resource.getValue());
-        }
-        this.resourceLimits = resourceLimits_;
+        this.resourceLimits = queryGroup.getResourceLimits();
         this.resiliencyMode = queryGroup.getResiliencyMode();
         this.updatedAtInMillis = queryGroup.getUpdatedAtInMillis();
     }
@@ -79,7 +74,7 @@ public CreateQueryGroupRequest(
         String name,
         String _id,
         ResiliencyMode mode,
-        Map<String, Object> resourceLimits,
+        Map<ResourceType, Object> resourceLimits,
         long updatedAtInMillis
     ) {
         this.name = name;
@@ -98,7 +93,7 @@ public CreateQueryGroupRequest(StreamInput in) throws IOException {
         name = in.readString();
         _id = in.readString();
         resiliencyMode = ResiliencyMode.fromName(in.readString());
-        resourceLimits = in.readMap();
+        resourceLimits = in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readGenericValue);
         updatedAtInMillis = in.readLong();
     }
 
@@ -124,12 +119,10 @@ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws
         XContentParser.Token token;
         String fieldName = "";
         String name = null;
-        String _id = UUIDs.randomBase64UUID();
         ResiliencyMode mode = null;
-        long updatedAtInMillis = Instant.now().getMillis();
 
         // Map to hold resources
-        final Map<String, Object> resourceLimits = new HashMap<>();
+        final Map<ResourceType, Object> resourceLimits = new HashMap<>();
         while ((token = parser.nextToken()) != null) {
             if (token == XContentParser.Token.FIELD_NAME) {
                 fieldName = parser.currentName();
@@ -151,12 +144,12 @@ public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws
                     if (token == XContentParser.Token.FIELD_NAME) {
                         fieldName = parser.currentName();
                     } else {
-                        resourceLimits.put(fieldName, parser.doubleValue());
+                        resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue());
                     }
                 }
             }
         }
-        return new CreateQueryGroupRequest(name, _id, mode, resourceLimits, updatedAtInMillis);
+        return new CreateQueryGroupRequest(name, UUIDs.randomBase64UUID(), mode, resourceLimits, Instant.now().getMillis());
     }
 
     @Override
@@ -182,15 +175,15 @@ public void setName(String name) {
     /**
      * resourceLimits getter
      */
-    public Map<String, Object> getResourceLimits() {
+    public Map<ResourceType, Object> getResourceLimits() {
         return resourceLimits;
     }
 
     /**
      * resourceLimits setter
      * @param resourceLimits - resourceLimit to be set
      */
-    public void setResourceLimits(Map<String, Object> resourceLimits) {
+    public void setResourceLimits(Map<ResourceType, Object> resourceLimits) {
         this.resourceLimits = resourceLimits;
     }
 
@@ -215,7 +208,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeString(name);
         out.writeString(_id);
         out.writeString(resiliencyMode.getName());
-        out.writeMap(resourceLimits);
+        out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue);
         out.writeLong(updatedAtInMillis);
     }
 

diff --git a/...ement/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java b/...ement/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java
@@ -14,14 +14,10 @@
 import org.opensearch.common.inject.Inject;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.plugin.wlm.action.service.Persistable;
-import org.opensearch.search.ResourceType;
 import org.opensearch.tasks.Task;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
 
-import java.util.HashMap;
-import java.util.Map;
-
 import static org.opensearch.cluster.metadata.QueryGroup.builder;
 
 /**
@@ -58,15 +54,10 @@ public TransportCreateQueryGroupAction(
 
     @Override
     protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener<CreateQueryGroupResponse> listener) {
-        Map<ResourceType, Object> resourceTypesMap = new HashMap<>();
-        Map<String, Object> resourceLimitsStringMap = request.getResourceLimits();
-        for (Map.Entry<String, Object> resource : resourceLimitsStringMap.entrySet()) {
-            resourceTypesMap.put(ResourceType.fromName(resource.getKey()), resource.getValue());
-        }
         QueryGroup queryGroup = builder().name(request.getName())
             ._id(request.get_id())
             .mode(request.getResiliencyMode().getName())
-            .resourceLimits(resourceTypesMap)
+            .resourceLimits(request.getResourceLimits())
             .updatedAt(request.getUpdatedAtInMillis())
             .build();
         threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> queryGroupPersistenceService.persist(queryGroup, listener));

diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java
@@ -7,6 +7,6 @@
  */
 
 /**
- * Base Package of CRUD API of QueryGroup
+ * Base package for CRUD API of QueryGroup
  */
 package org.opensearch.plugin.wlm.action;
diff --git a/...ement/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java b/...ement/src/main/java/org/opensearch/plugin/wlm/action/rest/RestCreateQueryGroupAction.java
@@ -50,7 +50,7 @@ public String getName() {
      */
     @Override
     public List<Route> routes() {
-        return List.of(new Route(POST, "_query_group/"), new Route(PUT, "_query_group/"));
+        return List.of(new Route(POST, "_wlm/_query_group/"), new Route(PUT, "_wlm/_query_group/"));
     }
 
     @Override

diff --git a/...rkload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java b/...rkload-management/src/test/java/org/opensearch/plugin/wlm/action/QueryGroupTestUtils.java
@@ -18,6 +18,7 @@
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.plugin.wlm.action.service.QueryGroupPersistenceService;
+import org.opensearch.search.ResourceType;
 import org.opensearch.threadpool.ThreadPool;
 
 import java.util.ArrayList;
@@ -119,7 +120,7 @@ public static List<Object> preparePersistenceServiceSetup(Map<String, QueryGroup
         return List.of(queryGroupPersistenceService, clusterState);
     }
 
-    public static void compareResourceLimits(Map<String, Object> resourceLimitMapOne, Map<String, Object> resourceLimitMapTwo) {
+    public static void compareResourceLimits(Map<ResourceType, Object> resourceLimitMapOne, Map<ResourceType, Object> resourceLimitMapTwo) {
         assertTrue(resourceLimitMapOne.keySet().containsAll(resourceLimitMapTwo.keySet()));
         assertTrue(resourceLimitMapOne.values().containsAll(resourceLimitMapTwo.values()));
     }