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..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,25 +20,33 @@ package org.apache.helix.guardrail; import org.apache.helix.HelixDataAccessor; +import org.apache.helix.model.IdealState; +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} 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() { @@ -54,6 +62,27 @@ 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; + } + + /** + * 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); } @@ -62,6 +91,8 @@ public static final class Builder { private final String clusterName; private ReadOnlyDataAccessor dataAccessor; private String instanceName; + private ResourceConfig proposedResourceConfig; + private IdealState proposedIdealState; private Builder(String clusterName) { this.clusterName = clusterName; @@ -77,6 +108,16 @@ public Builder instanceName(String instanceName) { return this; } + public Builder proposedResourceConfig(ResourceConfig proposedResourceConfig) { + this.proposedResourceConfig = 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 new file mode 100644 index 0000000000..acbe6284ef --- /dev/null +++ b/helix-core/src/main/java/org/apache/helix/guardrail/rules/PartitionWeightCapacityGuardrailRule.java @@ -0,0 +1,316 @@ +/* + * 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.ArrayList; +import java.util.Collections; +import java.util.Comparator; +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; +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.IdealState; +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 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}. + * 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. + *

+ * 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 + * {@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. + *

+ * 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"; + + // 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; + } + + @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(); + } + + 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. + return ValidationResult.feasible(); + } + + // 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()) { + maxInstanceCapacity.merge(entry.getKey(), entry.getValue(), Math::max); + } + } + + if (maxInstanceCapacity.isEmpty()) { + // 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(); + } + + 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(); + Set realPartitions = realPartitionNames(context.getProposedIdealState()); + + // 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<>(); + 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, + // 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); + effectiveWeight.putAll(partitionCapacityMap.get(partitionName)); + + // 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) { + continue; + } + Integer maxCapacity = maxInstanceCapacity.get(dimension); + if (maxCapacity == null) { + // 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 = + ResourceConfig.DEFAULT_PARTITION_KEY.equals(partitionName) ? null : partitionName; + // 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 or raise instance capacity.", weight, dimension, maxCapacity, + reportedPartition == null ? "every partition" : "partition " + reportedPartition)) + .build()); + } + } + } + + 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); + } + + /** + * 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/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 new file mode 100644 index 0000000000..5ed1fe9155 --- /dev/null +++ b/helix-core/src/test/java/org/apache/helix/guardrail/rules/TestPartitionWeightCapacityGuardrailRule.java @@ -0,0 +1,469 @@ +/* + * 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.HashMap; +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.constants.InstanceConstants; +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.IdealState; +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.doThrow; +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 and ideal state are 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 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); + 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()); + } + + @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"); + } + + @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()); + } + + @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")); + } + + @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. + 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(); + } + + 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) { + // 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()); + 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 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); + 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..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 @@ -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 { @@ -248,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) { @@ -296,10 +312,38 @@ 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); + 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. + GuardrailContext context = GuardrailContext.newBuilder(clusterId) + .dataAccessor(getDataAccssor(clusterId)) + .proposedResourceConfig(proposedResourceConfig) + .proposedIdealState(proposedIdealState) + .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)); + 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); @@ -318,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 2bb91d416c..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 @@ -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,266 @@ 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(dependsOnMethods = "testAddResourceWithWeight") + 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"; + String disabledResource = "guardrailDisabledWagedResource"; + + 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)); + + // 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()); + 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, disabledResource, + blockedResource)) { + try { + _gSetupTool.getClusterManagementTool().dropResource(CLUSTER_NAME, resource); + } catch (Exception ignored) { + } + } + // 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); + instanceConfig.setInstanceCapacityMap(originalInstanceCapacities.get(instance)); + _configAccessor.setInstanceConfig(CLUSTER_NAME, instance, instanceConfig); + } + } + 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); + 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; + } + + 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