Add partition-weight guard rail for addWagedResource

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.
  * <p>
  * 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 <em>proposed</em> 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 <em>count</em>
+   * ({@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);
     }