PartitionMovementSoftConstraint Implementation

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/PartitionMovementConstraint.java

@@ -0,0 +1,96 @@
+package org.apache.helix.controller.rebalancer.waged.constraints;
+
+/*
+ * 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.
+ */
+
+import java.util.Collections;
+import java.util.Map;
+
+import org.apache.helix.controller.rebalancer.waged.model.AssignableNode;
+import org.apache.helix.controller.rebalancer.waged.model.AssignableReplica;
+import org.apache.helix.controller.rebalancer.waged.model.ClusterContext;
+import org.apache.helix.model.Partition;
+import org.apache.helix.model.ResourceAssignment;
+
+/**
+ * Evaluate the proposed assignment according to the potential partition movements cost.
+ * The cost is evaluated based on the difference between the old assignment and the new assignment.
+ * In detail, we consider the following two previous assignments as the base.
+ * - Baseline assignment that is calculated regardless of the node state (online/offline).
+ * - Previous Best Possible assignment.
+ * Any change to these two assignments will increase the partition movements cost, so that the
+ * evaluated score will become lower.
+ */
+class PartitionMovementConstraint extends SoftConstraint {
+  private static final float MAX_SCORE = 1f;
+  private static final float MIN_SCORE = 0f;
+  //TODO: these factors will be tuned based on user's preference
+  // This factor indicates the default score that is evaluated if only partition allocation matches
+  // (states are different).
+  private static final float ALLOCATION_MATCH_FACTOR = 0.5f;
+  // This factor indicates the contribution of the Baseline assignment matching to the final score.
+  private static final float BASELINE_MATCH_FACTOR = 0.25f;
+
+  PartitionMovementConstraint() {
+    super(MAX_SCORE, MIN_SCORE);
+  }
+
+  @Override
+  protected float getAssignmentScore(AssignableNode node, AssignableReplica replica,
+      ClusterContext clusterContext) {
+    Map<String, String> bestPossibleStateMap =
+        getStateMap(replica, clusterContext.getBestPossibleAssignment());
+    Map<String, String> baselineStateMap =
+        getStateMap(replica, clusterContext.getBaselineAssignment());
+
+    // Prioritize the matching of the previous Best Possible assignment.
+    float scale = calculateAssignmentScale(node, replica, bestPossibleStateMap);
+    // If the baseline is also provided, adjust the final score accordingly.
+    scale = scale * (1 - BASELINE_MATCH_FACTOR)
+        + calculateAssignmentScale(node, replica, baselineStateMap) * BASELINE_MATCH_FACTOR;
+
+    return scale;
+  }
+
+  @Override
+  NormalizeFunction getNormalizeFunction() {
+    return score -> score * (getMaxScore() - getMinScore()) + getMinScore();
+  }
+
+  private Map<String, String> getStateMap(AssignableReplica replica,
+      Map<String, ResourceAssignment> assignment) {
+    String resourceName = replica.getResourceName();
+    String partitionName = replica.getPartitionName();
+    if (assignment == null || !assignment.containsKey(resourceName)) {
+      return Collections.emptyMap();
+    }
+    return assignment.get(resourceName).getReplicaMap(new Partition(partitionName));
+  }
+
+  private float calculateAssignmentScale(AssignableNode node, AssignableReplica replica,
+      Map<String, String> instanceToStateMap) {
+    String instanceName = node.getInstanceName();
+    if (!instanceToStateMap.containsKey(instanceName)) {
+      return 0;
+    } else {
+      return (instanceToStateMap.get(instanceName).equals(replica.getReplicaState()) ? 1
+          : ALLOCATION_MATCH_FACTOR);
+    }
+  }
+}

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterContext.java

@@ -28,6 +28,8 @@
 import java.util.stream.Collectors;
 
 import org.apache.helix.HelixException;
+import org.apache.helix.model.ResourceAssignment;
+
 
 /**
  * This class tracks the rebalance-related global cluster status.
@@ -44,30 +46,47 @@ public class ClusterContext {
 
   // map{zoneName : map{resourceName : set(partitionNames)}}
   private Map<String, Map<String, Set<String>>> _assignmentForFaultZoneMap = new HashMap<>();
+  // Records about the previous assignment
+  // <ResourceName, ResourceAssignment contains the baseline assignment>
+  private final Map<String, ResourceAssignment> _baselineAssignment;
+  // <ResourceName, ResourceAssignment contains the best possible assignment>
+  private final Map<String, ResourceAssignment> _bestPossibleAssignment;
 
   /**
    * Construct the cluster context based on the current instance status.
    * @param replicaSet All the partition replicas that are managed by the rebalancer
    * @param instanceCount The count of all the active instances that can be used to host partitions.
    */
-  ClusterContext(Set<AssignableReplica> replicaSet, int instanceCount) {
+  ClusterContext(Set<AssignableReplica> replicaSet, int instanceCount,
+      Map<String, ResourceAssignment> baselineAssignment, Map<String, ResourceAssignment> bestPossibleAssignment) {
     int totalReplicas = 0;
     int totalTopStateReplicas = 0;
 
     for (Map.Entry<String, List<AssignableReplica>> entry : replicaSet.stream()
-        .collect(Collectors.groupingBy(AssignableReplica::getResourceName)).entrySet()) {
+        .collect(Collectors.groupingBy(AssignableReplica::getResourceName))
+        .entrySet()) {
       int replicas = entry.getValue().size();
       totalReplicas += replicas;
 
       int replicaCnt = Math.max(1, estimateAvgReplicaCount(replicas, instanceCount));
       _estimatedMaxPartitionByResource.put(entry.getKey(), replicaCnt);
 
-      totalTopStateReplicas +=
-          entry.getValue().stream().filter(AssignableReplica::isReplicaTopState).count();
+      totalTopStateReplicas += entry.getValue().stream().filter(AssignableReplica::isReplicaTopState).count();
     }
 
     _estimatedMaxPartitionCount = estimateAvgReplicaCount(totalReplicas, instanceCount);
     _estimatedMaxTopStateCount = estimateAvgReplicaCount(totalTopStateReplicas, instanceCount);
+    _baselineAssignment = baselineAssignment;
+    _bestPossibleAssignment = bestPossibleAssignment;
+  }
+
+  public Map<String, ResourceAssignment> getBaselineAssignment() {
+    return _baselineAssignment == null || _baselineAssignment.isEmpty() ? Collections.emptyMap() : _baselineAssignment;
+  }
+
+  public Map<String, ResourceAssignment> getBestPossibleAssignment() {
+    return _bestPossibleAssignment == null || _bestPossibleAssignment.isEmpty() ? Collections.emptyMap()
+        : _bestPossibleAssignment;
   }
 
   public Map<String, Map<String, Set<String>>> getAssignmentForFaultZoneMap() {
@@ -93,25 +112,25 @@ public Set<String> getPartitionsForResourceAndFaultZone(String resourceName, Str
 
   void addPartitionToFaultZone(String faultZoneId, String resourceName, String partition) {
     if (!_assignmentForFaultZoneMap.computeIfAbsent(faultZoneId, k -> new HashMap<>())
-        .computeIfAbsent(resourceName, k -> new HashSet<>()).add(partition)) {
+        .computeIfAbsent(resourceName, k -> new HashSet<>())
+        .add(partition)) {
       throw new HelixException(
-          String.format("Resource %s already has a replica from partition %s in fault zone %s",
-              resourceName, partition, faultZoneId));
+          String.format("Resource %s already has a replica from partition %s in fault zone %s", resourceName, partition,
+              faultZoneId));
     }
   }
 
   boolean removePartitionFromFaultZone(String faultZoneId, String resourceName, String partition) {
     return _assignmentForFaultZoneMap.getOrDefault(faultZoneId, Collections.emptyMap())
-        .getOrDefault(resourceName, Collections.emptySet()).remove(partition);
+        .getOrDefault(resourceName, Collections.emptySet())
+        .remove(partition);
   }
 
-  void setAssignmentForFaultZoneMap(
-      Map<String, Map<String, Set<String>>> assignmentForFaultZoneMap) {
+  void setAssignmentForFaultZoneMap(Map<String, Map<String, Set<String>>> assignmentForFaultZoneMap) {
     _assignmentForFaultZoneMap = assignmentForFaultZoneMap;
   }
 
   private int estimateAvgReplicaCount(int replicaCount, int instanceCount) {
-    return (int) Math
-        .ceil((float) replicaCount / instanceCount * ERROR_MARGIN_FOR_ESTIMATED_MAX_COUNT);
+    return (int) Math.ceil((float) replicaCount / instanceCount * ERROR_MARGIN_FOR_ESTIMATED_MAX_COUNT);
   }
 }

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModel.java

@@ -19,14 +19,13 @@
  * under the License.
  */
 
-import org.apache.helix.HelixException;
-import org.apache.helix.model.ResourceAssignment;
-
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import org.apache.helix.HelixException;
+
 /**
  * This class wraps the required input for the rebalance algorithm.
  */
@@ -39,23 +38,14 @@ public class ClusterModel {
   private final Map<String, Map<String, AssignableReplica>> _assignableReplicaIndex;
   private final Map<String, AssignableNode> _assignableNodeMap;
 
-  // Records about the previous assignment
-  // <ResourceName, ResourceAssignment contains the baseline assignment>
-  private final Map<String, ResourceAssignment> _baselineAssignment;
-  // <ResourceName, ResourceAssignment contains the best possible assignment>
-  private final Map<String, ResourceAssignment> _bestPossibleAssignment;
-
   /**
    * @param clusterContext         The initialized cluster context.
    * @param assignableReplicas     The replicas to be assigned.
    *                               Note that the replicas in this list shall not be included while initializing the context and assignable nodes.
    * @param assignableNodes        The active instances.
-   * @param baselineAssignment     The recorded baseline assignment.
-   * @param bestPossibleAssignment The current best possible assignment.
    */
   ClusterModel(ClusterContext clusterContext, Set<AssignableReplica> assignableReplicas,
-      Set<AssignableNode> assignableNodes, Map<String, ResourceAssignment> baselineAssignment,
-      Map<String, ResourceAssignment> bestPossibleAssignment) {
+      Set<AssignableNode> assignableNodes) {
     _clusterContext = clusterContext;
 
     // Save all the to be assigned replication
@@ -70,9 +60,6 @@ public class ClusterModel {
 
     _assignableNodeMap = assignableNodes.stream()
         .collect(Collectors.toMap(AssignableNode::getInstanceName, node -> node));
-
-    _baselineAssignment = baselineAssignment;
-    _bestPossibleAssignment = bestPossibleAssignment;
   }
 
   public ClusterContext getContext() {
@@ -87,14 +74,6 @@ public Map<String, Set<AssignableReplica>> getAssignableReplicaMap() {
     return _assignableReplicaMap;
   }
 
-  public Map<String, ResourceAssignment> getBaseline() {
-    return _baselineAssignment;
-  }
-
-  public Map<String, ResourceAssignment> getBestPossibleAssignment() {
-    return _bestPossibleAssignment;
-  }
-
   /**
    * Assign the given replica to the specified instance and record the assignment in the cluster model.
    * The cluster usage information will be updated accordingly.

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModelProvider.java

@@ -86,12 +86,11 @@ public static ClusterModel generateClusterModel(ResourceControllerDataProvider d
     // Construct and initialize cluster context.
     ClusterContext context = new ClusterContext(
         replicaMap.values().stream().flatMap(Set::stream).collect(Collectors.toSet()),
-        activeInstances.size());
+        activeInstances.size(), baselineAssignment, bestPossibleAssignment);
     // Initial the cluster context with the allocated assignments.
     context.setAssignmentForFaultZoneMap(mapAssignmentToFaultZone(assignableNodes));
 
-    return new ClusterModel(context, toBeAssignedReplicas, assignableNodes, baselineAssignment,
-        bestPossibleAssignment);
+    return new ClusterModel(context, toBeAssignedReplicas, assignableNodes);
   }
 
   /**