diff --git a/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java index d2e14948c297e..92ff011422f8b 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java @@ -22,6 +22,7 @@ import org.opensearch.wlm.MutableWorkloadGroupFragment; import org.opensearch.wlm.MutableWorkloadGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupThrottleSettings; import org.joda.time.Instant; import java.io.IOException; @@ -74,15 +75,22 @@ public WorkloadGroup(String name, String _id, MutableWorkloadGroupFragment mutab throw new IllegalArgumentException("WorkloadGroup.updatedAtInMillis is not a valid epoch"); } - // Normalize null settings to empty Settings for storage - if (mutableWorkloadGroupFragment.getSettings() == null) { + // Drop null-valued "clear" keys before storage (meaningful only during an update merge, not on create). + Settings normalizedSettings = stripClearMarkers(mutableWorkloadGroupFragment.getSettings()); + Settings normalizedThrottling = stripClearMarkers(mutableWorkloadGroupFragment.getThrottling()); + if (normalizedSettings.equals(mutableWorkloadGroupFragment.getSettings()) == false + || normalizedThrottling.equals(mutableWorkloadGroupFragment.getThrottling()) == false) { mutableWorkloadGroupFragment = new MutableWorkloadGroupFragment( mutableWorkloadGroupFragment.getResiliencyMode(), mutableWorkloadGroupFragment.getResourceLimits(), - Settings.EMPTY + normalizedSettings, + normalizedThrottling ); } + // Cross-field checks on the merged throttling config (attribute required with a limit; ceiling must be >= 1). + WorkloadGroupThrottleSettings.validateMergedConfig(mutableWorkloadGroupFragment.getThrottling()); + this.name = name; this._id = _id; this.mutableWorkloadGroupFragment = mutableWorkloadGroupFragment; @@ -114,40 +122,69 @@ public static WorkloadGroup updateExistingWorkloadGroup( } final ResiliencyMode mode = Optional.ofNullable(mutableWorkloadGroupFragment.getResiliencyMode()) .orElse(existingGroup.getResiliencyMode()); - // Handle settings update with merge semantics: - // null settings = not specified in request (keep existing) - // empty Settings = clear all settings - // non-empty Settings = merge with existing; keys with null values are removed - final Settings mutableFragmentSettings = mutableWorkloadGroupFragment.getSettings(); - final Settings updatedSettings; - if (mutableFragmentSettings == null) { - // Not specified - keep existing - updatedSettings = Settings.builder().put(existingGroup.getSettings()).build(); - } else if (mutableFragmentSettings.isEmpty()) { - // Explicitly empty - clear all settings - updatedSettings = Settings.EMPTY; - } else { - // Merge: start with existing settings, overlay new values, remove null-valued keys - Settings.Builder builder = Settings.builder().put(existingGroup.getSettings()); - for (String key : mutableFragmentSettings.keySet()) { - String value = mutableFragmentSettings.get(key); - if (value == null) { - // null value means "clear this setting" - builder.remove(key); - } else { - builder.put(key, value); - } - } - updatedSettings = builder.build(); - } + final Settings updatedSettings = mergeSettings(existingGroup.getSettings(), mutableWorkloadGroupFragment.getSettings()); + final Settings updatedThrottling = mergeSettings( + existingGroup.getMutableWorkloadGroupFragment().getThrottling(), + mutableWorkloadGroupFragment.getThrottling() + ); return new WorkloadGroup( existingGroup.getName(), existingGroup.get_id(), - new MutableWorkloadGroupFragment(mode, updatedResourceLimits, updatedSettings), + new MutableWorkloadGroupFragment(mode, updatedResourceLimits, updatedSettings, updatedThrottling), Instant.now().getMillis() ); } + /** + * Drops null-valued keys from a settings bag before storage. A null value is the API gesture for "clear this key", + * which only carries meaning during an update merge (which consumes it); any that reach a persisted group, e.g. a + * null sent on create where there is nothing to clear, are dropped so stored config never contains a null value. + * + * @param s the settings to normalize (may be null) + * @return the settings with all null-valued keys removed, or empty if {@code s} is null + */ + private static Settings stripClearMarkers(Settings s) { + if (s == null) { + return Settings.EMPTY; + } + Settings.Builder builder = Settings.builder(); + for (String key : s.keySet()) { + String value = s.get(key); + if (value != null) { + builder.put(key, value); + } + } + return builder.build(); + } + + /** + * Merges an incoming settings bag from an update request onto the existing one: + * a null incoming bag (field absent) keeps existing, an empty incoming bag clears all, and a non-empty bag overlays + * its values with a per-key null value clearing that key. + * + * @param existing the currently stored settings + * @param incoming the settings from the update request (may be null) + * @return the merged settings + */ + private static Settings mergeSettings(Settings existing, Settings incoming) { + if (incoming == null) { + return Settings.builder().put(existing).build(); + } + if (incoming.isEmpty()) { + return Settings.EMPTY; + } + Settings.Builder builder = Settings.builder().put(existing); + for (String key : incoming.keySet()) { + String value = incoming.get(key); + if (value == null) { + builder.remove(key); + } else { + builder.put(key, value); + } + } + return builder.build(); + } + @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(name); @@ -298,8 +335,8 @@ public static Builder fromXContent(XContentParser parser) throws IOException { } mutableWorkloadGroupFragment1.parseField(parser, fieldName); } else if (token == XContentParser.Token.VALUE_NULL) { - if (fieldName.equals(MutableWorkloadGroupFragment.SETTINGS_STRING)) { - // "settings": null means clear all settings + if (fieldName.equals(MutableWorkloadGroupFragment.SETTINGS_STRING) + || fieldName.equals(MutableWorkloadGroupFragment.THROTTLING_STRING)) { mutableWorkloadGroupFragment1.parseField(parser, fieldName); } } diff --git a/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java b/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java index 0d767a5b9fa7b..1d6b6b332eac1 100644 --- a/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java +++ b/server/src/main/java/org/opensearch/wlm/MutableWorkloadGroupFragment.java @@ -36,11 +36,18 @@ public class MutableWorkloadGroupFragment extends AbstractDiffable resourceLimits; private Settings settings; + private Settings throttling; - public static final List acceptedFieldNames = List.of(RESILIENCY_MODE_STRING, RESOURCE_LIMITS_STRING, SETTINGS_STRING); + public static final List acceptedFieldNames = List.of( + RESILIENCY_MODE_STRING, + RESOURCE_LIMITS_STRING, + SETTINGS_STRING, + THROTTLING_STRING + ); public MutableWorkloadGroupFragment() {} @@ -52,11 +59,22 @@ public MutableWorkloadGroupFragment(ResiliencyMode resiliencyMode, Map resourceLimits, Settings settings) { + this(resiliencyMode, resourceLimits, settings, Settings.EMPTY); + } + + public MutableWorkloadGroupFragment( + ResiliencyMode resiliencyMode, + Map resourceLimits, + Settings settings, + Settings throttling + ) { validateResourceLimits(resourceLimits); WorkloadGroupSearchSettings.validate(settings); + WorkloadGroupThrottleSettings.validate(throttling); this.resiliencyMode = resiliencyMode; this.resourceLimits = resourceLimits; this.settings = settings != null ? settings : Settings.EMPTY; + this.throttling = throttling != null ? throttling : Settings.EMPTY; } public MutableWorkloadGroupFragment(StreamInput in) throws IOException { @@ -69,6 +87,7 @@ public MutableWorkloadGroupFragment(StreamInput in) throws IOException { resiliencyMode = updatedResiliencyMode == null ? null : ResiliencyMode.fromName(updatedResiliencyMode); if (in.getVersion().onOrAfter(Version.V_3_7_0)) { settings = Settings.readOptionalSettingsFromStream(in); + throttling = Settings.readOptionalSettingsFromStream(in); } else if (in.getVersion().onOrAfter(Version.V_3_6_0)) { // Legacy 3.6 format: read and discard (experimental API, no backward compat guarantee) boolean isNull = in.readBoolean(); @@ -76,8 +95,10 @@ public MutableWorkloadGroupFragment(StreamInput in) throws IOException { in.readMap(StreamInput::readString, StreamInput::readString); } settings = Settings.EMPTY; + throttling = Settings.EMPTY; } else { settings = Settings.EMPTY; + throttling = Settings.EMPTY; } } @@ -119,12 +140,25 @@ public Settings parseField(XContentParser parser) throws IOException { } } + static class ThrottlingParser implements FieldParser { + public Settings parseField(XContentParser parser) throws IOException { + // "throttling": null means clear all throttling (disable) + if (parser.currentToken() == XContentParser.Token.VALUE_NULL) { + return Settings.EMPTY; + } + Settings throttling = Settings.fromXContent(parser); + WorkloadGroupThrottleSettings.validate(throttling); + return throttling; + } + } + static class FieldParserFactory { static Optional> fieldParserFor(String fieldName) { return switch (fieldName) { case RESILIENCY_MODE_STRING -> Optional.of(new ResiliencyModeParser()); case RESOURCE_LIMITS_STRING -> Optional.of(new ResourceLimitsParser()); case SETTINGS_STRING -> Optional.of(new SearchSettingsParser()); + case THROTTLING_STRING -> Optional.of(new ThrottlingParser()); default -> Optional.empty(); }; } @@ -153,21 +187,46 @@ static Optional> fieldParserFor(String fieldName) { }, SETTINGS_STRING, (builder) -> { try { builder.startObject(SETTINGS_STRING); - Settings s = settings != null ? settings : Settings.EMPTY; - Map sortedSettingsMap = new TreeMap<>(); - for (String key : s.keySet()) { - sortedSettingsMap.put(key, s.get(key)); - } - for (Map.Entry e : sortedSettingsMap.entrySet()) { - builder.field(e.getKey(), e.getValue()); - } + writeSettingsFields(builder, settings); builder.endObject(); return null; } catch (IOException e) { throw new IllegalStateException("writing error encountered for the field " + SETTINGS_STRING); } + }, THROTTLING_STRING, (builder) -> { + try { + // Unlike settings (always emitted as {}), throttling is omitted entirely when unset. + Settings t = throttling != null ? throttling : Settings.EMPTY; + if (t.isEmpty() == false) { + builder.startObject(THROTTLING_STRING); + if (t.hasValue(WorkloadGroupThrottleSettings.ATTRIBUTE.getKey())) { + builder.field(WorkloadGroupThrottleSettings.ATTRIBUTE.getKey(), WorkloadGroupThrottleSettings.ATTRIBUTE.get(t)); + } + if (t.hasValue(WorkloadGroupThrottleSettings.NODE_LIMIT.getKey())) { + builder.field(WorkloadGroupThrottleSettings.NODE_LIMIT.getKey(), WorkloadGroupThrottleSettings.NODE_LIMIT.get(t)); + } + if (t.hasValue(WorkloadGroupThrottleSettings.SHARED_LIMIT.getKey())) { + builder.field(WorkloadGroupThrottleSettings.SHARED_LIMIT.getKey(), WorkloadGroupThrottleSettings.SHARED_LIMIT.get(t)); + } + builder.endObject(); + } + return null; + } catch (IOException e) { + throw new IllegalStateException("writing error encountered for the field " + THROTTLING_STRING); + } }); + private static void writeSettingsFields(XContentBuilder builder, Settings s) throws IOException { + Settings source = s != null ? s : Settings.EMPTY; + Map sorted = new TreeMap<>(); + for (String key : source.keySet()) { + sorted.put(key, source.get(key)); + } + for (Map.Entry e : sorted.entrySet()) { + builder.field(e.getKey(), e.getValue()); + } + } + public static boolean shouldParse(String field) { return FieldParserFactory.fieldParserFor(field).isPresent(); } @@ -181,6 +240,7 @@ public void parseField(XContentParser parser, String field) { case RESILIENCY_MODE_STRING -> setResiliencyMode((ResiliencyMode) value); case RESOURCE_LIMITS_STRING -> setResourceLimits((Map) value); case SETTINGS_STRING -> setSettings((Settings) value); + case THROTTLING_STRING -> setThrottling((Settings) value); } } catch (IllegalArgumentException e) { throw e; @@ -205,6 +265,8 @@ public void writeTo(StreamOutput out) throws IOException { out.writeOptionalString(resiliencyMode == null ? null : resiliencyMode.getName()); if (out.getVersion().onOrAfter(Version.V_3_7_0)) { Settings.writeOptionalSettingsToStream(settings, out); + // TODO: when upstreamed, gate throttling behind its own release version to support mixed-build clusters. + Settings.writeOptionalSettingsToStream(throttling, out); } else if (out.getVersion().onOrAfter(Version.V_3_6_0)) { // Legacy 3.6 format: write empty map (experimental API, settings not preserved across versions) out.writeBoolean(false); @@ -234,12 +296,13 @@ public boolean equals(Object o) { MutableWorkloadGroupFragment that = (MutableWorkloadGroupFragment) o; return Objects.equals(resiliencyMode, that.resiliencyMode) && Objects.equals(resourceLimits, that.resourceLimits) - && Objects.equals(settings, that.settings); + && Objects.equals(settings, that.settings) + && Objects.equals(throttling, that.throttling); } @Override public int hashCode() { - return Objects.hash(resiliencyMode, resourceLimits, settings); + return Objects.hash(resiliencyMode, resourceLimits, settings, throttling); } public ResiliencyMode getResiliencyMode() { @@ -254,6 +317,10 @@ public Settings getSettings() { return settings; } + public Settings getThrottling() { + return throttling; + } + /** * This enum models the different WorkloadGroup resiliency modes * SOFT - means that this workload group can consume more than workload group resource limits if node is not in duress @@ -299,4 +366,9 @@ void setSettings(Settings settings) { this.settings = settings != null ? settings : Settings.EMPTY; } + void setThrottling(Settings throttling) { + WorkloadGroupThrottleSettings.validate(throttling); + this.throttling = throttling != null ? throttling : Settings.EMPTY; + } + } diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java b/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java new file mode 100644 index 0000000000000..88ea8dd2e6a68 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WorkloadGroupThrottleSettings.java @@ -0,0 +1,136 @@ +/* + * 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.wlm; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; + +import java.util.Map; +import java.util.Set; + +/** + * Registry of valid workload group throttle settings with their validators. Throttle config is a nested + * {@code throttling} object (a {@link Settings} bag) like {@code settings}, so per-key null clears a field and an + * absent key keeps the existing value with no extra bookkeeping. + */ +@ExperimentalApi +public class WorkloadGroupThrottleSettings { + + /** Sentinel for an unset limit, matching the {@code -1 = not set} convention of {@code WLM_SEARCH_TIMEOUT}. */ + public static final int UNSET_LIMIT = -1; + + /** Dimension the limit is keyed by: {@code group} (whole group) or per {@code username} / {@code role}. No default: unset when absent. */ + public static final Setting ATTRIBUTE = Setting.simpleString("attribute"); + + /** Per-node in-flight allowance admitted locally with no coordination. {@code -1} means unset. */ + public static final Setting NODE_LIMIT = Setting.intSetting("node_limit", UNSET_LIMIT, UNSET_LIMIT); + + /** Central in-flight pool drawn on after a node exhausts its local allowance. {@code -1} means unset. */ + public static final Setting SHARED_LIMIT = Setting.intSetting("shared_limit", UNSET_LIMIT, UNSET_LIMIT); + + /** Allowed attribute values; {@code username} / {@code role} map to the security {@code principal.*} attributes at enforcement. */ + public static final Set ALLOWED_ATTRIBUTES = Set.of("group", "username", "role"); + + private static final Map> REGISTERED_SETTINGS = Map.of( + ATTRIBUTE.getKey(), + ATTRIBUTE, + NODE_LIMIT.getKey(), + NODE_LIMIT, + SHARED_LIMIT.getKey(), + SHARED_LIMIT + ); + + private WorkloadGroupThrottleSettings() { + throw new UnsupportedOperationException("Utility class"); + } + + /** + * Per-key validation: every key must be registered, {@code attribute} must be an allowed value, and each limit + * must be a non-negative integer ({@code -1} is the internal "unset" sentinel and is not user-settable). Safe to + * run on a partial fragment from an update request; the cross-field checks live in + * {@link #validateMergedConfig(Settings)}. + * + * @param throttling the throttling settings to validate + * @throws IllegalArgumentException if any key is unknown or any value is invalid + */ + public static void validate(Settings throttling) { + if (throttling == null) { + return; + } + for (String key : throttling.keySet()) { + String value = throttling.get(key); + if (REGISTERED_SETTINGS.containsKey(key) == false) { + throw new IllegalArgumentException("Unknown throttle setting: " + key); + } + // null value means "clear this key" — skip value validation + if (value == null) { + continue; + } + // Limits are stored with -1 as the internal "unset" sentinel, but a user may only send a non-negative integer. + if (NODE_LIMIT.getKey().equals(key) || SHARED_LIMIT.getKey().equals(key)) { + validateUserLimit(key, value); + } + } + String attribute = throttling.get(ATTRIBUTE.getKey()); + if (attribute != null && ALLOWED_ATTRIBUTES.contains(attribute) == false) { + throw new IllegalArgumentException( + "throttling.attribute must be one of " + ALLOWED_ATTRIBUTES + " but was '" + attribute + "'" + ); + } + } + + // Rejects a user-supplied limit that is not a non-negative integer; -1 is reserved as the internal "unset" sentinel. + private static void validateUserLimit(String key, String value) { + int parsed; + try { + parsed = Integer.parseInt(value); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("throttling." + key + " must be a non-negative integer but was '" + value + "'"); + } + if (parsed < 0) { + throw new IllegalArgumentException("throttling." + key + " must be non-negative but was " + parsed); + } + } + + /** + * Cross-field validation on a fully-merged throttling config. A limit may only be set alongside an attribute + * (a limit with no attribute is meaningless), and when throttling is configured the effective ceiling + * {@code max(0, node_limit) + max(0, shared_limit)} must be at least 1, since a ceiling of 0 rejects every + * request. Must be called on the merged result, not a partial update fragment. + * + * @param throttling the merged throttling settings + * @throws IllegalArgumentException if a limit is set without an attribute, or the effective ceiling is 0 + */ + public static void validateMergedConfig(Settings throttling) { + if (throttling == null || throttling.isEmpty()) { + return; + } + boolean hasAttribute = throttling.hasValue(ATTRIBUTE.getKey()); + boolean hasNode = throttling.hasValue(NODE_LIMIT.getKey()); + boolean hasShared = throttling.hasValue(SHARED_LIMIT.getKey()); + + if ((hasNode || hasShared) && hasAttribute == false) { + throw new IllegalArgumentException("throttling.attribute is required when a throttle limit is set"); + } + + int node = NODE_LIMIT.get(throttling); + int shared = SHARED_LIMIT.get(throttling); + if (Math.max(0, node) + Math.max(0, shared) < 1) { + throw new IllegalArgumentException( + "Effective throttle ceiling is 0 (node_limit=" + + (node == UNSET_LIMIT ? "unset" : node) + + ", shared_limit=" + + (shared == UNSET_LIMIT ? "unset" : shared) + + "); this would reject all requests. " + + "Set at least one limit to a positive value, or set throttling as null to disable throttling" + ); + } + } +} diff --git a/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java index 0e57b739a5cd2..a32315e8bc89d 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/WorkloadGroupTests.java @@ -19,6 +19,7 @@ import org.opensearch.wlm.MutableWorkloadGroupFragment; import org.opensearch.wlm.MutableWorkloadGroupFragment.ResiliencyMode; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadGroupThrottleSettings; import org.joda.time.Instant; import java.io.IOException; @@ -37,7 +38,29 @@ static WorkloadGroup createRandomWorkloadGroup(String _id) { String name = randomAlphaOfLength(10); Map resourceLimit = new HashMap<>(); resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false)); - return new WorkloadGroup(name, _id, new MutableWorkloadGroupFragment(randomMode(), resourceLimit), Instant.now().getMillis()); + // Generate a valid throttling config: either disabled (empty), or enabled with a required attribute plus + // at least one positive limit (so the effective ceiling is >= 1). + Settings.Builder throttling = Settings.builder(); + if (randomBoolean()) { + throttling.put("attribute", randomFrom("group", "username", "role")); + if (randomBoolean()) { + throttling.put("node_limit", randomIntBetween(1, 100)); + if (randomBoolean()) { + throttling.put("shared_limit", randomIntBetween(0, 100)); + } + } else { + throttling.put("shared_limit", randomIntBetween(1, 100)); + if (randomBoolean()) { + throttling.put("node_limit", randomIntBetween(0, 100)); + } + } + } + return new WorkloadGroup( + name, + _id, + new MutableWorkloadGroupFragment(randomMode(), resourceLimit, Settings.EMPTY, throttling.build()), + Instant.now().getMillis() + ); } private static ResiliencyMode randomMode() { @@ -376,6 +399,276 @@ public void testUpdateOverrideRequestValuesPersistsThroughMerge() { assertEquals("true", updated.getSettings().get("override_request_values")); } + public void testToXContentOmitsUnsetThrottling() throws IOException { + WorkloadGroup workloadGroup = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5), Settings.EMPTY), + System.currentTimeMillis() + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + workloadGroup.toXContent(builder, ToXContent.EMPTY_PARAMS); + assertFalse(builder.toString().contains("throttling")); + } + + public void testToXContentEmitsThrottling() throws IOException { + long currentTimeInMillis = Instant.now().getMillis(); + String workloadGroupId = UUIDs.randomBase64UUID(); + Settings throttling = Settings.builder().put("attribute", "username").put("node_limit", 10).put("shared_limit", 20).build(); + WorkloadGroup workloadGroup = new WorkloadGroup( + "TestWorkloadGroup", + workloadGroupId, + new MutableWorkloadGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.30), Settings.EMPTY, throttling), + currentTimeInMillis + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + workloadGroup.toXContent(builder, ToXContent.EMPTY_PARAMS); + String expected = String.format( + Locale.ROOT, + "{\"_id\":\"%s\",\"name\":\"TestWorkloadGroup\",\"resiliency_mode\":\"enforced\"," + + "\"resource_limits\":{\"cpu\":0.3}," + + "\"settings\":{}," + + "\"throttling\":{\"attribute\":\"username\",\"node_limit\":10,\"shared_limit\":20}," + + "\"updated_at\":%d}", + workloadGroupId, + currentTimeInMillis + ); + assertEquals(expected, builder.toString()); + } + + public void testNegativeThrottleLimitRejected() { + // -1 is the internal "unset" sentinel and, like any negative value, is not user-settable. + for (int badLimit : new int[] { -1, -2 }) { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("node_limit", badLimit).build() + ) + ); + assertTrue(exception.getMessage().contains("node_limit must be non-negative")); + } + } + + public void testInvalidThrottleAttributeRejected() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "index").put("node_limit", 5).build() + ) + ); + assertTrue(exception.getMessage().contains("throttling.attribute must be one of")); + } + + public void testUnknownThrottleKeyRejected() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("bogus_limit", 5).build() + ) + ); + assertTrue(exception.getMessage().contains("Unknown throttle setting")); + } + + public void testZeroEffectiveCeilingRejected() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 0).build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(exception.getMessage().contains("Effective throttle ceiling is 0")); + } + + public void testCeilingErrorDoesNotLeakUnsetSentinel() { + // attribute set but no limits: both resolve to the -1 sentinel; the message must show "unset", not -1. + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(exception.getMessage().contains("node_limit=unset")); + assertTrue(exception.getMessage().contains("shared_limit=unset")); + assertFalse(exception.getMessage().contains("-1")); + } + + public void testZeroLimitValidWhenOtherLimitPositive() { + // node_limit=0 is fine as long as the effective ceiling is >= 1 (shared_limit carries it). + WorkloadGroup workloadGroup = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 0).put("shared_limit", 10).build() + ), + System.currentTimeMillis() + ); + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals(Integer.valueOf(0), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + assertEquals(Integer.valueOf(10), WorkloadGroupThrottleSettings.SHARED_LIMIT.get(throttling)); + } + + public void testLimitWithoutAttributeRejected() { + // A throttle limit requires an attribute; a limit with no attribute is rejected. + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("node_limit", 5).build() + ), + System.currentTimeMillis() + ) + ); + assertTrue(exception.getMessage().contains("throttling.attribute is required")); + } + + public void testWholeGroupThrottleWithExplicitAttribute() { + // attribute has no default; whole-group throttling must be requested explicitly with attribute=group. + WorkloadGroup workloadGroup = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "group").put("node_limit", 5).build() + ), + System.currentTimeMillis() + ); + Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals("group", WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling)); + assertEquals(Integer.valueOf(5), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + } + + public void testUpdateMergesThrottling() { + WorkloadGroup existing = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 10).put("shared_limit", 20).build() + ), + System.currentTimeMillis() + ); + + // Update only node_limit — attribute and shared_limit should persist + MutableWorkloadGroupFragment updateFragment = new MutableWorkloadGroupFragment( + null, + Map.of(), + Settings.EMPTY, + Settings.builder().put("node_limit", 50).build() + ); + + WorkloadGroup updated = WorkloadGroup.updateExistingWorkloadGroup(existing, updateFragment); + Settings throttling = updated.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals("username", WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling)); + assertEquals(Integer.valueOf(50), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + assertEquals(Integer.valueOf(20), WorkloadGroupThrottleSettings.SHARED_LIMIT.get(throttling)); + } + + public void testUpdateWithNullClearsThrottleKey() throws IOException { + WorkloadGroup existing = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 10).put("shared_limit", 20).build() + ), + System.currentTimeMillis() + ); + + // Send explicit null for throttling.node_limit — should clear it, keep attribute and shared_limit + String json = "{\"resource_limits\":{\"memory\":0.5},\"throttling\":{\"node_limit\":null}}"; + XContentParser parser = createParser(JsonXContent.jsonXContent, json); + MutableWorkloadGroupFragment updateFragment = WorkloadGroup.Builder.fromXContent(parser).getMutableWorkloadGroupFragment(); + + WorkloadGroup updated = WorkloadGroup.updateExistingWorkloadGroup(existing, updateFragment); + Settings throttling = updated.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals("username", WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling)); + assertNull(throttling.get("node_limit")); // raw check: cleared key is absent (typed getter would return the -1 default) + assertEquals(Integer.valueOf(20), WorkloadGroupThrottleSettings.SHARED_LIMIT.get(throttling)); + } + + public void testUpdateWithNullThrottlingObjectDisables() throws IOException { + WorkloadGroup existing = new WorkloadGroup( + "test", + "test_id", + new MutableWorkloadGroupFragment( + ResiliencyMode.ENFORCED, + Map.of(ResourceType.MEMORY, 0.5), + Settings.EMPTY, + Settings.builder().put("attribute", "username").put("node_limit", 10).build() + ), + System.currentTimeMillis() + ); + + // "throttling": null disables throttling entirely + String json = "{\"resource_limits\":{\"memory\":0.5},\"throttling\":null}"; + XContentParser parser = createParser(JsonXContent.jsonXContent, json); + MutableWorkloadGroupFragment updateFragment = WorkloadGroup.Builder.fromXContent(parser).getMutableWorkloadGroupFragment(); + + WorkloadGroup updated = WorkloadGroup.updateExistingWorkloadGroup(existing, updateFragment); + assertTrue(updated.getMutableWorkloadGroupFragment().getThrottling().isEmpty()); + } + + public void testCreateDropsNullThrottleValues() throws IOException { + // On create there is nothing to clear, so null-valued keys are dropped rather than persisted; an + // all-null throttling object therefore collapses to empty (disabled) instead of hitting a ceiling error. + WorkloadGroup partial = parseCreate( + "{\"resiliency_mode\":\"enforced\",\"resource_limits\":{\"memory\":0.5}," + + "\"throttling\":{\"attribute\":\"username\",\"node_limit\":10,\"shared_limit\":null}}" + ); + Settings throttling = partial.getMutableWorkloadGroupFragment().getThrottling(); + assertEquals(Integer.valueOf(10), WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling)); + assertFalse(throttling.keySet().contains("shared_limit")); // raw check: null-valued key was dropped, not persisted + + WorkloadGroup allNull = parseCreate( + "{\"resiliency_mode\":\"enforced\",\"resource_limits\":{\"memory\":0.5}," + "\"throttling\":{\"node_limit\":null}}" + ); + assertTrue(allNull.getMutableWorkloadGroupFragment().getThrottling().isEmpty()); + } + + private WorkloadGroup parseCreate(String json) throws IOException { + XContentParser parser = createParser(JsonXContent.jsonXContent, json); + return WorkloadGroup.Builder.fromXContent(parser).name("test")._id("test_id").updatedAt(System.currentTimeMillis()).build(); + } + public void testSettingsNullFromXContentClearsSettings() throws IOException { // Simulate parsing {"settings": null} via XContent String json = "{\"_id\":\"test_id\",\"name\":\"test\",\"resiliency_mode\":\"enforced\","