From 73a59ef438c55092a8577b596b39f1d67bb211c2 Mon Sep 17 00:00:00 2001 From: Sarthak Jain Date: Fri, 17 Jul 2026 11:39:13 +0530 Subject: [PATCH 1/6] Add partition-weight guard rail for addWagedResource Block adding a WAGED resource whose per-partition weight exceeds the largest single instance's capacity in any dimension, which would make the partition permanently unplaceable. Existing addWagedResource validation only checks that weight keys are present, never their magnitude, so today such a resource is accepted into ZooKeeper and only fails later at rebalance time. This rule closes that gap by pre-validating the mutation on the REST endpoint. - New PartitionWeightCapacityGuardrailRule computes, per capacity dimension, the maximum capacity advertised by any single instance and fails the mutation when a partition's effective weight exceeds it. - GuardrailContext carries the proposed ResourceConfig so rules can read the to-be-written weights before the object exists in ZK. - ResourceAccessor.addResource wires the rule into the addWagedResource path with force/dryRun, mirroring the existing instance-drop guard rail. - Unit tests for the rule plus an integration test for the endpoint (enforce, dry-run, force bypass, within-capacity happy path). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../helix/guardrail/GuardrailContext.java | 26 ++- .../PartitionWeightCapacityGuardrailRule.java | 163 ++++++++++++++ ...tPartitionWeightCapacityGuardrailRule.java | 206 ++++++++++++++++++ .../resources/helix/ResourceAccessor.java | 27 ++- .../rest/server/TestResourceAccessor.java | 131 +++++++++++ 5 files changed, 548 insertions(+), 5 deletions(-) create mode 100644 helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java create mode 100644 helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java b/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java index c7e2d65411..c16d9c5a03 100644 --- a/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java +++ b/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java @@ -20,25 +20,30 @@ package org.apache.helix.guardrail; import org.apache.helix.HelixDataAccessor; +import org.apache.helix.model.ResourceConfig; /** * Immutable bundle of everything a {@link GuardrailRule} needs to evaluate a proposed mutation. *

* The context is intentionally small: it carries the cluster name, a narrow read-only view of * cluster state ({@link ReadOnlyDataAccessor}) for the target cluster, and the target instance name - * for instance-scoped operations. When rules for other object types (e.g. resources) are added, the - * corresponding field can be introduced here through the {@link Builder} without breaking existing - * rules. + * for instance-scoped operations. When rules need the actual object a mutation would write (rather + * than only current cluster state read through the accessor), that proposed object is + * supplied here as well; {@code proposedResourceConfig} is the first such field. New object types + * (e.g. a proposed instance config) are added the same way, through the {@link Builder}, without + * breaking existing rules. */ public class GuardrailContext { private final String clusterName; private final ReadOnlyDataAccessor dataAccessor; private final String instanceName; + private final ResourceConfig proposedResourceConfig; private GuardrailContext(Builder builder) { this.clusterName = builder.clusterName; this.dataAccessor = builder.dataAccessor; this.instanceName = builder.instanceName; + this.proposedResourceConfig = builder.proposedResourceConfig; } public String getClusterName() { @@ -54,6 +59,15 @@ public String getInstanceName() { return instanceName; } + /** + * The resource config a mutation proposes to write, or {@code null} if the operation is not + * resource-scoped. Rules read the to-be-written weights/settings from here rather than from ZK, + * since the object does not exist in ZK yet at pre-validation time. + */ + public ResourceConfig getProposedResourceConfig() { + return proposedResourceConfig; + } + public static Builder newBuilder(String clusterName) { return new Builder(clusterName); } @@ -62,6 +76,7 @@ public static final class Builder { private final String clusterName; private ReadOnlyDataAccessor dataAccessor; private String instanceName; + private ResourceConfig proposedResourceConfig; private Builder(String clusterName) { this.clusterName = clusterName; @@ -77,6 +92,11 @@ public Builder instanceName(String instanceName) { return this; } + public Builder proposedResourceConfig(ResourceConfig proposedResourceConfig) { + this.proposedResourceConfig = proposedResourceConfig; + return this; + } + public GuardrailContext build() { return new GuardrailContext(this); } diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java new file mode 100644 index 0000000000..5c0116fe01 --- /dev/null +++ b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java @@ -0,0 +1,163 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.helix.guardrail.rules; + +import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import org.apache.helix.PropertyKey; +import org.apache.helix.guardrail.GuardrailContext; +import org.apache.helix.guardrail.GuardrailRule; +import org.apache.helix.guardrail.ReadOnlyDataAccessor; +import org.apache.helix.guardrail.ValidationResult; +import org.apache.helix.guardrail.Violation; +import org.apache.helix.model.ClusterConfig; +import org.apache.helix.model.InstanceConfig; +import org.apache.helix.model.ResourceConfig; + +/** + * Guard rail that blocks adding a WAGED resource whose per-partition weight, in any capacity + * dimension, exceeds the largest capacity advertised by any single instance in that dimension. + *

+ * WAGED places a partition on exactly one instance per replica, so a partition can only ever be + * placed if, for every weight dimension {@code d}, some instance has {@code capacity_d >= weight_d}. + * If {@code weight_d} is larger than the maximum instance capacity in {@code d}, no arrangement of + * the cluster can ever host that partition — it is permanently unplaceable. Existing + * validation on {@code addWagedResource} only checks that the required weight keys are + * present; it never compares their magnitudes to instance capacity, so today such a resource is + * accepted into ZooKeeper and only fails later at rebalance time. This rule closes that gap by + * rejecting the mutation up front. + *

+ * The check is a necessary (not sufficient) condition for placeability: it compares each dimension + * independently against the best instance in that dimension. It is deliberately conservative so it + * never blocks a resource that could plausibly be placed — it only fails the cases that are + * provably impossible. + */ +public class PartitionWeightCapacityGuardrailRule implements GuardrailRule { + public static final String RULE_ID = "PARTITION_WEIGHT_EXCEEDS_INSTANCE_CAPACITY"; + + @Override + public String getId() { + return RULE_ID; + } + + @Override + public ValidationResult validate(GuardrailContext context) { + ResourceConfig proposedResourceConfig = context.getProposedResourceConfig(); + if (proposedResourceConfig == null) { + // Not a resource-scoped mutation; nothing for this rule to certify. + return ValidationResult.feasible(); + } + + ReadOnlyDataAccessor dataAccessor = context.getDataAccessor(); + PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder(); + ClusterConfig clusterConfig = dataAccessor.getProperty(keyBuilder.clusterConfig()); + if (clusterConfig == null) { + // No cluster config to interpret weights against; defer to downstream validation. + return ValidationResult.feasible(); + } + + List capacityKeys = clusterConfig.getInstanceCapacityKeys(); + if (capacityKeys.isEmpty()) { + // Cluster does not use the WAGED capacity/weight model, so weights carry no meaning here. + return ValidationResult.feasible(); + } + + // Largest capacity any single instance advertises, per dimension, folding in the cluster-level + // default instance capacity the same way the WAGED rebalancer does. + Map defaultInstanceCapacity = clusterConfig.getDefaultInstanceCapacityMap(); + List instanceConfigs = + dataAccessor.getChildValues(keyBuilder.instanceConfigs(), true); + Map maxInstanceCapacity = new HashMap<>(); + for (InstanceConfig instanceConfig : instanceConfigs) { + Map instanceCapacity = new HashMap<>(defaultInstanceCapacity); + instanceCapacity.putAll(instanceConfig.getInstanceCapacityMap()); + for (Map.Entry entry : instanceCapacity.entrySet()) { + maxInstanceCapacity.merge(entry.getKey(), entry.getValue(), Math::max); + } + } + + if (maxInstanceCapacity.isEmpty()) { + // No instance advertises any capacity yet, so there is nothing to compare against. Leave this + // to existing key-coverage validation rather than emit a misleading "unplaceable" verdict. + return ValidationResult.feasible(); + } + + Map> partitionCapacityMap; + try { + partitionCapacityMap = proposedResourceConfig.getPartitionCapacityMap(); + } catch (IOException e) { + // The weight map is malformed; we cannot certify the resource as placeable. + return ValidationResult.infeasible(Violation.newBuilder(RULE_ID) + .resource(proposedResourceConfig.getResourceName()) + .message(String.format("Could not parse partition weight map for resource %s: %s", + proposedResourceConfig.getResourceName(), e.getMessage())) + .build()); + } + + if (partitionCapacityMap.isEmpty()) { + // No explicit weights: the resource relies entirely on cluster defaults. Evaluate the DEFAULT + // partition so those defaults are still checked against instance capacity. + partitionCapacityMap = + Collections.singletonMap(ResourceConfig.DEFAULT_PARTITION_KEY, Collections.emptyMap()); + } + + Map defaultPartitionWeight = clusterConfig.getDefaultPartitionWeightMap(); + for (Map.Entry> partitionEntry : partitionCapacityMap.entrySet()) { + String partitionName = partitionEntry.getKey(); + // Effective weight = cluster default weight overridden by this partition's explicit weight, + // mirroring WagedValidationUtil#validateAndGetPartitionCapacity. + Map effectiveWeight = new HashMap<>(defaultPartitionWeight); + effectiveWeight.putAll(partitionEntry.getValue()); + + // Only the cluster's declared capacity dimensions are meaningful to WAGED placement. A + // required dimension missing from the weight is a key-coverage problem enforced separately by + // addResourceWithWeight, so it is skipped here rather than reported as an over-weight. + for (String dimension : capacityKeys) { + Integer weight = effectiveWeight.get(dimension); + if (weight == null) { + continue; + } + int maxCapacity = maxInstanceCapacity.getOrDefault(dimension, 0); + if (weight > maxCapacity) { + // DEFAULT_PARTITION_KEY is a placeholder for "every partition", not a real partition, so + // report it as unscoped for a clearer message. + String reportedPartition = + ResourceConfig.DEFAULT_PARTITION_KEY.equals(partitionName) ? null : partitionName; + return ValidationResult.infeasible(Violation.newBuilder(RULE_ID) + .resource(proposedResourceConfig.getResourceName()) + .partition(reportedPartition) + .message(String.format( + "Partition weight %d for dimension '%s' exceeds the largest single instance " + + "capacity %d in that dimension, making %s permanently unplaceable. Lower the " + + "weight, raise instance capacity, or use force=true to override.", weight, + dimension, maxCapacity, + reportedPartition == null ? "every partition" : "partition " + reportedPartition)) + .build()); + } + } + } + + return ValidationResult.feasible(); + } +} diff --git a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java new file mode 100644 index 0000000000..cc38acf05b --- /dev/null +++ b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java @@ -0,0 +1,206 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.helix.guardrail.rules; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Map; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.apache.helix.HelixDataAccessor; +import org.apache.helix.PropertyKey; +import org.apache.helix.guardrail.GuardrailContext; +import org.apache.helix.guardrail.ValidationResult; +import org.apache.helix.guardrail.Violation; +import org.apache.helix.model.ClusterConfig; +import org.apache.helix.model.InstanceConfig; +import org.apache.helix.model.ResourceConfig; +import org.testng.Assert; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link PartitionWeightCapacityGuardrailRule}. Cluster state (cluster config + + * instance configs) is supplied through a mocked {@link HelixDataAccessor}; the proposed resource + * config is passed directly through the {@link GuardrailContext}. + */ +public class TestPartitionWeightCapacityGuardrailRule { + private static final String CLUSTER = "testCluster"; + private static final String RESOURCE = "testResource"; + private static final PropertyKey.Builder BUILDER = new PropertyKey.Builder(CLUSTER); + + private final PartitionWeightCapacityGuardrailRule rule = + new PartitionWeightCapacityGuardrailRule(); + + @Test + public void testNullResourceConfigIsFeasible() { + GuardrailContext context = GuardrailContext.newBuilder(CLUSTER) + .dataAccessor(mock(HelixDataAccessor.class)) + .build(); + Assert.assertTrue(rule.validate(context).isFeasible()); + } + + @Test + public void testNullClusterConfigIsFeasible() throws IOException { + HelixDataAccessor dataAccessor = mock(HelixDataAccessor.class); + when(dataAccessor.keyBuilder()).thenReturn(BUILDER); + doReturn(null).when(dataAccessor).getProperty(BUILDER.clusterConfig()); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 1000))))); + Assert.assertTrue(result.isFeasible()); + } + + @Test + public void testNoCapacityKeysIsFeasible() throws IOException { + // Cluster does not use the WAGED capacity model, so weights are not interpreted. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 1000))))); + Assert.assertTrue(result.isFeasible()); + } + + @Test + public void testNoInstanceCapacityIsFeasible() throws IOException { + // Capacity keys are declared but no instance advertises capacity: nothing to compare against. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO", "BAR")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, ImmutableList.of()); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 1000))))); + Assert.assertTrue(result.isFeasible()); + } + + @Test + public void testWeightWithinCapacityIsFeasible() throws IOException { + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO", "BAR")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, ImmutableList.of( + instanceConfig("instance0", ImmutableMap.of("FOO", 100, "BAR", 100)))); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 100, "BAR", 100))))); + Assert.assertTrue(result.isFeasible()); + Assert.assertTrue(result.getViolations().isEmpty()); + } + + @Test + public void testWeightExceedsCapacityIsInfeasible() throws IOException { + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO", "BAR")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, ImmutableList.of( + instanceConfig("instance0", ImmutableMap.of("FOO", 100, "BAR", 100)), + instanceConfig("instance1", ImmutableMap.of("FOO", 100, "BAR", 100)))); + + // FOO weight 1000 exceeds the largest instance FOO capacity (100). + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 1000, "BAR", 100))))); + + Assert.assertFalse(result.isFeasible()); + Assert.assertEquals(result.getViolations().size(), 1); + Violation violation = result.getViolations().get(0); + Assert.assertEquals(violation.getRuleId(), PartitionWeightCapacityGuardrailRule.RULE_ID); + Assert.assertEquals(violation.getResourceName(), RESOURCE); + // A DEFAULT-scoped weight applies to every partition, so it is reported unscoped. + Assert.assertNull(violation.getPartitionName()); + Assert.assertTrue(violation.getMessage().contains("FOO")); + Assert.assertTrue(violation.getMessage().contains("1000")); + Assert.assertTrue(violation.getMessage().contains("100")); + } + + @Test + public void testPerPartitionOverrideExceedsIsInfeasible() throws IOException { + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + // DEFAULT weight is fine (50 <= 100), but the explicit override for testResource_0 is not. + ResourceConfig resourceConfig = resourceConfig(ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 50), + RESOURCE + "_0", ImmutableMap.of("FOO", 1000))); + ValidationResult result = rule.validate(contextWith(dataAccessor, resourceConfig)); + + Assert.assertFalse(result.isFeasible()); + Violation violation = result.getViolations().get(0); + Assert.assertEquals(violation.getRuleId(), PartitionWeightCapacityGuardrailRule.RULE_ID); + Assert.assertEquals(violation.getPartitionName(), RESOURCE + "_0"); + } + + @Test + public void testMaxCapacityAcrossInstancesUsed() throws IOException { + // The largest instance in each dimension is what matters, not the smallest: a weight of 500 is + // placeable as long as one instance has capacity >= 500. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, ImmutableList.of( + instanceConfig("instance0", ImmutableMap.of("FOO", 100)), + instanceConfig("instance1", ImmutableMap.of("FOO", 1000)))); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 500))))); + Assert.assertTrue(result.isFeasible()); + } + + private GuardrailContext contextWith(HelixDataAccessor dataAccessor, + ResourceConfig proposedResourceConfig) { + return GuardrailContext.newBuilder(CLUSTER) + .dataAccessor(dataAccessor) + .proposedResourceConfig(proposedResourceConfig) + .build(); + } + + private HelixDataAccessor mockAccessor(ClusterConfig clusterConfig, + List instanceConfigs) { + HelixDataAccessor dataAccessor = mock(HelixDataAccessor.class); + when(dataAccessor.keyBuilder()).thenReturn(BUILDER); + doReturn(clusterConfig).when(dataAccessor).getProperty(BUILDER.clusterConfig()); + doReturn(instanceConfigs).when(dataAccessor).getChildValues(BUILDER.instanceConfigs(), true); + return dataAccessor; + } + + private static InstanceConfig instanceConfig(String name, Map capacity) { + InstanceConfig instanceConfig = new InstanceConfig(name); + instanceConfig.setInstanceCapacityMap(capacity); + return instanceConfig; + } + + private static ResourceConfig resourceConfig(Map> partitionCapacity) + throws IOException { + ResourceConfig resourceConfig = new ResourceConfig(RESOURCE); + resourceConfig.setPartitionCapacityMap(partitionCapacity); + return resourceConfig; + } +} diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java index cc4988569f..d3fe51ea4b 100644 --- a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java +++ b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java @@ -26,6 +26,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import javax.ws.rs.DELETE; import javax.ws.rs.DefaultValue; @@ -47,6 +48,9 @@ import org.apache.helix.HelixAdmin; import org.apache.helix.HelixException; import org.apache.helix.PropertyPathBuilder; +import org.apache.helix.guardrail.GuardrailContext; +import org.apache.helix.guardrail.GuardrailPipeline; +import org.apache.helix.guardrail.rules.PartitionWeightCapacityGuardrailRule; import org.apache.helix.model.CustomizedView; import org.apache.helix.model.ExternalView; import org.apache.helix.model.HelixConfigScope; @@ -240,7 +244,9 @@ public Response addResource(@PathParam("clusterId") String clusterId, @DefaultValue("DEFAULT") @QueryParam("rebalanceStrategy") String rebalanceStrategy, @DefaultValue("0") @QueryParam("bucketSize") int bucketSize, @DefaultValue("-1") @QueryParam("maxPartitionsPerInstance") int maxPartitionsPerInstance, - @DefaultValue("addResource") @QueryParam("command") String command, String content) { + @DefaultValue("addResource") @QueryParam("command") String command, + @DefaultValue("false") @QueryParam("force") boolean force, + @DefaultValue("false") @QueryParam("dryRun") boolean dryRun, String content) { // Get the command. If not provided, the default would be "addResource" Command cmd; try { @@ -296,10 +302,27 @@ record = toZNRecord(content); _logger.error("Input does not contain both IdealState and ResourceConfig!"); return badRequest("Input does not contain both IdealState and ResourceConfig!"); } + + ResourceConfig proposedResourceConfig = new ResourceConfig(resourceConfigRecord); + + // Guard rail: block (or simulate) adding a resource whose partition weight exceeds the + // largest single instance's capacity in any dimension, which would make it permanently + // unplaceable. force=true overrides; dryRun=true only reports the verdict without writing. + GuardrailContext context = GuardrailContext.newBuilder(clusterId) + .dataAccessor(getDataAccssor(clusterId)) + .proposedResourceConfig(proposedResourceConfig) + .build(); + GuardrailPipeline pipeline = + new GuardrailPipeline(new PartitionWeightCapacityGuardrailRule()); + Optional preflightResponse = preflight(pipeline, context, force, dryRun); + if (preflightResponse.isPresent()) { + return preflightResponse.get(); + } + // Add using HelixAdmin API try { admin.addResourceWithWeight(clusterId, new IdealState(idealStateRecord), - new ResourceConfig(resourceConfigRecord)); + proposedResourceConfig); } catch (HelixException e) { String errMsg = String.format("Failed to add resource %s with weight in cluster %s!", idealStateRecord.getId(), clusterId); diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java index 2bb91d416c..858f5ace1c 100644 --- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java +++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java @@ -29,6 +29,7 @@ import java.util.Map; import java.util.Set; import javax.ws.rs.client.Entity; +import javax.ws.rs.client.WebTarget; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; @@ -43,6 +44,7 @@ import org.apache.helix.PropertyPathBuilder; import org.apache.helix.TestHelper; import org.apache.helix.controller.rebalancer.waged.WagedRebalancer; +import org.apache.helix.guardrail.rules.PartitionWeightCapacityGuardrailRule; import org.apache.helix.model.ClusterConfig; import org.apache.helix.model.CustomizedView; import org.apache.helix.model.ExternalView; @@ -604,6 +606,135 @@ public void testAddResourceWithWeight() throws IOException { Response.Status.BAD_REQUEST.getStatusCode()); } + /** + * Guard rail: adding a WAGED resource whose partition weight exceeds the largest single instance's + * capacity in any dimension is rejected before the resource is written to ZooKeeper, because such + * a resource is permanently unplaceable. Verifies enforcement (400 + verdict), dry-run (200 + + * verdict, no write), force bypass (created), and the within-capacity happy path (created). The + * cluster/instance capacity configuration is saved and restored so this test does not perturb the + * other resource tests that share {@value #CLUSTER_NAME}. + */ + @Test + public void testAddWagedResourceWeightGuardrail() throws Exception { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + + ClusterConfig clusterConfig = _configAccessor.getClusterConfig(CLUSTER_NAME); + List originalCapacityKeys = clusterConfig.getInstanceCapacityKeys(); + List instances = + _gSetupTool.getClusterManagementTool().getInstancesInCluster(CLUSTER_NAME); + Map> originalInstanceCapacities = new HashMap<>(); + for (String instance : instances) { + originalInstanceCapacities.put(instance, + _configAccessor.getInstanceConfig(CLUSTER_NAME, instance).getInstanceCapacityMap()); + } + + String blockedResource = "guardrailBlockedWagedResource"; + String forcedResource = "guardrailForcedWagedResource"; + String validResource = "guardrailValidWagedResource"; + + try { + // Declare two capacity dimensions and give every instance capacity 100 in each. + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO", "BAR")); + _configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig); + Map instanceCapacity = ImmutableMap.of("FOO", 100, "BAR", 100); + for (String instance : instances) { + InstanceConfig instanceConfig = _configAccessor.getInstanceConfig(CLUSTER_NAME, instance); + instanceConfig.setInstanceCapacityMap(instanceCapacity); + _configAccessor.setInstanceConfig(CLUSTER_NAME, instance, instanceConfig); + } + + // FOO weight 1000 exceeds the largest instance's FOO capacity (100): permanently unplaceable. + Map> overWeight = ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 1000, "BAR", 100)); + + // 1) Enforcement: blocked with 400 + verdict, and nothing written to ZK. + Response blocked = putWagedResource(blockedResource, + wagedResourceConfig(blockedResource, overWeight), Collections.emptyMap()); + Assert.assertEquals(blocked.getStatus(), Response.Status.BAD_REQUEST.getStatusCode()); + JsonNode blockedVerdict = OBJECT_MAPPER.readTree(blocked.readEntity(String.class)); + Assert.assertFalse(blockedVerdict.get("feasible").asBoolean()); + Assert.assertTrue( + blockedVerdict.toString().contains(PartitionWeightCapacityGuardrailRule.RULE_ID)); + Assert.assertFalse(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(blockedResource)); + + // 2) Dry-run: always 200 with the same infeasible verdict, and still nothing written. + Response dryRun = putWagedResource(blockedResource, + wagedResourceConfig(blockedResource, overWeight), ImmutableMap.of("dryRun", true)); + Assert.assertEquals(dryRun.getStatus(), Response.Status.OK.getStatusCode()); + JsonNode dryRunVerdict = OBJECT_MAPPER.readTree(dryRun.readEntity(String.class)); + Assert.assertFalse(dryRunVerdict.get("feasible").asBoolean()); + Assert.assertTrue( + dryRunVerdict.toString().contains(PartitionWeightCapacityGuardrailRule.RULE_ID)); + Assert.assertFalse(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(blockedResource)); + + // 3) force=true bypasses the guard rail: the over-weight resource is actually created. + Response forced = putWagedResource(forcedResource, + wagedResourceConfig(forcedResource, overWeight), ImmutableMap.of("force", true)); + Assert.assertEquals(forced.getStatus(), Response.Status.OK.getStatusCode()); + Assert.assertTrue(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(forcedResource)); + + // 4) A resource within capacity passes the guard rail and is created normally. + Map> withinCapacity = ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 100, "BAR", 100)); + Response valid = putWagedResource(validResource, + wagedResourceConfig(validResource, withinCapacity), Collections.emptyMap()); + Assert.assertEquals(valid.getStatus(), Response.Status.OK.getStatusCode()); + Assert.assertTrue(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(validResource)); + } finally { + // Drop any resources this test created (blockedResource was never created; ignore failures). + for (String resource : Arrays.asList(forcedResource, validResource, blockedResource)) { + try { + _gSetupTool.getClusterManagementTool().dropResource(CLUSTER_NAME, resource); + } catch (Exception ignored) { + } + } + // Restore cluster + instance capacity configuration to its original values. + ClusterConfig restore = _configAccessor.getClusterConfig(CLUSTER_NAME); + restore.setInstanceCapacityKeys(originalCapacityKeys); + _configAccessor.setClusterConfig(CLUSTER_NAME, restore); + for (String instance : instances) { + InstanceConfig instanceConfig = _configAccessor.getInstanceConfig(CLUSTER_NAME, instance); + instanceConfig.setInstanceCapacityMap(originalInstanceCapacities.get(instance)); + _configAccessor.setInstanceConfig(CLUSTER_NAME, instance, instanceConfig); + } + } + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + private Response putWagedResource(String resourceName, ResourceConfig resourceConfig, + Map flags) throws IOException { + IdealState idealState = new IdealState(resourceName); + idealState.getRecord().getSimpleFields().putAll(_gSetupTool.getClusterManagementTool() + .getResourceIdealState(CLUSTER_NAME, RESOURCE_NAME).getRecord().getSimpleFields()); + idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO); + idealState.setRebalancerClassName(WagedRebalancer.class.getName()); + idealState.setNumPartitions(1); + + Map inputMap = ImmutableMap.of( + ResourceAccessor.ResourceProperties.idealState.name(), idealState.getRecord(), + ResourceAccessor.ResourceProperties.resourceConfig.name(), resourceConfig.getRecord()); + Entity entity = + Entity.entity(OBJECT_MAPPER.writeValueAsString(inputMap), MediaType.APPLICATION_JSON_TYPE); + + WebTarget webTarget = target("clusters/" + CLUSTER_NAME + "/resources/" + resourceName) + .queryParam("command", "addWagedResource"); + for (Map.Entry flag : flags.entrySet()) { + webTarget = webTarget.queryParam(flag.getKey(), flag.getValue()); + } + return webTarget.request().put(entity); + } + + private static ResourceConfig wagedResourceConfig(String resourceName, + Map> partitionWeights) throws IOException { + ResourceConfig resourceConfig = new ResourceConfig(resourceName); + resourceConfig.setPartitionCapacityMap(partitionWeights); + return resourceConfig; + } + @Test(dependsOnMethods = "testAddResourceWithWeight") public void testValidateResource() throws IOException { // Define weight keys in ClusterConfig From 0e47b04bf1ea185bed3471f094c09b9c3942c2f2 Mon Sep 17 00:00:00 2001 From: Sarthak Jain Date: Fri, 14 Aug 2026 16:37:20 +0530 Subject: [PATCH 2/6] Ignore ghost partition keys in partition-weight guard rail The PARTITION_CAPACITY_MAP is operator-supplied and can carry stale or mistyped entries naming partitions the resource does not actually have (e.g. leftovers after lowering NUM_PARTITIONS). WAGED ignores such ghost entries at placement time and ZKHelixAdmin.validateWeightForResourceConfig tolerates them on the write path, so blocking on them made the guard rail stricter than the operation it fronts and produced false positives on valid resources. Thread the proposed IdealState through GuardrailContext and skip any weight-map key that is neither DEFAULT nor a real partition of the resource. Real partition names come from the ideal state's partition set when populated, otherwise from NUM_PARTITIONS via Helix's canonical _ naming (a freshly-proposed WAGED ideal state has no assignment yet, so its partition set is empty at pre-validation time). When no ideal state is supplied the rule falls back to evaluating every key, preserving prior behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../helix/guardrail/GuardrailContext.java | 27 ++++++++- .../PartitionWeightCapacityGuardrailRule.java | 52 ++++++++++++++++++ ...tPartitionWeightCapacityGuardrailRule.java | 55 ++++++++++++++++++- .../resources/helix/ResourceAccessor.java | 5 +- 4 files changed, 133 insertions(+), 6 deletions(-) diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java b/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java index c16d9c5a03..32bb27326d 100644 --- a/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java +++ b/helix-core/src/main/java/org/apache/helix/guardrail/GuardrailContext.java @@ -20,6 +20,7 @@ package org.apache.helix.guardrail; import org.apache.helix.HelixDataAccessor; +import org.apache.helix.model.IdealState; import org.apache.helix.model.ResourceConfig; /** @@ -29,21 +30,23 @@ * cluster state ({@link ReadOnlyDataAccessor}) for the target cluster, and the target instance name * for instance-scoped operations. When rules need the actual object a mutation would write (rather * than only current cluster state read through the accessor), that proposed object is - * supplied here as well; {@code proposedResourceConfig} is the first such field. New object types - * (e.g. a proposed instance config) are added the same way, through the {@link Builder}, without - * breaking existing rules. + * supplied here as well; {@code proposedResourceConfig} and {@code proposedIdealState} are the first + * such fields. New object types (e.g. a proposed instance config) are added the same way, through + * the {@link Builder}, without breaking existing rules. */ public class GuardrailContext { private final String clusterName; private final ReadOnlyDataAccessor dataAccessor; private final String instanceName; private final ResourceConfig proposedResourceConfig; + private final IdealState proposedIdealState; private GuardrailContext(Builder builder) { this.clusterName = builder.clusterName; this.dataAccessor = builder.dataAccessor; this.instanceName = builder.instanceName; this.proposedResourceConfig = builder.proposedResourceConfig; + this.proposedIdealState = builder.proposedIdealState; } public String getClusterName() { @@ -68,6 +71,18 @@ public ResourceConfig getProposedResourceConfig() { return proposedResourceConfig; } + /** + * The ideal state a mutation proposes to write, or {@code null} if the operation is not + * resource-scoped. Rules read the resource's structure (e.g. its partition count / names) from + * here rather than from ZK, since the object does not exist in ZK yet at pre-validation time. Note + * that a freshly-proposed ideal state has no computed assignment yet: its partition count + * ({@link IdealState#getNumPartitions()}) is set, but its per-partition preference lists are still + * empty, so {@link IdealState#getPartitionSet()} may be empty at this point. + */ + public IdealState getProposedIdealState() { + return proposedIdealState; + } + public static Builder newBuilder(String clusterName) { return new Builder(clusterName); } @@ -77,6 +92,7 @@ public static final class Builder { private ReadOnlyDataAccessor dataAccessor; private String instanceName; private ResourceConfig proposedResourceConfig; + private IdealState proposedIdealState; private Builder(String clusterName) { this.clusterName = clusterName; @@ -97,6 +113,11 @@ public Builder proposedResourceConfig(ResourceConfig proposedResourceConfig) { return this; } + public Builder proposedIdealState(IdealState proposedIdealState) { + this.proposedIdealState = proposedIdealState; + return this; + } + public GuardrailContext build() { return new GuardrailContext(this); } diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java index 5c0116fe01..9e20d0157c 100644 --- a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java @@ -22,8 +22,10 @@ import java.io.IOException; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import org.apache.helix.PropertyKey; import org.apache.helix.guardrail.GuardrailContext; @@ -32,6 +34,7 @@ import org.apache.helix.guardrail.ValidationResult; import org.apache.helix.guardrail.Violation; import org.apache.helix.model.ClusterConfig; +import org.apache.helix.model.IdealState; import org.apache.helix.model.InstanceConfig; import org.apache.helix.model.ResourceConfig; @@ -52,6 +55,15 @@ * independently against the best instance in that dimension. It is deliberately conservative so it * never blocks a resource that could plausibly be placed — it only fails the cases that are * provably impossible. + *

+ * Only weights for the resource's real partitions are evaluated. A resource's + * {@code PARTITION_CAPACITY_MAP} is operator-supplied and may carry stale or mistyped entries naming + * partitions the resource does not actually have (e.g. leftovers after lowering + * {@code NUM_PARTITIONS}). WAGED ignores such ghost entries at placement time and + * {@code ZKHelixAdmin.validateWeightForResourceConfig} tolerates them on the write path, so this + * rule skips any weight-map key that is neither {@code DEFAULT} nor a real partition of the proposed + * ideal state — blocking on a partition that will never exist would be a false positive + * stricter than the operation it fronts. */ public class PartitionWeightCapacityGuardrailRule implements GuardrailRule { public static final String RULE_ID = "PARTITION_WEIGHT_EXCEEDS_INSTANCE_CAPACITY"; @@ -123,8 +135,20 @@ public ValidationResult validate(GuardrailContext context) { } Map defaultPartitionWeight = clusterConfig.getDefaultPartitionWeightMap(); + Set realPartitions = realPartitionNames(context.getProposedIdealState()); for (Map.Entry> partitionEntry : partitionCapacityMap.entrySet()) { String partitionName = partitionEntry.getKey(); + + // Skip weights for partitions this resource does not actually have. The capacity map is + // operator-supplied and can carry stale/typo'd entries; WAGED ignores them at placement time, + // so blocking on them would be a false positive stricter than the write path we front. When + // the real partition list is unknown (no proposed ideal state) we cannot tell ghosts apart, + // so every entry is evaluated as before. + if (!ResourceConfig.DEFAULT_PARTITION_KEY.equals(partitionName) && realPartitions != null + && !realPartitions.contains(partitionName)) { + continue; + } + // Effective weight = cluster default weight overridden by this partition's explicit weight, // mirroring WagedValidationUtil#validateAndGetPartitionCapacity. Map effectiveWeight = new HashMap<>(defaultPartitionWeight); @@ -160,4 +184,32 @@ public ValidationResult validate(GuardrailContext context) { return ValidationResult.feasible(); } + + /** + * The names of the partitions the proposed resource actually has, or {@code null} if they cannot + * be determined (no proposed ideal state supplied). + *

+ * A freshly-proposed WAGED ideal state carries {@code NUM_PARTITIONS} but no computed assignment, + * so its preference lists — and therefore {@link IdealState#getPartitionSet()} — are + * still empty at pre-validation time. When that is the case the names are reconstructed from the + * partition count using Helix's canonical {@code _} scheme (the same naming the + * controller applies in {@code ResourceComputationStage}). If preference lists are already + * populated (e.g. a CUSTOMIZED ideal state), those partition names are used directly. + */ + private static Set realPartitionNames(IdealState idealState) { + if (idealState == null) { + return null; + } + Set declaredPartitions = idealState.getPartitionSet(); + if (declaredPartitions != null && !declaredPartitions.isEmpty()) { + return declaredPartitions; + } + int numPartitions = idealState.getNumPartitions(); + String resourceName = idealState.getResourceName(); + Set partitionNames = new HashSet<>(); + for (int i = 0; i < numPartitions; i++) { + partitionNames.add(resourceName + "_" + i); + } + return partitionNames; + } } diff --git a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java index cc38acf05b..9ab56d75b3 100644 --- a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java @@ -32,6 +32,7 @@ import org.apache.helix.guardrail.ValidationResult; import org.apache.helix.guardrail.Violation; import org.apache.helix.model.ClusterConfig; +import org.apache.helix.model.IdealState; import org.apache.helix.model.InstanceConfig; import org.apache.helix.model.ResourceConfig; import org.testng.Assert; @@ -44,7 +45,7 @@ /** * Unit tests for {@link PartitionWeightCapacityGuardrailRule}. Cluster state (cluster config + * instance configs) is supplied through a mocked {@link HelixDataAccessor}; the proposed resource - * config is passed directly through the {@link GuardrailContext}. + * config and ideal state are passed directly through the {@link GuardrailContext}. */ public class TestPartitionWeightCapacityGuardrailRule { private static final String CLUSTER = "testCluster"; @@ -174,11 +175,63 @@ public void testMaxCapacityAcrossInstancesUsed() throws IOException { Assert.assertTrue(result.isFeasible()); } + @Test + public void testGhostPartitionKeyIsIgnored() throws IOException { + // The capacity map names a partition (testResource_99999) the resource does not have: only _0 + // and _1 are real. WAGED ignores such stale/typo'd entries at placement time, so this rule must + // too, even though the ghost's weight (1000) far exceeds the largest instance capacity (100). + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + ResourceConfig resourceConfig = resourceConfig(ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 50), + RESOURCE + "_99999", ImmutableMap.of("FOO", 1000))); + ValidationResult result = rule.validate(contextWith(dataAccessor, resourceConfig, 2)); + + Assert.assertTrue(result.isFeasible()); + Assert.assertTrue(result.getViolations().isEmpty()); + } + + @Test + public void testRealPartitionStillFlaggedAlongsideGhost() throws IOException { + // Skipping ghosts must not mask a genuinely unplaceable real partition: testResource_99999 is + // ignored, but the real testResource_1 override (999 > 100) is still caught. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + ResourceConfig resourceConfig = resourceConfig(ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 50), + RESOURCE + "_99999", ImmutableMap.of("FOO", 1000), + RESOURCE + "_1", ImmutableMap.of("FOO", 999))); + ValidationResult result = rule.validate(contextWith(dataAccessor, resourceConfig, 2)); + + Assert.assertFalse(result.isFeasible()); + Violation violation = result.getViolations().get(0); + Assert.assertEquals(violation.getRuleId(), PartitionWeightCapacityGuardrailRule.RULE_ID); + Assert.assertEquals(violation.getPartitionName(), RESOURCE + "_1"); + } + private GuardrailContext contextWith(HelixDataAccessor dataAccessor, ResourceConfig proposedResourceConfig) { + // Default to a single-partition resource so the canonical testResource_0 partition is real. + return contextWith(dataAccessor, proposedResourceConfig, 1); + } + + private GuardrailContext contextWith(HelixDataAccessor dataAccessor, + ResourceConfig proposedResourceConfig, int numPartitions) { + // Mirror a freshly-proposed WAGED ideal state: partition count is set but the assignment (and + // thus getPartitionSet()) is still empty, so the rule reconstructs names from numPartitions. + IdealState idealState = new IdealState(RESOURCE); + idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO); + idealState.setNumPartitions(numPartitions); return GuardrailContext.newBuilder(CLUSTER) .dataAccessor(dataAccessor) .proposedResourceConfig(proposedResourceConfig) + .proposedIdealState(idealState) .build(); } diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java index d3fe51ea4b..69c52a270d 100644 --- a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java +++ b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java @@ -304,6 +304,7 @@ record = toZNRecord(content); } ResourceConfig proposedResourceConfig = new ResourceConfig(resourceConfigRecord); + IdealState proposedIdealState = new IdealState(idealStateRecord); // Guard rail: block (or simulate) adding a resource whose partition weight exceeds the // largest single instance's capacity in any dimension, which would make it permanently @@ -311,6 +312,7 @@ record = toZNRecord(content); GuardrailContext context = GuardrailContext.newBuilder(clusterId) .dataAccessor(getDataAccssor(clusterId)) .proposedResourceConfig(proposedResourceConfig) + .proposedIdealState(proposedIdealState) .build(); GuardrailPipeline pipeline = new GuardrailPipeline(new PartitionWeightCapacityGuardrailRule()); @@ -321,8 +323,7 @@ record = toZNRecord(content); // Add using HelixAdmin API try { - admin.addResourceWithWeight(clusterId, new IdealState(idealStateRecord), - proposedResourceConfig); + admin.addResourceWithWeight(clusterId, proposedIdealState, proposedResourceConfig); } catch (HelixException e) { String errMsg = String.format("Failed to add resource %s with weight in cluster %s!", idealStateRecord.getId(), clusterId); From 6990e089d597830194871152c2e60a517d43560f Mon Sep 17 00:00:00 2001 From: Sarthak Jain Date: Fri, 14 Aug 2026 16:52:34 +0530 Subject: [PATCH 3/6] Address weight guard rail review: honest dry-run, assignable-only capacity, no misblame Fixes three reviewer findings on the partition-weight guard rail: 1. force/dryRun were on the shared addResource dispatch method but only honored in the addWagedResource branch, so dryRun=true on a plain addResource silently performed a real write. Reject both flags with a 400 for any command other than addWagedResource so a "simulation" can never mutate ZK. 3. The rule used getOrDefault(dimension, 0), so when the cluster declared a capacity key the instances did not advertise, a resource weight in that dimension was reported as "exceeds capacity 0" and the resource author was told to lower a weight that cannot go below 0. That missing capacity key is an instance-side misconfiguration already reported by WagedValidationUtil, so skip the dimension instead (mirroring the existing weight == null skip) rather than misblaming the resource. 5. maxInstanceCapacity folded over every InstanceConfig, including non-assignable ones (EVACUATE / SWAP_IN / UNKNOWN). WAGED only places on getAssignableInstanceConfigMap() instances, so counting a decommissioning instance's capacity let the rule certify a resource WAGED can never place. Filter to InstanceConfig.isAssignable() instances, matching the rebalancer. Also documented the fail-closed behavior of the getChildValues(..., true) instance-config read. Adds unit tests for the assignable-only and missing-dimension cases and an integration test asserting force/dryRun are rejected for a non-waged command. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../PartitionWeightCapacityGuardrailRule.java | 39 +++++++++++--- ...tPartitionWeightCapacityGuardrailRule.java | 51 +++++++++++++++++++ .../resources/helix/ResourceAccessor.java | 10 ++++ .../rest/server/TestResourceAccessor.java | 31 +++++++++++ 4 files changed, 125 insertions(+), 6 deletions(-) diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java index 9e20d0157c..8dbf38fefa 100644 --- a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java @@ -40,7 +40,8 @@ /** * Guard rail that blocks adding a WAGED resource whose per-partition weight, in any capacity - * dimension, exceeds the largest capacity advertised by any single instance in that dimension. + * dimension, exceeds the largest capacity advertised by any single assignable instance in + * that dimension. *

* WAGED places a partition on exactly one instance per replica, so a partition can only ever be * placed if, for every weight dimension {@code d}, some instance has {@code capacity_d >= weight_d}. @@ -95,13 +96,28 @@ public ValidationResult validate(GuardrailContext context) { return ValidationResult.feasible(); } - // Largest capacity any single instance advertises, per dimension, folding in the cluster-level - // default instance capacity the same way the WAGED rebalancer does. + // Largest capacity any single ASSIGNABLE instance advertises, per dimension, folding in the + // cluster-level default instance capacity the same way the WAGED rebalancer does. Only + // assignable instances are counted: WAGED places exclusively on the instances in + // BaseControllerDataProvider#getAssignableInstanceConfigMap(), i.e. those where + // InstanceConfig#isAssignable() is true (this excludes EVACUATE / SWAP_IN / UNKNOWN operations). + // Counting capacity advertised by a non-assignable instance would let this rule certify a + // resource that WAGED can never actually place. + // + // getChildValues(..., true) reads instance configs fail-closed: a transient ZK read error or a + // single unreadable instance-config znode propagates out, and the guard rail pipeline then turns + // the add into a 400 rather than silently validating against partial cluster state. That is the + // safe default for a guard rail, at the cost of coupling addWagedResource availability to + // instance-config readability. Map defaultInstanceCapacity = clusterConfig.getDefaultInstanceCapacityMap(); List instanceConfigs = dataAccessor.getChildValues(keyBuilder.instanceConfigs(), true); Map maxInstanceCapacity = new HashMap<>(); for (InstanceConfig instanceConfig : instanceConfigs) { + if (instanceConfig == null || !instanceConfig.isAssignable()) { + // WAGED will not place on a non-assignable instance, so its capacity is irrelevant here. + continue; + } Map instanceCapacity = new HashMap<>(defaultInstanceCapacity); instanceCapacity.putAll(instanceConfig.getInstanceCapacityMap()); for (Map.Entry entry : instanceCapacity.entrySet()) { @@ -110,8 +126,9 @@ public ValidationResult validate(GuardrailContext context) { } if (maxInstanceCapacity.isEmpty()) { - // No instance advertises any capacity yet, so there is nothing to compare against. Leave this - // to existing key-coverage validation rather than emit a misleading "unplaceable" verdict. + // No assignable instance advertises any capacity yet, so there is nothing to compare against. + // Leave this to existing key-coverage validation rather than emit a misleading "unplaceable" + // verdict. return ValidationResult.feasible(); } @@ -162,7 +179,17 @@ public ValidationResult validate(GuardrailContext context) { if (weight == null) { continue; } - int maxCapacity = maxInstanceCapacity.getOrDefault(dimension, 0); + Integer maxCapacity = maxInstanceCapacity.get(dimension); + if (maxCapacity == null) { + // No assignable instance advertises capacity for this dimension. That is an instance-side + // misconfiguration (a cluster-declared capacity key missing from the instances), which + // WagedValidationUtil#validateAndGetInstanceCapacity already reports against the instances. + // Treating the absent dimension as capacity 0 here would blame the resource author for an + // instance problem and tell them to lower a weight that cannot go below 0. It is also + // inconsistent with the weight == null skip above, where missing key coverage is likewise + // deferred to that separate validation. So skip this dimension rather than over-block. + continue; + } if (weight > maxCapacity) { // DEFAULT_PARTITION_KEY is a placeholder for "every partition", not a real partition, so // report it as unscoped for a clearer message. diff --git a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java index 9ab56d75b3..b24467a33d 100644 --- a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java @@ -28,6 +28,7 @@ import com.google.common.collect.ImmutableMap; import org.apache.helix.HelixDataAccessor; import org.apache.helix.PropertyKey; +import org.apache.helix.constants.InstanceConstants; import org.apache.helix.guardrail.GuardrailContext; import org.apache.helix.guardrail.ValidationResult; import org.apache.helix.guardrail.Violation; @@ -215,6 +216,49 @@ public void testRealPartitionStillFlaggedAlongsideGhost() throws IOException { Assert.assertEquals(violation.getPartitionName(), RESOURCE + "_1"); } + @Test + public void testNonAssignableInstanceCapacityIgnored() throws IOException { + // WAGED only places on assignable instances. A large-capacity instance that is EVACUATE (being + // decommissioned) is not assignable, so its capacity must not count toward placeability: with + // only the assignable instance's 100 capacity, a DEFAULT weight of 5000 is unplaceable. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, ImmutableList.of( + instanceConfig("assignable", ImmutableMap.of("FOO", 100)), + instanceConfig("evacuating", ImmutableMap.of("FOO", 10000), + InstanceConstants.InstanceOperation.EVACUATE))); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 5000))))); + + Assert.assertFalse(result.isFeasible()); + Violation violation = result.getViolations().get(0); + Assert.assertEquals(violation.getRuleId(), PartitionWeightCapacityGuardrailRule.RULE_ID); + // The 10000-capacity evacuating instance is ignored, so the reported ceiling is the assignable + // instance's 100 rather than 10000. + Assert.assertTrue(violation.getMessage().contains("capacity 100")); + } + + @Test + public void testMissingCapacityDimensionNotBlamedOnResource() throws IOException { + // The cluster declares two capacity keys but the instances only advertise FOO. A BAR weight must + // not be blamed on the resource as "exceeds capacity 0"; a capacity key missing from the + // instances is an instance-side misconfiguration reported separately, so the rule defers on that + // dimension (mirroring the missing-weight skip) and only checks the dimensions instances cover. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO", "BAR")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + ResourceConfig resourceConfig = resourceConfig(ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 50, "BAR", 1))); + ValidationResult result = rule.validate(contextWith(dataAccessor, resourceConfig)); + + Assert.assertTrue(result.isFeasible()); + Assert.assertTrue(result.getViolations().isEmpty()); + } + private GuardrailContext contextWith(HelixDataAccessor dataAccessor, ResourceConfig proposedResourceConfig) { // Default to a single-partition resource so the canonical testResource_0 partition is real. @@ -250,6 +294,13 @@ private static InstanceConfig instanceConfig(String name, Map c return instanceConfig; } + private static InstanceConfig instanceConfig(String name, Map capacity, + InstanceConstants.InstanceOperation operation) { + InstanceConfig instanceConfig = instanceConfig(name, capacity); + instanceConfig.setInstanceOperation(operation); + return instanceConfig; + } + private static ResourceConfig resourceConfig(Map> partitionCapacity) throws IOException { ResourceConfig resourceConfig = new ResourceConfig(RESOURCE); diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java index 69c52a270d..907a40ae7b 100644 --- a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java +++ b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java @@ -254,6 +254,16 @@ public Response addResource(@PathParam("clusterId") String clusterId, } catch (Exception e) { return badRequest("Invalid command : " + command); } + // force and dryRun are only honored by commands that run a guard rail pipeline (currently only + // addWagedResource). For any other command they are silently ignored and, worse, dryRun=true on + // a plain addResource would still perform a real write — the opposite of a simulation. Reject + // them up front for unsupported commands so callers are never misled into thinking a mutation + // was simulated or its violations overridden. + if ((force || dryRun) && cmd != Command.addWagedResource) { + return badRequest(String.format( + "The 'force' and 'dryRun' flags are only supported for the 'addWagedResource' command, " + + "not '%s'.", command)); + } HelixAdmin admin = getHelixAdmin(); try { switch (cmd) { diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java index 858f5ace1c..2602467b3b 100644 --- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java +++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java @@ -705,6 +705,37 @@ public void testAddWagedResourceWeightGuardrail() throws Exception { System.out.println("End test :" + TestHelper.getTestMethodName()); } + /** + * force/dryRun are only meaningful for the addWagedResource command, which is the only one that + * runs a guard rail pipeline. On any other command they were silently ignored, so dryRun=true on a + * plain addResource still performed a real write. They must now be rejected with 400 and create + * nothing. + */ + @Test(dependsOnMethods = "testAddResourceWithWeight") + public void testDryRunAndForceRejectedForNonWagedCommand() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + + String dryRunResource = "dryRunRejectedResource"; + put("clusters/" + CLUSTER_NAME + "/resources/" + dryRunResource, + ImmutableMap.of("command", "addResource", "numPartitions", "1", "stateModelRef", + "OnlineOffline", "dryRun", "true"), + Entity.entity("", MediaType.APPLICATION_JSON_TYPE), + Response.Status.BAD_REQUEST.getStatusCode()); + Assert.assertFalse(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(dryRunResource)); + + String forceResource = "forceRejectedResource"; + put("clusters/" + CLUSTER_NAME + "/resources/" + forceResource, + ImmutableMap.of("command", "addResource", "numPartitions", "1", "stateModelRef", + "OnlineOffline", "force", "true"), + Entity.entity("", MediaType.APPLICATION_JSON_TYPE), + Response.Status.BAD_REQUEST.getStatusCode()); + Assert.assertFalse(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(forceResource)); + + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + private Response putWagedResource(String resourceName, ResourceConfig resourceConfig, Map flags) throws IOException { IdealState idealState = new IdealState(resourceName); From 1249ec5704c204e9184962b84b803bb26f06c84c Mon Sep 17 00:00:00 2001 From: Sarthak Jain Date: Fri, 14 Aug 2026 17:15:55 +0530 Subject: [PATCH 4/6] Address PR review tiers 2-3 for partition-weight guard rail Tier 2 (#4): Document the cluster-wide blast radius of an unplaceable WAGED resource (CAPACITY_DEFICIT stalls the global rebalance so nothing added afterward gets placed) and stop suggesting force=true, which is the exact action that triggers the deficit. The force capability itself is retained. Tier 3: - #6: Collect and report every over-capacity partition/dimension in a single verdict (deterministic order: DEFAULT first then natural), so a caller sees all problems at once instead of fixing one and resubmitting. - #7: Run cheap local structural checks (IdealState/ResourceConfig name match, non-negative weights) before building the guard rail context, so a dry-run reflects them and a structurally invalid request never reaches ZooKeeper. Negative-weight validation is explicit because the endpoint builds ResourceConfig from a raw ZNRecord that bypasses the setter's own check. - #8: Order the guard rail integration tests after testAddResourceWithWeight via dependsOnMethods. Adds a unit test for multi-violation aggregation/ordering and an integration test asserting structural checks are applied before the guard rail (including the raw-record negative-weight path). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../PartitionWeightCapacityGuardrailRule.java | 43 +++++++--- ...tPartitionWeightCapacityGuardrailRule.java | 24 ++++++ .../resources/helix/ResourceAccessor.java | 53 ++++++++++++ .../rest/server/TestResourceAccessor.java | 85 ++++++++++++++++++- 4 files changed, 194 insertions(+), 11 deletions(-) diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java index 8dbf38fefa..97702faad9 100644 --- a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java @@ -20,7 +20,9 @@ package org.apache.helix.guardrail.rules; import java.io.IOException; +import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -57,6 +59,14 @@ * never blocks a resource that could plausibly be placed — it only fails the cases that are * provably impossible. *

+ * Why this is enforced up front rather than left to rebalance. An unplaceable WAGED resource + * is not a resource-local failure. Once it exists, the WAGED global rebalance fails to compute a + * baseline assignment for the whole cluster (a {@code CAPACITY_DEFICIT} error), so no + * resource added after it gets placed anywhere until the offending resource is dropped. Resources + * already assigned keep their assignment, so the breakage is silent. That cluster-wide blast radius + * is why this is a hard pre-write guard rail; it is also why callers should not reach for + * {@code force=true} to bypass it, as forcing the resource in is exactly what triggers the deficit. + *

* Only weights for the resource's real partitions are evaluated. A resource's * {@code PARTITION_CAPACITY_MAP} is operator-supplied and may carry stale or mistyped entries naming * partitions the resource does not actually have (e.g. leftovers after lowering @@ -153,9 +163,17 @@ public ValidationResult validate(GuardrailContext context) { Map defaultPartitionWeight = clusterConfig.getDefaultPartitionWeightMap(); Set realPartitions = realPartitionNames(context.getProposedIdealState()); - for (Map.Entry> partitionEntry : partitionCapacityMap.entrySet()) { - String partitionName = partitionEntry.getKey(); + // Evaluate partitions in a deterministic order (the DEFAULT placeholder first, then the rest in + // natural order) rather than HashMap iteration order, so that when several partitions or + // dimensions are over capacity the set and order of reported violations is stable between runs. + List orderedPartitions = new ArrayList<>(partitionCapacityMap.keySet()); + orderedPartitions.sort( + Comparator.comparing((String p) -> !ResourceConfig.DEFAULT_PARTITION_KEY.equals(p)) + .thenComparing(Comparator.naturalOrder())); + + List violations = new ArrayList<>(); + for (String partitionName : orderedPartitions) { // Skip weights for partitions this resource does not actually have. The capacity map is // operator-supplied and can carry stale/typo'd entries; WAGED ignores them at placement time, // so blocking on them would be a false positive stricter than the write path we front. When @@ -169,11 +187,14 @@ public ValidationResult validate(GuardrailContext context) { // Effective weight = cluster default weight overridden by this partition's explicit weight, // mirroring WagedValidationUtil#validateAndGetPartitionCapacity. Map effectiveWeight = new HashMap<>(defaultPartitionWeight); - effectiveWeight.putAll(partitionEntry.getValue()); + effectiveWeight.putAll(partitionCapacityMap.get(partitionName)); - // Only the cluster's declared capacity dimensions are meaningful to WAGED placement. A - // required dimension missing from the weight is a key-coverage problem enforced separately by - // addResourceWithWeight, so it is skipped here rather than reported as an over-weight. + // Only the cluster's declared capacity dimensions are meaningful to WAGED placement, and + // capacityKeys is a List, so iterating it gives a fixed dimension order. A required dimension + // missing from the weight is a key-coverage problem enforced separately by + // addResourceWithWeight, so it is skipped here rather than reported as an over-weight. Every + // over-capacity dimension is collected (not just the first) so a caller sees all problems in + // one response instead of fixing one and resubmitting to discover the next. for (String dimension : capacityKeys) { Integer weight = effectiveWeight.get(dimension); if (weight == null) { @@ -195,21 +216,23 @@ public ValidationResult validate(GuardrailContext context) { // report it as unscoped for a clearer message. String reportedPartition = ResourceConfig.DEFAULT_PARTITION_KEY.equals(partitionName) ? null : partitionName; - return ValidationResult.infeasible(Violation.newBuilder(RULE_ID) + // Intentionally no force=true hint: forcing an unplaceable resource in is what triggers + // the cluster-wide CAPACITY_DEFICIT described in the class javadoc, so the message only + // points at the safe remedies. + violations.add(Violation.newBuilder(RULE_ID) .resource(proposedResourceConfig.getResourceName()) .partition(reportedPartition) .message(String.format( "Partition weight %d for dimension '%s' exceeds the largest single instance " + "capacity %d in that dimension, making %s permanently unplaceable. Lower the " - + "weight, raise instance capacity, or use force=true to override.", weight, - dimension, maxCapacity, + + "weight or raise instance capacity.", weight, dimension, maxCapacity, reportedPartition == null ? "every partition" : "partition " + reportedPartition)) .build()); } } } - return ValidationResult.feasible(); + return violations.isEmpty() ? ValidationResult.feasible() : ValidationResult.of(violations); } /** diff --git a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java index b24467a33d..39853b7b11 100644 --- a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java @@ -259,6 +259,30 @@ public void testMissingCapacityDimensionNotBlamedOnResource() throws IOException Assert.assertTrue(result.getViolations().isEmpty()); } + @Test + public void testMultipleViolationsAllReportedInFixedOrder() throws IOException { + // Both declared dimensions are over capacity. The rule must report both (not just the first) and + // in a stable order matching the cluster's capacity-key order, so a caller sees every problem in + // one response instead of fixing one, resubmitting, and discovering the next. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO", "BAR")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100, "BAR", 100)))); + + ResourceConfig resourceConfig = resourceConfig(ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 5000, "BAR", 9000))); + ValidationResult result = rule.validate(contextWith(dataAccessor, resourceConfig)); + + Assert.assertFalse(result.isFeasible()); + Assert.assertEquals(result.getViolations().size(), 2); + // Fixed order: FOO before BAR, matching the declared capacity-key order. + Assert.assertTrue(result.getViolations().get(0).getMessage().contains("'FOO'")); + Assert.assertTrue(result.getViolations().get(1).getMessage().contains("'BAR'")); + // No force=true suggestion: forcing an unplaceable resource is what triggers the cluster-wide + // capacity deficit the rule exists to prevent. + Assert.assertFalse(result.getViolations().get(0).getMessage().contains("force")); + } + private GuardrailContext contextWith(HelixDataAccessor dataAccessor, ResourceConfig proposedResourceConfig) { // Default to a single-partition resource so the canonical testResource_0 partition is real. diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java index 907a40ae7b..617503caab 100644 --- a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java +++ b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ResourceAccessor.java @@ -316,6 +316,16 @@ record = toZNRecord(content); ResourceConfig proposedResourceConfig = new ResourceConfig(resourceConfigRecord); IdealState proposedIdealState = new IdealState(idealStateRecord); + // Cheap, local structural validation before any ZK-backed guard rail work. Running it here + // means these failures are reflected by a dry-run (instead of a misleading feasible verdict) + // and are caught before the guard rail's instance-config scan, so a structurally invalid + // request never reaches ZooKeeper. + Optional structuralError = + validateWagedResourceStructure(proposedIdealState, proposedResourceConfig); + if (structuralError.isPresent()) { + return structuralError.get(); + } + // Guard rail: block (or simulate) adding a resource whose partition weight exceeds the // largest single instance's capacity in any dimension, which would make it permanently // unplaceable. force=true overrides; dryRun=true only reports the verdict without writing. @@ -352,6 +362,49 @@ record = toZNRecord(content); return OK(); } + /** + * Cheap, local (no ZooKeeper) structural checks for an addWagedResource request. Returns a + * {@code 400} response if the request is malformed, or {@link Optional#empty()} if it is + * structurally sound. These are validated before the guard rail pipeline so that a dry-run + * reflects them and a structurally invalid request never triggers the guard rail's instance-config + * read. + */ + private Optional validateWagedResourceStructure(IdealState idealState, + ResourceConfig resourceConfig) { + // IdealState and ResourceConfig must describe the same resource. addResourceWithWeight enforces + // this on the write path, but checking here means a dry-run reports it instead of returning a + // feasible verdict for a request that would then fail for real. + if (!idealState.getResourceName().equals(resourceConfig.getResourceName())) { + return Optional.of(badRequest(String.format( + "Resource names in IdealState (%s) and ResourceConfig (%s) are different!", + idealState.getResourceName(), resourceConfig.getResourceName()))); + } + + // Partition weights must be non-negative. ResourceConfig#setPartitionCapacityMap rejects + // negatives, but this endpoint constructs the ResourceConfig straight from a raw ZNRecord and + // bypasses that setter, so a negative weight would otherwise slip through (the guard rail's + // "weight > capacity" check does not catch it either). Validate it explicitly. + Map> partitionCapacityMap; + try { + partitionCapacityMap = resourceConfig.getPartitionCapacityMap(); + } catch (IOException e) { + return Optional.of(badRequest(String.format( + "Could not parse partition weight map for resource %s: %s", + resourceConfig.getResourceName(), e.getMessage()))); + } + for (Map.Entry> partitionEntry : partitionCapacityMap.entrySet()) { + for (Map.Entry dimensionEntry : partitionEntry.getValue().entrySet()) { + if (dimensionEntry.getValue() != null && dimensionEntry.getValue() < 0) { + return Optional.of(badRequest(String.format( + "Partition weight for resource %s, partition '%s', dimension '%s' is negative (%d); " + + "weights must be non-negative.", resourceConfig.getResourceName(), + partitionEntry.getKey(), dimensionEntry.getKey(), dimensionEntry.getValue()))); + } + } + } + return Optional.empty(); + } + @ResponseMetered(name = HttpConstants.WRITE_REQUEST) @Timed(name = HttpConstants.WRITE_REQUEST) @POST diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java index 2602467b3b..f33ff3a1c5 100644 --- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java +++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java @@ -614,7 +614,7 @@ public void testAddResourceWithWeight() throws IOException { * cluster/instance capacity configuration is saved and restored so this test does not perturb the * other resource tests that share {@value #CLUSTER_NAME}. */ - @Test + @Test(dependsOnMethods = "testAddResourceWithWeight") public void testAddWagedResourceWeightGuardrail() throws Exception { System.out.println("Start test :" + TestHelper.getTestMethodName()); @@ -766,6 +766,89 @@ private static ResourceConfig wagedResourceConfig(String resourceName, return resourceConfig; } + private static ResourceConfig rawWagedResourceConfig(String resourceName, + Map> partitionWeights) throws IOException { + // Build PARTITION_CAPACITY_MAP directly on the record, bypassing + // ResourceConfig#setPartitionCapacityMap so values it would reject (e.g. negatives) can be + // exercised through the raw-ZNRecord path the endpoint actually uses. + ResourceConfig resourceConfig = new ResourceConfig(resourceName); + Map rawCapacityRecord = new HashMap<>(); + for (Map.Entry> entry : partitionWeights.entrySet()) { + rawCapacityRecord.put(entry.getKey(), OBJECT_MAPPER.writeValueAsString(entry.getValue())); + } + resourceConfig.getRecord().setMapField( + ResourceConfig.ResourceConfigProperty.PARTITION_CAPACITY_MAP.name(), rawCapacityRecord); + return resourceConfig; + } + + /** + * Structural checks (IdealState/ResourceConfig name match, non-negative weights) run before the + * guard rail, so a dry-run reflects them and a structurally invalid request never reaches ZK. + * Capacity configuration is saved and restored so this test does not perturb the other resource + * tests that share {@value #CLUSTER_NAME}. + */ + @Test(dependsOnMethods = "testAddResourceWithWeight") + public void testWagedStructuralChecksAppliedBeforeGuardrail() throws Exception { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + + ClusterConfig clusterConfig = _configAccessor.getClusterConfig(CLUSTER_NAME); + List originalCapacityKeys = clusterConfig.getInstanceCapacityKeys(); + List instances = + _gSetupTool.getClusterManagementTool().getInstancesInCluster(CLUSTER_NAME); + Map> originalInstanceCapacities = new HashMap<>(); + for (String instance : instances) { + originalInstanceCapacities.put(instance, + _configAccessor.getInstanceConfig(CLUSTER_NAME, instance).getInstanceCapacityMap()); + } + + String resourceName = "structuralCheckResource"; + try { + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO", "BAR")); + _configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig); + Map instanceCapacity = ImmutableMap.of("FOO", 100, "BAR", 100); + for (String instance : instances) { + InstanceConfig instanceConfig = _configAccessor.getInstanceConfig(CLUSTER_NAME, instance); + instanceConfig.setInstanceCapacityMap(instanceCapacity); + _configAccessor.setInstanceConfig(CLUSTER_NAME, instance, instanceConfig); + } + + // 1) Name mismatch is a structural failure. Even a dry-run must report it (400) instead of + // returning a feasible verdict for a request that would then fail for real. + Map> withinCapacity = ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 100, "BAR", 100)); + Response mismatchDryRun = putWagedResource(resourceName, + wagedResourceConfig("someOtherName", withinCapacity), ImmutableMap.of("dryRun", true)); + Assert.assertEquals(mismatchDryRun.getStatus(), Response.Status.BAD_REQUEST.getStatusCode()); + Assert.assertFalse(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(resourceName)); + + // 2) Negative weights are rejected (400) even though they do not exceed capacity, and even + // though the endpoint builds the ResourceConfig from a raw ZNRecord that bypasses + // ResourceConfig#setPartitionCapacityMap's own negative check. + Response negative = putWagedResource(resourceName, + rawWagedResourceConfig(resourceName, ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", -5, "BAR", 100))), + Collections.emptyMap()); + Assert.assertEquals(negative.getStatus(), Response.Status.BAD_REQUEST.getStatusCode()); + Assert.assertFalse(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(resourceName)); + } finally { + try { + _gSetupTool.getClusterManagementTool().dropResource(CLUSTER_NAME, resourceName); + } catch (Exception ignored) { + } + ClusterConfig restore = _configAccessor.getClusterConfig(CLUSTER_NAME); + restore.setInstanceCapacityKeys(originalCapacityKeys); + _configAccessor.setClusterConfig(CLUSTER_NAME, restore); + for (String instance : instances) { + InstanceConfig instanceConfig = _configAccessor.getInstanceConfig(CLUSTER_NAME, instance); + instanceConfig.setInstanceCapacityMap(originalInstanceCapacities.get(instance)); + _configAccessor.setInstanceConfig(CLUSTER_NAME, instance, instanceConfig); + } + } + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + @Test(dependsOnMethods = "testAddResourceWithWeight") public void testValidateResource() throws IOException { // Define weight keys in ClusterConfig From 3512d345883a97d7847f74bdede3238916c67a1e Mon Sep 17 00:00:00 2001 From: Sarthak Jain Date: Sun, 16 Aug 2026 11:48:21 +0530 Subject: [PATCH 5/6] Address PR review round 2: skip-comment accuracy, violation cap, coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - A (rule): Reword the maxCapacity==null skip comment. It falsely claimed WagedValidationUtil#validateAndGetInstanceCapacity reports the missing-instance-dimension case at add time; that method only runs in the rebalancer and in the separate validateInstancesForWagedRebalance admin call, neither on the addResourceWithWeight path. The comment now states the gap is deliberate and uncovered at add time (surfaces only later as a WAGED placement failure). Behavior (skip) is unchanged. - C (rule): Cap enumerated violations at MAX_REPORTED_VIOLATIONS (100) and append a single summary entry recording how many were omitted, so a resource with explicit per-partition weights breaching capacity on many partitions/dimensions cannot produce a multi-megabyte 400 body. - D (tests): Cover the two realPartitionNames branches that had no test — a CUSTOMIZED IdealState whose populated partition set is used directly, and NUM_PARTITIONS=0 where the reconstructed set is empty and every explicit per-partition weight is skipped. Adds a test for the C cap too. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../PartitionWeightCapacityGuardrailRule.java | 51 +++++++++-- ...tPartitionWeightCapacityGuardrailRule.java | 89 +++++++++++++++++++ 2 files changed, 132 insertions(+), 8 deletions(-) diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java index 97702faad9..1e7bded96a 100644 --- a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java @@ -79,6 +79,13 @@ public class PartitionWeightCapacityGuardrailRule implements GuardrailRule { public static final String RULE_ID = "PARTITION_WEIGHT_EXCEEDS_INSTANCE_CAPACITY"; + // Upper bound on the number of individual weight violations enumerated in a single verdict. A + // resource that sets explicit per-partition weights can breach capacity on every partition and + // dimension at once (e.g. a 10k-partition, 3-dimension resource is ~30k violations), which would + // otherwise produce a multi-megabyte 400 response. Beyond this cap the extra violations are + // summarized in a single trailing entry that records how many were omitted. + private static final int MAX_REPORTED_VIOLATIONS = 100; + @Override public String getId() { return RULE_ID; @@ -173,6 +180,7 @@ public ValidationResult validate(GuardrailContext context) { .thenComparing(Comparator.naturalOrder())); List violations = new ArrayList<>(); + int totalViolations = 0; for (String partitionName : orderedPartitions) { // Skip weights for partitions this resource does not actually have. The capacity map is // operator-supplied and can carry stale/typo'd entries; WAGED ignores them at placement time, @@ -202,16 +210,27 @@ public ValidationResult validate(GuardrailContext context) { } Integer maxCapacity = maxInstanceCapacity.get(dimension); if (maxCapacity == null) { - // No assignable instance advertises capacity for this dimension. That is an instance-side - // misconfiguration (a cluster-declared capacity key missing from the instances), which - // WagedValidationUtil#validateAndGetInstanceCapacity already reports against the instances. - // Treating the absent dimension as capacity 0 here would blame the resource author for an - // instance problem and tell them to lower a weight that cannot go below 0. It is also - // inconsistent with the weight == null skip above, where missing key coverage is likewise - // deferred to that separate validation. So skip this dimension rather than over-block. + // No assignable instance advertises capacity for this dimension: a cluster-declared + // capacity key that is missing from every instance. This is an instance-side + // misconfiguration, not a fault of the resource being added, so we deliberately skip it + // rather than fail the add. Treating the absent dimension as capacity 0 would blame the + // resource author and tell them to lower a weight that cannot go below 0, and in this + // state every WAGED resource is already unplaceable, not just this one. + // + // This gap is intentionally left uncovered here: nothing on the addResourceWithWeight + // path validates instance-side capacity coverage. WagedValidationUtil# + // validateAndGetInstanceCapacity runs only inside the rebalancer and from the separate + // validateInstancesForWagedRebalance admin call, neither of which is on this path, so such + // a resource is accepted at add time and only surfaces later as a WAGED placement failure. continue; } if (weight > maxCapacity) { + totalViolations++; + // Enumerate at most MAX_REPORTED_VIOLATIONS; any overflow is summarized after the loop so + // a pathological resource cannot return a multi-megabyte body. + if (violations.size() >= MAX_REPORTED_VIOLATIONS) { + continue; + } // DEFAULT_PARTITION_KEY is a placeholder for "every partition", not a real partition, so // report it as unscoped for a clearer message. String reportedPartition = @@ -232,7 +251,23 @@ public ValidationResult validate(GuardrailContext context) { } } - return violations.isEmpty() ? ValidationResult.feasible() : ValidationResult.of(violations); + if (violations.isEmpty()) { + return ValidationResult.feasible(); + } + if (totalViolations > violations.size()) { + // More partitions/dimensions breached capacity than we enumerated. Record the overflow so the + // caller knows the list is truncated and by how much, instead of silently dropping them. + int reported = violations.size(); + violations.add(Violation.newBuilder(RULE_ID) + .resource(proposedResourceConfig.getResourceName()) + .message(String.format( + "Showing the first %d of %d partition-weight violations; %d were omitted to bound the " + + "response size. The omitted violations are further partitions breaching the same " + + "dimension(s); fix the reported dimensions and resubmit.", + reported, totalViolations, totalViolations - reported)) + .build()); + } + return ValidationResult.of(violations); } /** diff --git a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java index 39853b7b11..ac35be2b69 100644 --- a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.Arrays; +import java.util.HashMap; import java.util.List; import java.util.Map; @@ -283,6 +284,83 @@ public void testMultipleViolationsAllReportedInFixedOrder() throws IOException { Assert.assertFalse(result.getViolations().get(0).getMessage().contains("force")); } + @Test + public void testCustomizedIdealStatePartitionsUsedDirectly() throws IOException { + // A CUSTOMIZED ideal state already carries its partitions (in its map fields), so + // getPartitionSet() is non-empty and realPartitionNames returns those names directly rather than + // reconstructing _ from the partition count. A weight on one of those real + // partitions is evaluated, while a canonical-scheme name that is NOT among them is a ghost. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + IdealState idealState = new IdealState(RESOURCE); + idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED); + idealState.setPartitionState("realPartition", "instance0", "MASTER"); + + ResourceConfig resourceConfig = resourceConfig(ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 1), + "realPartition", ImmutableMap.of("FOO", 5000), + RESOURCE + "_0", ImmutableMap.of("FOO", 5000))); + + ValidationResult result = rule.validate(contextWith(dataAccessor, resourceConfig, idealState)); + + // Only the real (declared) partition is reported; RESOURCE_0 is a ghost under this ideal state. + Assert.assertFalse(result.isFeasible()); + Assert.assertEquals(result.getViolations().size(), 1); + Assert.assertEquals(result.getViolations().get(0).getPartitionName(), "realPartition"); + } + + @Test + public void testZeroPartitionsSkipsAllExplicitWeights() throws IOException { + // NUM_PARTITIONS=0 with an empty assignment reconstructs an empty partition set, so every + // explicit per-partition weight names a partition the resource does not have and is skipped as a + // ghost. The verdict is feasible even though the weight would exceed capacity for a real + // partition. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + ResourceConfig resourceConfig = resourceConfig(ImmutableMap.of( + ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 1), + RESOURCE + "_0", ImmutableMap.of("FOO", 5000))); + + ValidationResult result = rule.validate(contextWith(dataAccessor, resourceConfig, 0)); + Assert.assertTrue(result.isFeasible()); + } + + @Test + public void testViolationsCappedWithSuppressedCount() throws IOException { + // A resource that breaches capacity on many partitions at once must not enumerate every + // violation; the list is capped and a trailing summary records how many were omitted. + int partitionCount = 150; + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + HelixDataAccessor dataAccessor = mockAccessor(clusterConfig, + ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))); + + Map> weights = new HashMap<>(); + weights.put(ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 1)); + for (int i = 0; i < partitionCount; i++) { + weights.put(RESOURCE + "_" + i, ImmutableMap.of("FOO", 5000)); + } + ResourceConfig resourceConfig = resourceConfig(weights); + + ValidationResult result = + rule.validate(contextWith(dataAccessor, resourceConfig, partitionCount)); + + Assert.assertFalse(result.isFeasible()); + // 100 enumerated violations plus a single trailing summary entry. + Assert.assertEquals(result.getViolations().size(), 101); + Violation summary = result.getViolations().get(result.getViolations().size() - 1); + Assert.assertTrue(summary.getMessage().contains("150")); + Assert.assertTrue(summary.getMessage().contains("omitted")); + // The summary is not scoped to any single partition. + Assert.assertNull(summary.getPartitionName()); + } + private GuardrailContext contextWith(HelixDataAccessor dataAccessor, ResourceConfig proposedResourceConfig) { // Default to a single-partition resource so the canonical testResource_0 partition is real. @@ -303,6 +381,17 @@ private GuardrailContext contextWith(HelixDataAccessor dataAccessor, .build(); } + private GuardrailContext contextWith(HelixDataAccessor dataAccessor, + ResourceConfig proposedResourceConfig, IdealState proposedIdealState) { + // Use a caller-supplied ideal state (e.g. a CUSTOMIZED one with an already-populated partition + // set) instead of one reconstructed from a partition count. + return GuardrailContext.newBuilder(CLUSTER) + .dataAccessor(dataAccessor) + .proposedResourceConfig(proposedResourceConfig) + .proposedIdealState(proposedIdealState) + .build(); + } + private HelixDataAccessor mockAccessor(ClusterConfig clusterConfig, List instanceConfigs) { HelixDataAccessor dataAccessor = mock(HelixDataAccessor.class); From 5c45e20566d2cad5adf12a6fae3b3ed639dd5575 Mon Sep 17 00:00:00 2001 From: Sarthak Jain Date: Sun, 16 Aug 2026 12:07:23 +0530 Subject: [PATCH 6/6] Make partition-weight guard rail opt-in via ClusterConfig kill switch Addresses review comment on the addWagedResource guard rail: the pipeline was constructed inline and ran unconditionally for every caller, with no cluster- or server-level way to disable it. Combined with the fail-closed instance-config scan, a single unreadable instance-config znode could take addWagedResource down cluster-wide, and force=true only bypasses the rule per request. Add an opt-in ClusterConfig flag, PARTITION_WEIGHT_GUARDRAIL_ENABLED (default OFF), checked inside the rule after the null-config check and BEFORE the instance-config scan: - Enabling the guard rail is now a deliberate, per-cluster decision, so it can be rolled out in stages instead of turning on everywhere at once. - Disabling it is a single ClusterConfig change that backs the rule out for every caller with no client change and no helix-rest redeploy - a proper kill switch for a false positive. - Because the flag is checked before the instance read, a disabled cluster never runs the fail-closed scan, so an unreadable znode cannot break addWagedResource on clusters that have not opted in. The fail-closed read is kept for clusters that do opt in; the flag is the escape hatch. Tests: - ClusterConfig: add isPartitionWeightGuardrailEnabled/ setPartitionWeightGuardrailEnabled (getBooleanField default false). - Unit: enable the flag centrally in the mockAccessor helper so the existing enforcement tests keep exercising the rule; add testGuardrailDisabledByDefaultShortCircuits (stubs the instance scan to throw and asserts it is never reached) and testGuardrailExplicitlyDisabledAllowsOverCapacity. - Integration: testAddWagedResourceWeightGuardrail now asserts an over-capacity resource is allowed through with the flag off (default), enables the flag for the enforcement/dry-run/force/within-capacity steps, and disables it again on restore. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../PartitionWeightCapacityGuardrailRule.java | 16 +++++++ .../org/apache/helix/model/ClusterConfig.java | 31 +++++++++++++ ...tPartitionWeightCapacityGuardrailRule.java | 46 +++++++++++++++++++ .../rest/server/TestResourceAccessor.java | 21 ++++++++- 4 files changed, 112 insertions(+), 2 deletions(-) diff --git a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java index 1e7bded96a..acbe6284ef 100644 --- a/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java @@ -75,6 +75,14 @@ * rule skips any weight-map key that is neither {@code DEFAULT} nor a real partition of the proposed * ideal state — blocking on a partition that will never exist would be a false positive * stricter than the operation it fronts. + *

+ * Opt-in. This guard rail runs only when the cluster explicitly enables it via + * {@link ClusterConfig#setPartitionWeightGuardrailEnabled(boolean)}; it is disabled by default. That + * makes turning it on a deliberate per-cluster decision and, just as importantly, gives operators a + * single-config-change kill switch: if the rule ever produces a false positive, disabling it via + * ClusterConfig immediately backs it out for every caller with no client change and no helix-rest + * redeploy. When the cluster has it disabled the rule returns feasible before reading any instance + * config, so a disabled cluster is never exposed to the fail-closed instance-config scan below. */ public class PartitionWeightCapacityGuardrailRule implements GuardrailRule { public static final String RULE_ID = "PARTITION_WEIGHT_EXCEEDS_INSTANCE_CAPACITY"; @@ -107,6 +115,14 @@ public ValidationResult validate(GuardrailContext context) { return ValidationResult.feasible(); } + if (!clusterConfig.isPartitionWeightGuardrailEnabled()) { + // Opt-in guard rail, disabled by default. Returning here (before the instance-config scan + // below) is also the kill switch: disabling the rule via ClusterConfig backs it out for every + // caller with a single config change, and a disabled cluster never runs the fail-closed + // instance-config read, so one unreadable znode cannot take addWagedResource down. + return ValidationResult.feasible(); + } + List capacityKeys = clusterConfig.getInstanceCapacityKeys(); if (capacityKeys.isEmpty()) { // Cluster does not use the WAGED capacity/weight model, so weights carry no meaning here. diff --git a/helix-core/src/main/java/org/apache/helix/model/ClusterConfig.java b/helix-core/src/main/java/org/apache/helix/model/ClusterConfig.java index 2c8dfcf4ac..e6635d5298 100644 --- a/helix-core/src/main/java/org/apache/helix/model/ClusterConfig.java +++ b/helix-core/src/main/java/org/apache/helix/model/ClusterConfig.java @@ -130,6 +130,12 @@ public enum ClusterConfigProperty { DEFAULT_INSTANCE_CAPACITY_MAP, // The default partition weights if no weight is configured in the Resource Config node. DEFAULT_PARTITION_WEIGHT_MAP, + // Opt-in toggle for the helix-rest PartitionWeightCapacityGuardrailRule, which pre-validates an + // addWagedResource request and rejects a resource whose partition weight exceeds the largest + // single instance capacity (making it permanently unplaceable). Disabled by default so enabling + // it is a deliberate per-cluster decision; it can be turned off again with a single ClusterConfig + // change (no client change or helix-rest redeploy) to back out a false positive. + PARTITION_WEIGHT_GUARDRAIL_ENABLED, // The preference of the rebalance result. // EVENNESS - Evenness of the resource utilization, partition, and top state distribution. // LESS_MOVEMENT - the tendency of keeping the current assignment instead of moving the partition for optimal assignment. @@ -1192,6 +1198,31 @@ public void setDefaultPartitionWeightMap(Map weightDataMap) setDefaultCapacityMap(ClusterConfigProperty.DEFAULT_PARTITION_WEIGHT_MAP, weightDataMap); } + /** + * Whether the helix-rest partition-weight capacity guard rail is enabled for this cluster. When + * enabled, an addWagedResource request is pre-validated and rejected before any ZooKeeper write if + * a partition weight exceeds the largest single instance capacity in any dimension (which would + * make the resource permanently unplaceable and stall the WAGED global rebalance cluster-wide). + *

+ * Disabled by default: enabling the guard rail is an opt-in, per-cluster decision, and it can be + * turned off again with a single ClusterConfig change to back out a false positive without + * changing any client or redeploying helix-rest. + * @return true if the guard rail is enabled; false (the default) otherwise. + */ + public boolean isPartitionWeightGuardrailEnabled() { + return _record.getBooleanField( + ClusterConfigProperty.PARTITION_WEIGHT_GUARDRAIL_ENABLED.name(), false); + } + + /** + * Enable or disable the helix-rest partition-weight capacity guard rail for this cluster. + * @param enabled true to enable the guard rail, false to disable it. + */ + public void setPartitionWeightGuardrailEnabled(boolean enabled) { + _record.setBooleanField( + ClusterConfigProperty.PARTITION_WEIGHT_GUARDRAIL_ENABLED.name(), enabled); + } + private Map getDefaultCapacityMap(ClusterConfigProperty capacityPropertyType) { Map capacityData = _record.getMapField(capacityPropertyType.name()); if (capacityData != null) { diff --git a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java index ac35be2b69..5ed1fe9155 100644 --- a/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java +++ b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java @@ -41,6 +41,7 @@ import org.testng.annotations.Test; import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -103,6 +104,46 @@ public void testNoInstanceCapacityIsFeasible() throws IOException { Assert.assertTrue(result.isFeasible()); } + @Test + public void testGuardrailDisabledByDefaultShortCircuits() throws IOException { + // The guard rail is opt-in: with the flag left unset (its default), an over-capacity weight that + // would otherwise be flagged is allowed through, AND the fail-closed instance-config scan is + // never performed. getChildValues is stubbed to throw so this test fails loudly if the + // short-circuit ever regresses and the rule reaches the scan on a disabled cluster. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + // Flag intentionally left unset -> defaults to false. + HelixDataAccessor dataAccessor = mock(HelixDataAccessor.class); + when(dataAccessor.keyBuilder()).thenReturn(BUILDER); + doReturn(clusterConfig).when(dataAccessor).getProperty(BUILDER.clusterConfig()); + doThrow(new RuntimeException("instance-config scan must not run when the guard rail is disabled")) + .when(dataAccessor).getChildValues(BUILDER.instanceConfigs(), true); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 5000))))); + Assert.assertTrue(result.isFeasible()); + } + + @Test + public void testGuardrailExplicitlyDisabledAllowsOverCapacity() throws IOException { + // Explicitly setting the flag to false is equivalent to leaving it unset: an over-capacity + // weight that the enabled rule would flag is allowed through. + ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); + clusterConfig.setInstanceCapacityKeys(Arrays.asList("FOO")); + clusterConfig.setPartitionWeightGuardrailEnabled(false); + HelixDataAccessor dataAccessor = mock(HelixDataAccessor.class); + when(dataAccessor.keyBuilder()).thenReturn(BUILDER); + doReturn(clusterConfig).when(dataAccessor).getProperty(BUILDER.clusterConfig()); + doReturn(ImmutableList.of(instanceConfig("instance0", ImmutableMap.of("FOO", 100)))) + .when(dataAccessor).getChildValues(BUILDER.instanceConfigs(), true); + + ValidationResult result = rule.validate(contextWith(dataAccessor, + resourceConfig(ImmutableMap.of(ResourceConfig.DEFAULT_PARTITION_KEY, + ImmutableMap.of("FOO", 5000))))); + Assert.assertTrue(result.isFeasible()); + } + @Test public void testWeightWithinCapacityIsFeasible() throws IOException { ClusterConfig clusterConfig = new ClusterConfig(CLUSTER); @@ -394,6 +435,11 @@ private GuardrailContext contextWith(HelixDataAccessor dataAccessor, private HelixDataAccessor mockAccessor(ClusterConfig clusterConfig, List instanceConfigs) { + // The guard rail is opt-in (disabled by default) on a real cluster, but these unit tests exist + // to exercise its enforcement, which only runs when enabled. Enable it here so each enforcement + // test does not have to repeat it; the disabled-cluster behavior is covered explicitly by the + // testGuardrail*Disabled* cases, which build their accessor without this helper. + clusterConfig.setPartitionWeightGuardrailEnabled(true); HelixDataAccessor dataAccessor = mock(HelixDataAccessor.class); when(dataAccessor.keyBuilder()).thenReturn(BUILDER); doReturn(clusterConfig).when(dataAccessor).getProperty(BUILDER.clusterConfig()); diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java index f33ff3a1c5..0f7e417416 100644 --- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java +++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestResourceAccessor.java @@ -631,6 +631,7 @@ public void testAddWagedResourceWeightGuardrail() throws Exception { String blockedResource = "guardrailBlockedWagedResource"; String forcedResource = "guardrailForcedWagedResource"; String validResource = "guardrailValidWagedResource"; + String disabledResource = "guardrailDisabledWagedResource"; try { // Declare two capacity dimensions and give every instance capacity 100 in each. @@ -647,6 +648,19 @@ public void testAddWagedResourceWeightGuardrail() throws Exception { Map> overWeight = ImmutableMap.of( ResourceConfig.DEFAULT_PARTITION_KEY, ImmutableMap.of("FOO", 1000, "BAR", 100)); + // 0) Opt-in: the guard rail is disabled by default, so an over-capacity resource is allowed + // through and actually created even without force=true. + Response disabled = putWagedResource(disabledResource, + wagedResourceConfig(disabledResource, overWeight), Collections.emptyMap()); + Assert.assertEquals(disabled.getStatus(), Response.Status.OK.getStatusCode()); + Assert.assertTrue(_gSetupTool.getClusterManagementTool().getResourcesInCluster(CLUSTER_NAME) + .contains(disabledResource)); + + // Enable the guard rail for the remainder of the test (opt-in per cluster). + clusterConfig = _configAccessor.getClusterConfig(CLUSTER_NAME); + clusterConfig.setPartitionWeightGuardrailEnabled(true); + _configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig); + // 1) Enforcement: blocked with 400 + verdict, and nothing written to ZK. Response blocked = putWagedResource(blockedResource, wagedResourceConfig(blockedResource, overWeight), Collections.emptyMap()); @@ -686,15 +700,18 @@ public void testAddWagedResourceWeightGuardrail() throws Exception { .contains(validResource)); } finally { // Drop any resources this test created (blockedResource was never created; ignore failures). - for (String resource : Arrays.asList(forcedResource, validResource, blockedResource)) { + for (String resource : Arrays.asList(forcedResource, validResource, disabledResource, + blockedResource)) { try { _gSetupTool.getClusterManagementTool().dropResource(CLUSTER_NAME, resource); } catch (Exception ignored) { } } - // Restore cluster + instance capacity configuration to its original values. + // Restore cluster + instance capacity configuration to its original values, and disable the + // opt-in guard rail again so it does not leak into other tests sharing this cluster. ClusterConfig restore = _configAccessor.getClusterConfig(CLUSTER_NAME); restore.setInstanceCapacityKeys(originalCapacityKeys); + restore.setPartitionWeightGuardrailEnabled(false); _configAccessor.setClusterConfig(CLUSTER_NAME, restore); for (String instance : instances) { InstanceConfig instanceConfig = _configAccessor.getInstanceConfig(CLUSTER_NAME, instance);