add QueryGroup.toXContent test case

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java

@@ -10,20 +10,18 @@
 
 import org.opensearch.cluster.AbstractDiffable;
 import org.opensearch.cluster.Diff;
-import org.opensearch.common.ResourceType;
 import org.opensearch.common.UUIDs;
-import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
 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.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 
@@ -39,23 +37,22 @@
  *              "updatedAt": 4513232415
  * }
  */
-@ExperimentalApi
+@PublicApi(since = "2.15")
 public class QueryGroup extends AbstractDiffable<QueryGroup> implements ToXContentObject {
 
     public static final int MAX_CHARS_ALLOWED_IN_NAME = 50;
     private final String name;
     private final String _id;
-    private final QueryGroupMode resiliencyMode;
+    private final ResiliencyMode resiliencyMode;
     // It is an epoch in millis
     private final long updatedAtInMillis;
     private final Map<ResourceType, Object> resourceLimits;
 
-
-    public QueryGroup(String name, QueryGroupMode resiliencyMode, Map<ResourceType, Object> resourceLimits) {
+    public QueryGroup(String name, ResiliencyMode resiliencyMode, Map<ResourceType, Object> resourceLimits) {
         this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis());
     }
 
-    public QueryGroup(String name, String _id, QueryGroupMode resiliencyMode, Map<ResourceType, Object> resourceLimits, long updatedAt) {
+    public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map<ResourceType, Object> resourceLimits, long updatedAt) {
         Objects.requireNonNull(name, "QueryGroup.name can't be null");
         Objects.requireNonNull(resourceLimits, "QueryGroup.resourceLimits can't be null");
         Objects.requireNonNull(resiliencyMode, "QueryGroup.resiliencyMode can't be null");
@@ -94,7 +91,7 @@ public QueryGroup(StreamInput in) throws IOException {
         this(
             in.readString(),
             in.readString(),
-            QueryGroupMode.fromName(in.readString()),
+            ResiliencyMode.fromName(in.readString()),
             in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readGenericValue),
             in.readLong()
         );
@@ -141,8 +138,10 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa
         builder.field("updatedAt", updatedAtInMillis);
         // write resource limits
         builder.startObject("resourceLimits");
-        for (Map.Entry<ResourceType, Object> resourceLimit : resourceLimits.entrySet()) {
-            builder.field(resourceLimit.getKey().getName(), resourceLimit.getValue());
+        for (ResourceType resourceType : ResourceType.values()) {
+            if (resourceLimits.containsKey(resourceType)) {
+                builder.field(resourceType.getName(), resourceLimits.get(resourceType));
+            }
         }
         builder.endObject();
 
@@ -181,10 +180,12 @@ public static QueryGroup fromXContent(final XContentParser parser) throws IOExce
                 } else {
                     throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup");
                 }
-            } else if (token == XContentParser.Token.START_OBJECT ) {
+            } else if (token == XContentParser.Token.START_OBJECT) {
 
                 if (!fieldName.equals("resourceLimits")) {
-                    throw new IllegalArgumentException("was expecting the { " + " but found " + token);
+                    throw new IllegalArgumentException(
+                        "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token
+                    );
                 }
 
                 while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
@@ -225,7 +226,7 @@ public String getName() {
         return name;
     }
 
-    public QueryGroupMode getResiliencyMode() {
+    public ResiliencyMode getResiliencyMode() {
         return resiliencyMode;
     }
 
@@ -250,26 +251,29 @@ public static Builder builder() {
     }
 
     /**
-     * This enum models the different sandbox modes
+     * 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 QueryGroupMode {
+    @PublicApi(since = "2.15")
+    public enum ResiliencyMode {
         SOFT("soft"),
         ENFORCED("enforced"),
         MONITOR("monitor");
 
         private final String name;
 
-        QueryGroupMode(String mode) {
+        ResiliencyMode(String mode) {
             this.name = mode;
         }
 
         public String getName() {
             return name;
         }
 
-        public static QueryGroupMode fromName(String s) {
-            for (QueryGroupMode mode : values()) {
+        public static ResiliencyMode fromName(String s) {
+            for (ResiliencyMode mode : values()) {
                 if (mode.getName().equalsIgnoreCase(s)) return mode;
 
             }
@@ -281,11 +285,11 @@ public static QueryGroupMode fromName(String s) {
     /**
      * Builder class for {@link QueryGroup}
      */
-    @ExperimentalApi
+    @PublicApi(since = "2.15")
     public static class Builder {
         private String name;
         private String _id;
-        private QueryGroupMode resiliencyMode;
+        private ResiliencyMode resiliencyMode;
         private long updatedAt;
         private Map<ResourceType, Object> resourceLimits;
 
@@ -302,7 +306,7 @@ public Builder _id(String _id) {
         }
 
         public Builder mode(String mode) {
-            this.resiliencyMode = QueryGroupMode.fromName(mode);
+            this.resiliencyMode = ResiliencyMode.fromName(mode);
             return this;
         }
 

server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java

@@ -12,7 +12,6 @@
 import org.opensearch.cluster.Diff;
 import org.opensearch.cluster.DiffableUtils;
 import org.opensearch.cluster.NamedDiff;
-import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.io.stream.StreamInput;
@@ -42,7 +41,6 @@
  *     }
  * }
  */
-@ExperimentalApi
 public class QueryGroupMetadata implements Metadata.Custom {
     public static final String TYPE = "queryGroups";
     private static final ParseField QUERY_GROUP_FIELD = new ParseField("queryGroups");

server/src/main/java/org/opensearch/search/ResourceType.java

@@ -8,12 +8,18 @@
 
 package org.opensearch.search;
 
+import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
 /**
  * Enum to hold the resource type
  */
+@PublicApi(since = "1.3")
 public enum ResourceType {
     CPU("cpu"),
-    JVM("jvm");
+    MEMORY("memory");
 
     private final String name;
 
@@ -35,7 +41,11 @@ public static ResourceType fromName(String s) {
         throw new IllegalArgumentException("Unknown resource type: [" + s + "]");
     }
 
-    private String getName() {
+    public static void writeTo(StreamOutput out, ResourceType resourceType) throws IOException {
+        out.writeString(resourceType.getName());
+    }
+
+    public String getName() {
         return name;
     }
 }

server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java

@@ -71,7 +71,7 @@ public class SearchBackpressureService extends AbstractLifecycleComponent implem
         TaskResourceUsageTrackerType.CPU_USAGE_TRACKER,
         (nodeDuressTrackers) -> nodeDuressTrackers.isResourceInDuress(ResourceType.CPU),
         TaskResourceUsageTrackerType.HEAP_USAGE_TRACKER,
-        (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.JVM),
+        (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.MEMORY),
         TaskResourceUsageTrackerType.ELAPSED_TIME_TRACKER,
         (nodeDuressTrackers) -> true
     );
@@ -105,7 +105,7 @@ public SearchBackpressureService(
                     )
                 );
                 put(
-                    ResourceType.JVM,
+                    ResourceType.MEMORY,
                     new NodeDuressTracker(
                         () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings()
                             .getHeapThreshold(),

server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java

@@ -331,9 +331,10 @@ public void testRejectsDuplicateExistingShardsAllocatorName() {
 
     public void testQueryGroupMetadataRegister() {
         List<NamedWriteableRegistry.Entry> customEntries = ClusterModule.getNamedWriteables();
-        assertTrue(customEntries.stream().anyMatch(entry ->
-            entry.categoryClass == Metadata.Custom.class && entry.name.equals(QueryGroupMetadata.TYPE)
-        ));
+        assertTrue(
+            customEntries.stream()
+                .anyMatch(entry -> entry.categoryClass == Metadata.Custom.class && entry.name.equals(QueryGroupMetadata.TYPE))
+        );
     }
 
     private static ClusterPlugin existingShardsAllocatorPlugin(final String allocatorName) {

server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java

@@ -8,9 +8,13 @@
 
 package org.opensearch.cluster.metadata;
 
-import org.opensearch.common.ResourceType;
+import org.opensearch.common.UUIDs;
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.search.ResourceType;
 import org.opensearch.test.AbstractSerializingTestCase;
 import org.joda.time.Instant;
 
@@ -22,20 +26,20 @@
 
 public class QueryGroupTests extends AbstractSerializingTestCase<QueryGroup> {
 
-    private static final List<QueryGroup.QueryGroupMode> allowedModes = List.of(
-        QueryGroup.QueryGroupMode.SOFT,
-        QueryGroup.QueryGroupMode.ENFORCED,
-        QueryGroup.QueryGroupMode.MONITOR
+    private static final List<QueryGroup.ResiliencyMode> allowedModes = List.of(
+        QueryGroup.ResiliencyMode.SOFT,
+        QueryGroup.ResiliencyMode.ENFORCED,
+        QueryGroup.ResiliencyMode.MONITOR
     );
 
     static QueryGroup createRandomQueryGroup(String _id) {
         String name = randomAlphaOfLength(10);
         Map<ResourceType, Object> resourceLimit = new HashMap<>();
-        resourceLimit.put(ResourceType.HEAP_ALLOCATIONS, randomDoubleBetween(0.0, 0.80, false));
+        resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false));
         return new QueryGroup(name, _id, randomMode(), resourceLimit, Instant.now().getMillis());
     }
 
-    private static QueryGroup.QueryGroupMode randomMode() {
+    private static QueryGroup.ResiliencyMode randomMode() {
         return allowedModes.get(randomIntBetween(0, allowedModes.size() - 1));
     }
 
@@ -95,7 +99,7 @@ public void testEmptyResourceLimits() {
     public void testIllegalQueryGroupMode() {
         assertThrows(
             NullPointerException.class,
-            () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.HEAP_ALLOCATIONS, (Object) 0.4), Instant.now().getMillis())
+            () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, (Object) 0.4), Instant.now().getMillis())
         );
     }
 
@@ -106,7 +110,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() {
                 "analytics",
                 "_id",
                 randomMode(),
-                Map.of(ResourceType.HEAP_ALLOCATIONS, (Object) randomDoubleBetween(1.1, 1.8, false)),
+                Map.of(ResourceType.MEMORY, (Object) randomDoubleBetween(1.1, 1.8, false)),
                 Instant.now().getMillis()
             )
         );
@@ -117,7 +121,7 @@ public void testValidQueryGroup() {
             "analytics",
             "_id",
             randomMode(),
-            Map.of(ResourceType.HEAP_ALLOCATIONS, randomDoubleBetween(0.01, 0.8, false)),
+            Map.of(ResourceType.MEMORY, randomDoubleBetween(0.01, 0.8, false)),
             Instant.ofEpochMilli(1717187289).getMillis()
         );
 
@@ -129,4 +133,26 @@ public void testValidQueryGroup() {
         assertTrue(allowedModes.contains(queryGroup.getResiliencyMode()));
         assertEquals(1717187289, queryGroup.getUpdatedAtInMillis());
     }
+
+    public void testToXContent() throws IOException {
+        long currentTimeInMillis = Instant.now().getMillis();
+        String queryGroupId = UUIDs.randomBase64UUID();
+        QueryGroup queryGroup = new QueryGroup(
+            "TestQueryGroup",
+            queryGroupId,
+            QueryGroup.ResiliencyMode.ENFORCED,
+            Map.of(ResourceType.CPU, 0.30, ResourceType.MEMORY, 0.40),
+            currentTimeInMillis
+        );
+        XContentBuilder builder = JsonXContent.contentBuilder();
+        queryGroup.toXContent(builder, ToXContent.EMPTY_PARAMS);
+        assertEquals(
+            String.format(
+                "{\"_id\":\"%s\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"updatedAt\":%d,\"resourceLimits\":{\"cpu\":0.3,\"memory\":0.4}}",
+                queryGroupId,
+                currentTimeInMillis
+            ),
+            builder.toString()
+        );
+    }
 }