You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@helix.apache.org by ji...@apache.org on 2019/10/28 22:32:55 UTC
[helix] 17/50: Implement the basic constraint based algorithm (#381)
This is an automated email from the ASF dual-hosted git repository.
jiajunwang pushed a commit to branch wagedRebalancer
in repository https://gitbox.apache.org/repos/asf/helix.git
commit 2659f68b21cf1dbbe6c091a9e4cc472b1dc50258
Author: Yi Wang <i3...@gmail.com>
AuthorDate: Fri Sep 6 16:57:25 2019 -0700
Implement the basic constraint based algorithm (#381)
Implement basic constraint algorithm: Greedy based, each time it picks the best scores given each replica and assigns the replica to the node. It doesn't guarantee to achieve global optimal but local optimal result
The algorithm is based on a given set of constraints
* HardConstraint: Approve or deny the assignment given its condition, any assignment cannot bypass any "hard constraint"
* SoftConstraint: Evaluate the assignment by points/rewards/scores, a higher point means a better assignment
The goal is to avoid all "hard constraints" while accumulating the most points(rewards) from "soft constraints"
---
.../org/apache/helix/HelixRebalanceException.java | 13 ++-
.../{constraints => }/RebalanceAlgorithm.java | 27 +++--
.../rebalancer/waged/WagedRebalancer.java | 13 +--
.../constraints/ConstraintBasedAlgorithm.java | 129 +++++++++++++++++++++
.../ConstraintBasedAlgorithmFactory.java | 41 +++++++
.../constraints/ConstraintsRebalanceAlgorithm.java | 48 --------
.../constraints/SoftConstraintWeightModel.java | 6 +-
.../rebalancer/waged/model/OptimalAssignment.java | 69 +++++++++++
.../constraints/TestConstraintBasedAlgorithm.java | 74 ++++++++++++
.../waged/model/ClusterModelTestHelper.java | 50 ++++++++
10 files changed, 396 insertions(+), 74 deletions(-)
diff --git a/helix-core/src/main/java/org/apache/helix/HelixRebalanceException.java b/helix-core/src/main/java/org/apache/helix/HelixRebalanceException.java
index c01b173..d01fc60 100644
--- a/helix-core/src/main/java/org/apache/helix/HelixRebalanceException.java
+++ b/helix-core/src/main/java/org/apache/helix/HelixRebalanceException.java
@@ -23,21 +23,26 @@ package org.apache.helix;
* Exception thrown by Helix due to rebalance failures.
*/
public class HelixRebalanceException extends Exception {
- enum RebalanceFailureType {
+ public enum Type {
INVALID_CLUSTER_STATUS,
INVALID_REBALANCER_STATUS,
FAILED_TO_CALCULATE,
UNKNOWN_FAILURE
}
- private final RebalanceFailureType _type;
+ private final Type _type;
- public HelixRebalanceException(String message, RebalanceFailureType type, Throwable cause) {
+ public HelixRebalanceException(String message, Type type, Throwable cause) {
super(String.format("%s. Failure Type: %s", message, type.name()), cause);
_type = type;
}
- public RebalanceFailureType getFailureType() {
+ public HelixRebalanceException(String message, Type type) {
+ super(String.format("%s. Failure Type: %s", message, type.name()));
+ _type = type;
+ }
+
+ public Type getFailureType() {
return _type;
}
}
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/RebalanceAlgorithm.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java
similarity index 53%
rename from helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/RebalanceAlgorithm.java
rename to helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java
index b652836..1374162 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/RebalanceAlgorithm.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java
@@ -1,4 +1,4 @@
-package org.apache.helix.controller.rebalancer.waged.constraints;
+package org.apache.helix.controller.rebalancer.waged;
/*
* Licensed to the Apache Software Foundation (ASF) under one
@@ -9,7 +9,7 @@ package org.apache.helix.controller.rebalancer.waged.constraints;
* "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
+ * 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
@@ -19,26 +19,25 @@ package org.apache.helix.controller.rebalancer.waged.constraints;
* under the License.
*/
+import org.apache.helix.HelixRebalanceException;
import org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
-import org.apache.helix.model.ResourceAssignment;
-
-import java.util.Map;
+import org.apache.helix.controller.rebalancer.waged.model.OptimalAssignment;
/**
- * A generic rebalance algorithm interface for the WAGED rebalancer.
+ * A generic interface to generate the optimal assignment given the runtime cluster environment.
*
- * @see <a href="Rebalance Algorithm">https://github.com/apache/helix/wiki/Design-Proposal---Weight-Aware-Globally-Even-Distribute-Rebalancer#rebalance-algorithm-adapter</a>
+ * <pre>
+ * @see <a href="https://github.com/apache/helix/wiki/
+ * Design-Proposal---Weight-Aware-Globally-Even-Distribute-Rebalancer
+ * #rebalance-algorithm-adapter">Rebalance Algorithm</a>
+ * </pre>
*/
public interface RebalanceAlgorithm {
/**
* Rebalance the Helix resource partitions based on the input cluster model.
- *
- * @param clusterModel
- * @param failureReasons Return the failures <ResourceName, <FailureReason, Count>> that happen during the rebalance calculation.
- * If the map is null, no failure will be returned.
- * @return A map of <ResourceName, ResourceAssignment>.
+ * @param clusterModel The run time cluster model that contains all necessary information
+ * @return An instance of {@link OptimalAssignment}
*/
- Map<String, ResourceAssignment> rebalance(ClusterModel clusterModel,
- Map<String, Map<HardConstraint, Integer>> failureReasons);
+ OptimalAssignment calculate(ClusterModel clusterModel) throws HelixRebalanceException;
}
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java
index 43b2564..5b9634e 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java
@@ -19,23 +19,22 @@ package org.apache.helix.controller.rebalancer.waged;
* under the License.
*/
+import java.util.HashMap;
+import java.util.Map;
+
import org.apache.helix.HelixManager;
import org.apache.helix.HelixRebalanceException;
import org.apache.helix.controller.changedetector.ResourceChangeDetector;
import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
import org.apache.helix.controller.rebalancer.DelayedAutoRebalancer;
import org.apache.helix.controller.rebalancer.internal.MappingCalculator;
-import org.apache.helix.controller.rebalancer.waged.constraints.ConstraintsRebalanceAlgorithm;
-import org.apache.helix.controller.rebalancer.waged.constraints.RebalanceAlgorithm;
+import org.apache.helix.controller.rebalancer.waged.constraints.ConstraintBasedAlgorithmFactory;
import org.apache.helix.controller.stages.CurrentStateOutput;
import org.apache.helix.model.IdealState;
import org.apache.helix.model.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import java.util.HashMap;
-import java.util.Map;
-
/**
* A placeholder before we have the implementation.
* Weight-Aware Globally-Even Distribute Rebalancer.
@@ -69,8 +68,8 @@ public class WagedRebalancer {
public WagedRebalancer(HelixManager helixManager) {
// TODO init the metadata store according to their requirement when integrate, or change to final static method if possible.
_assignmentMetadataStore = new AssignmentMetadataStore();
- // TODO init the algorithm according to the requirement when integrate.
- _rebalanceAlgorithm = new ConstraintsRebalanceAlgorithm();
+ // TODO parse the cluster setting
+ _rebalanceAlgorithm = ConstraintBasedAlgorithmFactory.getInstance();
// Use the mapping calculator in DelayedAutoRebalancer for calculating the final assignment
// output.
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintBasedAlgorithm.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintBasedAlgorithm.java
new file mode 100644
index 0000000..99d8d2a
--- /dev/null
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintBasedAlgorithm.java
@@ -0,0 +1,129 @@
+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.ArrayList;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import org.apache.helix.HelixRebalanceException;
+import org.apache.helix.controller.rebalancer.waged.RebalanceAlgorithm;
+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.controller.rebalancer.waged.model.ClusterModel;
+import org.apache.helix.controller.rebalancer.waged.model.OptimalAssignment;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.google.common.collect.Maps;
+
+/**
+ * The algorithm is based on a given set of constraints
+ * - HardConstraint: Approve or deny the assignment given its condition, any assignment cannot
+ * bypass any "hard constraint"
+ * - SoftConstraint: Evaluate the assignment by points/rewards/scores, a higher point means a better
+ * assignment
+ * The goal is to accumulate the most points(rewards) from "soft constraints" while avoiding any
+ * "hard constraints"
+ */
+class ConstraintBasedAlgorithm implements RebalanceAlgorithm {
+ private static final Logger LOG = LoggerFactory.getLogger(ConstraintBasedAlgorithm.class);
+ private final List<HardConstraint> _hardConstraints;
+ private final List<SoftConstraint> _softConstraints;
+ private final SoftConstraintWeightModel _softConstraintsWeightModel;
+
+ ConstraintBasedAlgorithm(List<HardConstraint> hardConstraints,
+ List<SoftConstraint> softConstraints, SoftConstraintWeightModel softConstraintWeightModel) {
+ _hardConstraints = hardConstraints;
+ _softConstraints = softConstraints;
+ _softConstraintsWeightModel = softConstraintWeightModel;
+ }
+
+ @Override
+ public OptimalAssignment calculate(ClusterModel clusterModel) throws HelixRebalanceException {
+ OptimalAssignment optimalAssignment = new OptimalAssignment();
+ Map<String, Set<AssignableReplica>> replicasByResource = clusterModel.getAssignableReplicaMap();
+ List<AssignableNode> nodes = new ArrayList<>(clusterModel.getAssignableNodes().values());
+
+ // TODO: different orders of resource/replica could lead to different greedy assignments, will
+ // revisit and improve the performance
+ for (String resource : replicasByResource.keySet()) {
+ for (AssignableReplica replica : replicasByResource.get(resource)) {
+ Optional<AssignableNode> maybeBestNode =
+ getNodeWithHighestPoints(replica, nodes, clusterModel.getContext(), optimalAssignment);
+ // stop immediately if any replica cannot find best assignable node
+ if (optimalAssignment.hasAnyFailure()) {
+ String errorMessage = String.format(
+ "Unable to find any available candidate node for partition %s; Fail reasons: %s",
+ replica.getPartitionName(), optimalAssignment.getFailures());
+ throw new HelixRebalanceException(errorMessage,
+ HelixRebalanceException.Type.FAILED_TO_CALCULATE);
+ }
+ maybeBestNode.ifPresent(node -> clusterModel.assign(replica.getResourceName(),
+ replica.getPartitionName(), replica.getReplicaState(), node.getInstanceName()));
+ }
+ }
+
+ return optimalAssignment.convertFrom(clusterModel);
+ }
+
+ private Optional<AssignableNode> getNodeWithHighestPoints(AssignableReplica replica,
+ List<AssignableNode> assignableNodes, ClusterContext clusterContext,
+ OptimalAssignment optimalAssignment) {
+ Map<AssignableNode, List<HardConstraint>> hardConstraintFailures = new HashMap<>();
+ List<AssignableNode> candidateNodes = assignableNodes.stream().filter(candidateNode -> {
+ boolean isValid = true;
+ // need to record all the failure reasons and it gives us the ability to debug/fix the runtime
+ // cluster environment
+ for (HardConstraint hardConstraint : _hardConstraints) {
+ if (!hardConstraint.isAssignmentValid(candidateNode, replica, clusterContext)) {
+ hardConstraintFailures.computeIfAbsent(candidateNode, node -> new ArrayList<>())
+ .add(hardConstraint);
+ isValid = false;
+ }
+ }
+ return isValid;
+ }).collect(Collectors.toList());
+ if (candidateNodes.isEmpty()) {
+ optimalAssignment.recordAssignmentFailure(replica,
+ Maps.transformValues(hardConstraintFailures, this::convertFailureReasons));
+ return Optional.empty();
+ }
+
+ Function<AssignableNode, Float> calculatePoints =
+ (candidateNode) -> _softConstraintsWeightModel.getSumOfScores(_softConstraints.stream()
+ .collect(Collectors.toMap(Function.identity(), softConstraint -> softConstraint
+ .getAssignmentOriginScore(candidateNode, replica, clusterContext))));
+
+ return candidateNodes.stream().max(Comparator.comparing(calculatePoints));
+ }
+
+ private List<String> convertFailureReasons(List<HardConstraint> hardConstraints) {
+ return hardConstraints.stream().map(HardConstraint::getDescription)
+ .collect(Collectors.toList());
+ }
+}
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintBasedAlgorithmFactory.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintBasedAlgorithmFactory.java
new file mode 100644
index 0000000..895fa61
--- /dev/null
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintBasedAlgorithmFactory.java
@@ -0,0 +1,41 @@
+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.ArrayList;
+import java.util.List;
+
+import org.apache.helix.controller.rebalancer.waged.RebalanceAlgorithm;
+import org.apache.helix.model.ClusterConfig;
+
+public class ConstraintBasedAlgorithmFactory {
+
+ // TODO: the parameter comes from cluster config, will tune how these 2 integers will change the
+ // soft constraint weight model
+ public static RebalanceAlgorithm getInstance() {
+ // TODO initialize constraints, depending on constraints implementations PRs
+ List<HardConstraint> hardConstraints = new ArrayList<>();
+ List<SoftConstraint> softConstraints = new ArrayList<>();
+ SoftConstraintWeightModel softConstraintWeightModel = new SoftConstraintWeightModel();
+
+ return new ConstraintBasedAlgorithm(hardConstraints, softConstraints,
+ softConstraintWeightModel);
+ }
+}
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
deleted file mode 100644
index 8adaa73..0000000
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/ConstraintsRebalanceAlgorithm.java
+++ /dev/null
@@ -1,48 +0,0 @@
-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 org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
-import org.apache.helix.model.ResourceAssignment;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * A placeholder before we have the implementation.
- * The constraint-based rebalance algorithm that is used in the WAGED rebalancer.
- */
-public class ConstraintsRebalanceAlgorithm implements RebalanceAlgorithm {
- private static final Logger LOG = LoggerFactory.getLogger(ConstraintsRebalanceAlgorithm.class);
-
- private Map<HardConstraint, Integer> _failureReasonCounterMap = new HashMap<>();
-
- public ConstraintsRebalanceAlgorithm() {
- // TODO Constraints initialization
- }
-
- @Override
- public Map<String, ResourceAssignment> rebalance(ClusterModel clusterModel,
- Map<String, Map<HardConstraint, Integer>> failureReasons) {
- return new HashMap<>();
- }
-}
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/SoftConstraintWeightModel.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/SoftConstraintWeightModel.java
index 10201ce..41e4334 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/SoftConstraintWeightModel.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/constraints/SoftConstraintWeightModel.java
@@ -29,8 +29,12 @@ import com.google.common.collect.ImmutableMap;
class SoftConstraintWeightModel {
private static Map<? extends SoftConstraint, Float> MODEL;
+ // TODO either define the weights in property files or zookeeper node or static human input
+ SoftConstraintWeightModel() {
+
+ }
+
static {
- // TODO either define the weights in property files or zookeeper node or static human input
MODEL = ImmutableMap.<SoftConstraint, Float> builder()
.put(LeastPartitionCountConstraint.INSTANCE, 1.0f).build();
}
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/OptimalAssignment.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/OptimalAssignment.java
new file mode 100644
index 0000000..31cb181
--- /dev/null
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/OptimalAssignment.java
@@ -0,0 +1,69 @@
+package org.apache.helix.controller.rebalancer.waged.model;
+
+/*
+ * 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.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.helix.model.ResourceAssignment;
+
+/**
+ * The data model represents the optimal assignment of N replicas assigned to M instances;
+ * It's mostly used as the return parameter of an assignment calculation algorithm; If the algorithm
+ * failed to find optimal assignment given the endeavor, the user could check the failure reasons
+ */
+public class OptimalAssignment {
+ private Map<AssignableNode, List<AssignableReplica>> _optimalAssignment = new HashMap<>();
+ private Map<AssignableReplica, Map<AssignableNode, List<String>>> _failedAssignments =
+ new HashMap<>();
+
+ public OptimalAssignment() {
+
+ }
+
+ public void updateAssignments(ClusterModel clusterModel) {
+
+ }
+
+ // TODO: determine the output of final assignment format
+ public Map<String, ResourceAssignment> getOptimalResourceAssignment() {
+ throw new UnsupportedOperationException("Not implemented yet");
+ }
+
+ // TODO: the convert method is not the best choice so far, will revisit the data model
+ public OptimalAssignment convertFrom(ClusterModel clusterModel) {
+ return this;
+ }
+
+ public void recordAssignmentFailure(AssignableReplica replica,
+ Map<AssignableNode, List<String>> failedReasons) {
+ _failedAssignments.put(replica, failedReasons);
+ }
+
+ public boolean hasAnyFailure() {
+ return !_failedAssignments.isEmpty();
+ }
+
+ public String getFailures() {
+ // TODO: format the error string
+ return _failedAssignments.toString();
+ }
+}
diff --git a/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/constraints/TestConstraintBasedAlgorithm.java b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/constraints/TestConstraintBasedAlgorithm.java
new file mode 100644
index 0000000..d06cc5f
--- /dev/null
+++ b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/constraints/TestConstraintBasedAlgorithm.java
@@ -0,0 +1,74 @@
+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 static org.mockito.Matchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+
+import org.apache.helix.HelixRebalanceException;
+import org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
+import org.apache.helix.controller.rebalancer.waged.model.ClusterModelTestHelper;
+import org.apache.helix.controller.rebalancer.waged.model.OptimalAssignment;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.google.common.collect.ImmutableList;
+
+public class TestConstraintBasedAlgorithm {
+ private ConstraintBasedAlgorithm _algorithm;
+
+ @BeforeMethod
+ public void beforeMethod() {
+ HardConstraint mockHardConstraint = mock(HardConstraint.class);
+ SoftConstraint mockSoftConstraint = mock(SoftConstraint.class);
+ SoftConstraintWeightModel mockSoftConstraintWeightModel = mock(SoftConstraintWeightModel.class);
+ when(mockHardConstraint.isAssignmentValid(any(), any(), any())).thenReturn(false);
+ when(mockSoftConstraint.getAssignmentOriginScore(any(), any(), any())).thenReturn(1.0f);
+
+ _algorithm = new ConstraintBasedAlgorithm(ImmutableList.of(mockHardConstraint),
+ ImmutableList.of(mockSoftConstraint), mockSoftConstraintWeightModel);
+ }
+
+ @Test(expectedExceptions = HelixRebalanceException.class)
+ public void testCalculateNoValidAssignment() throws IOException, HelixRebalanceException {
+ ClusterModel clusterModel = new ClusterModelTestHelper().getDefaultClusterModel();
+ _algorithm.calculate(clusterModel);
+ }
+
+ @Test
+ public void testCalculateWithValidAssignment() throws IOException, HelixRebalanceException {
+ HardConstraint mockHardConstraint = mock(HardConstraint.class);
+ SoftConstraint mockSoftConstraint = mock(SoftConstraint.class);
+ SoftConstraintWeightModel mockSoftConstraintWeightModel = mock(SoftConstraintWeightModel.class);
+ when(mockHardConstraint.isAssignmentValid(any(), any(), any())).thenReturn(true);
+ when(mockSoftConstraint.getAssignmentOriginScore(any(), any(), any())).thenReturn(1.0f);
+ when(mockSoftConstraintWeightModel.getSumOfScores(any())).thenReturn(1.0f);
+ _algorithm = new ConstraintBasedAlgorithm(ImmutableList.of(mockHardConstraint),
+ ImmutableList.of(mockSoftConstraint), mockSoftConstraintWeightModel);
+ ClusterModel clusterModel = new ClusterModelTestHelper().getDefaultClusterModel();
+ OptimalAssignment optimalAssignment = _algorithm.calculate(clusterModel);
+
+ Assert.assertFalse(optimalAssignment.hasAnyFailure());
+ }
+}
diff --git a/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModelTestHelper.java b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModelTestHelper.java
new file mode 100644
index 0000000..76f1141
--- /dev/null
+++ b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModelTestHelper.java
@@ -0,0 +1,50 @@
+package org.apache.helix.controller.rebalancer.waged.model;
+
+/*
+ * 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.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+
+import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
+
+public class ClusterModelTestHelper extends AbstractTestClusterModel {
+
+ public ClusterModel getDefaultClusterModel() throws IOException {
+ initialize();
+ ResourceControllerDataProvider testCache = setupClusterDataCache();
+ Set<AssignableReplica> assignableReplicas = generateReplicas(testCache);
+ Set<AssignableNode> assignableNodes = generateNodes(testCache);
+
+ ClusterContext context = new ClusterContext(assignableReplicas, 2);
+ return new ClusterModel(context, assignableReplicas, assignableNodes, Collections.emptyMap(),
+ Collections.emptyMap());
+ }
+
+ private Set<AssignableNode> generateNodes(ResourceControllerDataProvider testCache) {
+ Set<AssignableNode> nodeSet = new HashSet<>();
+ testCache.getInstanceConfigMap().values().stream()
+ .forEach(config -> nodeSet.add(new AssignableNode(testCache.getClusterConfig(),
+ testCache.getInstanceConfigMap().get(_testInstanceId), config.getInstanceName(),
+ Collections.emptyList())));
+ return nodeSet;
+ }
+}