diff --git a/CHANGELOG.md b/CHANGELOG.md index cfb12ee853a8f..4d0990db31d20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - [Remote Store] Rate limiter for remote store low priority uploads ([#14374](https://github.com/opensearch-project/OpenSearch/pull/14374/)) - Apply the date histogram rewrite optimization to range aggregation ([#13865](https://github.com/opensearch-project/OpenSearch/pull/13865)) - [Writable Warm] Add composite directory implementation and integrate it with FileCache ([12782](https://github.com/opensearch-project/OpenSearch/pull/12782)) +- [Workload Management] Add QueryGroup schema ([13669](https://github.com/opensearch-project/OpenSearch/pull/13669)) - Add batching supported processor base type AbstractBatchingProcessor ([#14554](https://github.com/opensearch-project/OpenSearch/pull/14554)) - Fix race condition while parsing derived fields from search definition ([14445](https://github.com/opensearch-project/OpenSearch/pull/14445)) - Add allowlist setting for ingest-common and search-pipeline-common processors ([#14439](https://github.com/opensearch-project/OpenSearch/issues/14439)) diff --git a/server/src/main/java/org/opensearch/cluster/ClusterModule.java b/server/src/main/java/org/opensearch/cluster/ClusterModule.java index c7fd263bda56a..bb51c42252448 100644 --- a/server/src/main/java/org/opensearch/cluster/ClusterModule.java +++ b/server/src/main/java/org/opensearch/cluster/ClusterModule.java @@ -48,6 +48,7 @@ import org.opensearch.cluster.metadata.MetadataIndexTemplateService; import org.opensearch.cluster.metadata.MetadataMappingService; import org.opensearch.cluster.metadata.MetadataUpdateSettingsService; +import org.opensearch.cluster.metadata.QueryGroupMetadata; import org.opensearch.cluster.metadata.RepositoriesMetadata; import org.opensearch.cluster.metadata.ViewMetadata; import org.opensearch.cluster.metadata.WeightedRoutingMetadata; @@ -214,6 +215,8 @@ public static List getNamedWriteables() { DecommissionAttributeMetadata::new, DecommissionAttributeMetadata::readDiffFrom ); + + registerMetadataCustom(entries, QueryGroupMetadata.TYPE, QueryGroupMetadata::new, QueryGroupMetadata::readDiffFrom); // Task Status (not Diffable) entries.add(new Entry(Task.Status.class, PersistentTasksNodeService.Status.NAME, PersistentTasksNodeService.Status::new)); return entries; diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java index e3f63b1c27b83..2a54f6444ffda 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -1368,6 +1368,25 @@ public Builder removeDataStream(String name) { return this; } + public Builder queryGroups(final Map queryGroups) { + this.customs.put(QueryGroupMetadata.TYPE, new QueryGroupMetadata(queryGroups)); + return this; + } + + public Builder put(final QueryGroup queryGroup) { + Objects.requireNonNull(queryGroup, "queryGroup should not be null"); + Map existing = new HashMap<>(getQueryGroups()); + existing.put(queryGroup.get_id(), queryGroup); + return queryGroups(existing); + } + + private Map getQueryGroups() { + return Optional.ofNullable(this.customs.get(QueryGroupMetadata.TYPE)) + .map(o -> (QueryGroupMetadata) o) + .map(QueryGroupMetadata::queryGroups) + .orElse(Collections.emptyMap()); + } + private Map getViews() { return Optional.ofNullable(customs.get(ViewMetadata.TYPE)) .map(o -> (ViewMetadata) o) diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java new file mode 100644 index 0000000000000..beaab198073df --- /dev/null +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -0,0 +1,317 @@ +/* + * 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.cluster.metadata; + +import org.opensearch.cluster.AbstractDiffable; +import org.opensearch.cluster.Diff; +import org.opensearch.common.UUIDs; +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.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.Map; +import java.util.Objects; + +/** + * Class to define the QueryGroup schema + * { + * "_id": "fafjafjkaf9ag8a9ga9g7ag0aagaga", + * "resourceLimits": { + * "jvm": 0.4 + * }, + * "resiliency_mode": "enforced", + * "name": "analytics", + * "updatedAt": 4513232415 + * } + */ +@ExperimentalApi +public class QueryGroup extends AbstractDiffable implements ToXContentObject { + + 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; + + public QueryGroup(String name, ResiliencyMode resiliencyMode, Map resourceLimits) { + this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis()); + } + + public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits, long updatedAt) { + 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(_id, "QueryGroup._id can't be null"); + + if (name.length() > MAX_CHARS_ALLOWED_IN_NAME) { + throw new IllegalArgumentException("QueryGroup.name shouldn't be more than 50 chars long"); + } + + if (resourceLimits.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.updatedAtInMillis = updatedAt; + } + + private static boolean isValid(long updatedAt) { + long minValidTimestamp = Instant.ofEpochMilli(0L).getMillis(); + + // Use Instant.now() to get the current time in seconds since epoch + long currentSeconds = Instant.now().getMillis(); + + // Check if the timestamp is within a reasonable range + return minValidTimestamp <= updatedAt && updatedAt <= currentSeconds; + } + + public QueryGroup(StreamInput in) throws IOException { + this( + in.readString(), + in.readString(), + ResiliencyMode.fromName(in.readString()), + in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readGenericValue), + in.readLong() + ); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(name); + out.writeString(_id); + out.writeString(resiliencyMode.getName()); + out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue); + out.writeLong(updatedAtInMillis); + } + + private void validateResourceLimits(Map resourceLimits) { + for (Map.Entry resource : resourceLimits.entrySet()) { + Double threshold = (Double) 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, 1.0) > 0) { + throw new IllegalArgumentException("resource value should be less than 1.0"); + } + } + } + + @Override + public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { + builder.startObject(); + builder.field("_id", _id); + builder.field("name", name); + builder.field("resiliency_mode", resiliencyMode.getName()); + builder.field("updatedAt", updatedAtInMillis); + // write resource limits + builder.startObject("resourceLimits"); + for (ResourceType resourceType : ResourceType.values()) { + if (resourceLimits.containsKey(resourceType)) { + builder.field(resourceType.getName(), resourceLimits.get(resourceType)); + } + } + builder.endObject(); + + builder.endObject(); + return builder; + } + + public static QueryGroup fromXContent(final XContentParser parser) throws IOException { + if (parser.currentToken() == null) { // fresh parser? move to the first token + parser.nextToken(); + } + + Builder builder = builder(); + + XContentParser.Token token = parser.currentToken(); + + if (token != XContentParser.Token.START_OBJECT) { + throw new IllegalArgumentException("Expected START_OBJECT token but found [" + parser.currentName() + "]"); + } + + String fieldName = ""; + // Map to hold resources + final Map resourceLimits = new HashMap<>(); + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else if (token.isValue()) { + if (fieldName.equals("_id")) { + builder._id(parser.text()); + } else if (fieldName.equals("name")) { + builder.name(parser.text()); + } else if (fieldName.equals("resiliency_mode")) { + builder.mode(parser.text()); + } else if (fieldName.equals("updatedAt")) { + builder.updatedAt(parser.longValue()); + } else { + throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup"); + } + } else if (token == XContentParser.Token.START_OBJECT) { + + if (!fieldName.equals("resourceLimits")) { + throw new IllegalArgumentException( + "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token + ); + } + + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else { + resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); + } + } + + } + } + builder.resourceLimits(resourceLimits); + return builder.build(); + } + + public static Diff readDiff(final StreamInput in) throws IOException { + return readDiffFrom(QueryGroup::new, in); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + QueryGroup that = (QueryGroup) o; + return Objects.equals(name, that.name) + && Objects.equals(resourceLimits, that.resourceLimits) + && Objects.equals(_id, that._id) + && updatedAtInMillis == that.updatedAtInMillis; + } + + @Override + public int hashCode() { + return Objects.hash(name, resourceLimits, updatedAtInMillis, _id); + } + + public String getName() { + return name; + } + + public ResiliencyMode getResiliencyMode() { + return resiliencyMode; + } + + public Map getResourceLimits() { + return resourceLimits; + } + + public String get_id() { + return _id; + } + + public long getUpdatedAtInMillis() { + return updatedAtInMillis; + } + + /** + * builder method for the {@link QueryGroup} + * @return Builder object + */ + 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} + */ + @ExperimentalApi + public static class Builder { + private String name; + private String _id; + private ResiliencyMode resiliencyMode; + private long updatedAt; + private Map resourceLimits; + + private Builder() {} + + public Builder name(String name) { + this.name = name; + return this; + } + + public Builder _id(String _id) { + this._id = _id; + return this; + } + + public Builder mode(String mode) { + this.resiliencyMode = ResiliencyMode.fromName(mode); + return this; + } + + public Builder updatedAt(long updatedAt) { + this.updatedAt = updatedAt; + return this; + } + + public Builder resourceLimits(Map resourceLimits) { + this.resourceLimits = resourceLimits; + return this; + } + + public QueryGroup build() { + return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); + } + + } +} diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java new file mode 100644 index 0000000000000..79732bc505ee2 --- /dev/null +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java @@ -0,0 +1,185 @@ +/* + * 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.cluster.metadata; + +import org.opensearch.Version; +import org.opensearch.cluster.Diff; +import org.opensearch.cluster.DiffableUtils; +import org.opensearch.cluster.NamedDiff; +import org.opensearch.core.ParseField; +import org.opensearch.core.common.Strings; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParser; + +import java.io.IOException; +import java.util.EnumSet; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import static org.opensearch.cluster.metadata.Metadata.ALL_CONTEXTS; + +/** + * This class holds the QueryGroupMetadata + * sample schema + * { + * "queryGroups": { + * "_id": { + * {@link QueryGroup} + * }, + * ... + * } + * } + */ +public class QueryGroupMetadata implements Metadata.Custom { + public static final String TYPE = "queryGroups"; + private static final ParseField QUERY_GROUP_FIELD = new ParseField("queryGroups"); + + private final Map queryGroups; + + public QueryGroupMetadata(Map queryGroups) { + this.queryGroups = queryGroups; + } + + public QueryGroupMetadata(StreamInput in) throws IOException { + this.queryGroups = in.readMap(StreamInput::readString, QueryGroup::new); + } + + public Map queryGroups() { + return this.queryGroups; + } + + /** + * Returns the name of the writeable object + */ + @Override + public String getWriteableName() { + return TYPE; + } + + /** + * The minimal version of the recipient this object can be sent to + */ + @Override + public Version getMinimalSupportedVersion() { + return Version.V_3_0_0; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeMap(queryGroups, StreamOutput::writeString, (stream, val) -> val.writeTo(stream)); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + for (Map.Entry entry : queryGroups.entrySet()) { + builder.field(entry.getKey(), entry.getValue()); + } + return builder; + } + + public static QueryGroupMetadata fromXContent(XContentParser parser) throws IOException { + Map queryGroupMap = new HashMap<>(); + + if (parser.currentToken() == null) { + parser.nextToken(); + } + + if (parser.currentToken() != XContentParser.Token.START_OBJECT) { + throw new IllegalArgumentException( + "QueryGroupMetadata.fromXContent was expecting a { token but found : " + parser.currentToken() + ); + } + XContentParser.Token token = parser.currentToken(); + String fieldName = parser.currentName(); + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else { + QueryGroup queryGroup = QueryGroup.fromXContent(parser); + queryGroupMap.put(fieldName, queryGroup); + } + } + + return new QueryGroupMetadata(queryGroupMap); + } + + @Override + public Diff diff(final Metadata.Custom previousState) { + return new QueryGroupMetadataDiff((QueryGroupMetadata) previousState, this); + } + + public static NamedDiff readDiffFrom(StreamInput in) throws IOException { + return new QueryGroupMetadataDiff(in); + } + + @Override + public EnumSet context() { + return ALL_CONTEXTS; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + QueryGroupMetadata that = (QueryGroupMetadata) o; + return Objects.equals(queryGroups, that.queryGroups); + } + + @Override + public int hashCode() { + return Objects.hash(queryGroups); + } + + @Override + public String toString() { + return Strings.toString(MediaTypeRegistry.JSON, this); + } + + /** + * QueryGroupMetadataDiff + */ + static class QueryGroupMetadataDiff implements NamedDiff { + final Diff> dataStreamDiff; + + QueryGroupMetadataDiff(final QueryGroupMetadata before, final QueryGroupMetadata after) { + dataStreamDiff = DiffableUtils.diff(before.queryGroups, after.queryGroups, DiffableUtils.getStringKeySerializer()); + } + + QueryGroupMetadataDiff(final StreamInput in) throws IOException { + this.dataStreamDiff = DiffableUtils.readJdkMapDiff( + in, + DiffableUtils.getStringKeySerializer(), + QueryGroup::new, + QueryGroup::readDiff + ); + } + + /** + * Returns the name of the writeable object + */ + @Override + public String getWriteableName() { + return TYPE; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + dataStreamDiff.writeTo(out); + } + + @Override + public Metadata.Custom apply(Metadata.Custom part) { + return new QueryGroupMetadata(new HashMap<>(dataStreamDiff.apply(((QueryGroupMetadata) part).queryGroups))); + } + } +} diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index 5bbcd7de1c2ce..fe5ce4dd2bb50 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 = "2.x") 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 f2d99a51f1c9a..97706927ba857 100644 --- a/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java +++ b/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java @@ -33,6 +33,7 @@ package org.opensearch.cluster; import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroupMetadata; import org.opensearch.cluster.metadata.RepositoriesMetadata; import org.opensearch.cluster.routing.ShardRouting; import org.opensearch.cluster.routing.allocation.ExistingShardsAllocator; @@ -69,6 +70,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsModule; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.gateway.GatewayAllocator; import org.opensearch.plugins.ClusterPlugin; import org.opensearch.telemetry.metrics.noop.NoopMetricsRegistry; @@ -327,6 +329,14 @@ 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)) + ); + } + private static ClusterPlugin existingShardsAllocatorPlugin(final String allocatorName) { return new ClusterPlugin() { @Override diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java new file mode 100644 index 0000000000000..d70a9ce5e10cd --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -0,0 +1,93 @@ +/* + * 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.cluster.metadata; + +import org.opensearch.cluster.Diff; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.search.ResourceType; +import org.opensearch.test.AbstractDiffableSerializationTestCase; + +import java.io.IOException; +import java.util.Collections; +import java.util.Map; + +import static org.opensearch.cluster.metadata.QueryGroupTests.createRandomQueryGroup; + +public class QueryGroupMetadataTests extends AbstractDiffableSerializationTestCase { + + public void testToXContent() throws IOException { + long updatedAt = 1720047207; + QueryGroupMetadata queryGroupMetadata = new QueryGroupMetadata( + Map.of( + "ajakgakg983r92_4242", + new QueryGroup( + "test", + "ajakgakg983r92_4242", + QueryGroup.ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + updatedAt + ) + ) + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + builder.startObject(); + queryGroupMetadata.toXContent(builder, null); + builder.endObject(); + assertEquals( + "{\"ajakgakg983r92_4242\":{\"_id\":\"ajakgakg983r92_4242\",\"name\":\"test\",\"resiliency_mode\":\"enforced\",\"updatedAt\":1720047207,\"resourceLimits\":{\"memory\":0.5}}}", + builder.toString() + ); + } + + @Override + protected NamedWriteableRegistry getNamedWriteableRegistry() { + return new NamedWriteableRegistry( + Collections.singletonList( + new NamedWriteableRegistry.Entry(QueryGroupMetadata.class, QueryGroupMetadata.TYPE, QueryGroupMetadata::new) + ) + ); + } + + @Override + protected Metadata.Custom makeTestChanges(Metadata.Custom testInstance) { + final QueryGroup queryGroup = createRandomQueryGroup("asdfakgjwrir23r25"); + final QueryGroupMetadata queryGroupMetadata = new QueryGroupMetadata(Map.of(queryGroup.get_id(), queryGroup)); + return queryGroupMetadata; + } + + @Override + protected Writeable.Reader> diffReader() { + return QueryGroupMetadata::readDiffFrom; + } + + @Override + protected Metadata.Custom doParseInstance(XContentParser parser) throws IOException { + return QueryGroupMetadata.fromXContent(parser); + } + + @Override + protected Writeable.Reader instanceReader() { + return QueryGroupMetadata::new; + } + + @Override + protected QueryGroupMetadata createTestInstance() { + return new QueryGroupMetadata(getRandomQueryGroups()); + } + + private Map getRandomQueryGroups() { + QueryGroup qg1 = createRandomQueryGroup("1243gsgsdgs"); + QueryGroup qg2 = createRandomQueryGroup("lkajga8080"); + return Map.of(qg1.get_id(), qg1, qg2.get_id(), qg2); + } +} diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java new file mode 100644 index 0000000000000..c564f0778e6f0 --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -0,0 +1,158 @@ +/* + * 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.cluster.metadata; + +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; + +import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class QueryGroupTests extends AbstractSerializingTestCase { + + 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.MEMORY, randomDoubleBetween(0.0, 0.80, false)); + return new QueryGroup(name, _id, randomMode(), resourceLimit, Instant.now().getMillis()); + } + + private static QueryGroup.ResiliencyMode randomMode() { + return allowedModes.get(randomIntBetween(0, allowedModes.size() - 1)); + } + + /** + * Parses to a new instance using the provided {@link XContentParser} + * + * @param parser + */ + @Override + protected QueryGroup doParseInstance(XContentParser parser) throws IOException { + return QueryGroup.fromXContent(parser); + } + + /** + * Returns a {@link Writeable.Reader} that can be used to de-serialize the instance + */ + @Override + protected Writeable.Reader instanceReader() { + return QueryGroup::new; + } + + /** + * Creates a random test instance to use in the tests. This method will be + * called multiple times during test execution and should return a different + * random instance each time it is called. + */ + @Override + protected QueryGroup createTestInstance() { + return createRandomQueryGroup("1232sfraeradf_"); + } + + public void testNullName() { + assertThrows( + NullPointerException.class, + () -> new QueryGroup(null, "_id", randomMode(), Collections.emptyMap(), Instant.now().getMillis()) + ); + } + + public void testNullId() { + assertThrows( + NullPointerException.class, + () -> new QueryGroup("Dummy", null, randomMode(), Collections.emptyMap(), Instant.now().getMillis()) + ); + } + + public void testNullResourceLimits() { + assertThrows(NullPointerException.class, () -> new QueryGroup("analytics", "_id", randomMode(), null, Instant.now().getMillis())); + } + + public void testEmptyResourceLimits() { + assertThrows( + IllegalArgumentException.class, + () -> new QueryGroup("analytics", "_id", randomMode(), Collections.emptyMap(), Instant.now().getMillis()) + ); + } + + public void testIllegalQueryGroupMode() { + assertThrows( + NullPointerException.class, + () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, (Object) 0.4), Instant.now().getMillis()) + ); + } + + public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { + assertThrows( + IllegalArgumentException.class, + () -> new QueryGroup( + "analytics", + "_id", + randomMode(), + Map.of(ResourceType.MEMORY, (Object) randomDoubleBetween(1.1, 1.8, false)), + Instant.now().getMillis() + ) + ); + } + + public void testValidQueryGroup() { + QueryGroup queryGroup = new QueryGroup( + "analytics", + "_id", + randomMode(), + Map.of(ResourceType.MEMORY, randomDoubleBetween(0.01, 0.8, false)), + Instant.ofEpochMilli(1717187289).getMillis() + ); + + assertNotNull(queryGroup.getName()); + assertEquals("analytics", queryGroup.getName()); + assertNotNull(queryGroup.getResourceLimits()); + assertFalse(queryGroup.getResourceLimits().isEmpty()); + assertEquals(1, queryGroup.getResourceLimits().size()); + 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( + "{\"_id\":\"" + + queryGroupId + + "\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"updatedAt\":" + + currentTimeInMillis + + ",\"resourceLimits\":{\"cpu\":0.3,\"memory\":0.4}}", + 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)); } };