From dfb5d01a414426ca3700c03491a68841b9ba8c83 Mon Sep 17 00:00:00 2001
From: jiajunwang <jjwang@linkedin.com>
Date: Fri, 9 Aug 2019 15:22:51 -0700
Subject: [PATCH] Change the rebalancer assignment record to be
 ResourceAssignment instead of IdealState.

ResourceAssignment fit the usage better. And there will be no unnecessary information to be recorded or read during the rebalance calculation.
---
 .../rebalancer/waged/RebalanceAlgorithm.java    |  3 ++-
 .../ConstraintsRebalanceAlgorithm.java          |  3 ++-
 .../rebalancer/waged/model/ClusterModel.java    | 17 +++++++++--------
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java
index 0e6c89122c..ae258ca637 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java
@@ -22,6 +22,7 @@
 import org.apache.helix.controller.rebalancer.waged.constraints.HardConstraint;
 import org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
 import org.apache.helix.model.IdealState;
+import org.apache.helix.model.ResourceAssignment;
 
 import java.util.Map;
 
@@ -40,6 +41,6 @@ public interface RebalanceAlgorithm {
    *                       If the map is null, no failure will be returned.
    * @return A map <ResourceName, FailureReason> of the rebalanced resource assignments that are saved in the IdeaStates.
    */
-  Map<String, IdealState> rebalance(ClusterModel clusterModel,
+  Map<String, ResourceAssignment> rebalance(ClusterModel clusterModel,
       Map<String, Map<HardConstraint.FailureReason, Integer>> failureReasons);
 }
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintsRebalanceAlgorithm.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintsRebalanceAlgorithm.java
index 292d90383f..a75854a4cc 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintsRebalanceAlgorithm.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintsRebalanceAlgorithm.java
@@ -22,6 +22,7 @@
 import org.apache.helix.controller.rebalancer.waged.RebalanceAlgorithm;
 import org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
 import org.apache.helix.model.IdealState;
+import org.apache.helix.model.ResourceAssignment;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -42,7 +43,7 @@ public ConstraintsRebalanceAlgorithm() {
   }
 
   @Override
-  public Map<String, IdealState> rebalance(ClusterModel clusterModel,
+  public Map<String, ResourceAssignment> rebalance(ClusterModel clusterModel,
       Map<String, Map<HardConstraint.FailureReason, Integer>> failureReasons) {
     return new HashMap<>();
   }
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModel.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModel.java
index 2908939032..1be527a80f 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModel.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModel.java
@@ -21,6 +21,7 @@
 
 import org.apache.helix.HelixException;
 import org.apache.helix.model.IdealState;
+import org.apache.helix.model.ResourceAssignment;
 
 import java.util.Collections;
 import java.util.Map;
@@ -40,10 +41,10 @@ public class ClusterModel {
   private final Map<String, AssignableNode> _assignableNodeMap;
 
   // Records about the previous assignment
-  // <ResourceName, IdealState contains the baseline assignment>
-  private final Map<String, IdealState> _baselineAssignment;
-  // <ResourceName, IdealState contains the best possible assignment>
-  private final Map<String, IdealState> _bestPossibleAssignment;
+  // <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.
@@ -54,8 +55,8 @@ public class ClusterModel {
    * @param bestPossibleAssignment The current best possible assignment.
    */
   ClusterModel(ClusterContext clusterContext, Set<AssignableReplica> assignableReplicas,
-      Set<AssignableNode> assignableNodes, Map<String, IdealState> baselineAssignment,
-      Map<String, IdealState> bestPossibleAssignment) {
+      Set<AssignableNode> assignableNodes, Map<String, ResourceAssignment> baselineAssignment,
+      Map<String, ResourceAssignment> bestPossibleAssignment) {
     _clusterContext = clusterContext;
 
     // Save all the to be assigned replication
@@ -87,11 +88,11 @@ public Map<String, Set<AssignableReplica>> getAssignableReplicaMap() {
     return _assignableReplicaMap;
   }
 
-  public Map<String, IdealState> getBaseline() {
+  public Map<String, ResourceAssignment> getBaseline() {
     return _baselineAssignment;
   }
 
-  public Map<String, IdealState> getBestPossibleAssignment() {
+  public Map<String, ResourceAssignment> getBestPossibleAssignment() {
     return _bestPossibleAssignment;
   }