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 2a52c18e3a26c..abcc2c51d03d1 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -48,8 +48,6 @@ public class QueryGroup extends AbstractDiffable implements ToXConte private final long updatedAtInMillis; private final Map resourceLimits; - // list of resources that are allowed to be present in the QueryGroup schema - public static final List ALLOWED_RESOURCES = List.of("heap_allocations", "cpu"); public QueryGroup(String name, QueryGroupMode resiliencyMode, Map resourceLimits) { this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis()); @@ -116,20 +114,13 @@ public void writeTo(StreamOutput out) throws IOException { private void validateResourceLimits(Map resourceLimits) { for (Map.Entry resource : resourceLimits.entrySet()) { - String resourceName = resource.getKey().getName(); Double threshold = (Double) resource.getValue(); - Objects.requireNonNull(resourceName, "resourceName can't be null"); - Objects.requireNonNull(threshold, "resource limit threshold for" + resourceName + " : can't be null"); + 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"); } - - if (!ALLOWED_RESOURCES.contains(resourceName.toLowerCase(Locale.ROOT))) { - throw new IllegalArgumentException( - "resource has to be valid, valid resources " + ALLOWED_RESOURCES.stream().reduce((x, e) -> x + ", " + e).get() - ); - } } } @@ -181,10 +172,8 @@ public static QueryGroup fromXContent(final XContentParser parser) throws IOExce builder.mode(parser.text()); } else if (fieldName.equals("updatedAt")) { builder.updatedAt(parser.longValue()); - } else if (ALLOWED_RESOURCES.contains(fieldName)) { - resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); } else { - throw new IllegalArgumentException("unrecognised [field=" + fieldName + "] in QueryGroup"); + resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); } } } diff --git a/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java b/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java index f2d99a51f1c9a..5e9a91b69569a 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,13 @@ 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