Refine the WAGED rebalancer to minimize the partial rebalance workload. (apache#639)

* Refine the WAGED rebalancer to minimize the partial rebalance workload.

Split the cluster module calculation method so that different rebalance logic can have different rebalance scope calculation logic.
Also, refine the WAGED rebalancer logic to reduce duplicate code.

Author: jiajunwang

helix-core/src/main/java/org/apache/helix/controller/changedetector/ResourceChangeDetector.java

@@ -24,6 +24,8 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 import com.google.common.collect.Sets;
 import org.apache.helix.HelixConstants;
@@ -173,4 +175,20 @@ public Collection<String> getRemovalsByType(HelixConstants.ChangeType changeType
         changedItems -> getRemovedItems(determinePropertyMapByType(changeType, _oldSnapshot),
             determinePropertyMapByType(changeType, _newSnapshot)));
   }
+
+  /**
+   * @return A map contains all the changed items that are categorized by the change types.
+   */
+  public Map<HelixConstants.ChangeType, Set<String>> getAllChanges() {
+    return getChangeTypes().stream()
+        .collect(Collectors.toMap(changeType -> changeType, changeType -> {
+          Set<String> itemKeys = new HashSet<>();
+          itemKeys.addAll(getAdditionsByType(changeType));
+          itemKeys.addAll(getChangesByType(changeType));
+          itemKeys.addAll(getRemovalsByType(changeType));
+          return itemKeys;
+        })).entrySet().stream().filter(changeEntry -> !changeEntry.getValue().isEmpty()).collect(
+            Collectors
+                .toMap(changeEntry -> changeEntry.getKey(), changeEntry -> changeEntry.getValue()));
+  }
 }

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

@@ -254,7 +254,7 @@ private void reportInstanceCapacityMetrics(ClusterStatusMonitor clusterStatusMon
 
         Map<String, ResourceAssignment> currentStateAssignment =
             currentStateOutput.getAssignment(resourceToMonitorMap.keySet());
-        ClusterModel clusterModel = ClusterModelProvider.generateClusterModelFromCurrentState(
+        ClusterModel clusterModel = ClusterModelProvider.generateClusterModelFromExistingAssignment(
             dataProvider, resourceToMonitorMap, currentStateAssignment);
 
         Map<String, Double> maxUsageMap = new HashMap<>();

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

@@ -37,6 +37,7 @@
 import org.apache.helix.controller.pipeline.Stage;
 import org.apache.helix.controller.pipeline.StageContext;
 import org.apache.helix.controller.rebalancer.waged.AssignmentMetadataStore;
+import org.apache.helix.controller.rebalancer.waged.RebalanceAlgorithm;
 import org.apache.helix.controller.rebalancer.waged.WagedRebalancer;
 import org.apache.helix.controller.rebalancer.waged.constraints.ConstraintBasedAlgorithmFactory;
 import org.apache.helix.controller.stages.AttributeName;
@@ -413,45 +414,45 @@ public String toString() {
     return verifierName + "(" + _clusterName + "@" + _zkClient + "@resources["
        + (_resources != null ? Arrays.toString(_resources.toArray()) : "") + "])";
   }
-}
 
-/**
- * A Dryrun WAGED rebalancer that only calculates the assignment based on the cluster status but
- * never update the rebalancer assignment metadata.
- * This rebalacer is used in the verifiers or tests.
- */
-class DryrunWagedRebalancer extends WagedRebalancer {
-  DryrunWagedRebalancer(String metadataStoreAddrs, String clusterName,
-      Map<ClusterConfig.GlobalRebalancePreferenceKey, Integer> preferences) {
-    super(new ReadOnlyAssignmentMetadataStore(metadataStoreAddrs, clusterName),
-        ConstraintBasedAlgorithmFactory.getInstance(preferences));
-  }
+  /**
+   * A Dryrun WAGED rebalancer that only calculates the assignment based on the cluster status but
+   * never update the rebalancer assignment metadata.
+   * This rebalacer is used in the verifiers or tests.
+   */
+  private class DryrunWagedRebalancer extends WagedRebalancer {
+    DryrunWagedRebalancer(String metadataStoreAddrs, String clusterName,
+        Map<ClusterConfig.GlobalRebalancePreferenceKey, Integer> preferences) {
+      super(new ReadOnlyAssignmentMetadataStore(metadataStoreAddrs, clusterName),
+          ConstraintBasedAlgorithmFactory.getInstance(preferences));
+    }
 
-  @Override
-  protected Map<String, ResourceAssignment> computeBestPossibleAssignment(
-      ResourceControllerDataProvider clusterData, Map<String, Resource> resourceMap,
-      Set<String> activeNodes, CurrentStateOutput currentStateOutput)
-      throws HelixRebalanceException {
-    return getBestPossibleAssignment(getAssignmentMetadataStore(), currentStateOutput,
-        resourceMap.keySet());
+    @Override
+    protected Map<String, ResourceAssignment> computeBestPossibleAssignment(
+        ResourceControllerDataProvider clusterData, Map<String, Resource> resourceMap,
+        Set<String> activeNodes, CurrentStateOutput currentStateOutput, RebalanceAlgorithm algorithm)
+        throws HelixRebalanceException {
+      return getBestPossibleAssignment(getAssignmentMetadataStore(), currentStateOutput,
+          resourceMap.keySet());
+    }
   }
-}
 
-class ReadOnlyAssignmentMetadataStore extends AssignmentMetadataStore {
-  ReadOnlyAssignmentMetadataStore(String metadataStoreAddrs, String clusterName) {
-    super(new ZkBucketDataAccessor(metadataStoreAddrs), clusterName);
-  }
+  private class ReadOnlyAssignmentMetadataStore extends AssignmentMetadataStore {
+    ReadOnlyAssignmentMetadataStore(String metadataStoreAddrs, String clusterName) {
+      super(new ZkBucketDataAccessor(metadataStoreAddrs), clusterName);
+    }
 
-  @Override
-  public void persistBaseline(Map<String, ResourceAssignment> globalBaseline) {
-    // Update the in-memory reference only
-    _globalBaseline = globalBaseline;
-  }
+    @Override
+    public void persistBaseline(Map<String, ResourceAssignment> globalBaseline) {
+      // Update the in-memory reference only
+      _globalBaseline = globalBaseline;
+    }
 
-  @Override
-  public void persistBestPossibleAssignment(
-      Map<String, ResourceAssignment> bestPossibleAssignment) {
-    // Update the in-memory reference only
-    _bestPossibleAssignment = bestPossibleAssignment;
+    @Override
+    public void persistBestPossibleAssignment(
+        Map<String, ResourceAssignment> bestPossibleAssignment) {
+      // Update the in-memory reference only
+      _bestPossibleAssignment = bestPossibleAssignment;
+    }
   }
 }

helix-core/src/main/java/org/apache/helix/controller/stages/CurrentStateComputationStage.java

@@ -114,8 +114,7 @@ protected ResourceControllerDataProvider setupClusterDataCache() throws IOExcept
   }
 
   @Test
-  public void testRebalance()
-      throws IOException, HelixRebalanceException {
+  public void testRebalance() throws IOException, HelixRebalanceException {
     _metadataStore.clearMetadataStore();
     WagedRebalancer rebalancer = new WagedRebalancer(_metadataStore, _algorithm);
 
@@ -128,6 +127,10 @@ public void testRebalance()
               .forEach(partition -> resource.addPartition(partition));
           return resource;
         }));
+    // Mocking the change types for triggering a baseline rebalance.
+    when(clusterData.getRefreshedChangeTypes())
+        .thenReturn(Collections.singleton(HelixConstants.ChangeType.CLUSTER_CONFIG));
+
     Map<String, IdealState> newIdealStates =
         rebalancer.computeNewIdealStates(clusterData, resourceMap, new CurrentStateOutput());
     Map<String, ResourceAssignment> algorithmResult = _algorithm.getRebalanceResult();
@@ -150,6 +153,9 @@ public void testPartialRebalance() throws IOException, HelixRebalanceException {
               .forEach(partition -> resource.addPartition(partition));
           return resource;
         }));
+    // Mocking the change types for triggering a baseline rebalance.
+    when(clusterData.getRefreshedChangeTypes())
+        .thenReturn(Collections.singleton(HelixConstants.ChangeType.CLUSTER_CONFIG));
 
     // Test with partial resources listed in the resourceMap input.
     // Remove the first resource from the input. Note it still exists in the cluster data cache.
@@ -175,6 +181,9 @@ public void testRebalanceWithCurrentState() throws IOException, HelixRebalanceEx
               .forEach(partition -> resource.addPartition(partition));
           return resource;
         }));
+    // Mocking the change types for triggering a baseline rebalance.
+    when(clusterData.getRefreshedChangeTypes())
+        .thenReturn(Collections.singleton(HelixConstants.ChangeType.CLUSTER_CONFIG));
 
     // Test with current state exists, so the rebalancer should calculate for the intermediate state
     // Create current state based on the cluster data cache.
@@ -256,12 +265,12 @@ public void testInvalidClusterStatus() throws IOException, HelixRebalanceExcepti
         Collectors.toMap(resourceName -> resourceName, resourceName -> new Resource(resourceName)));
     try {
       rebalancer.computeBestPossibleAssignment(clusterData, resourceMap,
-          clusterData.getEnabledLiveInstances(), new CurrentStateOutput());
+          clusterData.getEnabledLiveInstances(), new CurrentStateOutput(), _algorithm);
       Assert.fail("Rebalance shall fail.");
     } catch (HelixRebalanceException ex) {
       Assert.assertEquals(ex.getFailureType(), HelixRebalanceException.Type.INVALID_CLUSTER_STATUS);
       Assert.assertEquals(ex.getMessage(),
-          "Failed to generate cluster model. Failure Type: INVALID_CLUSTER_STATUS");
+          "Failed to generate cluster model for partial rebalance. Failure Type: INVALID_CLUSTER_STATUS");
     }
 
     // The rebalance will be done with empty mapping result since there is no previously calculated
@@ -321,7 +330,7 @@ public void testAlgorithmException()
     // Calculation will fail
     try {
       rebalancer.computeBestPossibleAssignment(clusterData, resourceMap,
-          clusterData.getEnabledLiveInstances(), new CurrentStateOutput());
+          clusterData.getEnabledLiveInstances(), new CurrentStateOutput(), badAlgorithm);
       Assert.fail("Rebalance shall fail.");
     } catch (HelixRebalanceException ex) {
       Assert.assertEquals(ex.getFailureType(), HelixRebalanceException.Type.FAILED_TO_CALCULATE);