Implement POC work greedy based constraint algorithm

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/algorithm/ConstraintBasedAlgorithm.java

@@ -0,0 +1,120 @@
+package org.apache.helix.controller.rebalancer.waged.algorithm;
+
+/*
+ * 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.controller.rebalancer.waged.constraints.HardConstraint;
+import org.apache.helix.controller.rebalancer.waged.constraints.SoftConstraint;
+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;
+
+/**
+ * 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
+ * <p>
+ * The goal is to accumulate the most points(rewards) from "soft constraints" while avoiding all
+ * "hard constraints"
+ */
+public class ConstraintBasedAlgorithm implements RebalanceAlgorithm {
+  private static final Logger LOG = LoggerFactory.getLogger(ConstraintBasedAlgorithm.class);
+  private final List<HardConstraint> _hardConstraints;
+  private final List<SoftConstraint> _softConstraints;
+
+  public ConstraintBasedAlgorithm(List<HardConstraint> hardConstraints,
+      List<SoftConstraint> softConstraints) {
+    _hardConstraints = hardConstraints;
+    _softConstraints = softConstraints;
+  }
+
+  @Override
+  public OptimalAssignment calculate(ClusterModel clusterModel) {
+    ClusterContext clusterContext = clusterModel.getContext();
+    OptimalAssignment optimalAssignment = new OptimalAssignment(clusterContext);
+    Map<String, Set<AssignableReplica>> replicasByResource = clusterModel.getAssignableReplicaMap();
+    List<AssignableNode> nodes = new ArrayList<>(clusterModel.getAssignableNodes().values());
+
+    for (String resource : replicasByResource.keySet()) {
+      for (AssignableReplica replica : replicasByResource.get(resource)) {
+        Optional<AssignableNode> maybeBestNode =
+            getNodeWithHighestPoints(replica, nodes, optimalAssignment);
+        // stop immediately if any replica cannot find best assignable node
+        if (optimalAssignment.hasAnyFailure()) {
+          LOG.error(
+              "Unable to find any available candidate node for partition {}; Fail reasons: {}",
+              replica.getPartitionName(), optimalAssignment.getFailures());
+          return optimalAssignment;
+        }
+        maybeBestNode.ifPresent(node -> {
+          optimalAssignment.addAssignment(replica, node);
+          clusterModel.assign(replica.getResourceName(), replica.getPartitionName(),
+              replica.getReplicaState(), node.getInstanceName());
+        });
+      }
+    }
+
+    return optimalAssignment;
+  }
+
+  private Optional<AssignableNode> getNodeWithHighestPoints(AssignableReplica replica,
+      List<AssignableNode> assignableNodes, OptimalAssignment optimalAssignment) {
+    Map<AssignableNode, List<HardConstraint>> hardConstraintFailures = new HashMap<>();
+    List<AssignableNode> candidateNodes = assignableNodes.stream().filter(candidateNode -> {
+      boolean isValid = true;
+      // evaluate all hard constraints and record all the failure reasons why one assignment fails
+      for (HardConstraint hardConstraint : _hardConstraints) {
+        if (!hardConstraint.isAssignmentValid(candidateNode, replica,
+            optimalAssignment.getClusterContext())) {
+          hardConstraintFailures.computeIfAbsent(candidateNode, node -> new ArrayList<>())
+              .add(hardConstraint);
+          isValid = false;
+        }
+      }
+      return isValid;
+    }).collect(Collectors.toList());
+    if (candidateNodes.isEmpty()) {
+      optimalAssignment.trackAssignmentFailure(replica, hardConstraintFailures);
+      return Optional.empty();
+    }
+
+    Function<AssignableNode, Double> calculatePoints =
+        (candidateNode) -> _softConstraints.stream().map(softConstraint -> softConstraint
+            .assignmentScore(candidateNode, replica, optimalAssignment.getClusterContext()))
+            .mapToDouble(score -> score).sum();
+
+    return candidateNodes.stream().max(Comparator.comparing(calculatePoints));
+  }
+}

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/RebalanceAlgorithm.java → helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/algorithm/RebalanceAlgorithm.java

@@ -1,4 +1,4 @@
-package org.apache.helix.controller.rebalancer.waged;
+package org.apache.helix.controller.rebalancer.waged.algorithm;
 
 /*
  * Licensed to the Apache Software Foundation (ASF) under one
@@ -19,27 +19,21 @@
  * under the License.
  */
 
-import org.apache.helix.controller.rebalancer.waged.constraints.HardConstraint;
 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.
- *
- * @see <a href="Rebalance Algorithm">https://github.com/apache/helix/wiki/Design-Proposal---Weight-Aware-Globally-Even-Distribute-Rebalancer#rebalance-algorithm-adapter</a>
+ * @see <a
+ *      href="Rebalance Algorithm">https://github.com/apache/helix/wiki/Design-Proposal---Weight-Aware-Globally-Even-Distribute-Rebalancer#rebalance-algorithm-adapter</a>
  */
 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.FailureReason, Integer>> failureReasons);
+  OptimalAssignment calculate(ClusterModel clusterModel);
 }

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

@@ -71,6 +71,7 @@ public class ClusterContext {
     _estimatedMaxTopStateCount = estimateAvgReplicaCount(totalTopStateReplicas, instanceCount);
   }
 
+
   public Map<String, Map<String, Set<String>>> getAssignmentForFaultZoneMap() {
     return _assignmentForFaultZoneMap;
   }
@@ -115,4 +116,4 @@ private int estimateAvgReplicaCount(int replicaCount, int instanceCount) {
     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/OptimalAssignment.java

@@ -0,0 +1,73 @@
+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.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.helix.controller.rebalancer.waged.constraints.HardConstraint;
+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<HardConstraint>>> _failedAssignments =
+      new HashMap<>();
+  private final ClusterContext _clusterContext;
+
+  public OptimalAssignment(ClusterContext clusterContext) {
+    _clusterContext = clusterContext;
+  }
+
+  public ClusterContext getClusterContext() {
+    return _clusterContext;
+  }
+
+  public Map<String, ResourceAssignment> getOptimalResourceAssignment() {
+    // TODO: convert the optimal assignment to map
+    return Collections.emptyMap();
+  }
+
+  public void trackAssignmentFailure(AssignableReplica replica,
+      Map<AssignableNode, List<HardConstraint>> failedAssignments) {
+    _failedAssignments.put(replica, failedAssignments);
+  }
+
+  public boolean hasAnyFailure() {
+    return !_failedAssignments.isEmpty();
+  }
+
+  public String getFailures() {
+    // TODO: format the error string
+    return _failedAssignments.toString();
+  }
+
+  //TODO: delegate the add/release operations from clusterModel or clusterContext
+  public void addAssignment(AssignableReplica replica, AssignableNode node) {
+    _optimalAssignment.computeIfAbsent(node, key -> new ArrayList<>()).add(replica);
+  }
+}

helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/algorithm/TestConstraintBasedAlgorithm.java

@@ -0,0 +1,72 @@
+package org.apache.helix.controller.rebalancer.waged.algorithm;
+
+/*
+ * 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.controller.rebalancer.waged.constraints.HardConstraint;
+import org.apache.helix.controller.rebalancer.waged.constraints.SoftConstraint;
+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);
+    when(mockHardConstraint.isAssignmentValid(any(), any(), any())).thenReturn(false);
+    when(mockSoftConstraint.assignmentScore(any(), any(), any())).thenReturn(1.0f);
+
+    _algorithm = new ConstraintBasedAlgorithm(ImmutableList.of(mockHardConstraint), ImmutableList.of(mockSoftConstraint));
+  }
+
+  @Test
+  public void testCalculateNoValidAssignment() throws IOException  {
+    ClusterModel clusterModel = new ClusterModelTestHelper().getDefaultClusterModel();
+    OptimalAssignment optimalAssignment = _algorithm.calculate(clusterModel);
+
+    Assert.assertTrue(optimalAssignment.hasAnyFailure());
+  }
+
+  @Test
+  public void testCalculateWithValidAssignment() throws IOException {
+    HardConstraint mockHardConstraint = mock(HardConstraint.class);
+    SoftConstraint mockSoftConstraint = mock(SoftConstraint.class);
+    when(mockHardConstraint.isAssignmentValid(any(), any(), any())).thenReturn(true);
+    when(mockSoftConstraint.assignmentScore(any(), any(), any())).thenReturn(1.0f);
+    _algorithm = new ConstraintBasedAlgorithm(ImmutableList.of(mockHardConstraint), ImmutableList.of(mockSoftConstraint));
+    ClusterModel clusterModel = new ClusterModelTestHelper().getDefaultClusterModel();
+    OptimalAssignment optimalAssignment = _algorithm.calculate(clusterModel);
+
+    Assert.assertFalse(optimalAssignment.hasAnyFailure());
+  }
+}