From 769a46fb03930bb841a65e8e8d9b2da77df54175 Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Wed, 26 Jun 2024 16:44:00 -0700 Subject: [PATCH] add QueryGroup.toXContent test case Signed-off-by: Kaushal Kumar --- .../cluster/metadata/QueryGroup.java | 52 ++++++++++--------- .../cluster/metadata/QueryGroupMetadata.java | 2 - .../org/opensearch/common/ResourceType.java | 46 ---------------- .../org/opensearch/search/ResourceType.java | 14 ++++- .../SearchBackpressureService.java | 4 +- .../cluster/ClusterModuleTests.java | 7 +-- .../cluster/metadata/QueryGroupTests.java | 46 ++++++++++++---- .../SearchBackpressureServiceTests.java | 12 ++--- .../trackers/NodeDuressTrackersTests.java | 8 +-- 9 files changed, 92 insertions(+), 99 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/common/ResourceType.java 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 4bef98af6cd1f..ca15eb2446759 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/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 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 resourceLimits; - - public QueryGroup(String name, QueryGroupMode resiliencyMode, Map resourceLimits) { + public QueryGroup(String name, ResiliencyMode resiliencyMode, Map resourceLimits) { this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis()); } - public QueryGroup(String name, String _id, QueryGroupMode resiliencyMode, Map resourceLimits, long updatedAt) { + public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits, long updatedAt) { Objects.requireNonNull(name, "QueryGroup.name can't be null"); Objects.requireNonNull(resourceLimits, "QueryGroup.resourceLimits can't be null"); Objects.requireNonNull(resiliencyMode, "QueryGroup.resiliencyMode can't be null"); @@ -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 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,17 +251,20 @@ 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; } @@ -268,8 +272,8 @@ 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 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; } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java index 342708bdf73c4..ae6efd25a43f0 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java +++ b/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"); diff --git a/server/src/main/java/org/opensearch/common/ResourceType.java b/server/src/main/java/org/opensearch/common/ResourceType.java deleted file mode 100644 index 4776224219a32..0000000000000 --- a/server/src/main/java/org/opensearch/common/ResourceType.java +++ /dev/null @@ -1,46 +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.common; - -import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.core.common.io.stream.StreamOutput; - -import java.io.IOException; - -/** - * Enum to specify the system level resources - */ -@ExperimentalApi -public enum ResourceType { - CPU("cpu"), - HEAP_ALLOCATIONS("heap_allocations"); - - final String name; - - ResourceType(String s) { - this.name = s; - } - - public static ResourceType fromName(String s) { - for (ResourceType resourceType : values()) { - if (resourceType.getName().equals(s)) { - return resourceType; - } - } - throw new IllegalArgumentException(s + " is not a valid ResourceType"); - } - - public String getName() { - return name; - } - - public static void writeTo(StreamOutput out, ResourceType val) throws IOException { - out.writeString(val.getName()); - } -} diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index 5bbcd7de1c2ce..57140405ccf0b 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/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; } } diff --git a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java index 3e8ed3070e4ef..c26c5d63a3573 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java +++ b/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(), diff --git a/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java b/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java index 5e9a91b69569a..97706927ba857 100644 --- a/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java +++ b/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java @@ -331,9 +331,10 @@ public void testRejectsDuplicateExistingShardsAllocatorName() { public void testQueryGroupMetadataRegister() { List 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) { 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 776be05335737..170288b299e5b 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/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 { - private static final List allowedModes = List.of( - QueryGroup.QueryGroupMode.SOFT, - QueryGroup.QueryGroupMode.ENFORCED, - QueryGroup.QueryGroupMode.MONITOR + private static final List allowedModes = List.of( + QueryGroup.ResiliencyMode.SOFT, + QueryGroup.ResiliencyMode.ENFORCED, + QueryGroup.ResiliencyMode.MONITOR ); static QueryGroup createRandomQueryGroup(String _id) { String name = randomAlphaOfLength(10); Map 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() + ); + } } diff --git a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java index 43df482fcc2ae..15d0fcd10d701 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java @@ -57,7 +57,7 @@ import java.util.function.LongSupplier; import static org.opensearch.search.ResourceType.CPU; -import static org.opensearch.search.ResourceType.JVM; +import static org.opensearch.search.ResourceType.MEMORY; import static org.opensearch.search.backpressure.SearchBackpressureTestHelpers.createMockTaskWithResourceStats; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; @@ -102,7 +102,7 @@ public void testIsNodeInDuress() { EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(ResourceType.JVM, heapUsageTracker); + put(ResourceType.MEMORY, heapUsageTracker); put(ResourceType.CPU, cpuUsageTracker); } }; @@ -233,7 +233,7 @@ public void testSearchTaskInFlightCancellation() { EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(JVM, heapUsageTracker); + put(MEMORY, heapUsageTracker); put(CPU, mockNodeDuressTracker); } }; @@ -308,7 +308,7 @@ public void testSearchShardTaskInFlightCancellation() { EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(JVM, new NodeDuressTracker(() -> false, () -> 3)); + put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); put(CPU, mockNodeDuressTracker); } }; @@ -401,7 +401,7 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(JVM, new NodeDuressTracker(() -> false, () -> 3)); + put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); put(CPU, new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -495,7 +495,7 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(JVM, new NodeDuressTracker(() -> false, () -> 3)); + put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); put(CPU, new NodeDuressTracker(() -> true, () -> 3)); } }; diff --git a/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java b/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java index 2db251ee461db..801576bdf89d4 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java @@ -19,7 +19,7 @@ public class NodeDuressTrackersTests extends OpenSearchTestCase { public void testNodeNotInDuress() { EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.JVM, new NodeDuressTracker(() -> false, () -> 2)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> false, () -> 2)); put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 2)); } }; @@ -34,7 +34,7 @@ public void testNodeNotInDuress() { public void testNodeInDuressWhenHeapInDuress() { EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.JVM, new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> true, () -> 3)); put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 1)); } }; @@ -51,7 +51,7 @@ public void testNodeInDuressWhenHeapInDuress() { public void testNodeInDuressWhenCPUInDuress() { EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.JVM, new NodeDuressTracker(() -> false, () -> 1)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> false, () -> 1)); put(ResourceType.CPU, new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -68,7 +68,7 @@ public void testNodeInDuressWhenCPUInDuress() { public void testNodeInDuressWhenCPUAndHeapInDuress() { EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.JVM, new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> true, () -> 3)); put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 3)); } };