pkuwm commented on a change in pull request #362: The WAGED rebalancer cluster model implementation URL: https://github.com/apache/helix/pull/362#discussion_r307943962
########## File path: helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModel.java ########## @@ -19,9 +19,129 @@ * under the License. */ +import org.apache.helix.HelixException; +import org.apache.helix.model.IdealState; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + /** - * A placeholder before we have the implementation. - * * This class wraps the required input for the rebalance algorithm. */ -public class ClusterModel { } +public class ClusterModel { + private final ClusterContext _clusterContext; + // <Resource, <Replica Key, Set<Replica>>> + private final Map<String, Map<String, AssignableReplica>> _assignableReplicaIndex; + private final Map<String, AssignableNode> _assignableNodeMap; + + // Records about the previous assignment + private final Map<String, IdealState> _baselineAssignment; + private final Map<String, IdealState> _bestPossibleAssignment; + + /** + * @param clusterContext The initialized cluster context. + * @param assignableReplicas The replications 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, IdealState> baselineAssignment, + Map<String, IdealState> bestPossibleAssignment) { + _clusterContext = clusterContext; + + // Index all the replicas to be assigned. + _assignableReplicaIndex = new HashMap<>(); + assignableReplicas.stream().forEach(rep -> { + Map<String, AssignableReplica> replicaMap = + _assignableReplicaIndex.getOrDefault(rep.getResourceName(), new HashMap<>()); + replicaMap.put(rep.getReplicaKey(), rep); Review comment: `computeIfAbsent`? ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services