[WAGED Tag Isolation 7/16] Hook into the assignment loop

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

@@ -130,64 +130,93 @@ private OptimalAssignment calculateInternal(ClusterModel clusterModel,
     List<AssignableNode> nodes = new ArrayList<>(clusterModel.getAssignableNodes().values());
     Set<String> busyInstances =
         getBusyInstances(clusterModel.getContext().getBestPossibleAssignment().values());
+    List<AssignableReplica> allReplicas =
+        clusterModel.getAssignableReplicaMap().values().stream().flatMap(Collection::stream)
+            .collect(Collectors.toList());
+    // A no-op unless WAGED_INSTANCE_TAG_ISOLATION_ENABLED is set on the cluster config.
+    InstanceTagIsolation isolation =
+        new InstanceTagIsolation(clusterModel, allReplicas, nodes);
 
-    // create a always >0 capacity map to avoid divide by 0.
-    Map<String, Float> positiveEstimateClusterRemainCap = new HashMap<>();
-    for (Map.Entry<String, Long> clusterRemainingCap : clusterModel.getContext()
-        .getEstimateUtilizationMap().entrySet()) {
-      String capacityKey = clusterRemainingCap.getKey();
-      if (clusterRemainingCap.getValue() < 0) {
-        // all replicas' assignment will fail if there is one dimension's remain capacity <0.
-        long totalCapacity = clusterModel.getContext().getClusterCapacityMap().get(capacityKey);
-        long remainingCapacity = clusterRemainingCap.getValue();
-        long totalUsage = totalCapacity - remainingCapacity;
-        throw new HelixRebalanceException(String
-            .format("The cluster '%s' does not have enough %s capacity for all partitions. Total capacity: %d, Required: %d, Deficit: %d",
-                clusterModel.getContext().getClusterName(), capacityKey, totalCapacity, totalUsage, Math.abs(remainingCapacity)),
-            HelixRebalanceException.Type.FAILED_TO_CALCULATE,
-            HelixRebalanceException.FailureCategory.CAPACITY_DEFICIT);
-      }
-      // estimate remain capacity after assignment + %1 of current cluster capacity before assignment
-      positiveEstimateClusterRemainCap.put(capacityKey,
-           (float) clusterModel.getContext().getEstimateUtilizationMap().get(capacityKey) +
-          (clusterModel.getContext().getClusterCapacityMap().get(capacityKey) * DIV_GUARD));
-    }
+    Map<String, Float> positiveEstimateClusterRemainCap = computeScoringCapacities(clusterModel);
 
     // Create a wrapper for each AssignableReplica.
     List<AssignableReplicaWithScore> toBeAssignedReplicas =
-        clusterModel.getAssignableReplicaMap().values().stream().flatMap(Collection::stream).map(
+        allReplicas.stream().map(
             replica -> new AssignableReplicaWithScore(replica, clusterModel, positiveEstimateClusterRemainCap)).sorted()
             .collect(Collectors.toList());
 
     for (AssignableReplicaWithScore replicaWithScore : toBeAssignedReplicas) {
       AssignableReplica replica = replicaWithScore.getAssignableReplica();
+      if (isolation.shouldSkip(replica)) {
+        continue;
+      }
+      OptimalAssignment failureSink = isolation.failureSink(optimalAssignment);
       Optional<AssignableNode> maybeBestNode =
           getNodeWithHighestPoints(replica, nodes, clusterModel.getContext(), busyInstances,
-              optimalAssignment, blockingTypes);
+              failureSink, blockingTypes);
       // stop immediately if any replica cannot find best assignable node
-      if (!maybeBestNode.isPresent() || optimalAssignment.hasAnyFailure()) {
+      if (!maybeBestNode.isPresent() || failureSink.hasAnyFailure()) {
         String errorMessage = String.format(
             "Unable to find any available candidate node for partition %s (resource: %s, cluster: %s); "
                 + "Failure summary: %s; Fail reasons: %s",
             replica.getPartitionName(),
             replica.getResourceName(),
             clusterModel.getContext().getClusterName(),
-            optimalAssignment.getFailureSummary(),
-            optimalAssignment.getFailures());
-        throw new HelixRebalanceException(errorMessage,
+            failureSink.getFailureSummary(),
+            failureSink.getFailures());
+        HelixRebalanceException failure = new HelixRebalanceException(errorMessage,
             HelixRebalanceException.Type.FAILED_TO_CALCULATE,
             HelixRebalanceException.FailureCategory.NO_CANDIDATE_NODE);
+        // ... unless isolation can absorb the failure into the replica's own instance tag group.
+        if (!isolation.tryIsolate(replica, failure)) {
+          throw failure;
+        }
+        continue;
       }
       AssignableNode bestNode = maybeBestNode.get();
       // Assign the replica and update the cluster model.
       clusterModel
           .assign(replica.getResourceName(), replica.getPartitionName(), replica.getReplicaState(),
               bestNode.getInstanceName());
+      isolation.recordPlacement(replica, bestNode);
     }
+
+    isolation.finish(optimalAssignment);
     optimalAssignment.updateAssignments(clusterModel);
     return optimalAssignment;
   }
 
+  /**
+   * The per-capacity-key denominators used to score how scarce each dimension is, always kept above
+   * zero to avoid a divide by zero.
+   *
+   * @throws HelixRebalanceException when the cluster as a whole cannot hold all the partitions.
+   */
+  private static Map<String, Float> computeScoringCapacities(ClusterModel clusterModel)
+      throws HelixRebalanceException {
+    Map<String, Float> positiveEstimateClusterRemainCap = new HashMap<>();
+    for (Map.Entry<String, Long> clusterRemainingCap : clusterModel.getContext()
+        .getEstimateUtilizationMap().entrySet()) {
+      String capacityKey = clusterRemainingCap.getKey();
+      if (clusterRemainingCap.getValue() < 0) {
+        // all replicas' assignment will fail if there is one dimension's remain capacity <0.
+        long totalCapacity = clusterModel.getContext().getClusterCapacityMap().get(capacityKey);
+        long remainingCapacity = clusterRemainingCap.getValue();
+        long totalUsage = totalCapacity - remainingCapacity;
+        throw new HelixRebalanceException(String
+            .format("The cluster '%s' does not have enough %s capacity for all partitions. Total capacity: %d, Required: %d, Deficit: %d",
+                clusterModel.getContext().getClusterName(), capacityKey, totalCapacity, totalUsage, Math.abs(remainingCapacity)),
+            HelixRebalanceException.Type.FAILED_TO_CALCULATE,
+            HelixRebalanceException.FailureCategory.CAPACITY_DEFICIT);
+      }
+      // estimate remain capacity after assignment + %1 of current cluster capacity before assignment
+      positiveEstimateClusterRemainCap.put(capacityKey,
+           (float) clusterModel.getContext().getEstimateUtilizationMap().get(capacityKey) +
+          (clusterModel.getContext().getClusterCapacityMap().get(capacityKey) * DIV_GUARD));
+    }
+    return positiveEstimateClusterRemainCap;
+  }
+
   private Optional<AssignableNode> getNodeWithHighestPoints(AssignableReplica replica,
       List<AssignableNode> assignableNodes, ClusterContext clusterContext, Set<String> busyInstances,
       OptimalAssignment optimalAssignment, Set<HardConstraint.Type> blockingTypes)