Add WAGED rebalance build/solve latency split and partial movement counter

Improve WAGED rebalancer observability on the global (baseline) and partial
(best possible) paths:

- Split the existing per-phase latency into a model-build span and an
  algorithm-solve span for each path, via four new latency gauges
  (GlobalBaselineCalcBuild/SolveLatencyGauge and
  PartialRebalanceBuild/SolveLatencyGauge). This lets a slow rebalance be
  attributed to cluster-model construction versus constraint solving. The
  existing overall latency gauges are retained for backward compatibility.

- Count replica placements in the new best possible assignment that changed
  relative to the previous best possible assignment
  (PartialRebalanceReplicaMovementCounter) using a new
  ResourceUsageCalculator.countReplicaMovements helper. The best possible
  assignment is what actually drives state transitions, so this surfaces the
  churn / state-transition blast radius of each partial rebalance at
  intent-time (before throttling), attributable to the partial phase --
  complementing the existing downstream Pending*RebalanceReplica gauges.

Add a unit test for countReplicaMovements and extend TestWagedRebalancerMetrics
to assert the new latency gauges and movement counter are populated.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: sjainit

helix-core/src/main/java/org/apache/helix/controller/rebalancer/util/ResourceUsageCalculator.java

@@ -19,6 +19,7 @@
  * under the License.
  */
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -143,6 +144,49 @@ public static double measureBaselineDivergence(Map<String, ResourceAssignment> b
         : (1.0d - (double) numMatchedReplicas / (double) numTotalBestPossibleReplicas);
   }
 
+  /**
+   * Counts the number of replica placements in {@code newAssignment} that are new or changed
+   * relative to {@code oldAssignment}. A replica placement (identified by partition name and
+   * instance name) counts as a movement when the instance does not hold that partition's replica in
+   * the previous assignment, or holds it in a different state. This approximates the amount of churn
+   * a rebalance introduces (i.e. the number of state-transition messages it will generate), which is
+   * a useful blast-radius signal independent of how long the rebalance took.
+   *
+   * Note: replicas that were dropped (present in {@code oldAssignment} but absent in
+   * {@code newAssignment}) are intentionally not counted; the metric measures work introduced by the
+   * new assignment. A replica whose instance is unchanged but whose state changed (e.g. a
+   * SLAVE promoted to MASTER) is counted, because it still produces a state transition.
+   *
+   * @param oldAssignment the previous assignment (e.g. previous baseline or best possible)
+   * @param newAssignment the newly computed assignment
+   * @return the number of new or changed replica placements introduced by {@code newAssignment}
+   */
+  public static long countReplicaMovements(Map<String, ResourceAssignment> oldAssignment,
+      Map<String, ResourceAssignment> newAssignment) {
+    long movements = 0L;
+    for (Map.Entry<String, ResourceAssignment> resourceEntry : newAssignment.entrySet()) {
+      String resourceName = resourceEntry.getKey();
+      Map<String, Map<String, String>> newPartitions =
+          resourceEntry.getValue().getRecord().getMapFields();
+      ResourceAssignment oldResource = oldAssignment.get(resourceName);
+      Map<String, Map<String, String>> oldPartitions =
+          oldResource == null ? Collections.emptyMap() : oldResource.getRecord().getMapFields();
+
+      for (Map.Entry<String, Map<String, String>> partitionEntry : newPartitions.entrySet()) {
+        Map<String, String> newReplicas = partitionEntry.getValue();
+        Map<String, String> oldReplicas =
+            oldPartitions.getOrDefault(partitionEntry.getKey(), Collections.emptyMap());
+        for (Map.Entry<String, String> replicaEntry : newReplicas.entrySet()) {
+          // New or changed placement for this (partition, instance) => a movement.
+          if (!replicaEntry.getValue().equals(oldReplicas.get(replicaEntry.getKey()))) {
+            movements++;
+          }
+        }
+      }
+    }
+    return movements;
+  }
+
   /**
    * Calculates average partition weight per capacity key for a resource config. Example as below:
    * Input =

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

@@ -80,6 +80,10 @@ class GlobalRebalanceRunner implements AutoCloseable {
   private final LatencyMetric _writeLatency;
   private final CountMetric _baselineCalcCounter;
   private final LatencyMetric _baselineCalcLatency;
+  // Sub-phase latencies of _baselineCalcLatency: time spent building the cluster model vs. solving
+  // it with the assignment algorithm. Together they let a slow baseline be attributed to either half.
+  private final LatencyMetric _baselineCalcBuildLatency;
+  private final LatencyMetric _baselineCalcSolveLatency;
   // Reporter that ticks RebalanceFailureCounter plus the per-FailureCategory counters on both
   // the Rebalancer-domain and ClusterStatus-domain MBeans. Owned by WagedRebalancer; injected
   // so the runner doesn't need a direct ClusterStatusMonitor reference. For the Baseline phase this
@@ -114,6 +118,12 @@ public GlobalRebalanceRunner(AssignmentManager assignmentManager,
     _baselineCalcLatency = metricCollector.getMetric(
         WagedRebalancerMetricCollector.WagedRebalancerMetricNames.GlobalBaselineCalcLatencyGauge.name(),
         LatencyMetric.class);
+    _baselineCalcBuildLatency = metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.GlobalBaselineCalcBuildLatencyGauge.name(),
+        LatencyMetric.class);
+    _baselineCalcSolveLatency = metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.GlobalBaselineCalcSolveLatencyGauge.name(),
+        LatencyMetric.class);
     _asyncFailureReporter = asyncFailureReporter;
     _baselineComputeStatusReporter = baselineComputeStatusReporter;
     _asyncGlobalRebalanceEnabled = isAsyncGlobalRebalanceEnabled;
@@ -217,16 +227,25 @@ private void doGlobalRebalance(ResourceControllerDataProvider clusterData,
     Map<String, ResourceAssignment> currentBaseline =
         _assignmentManager.getBaselineAssignment(_assignmentMetadataStore, currentStateOutput, resourceMap.keySet());
     ClusterModel clusterModel;
+    _baselineCalcBuildLatency.startMeasuringLatency();
     try {
       clusterModel = ClusterModelProvider.generateClusterModelForBaseline(clusterData, resourceMap,
           allAssignableInstances, clusterChanges, currentBaseline);
     } catch (Exception ex) {
       throw new HelixRebalanceException("Failed to generate cluster model for global rebalance.",
           HelixRebalanceException.Type.INVALID_CLUSTER_STATUS,
           HelixRebalanceException.FailureCategory.INVALID_CLUSTER_CONFIG, ex);
+    } finally {
+      _baselineCalcBuildLatency.endMeasuringLatency();
     }
 
-    Map<String, ResourceAssignment> newBaseline = WagedRebalanceUtil.calculateAssignment(clusterModel, algorithm);
+    _baselineCalcSolveLatency.startMeasuringLatency();
+    Map<String, ResourceAssignment> newBaseline;
+    try {
+      newBaseline = WagedRebalanceUtil.calculateAssignment(clusterModel, algorithm);
+    } finally {
+      _baselineCalcSolveLatency.endMeasuringLatency();
+    }
     boolean isBaselineChanged =
         _assignmentMetadataStore != null && _assignmentMetadataStore.isBaselineChanged(newBaseline);
     // Write the new baseline to metadata store

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

@@ -30,6 +30,7 @@
 import java.util.function.Consumer;
 import org.apache.helix.HelixRebalanceException;
 import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
+import org.apache.helix.controller.rebalancer.util.ResourceUsageCalculator;
 import org.apache.helix.controller.rebalancer.util.WagedRebalanceUtil;
 import org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
 import org.apache.helix.controller.rebalancer.waged.model.ClusterModelProvider;
@@ -78,6 +79,14 @@ class PartialRebalanceRunner implements AutoCloseable {
   private final Runnable _partialRebalanceSuccessReporter;
   private final CountMetric _partialRebalanceCounter;
   private final LatencyMetric _partialRebalanceLatency;
+  // Sub-phase latencies of _partialRebalanceLatency: time spent building the cluster model vs.
+  // solving it with the assignment algorithm. Together they let a slow partial rebalance be
+  // attributed to either half.
+  private final LatencyMetric _partialRebalanceBuildLatency;
+  private final LatencyMetric _partialRebalanceSolveLatency;
+  // Count of replica placements in the new best possible assignment that differ from the previous
+  // one, i.e. the churn (state-transition blast radius) this partial rebalance introduces.
+  private final CountMetric _partialRebalanceReplicaMovementCounter;
 
   private boolean _asyncPartialRebalanceEnabled;
   private Future<Boolean> _asyncPartialRebalanceResult;
@@ -105,6 +114,18 @@ public PartialRebalanceRunner(AssignmentManager assignmentManager,
         WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceLatencyGauge
             .name(),
         LatencyMetric.class);
+    _partialRebalanceBuildLatency = metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceBuildLatencyGauge
+            .name(),
+        LatencyMetric.class);
+    _partialRebalanceSolveLatency = metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceSolveLatencyGauge
+            .name(),
+        LatencyMetric.class);
+    _partialRebalanceReplicaMovementCounter = metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceReplicaMovementCounter
+            .name(),
+        CountMetric.class);
     _baselineDivergenceGauge = metricCollector.getMetric(
         WagedRebalancerMetricCollector.WagedRebalancerMetricNames.BaselineDivergenceGauge.name(),
         BaselineDivergenceGauge.class);
@@ -199,6 +220,7 @@ private void doPartialRebalance(ResourceControllerDataProvider clusterData, Map<
         _assignmentManager.getBestPossibleAssignment(_assignmentMetadataStore, currentStateOutput,
             resourceMap.keySet());
     ClusterModel clusterModel;
+    _partialRebalanceBuildLatency.startMeasuringLatency();
     try {
       clusterModel = ClusterModelProvider
           .generateClusterModelForPartialRebalance(clusterData, resourceMap, activeNodes,
@@ -207,8 +229,25 @@ private void doPartialRebalance(ResourceControllerDataProvider clusterData, Map<
       throw new HelixRebalanceException("Failed to generate cluster model for partial rebalance.",
           HelixRebalanceException.Type.INVALID_CLUSTER_STATUS,
           HelixRebalanceException.FailureCategory.INVALID_CLUSTER_CONFIG, ex);
+    } finally {
+      _partialRebalanceBuildLatency.endMeasuringLatency();
+    }
+    _partialRebalanceSolveLatency.startMeasuringLatency();
+    Map<String, ResourceAssignment> newAssignment;
+    try {
+      newAssignment = WagedRebalanceUtil.calculateAssignment(clusterModel, algorithm);
+    } finally {
+      _partialRebalanceSolveLatency.endMeasuringLatency();
+    }
+    boolean isBestPossibleChanged = _assignmentMetadataStore != null
+        && _assignmentMetadataStore.isBestPossibleChanged(newAssignment);
+    // Report how many replica placements changed relative to the previous best possible assignment.
+    // This best possible assignment is what actually drives state transitions, so this is the churn
+    // magnitude the cluster will experience. Skip the diff pass when nothing changed.
+    if (isBestPossibleChanged) {
+      _partialRebalanceReplicaMovementCounter.increment(ResourceUsageCalculator
+          .countReplicaMovements(currentBestPossibleAssignment, newAssignment));
     }
-    Map<String, ResourceAssignment> newAssignment = WagedRebalanceUtil.calculateAssignment(clusterModel, algorithm);
 
     // Asynchronously report baseline divergence metric before persisting to metadata store,
     // just in case if persisting fails, we still have the metric.
@@ -223,7 +262,7 @@ private void doPartialRebalance(ResourceControllerDataProvider clusterData, Map<
         currentBaseline, newAssignmentCopy);
 
     boolean bestPossibleUpdateSuccessful = false;
-    if (_assignmentMetadataStore != null && _assignmentMetadataStore.isBestPossibleChanged(newAssignment)) {
+    if (isBestPossibleChanged) {
       // This will not persist the new Best Possible Assignment into ZK. It will only update the in-memory cache.
       // If this is done successfully, the new Best Possible Assignment will be persisted into ZK the next time that
       // the pipeline is triggered. We schedule the pipeline to run below.

helix-core/src/main/java/org/apache/helix/monitoring/metrics/WagedRebalancerMetricCollector.java

@@ -46,6 +46,15 @@ public enum WagedRebalancerMetricNames {
     EmergencyRebalanceLatencyGauge,
     RebalanceOverwriteLatencyGauge,
 
+    // Sub-phase breakdown of the Global Baseline and Partial rebalance latencies above. Each of the
+    // two phases is timed separately so a slow rebalance can be attributed to either building the
+    // cluster model (constructing assignable nodes/replicas and pinning existing placement) or
+    // solving it (running the constraint-based assignment algorithm).
+    GlobalBaselineCalcBuildLatencyGauge,
+    GlobalBaselineCalcSolveLatencyGauge,
+    PartialRebalanceBuildLatencyGauge,
+    PartialRebalanceSolveLatencyGauge,
+
     // The following latency metrics are related to AssignmentMetadataStore
     StateReadLatencyGauge,
     StateWriteLatencyGauge,
@@ -87,7 +96,13 @@ public enum WagedRebalancerMetricNames {
     GlobalBaselineCalcCounter,
     PartialRebalanceCounter,
     EmergencyRebalanceCounter,
-    RebalanceOverwriteCounter
+    RebalanceOverwriteCounter,
+
+    // Count of replica placements in the new best possible assignment that changed relative to the
+    // previous best possible assignment, i.e. the churn (state-transition blast radius) that the
+    // partial rebalance introduces (the best possible assignment is what actually drives state
+    // transitions in the cluster).
+    PartialRebalanceReplicaMovementCounter
   }
 
   public WagedRebalancerMetricCollector(String clusterName) {
@@ -129,6 +144,20 @@ private void createMetrics() {
     LatencyMetric rebalanceOverwriteLatencyGauge =
         new RebalanceLatencyGauge(WagedRebalancerMetricNames.RebalanceOverwriteLatencyGauge.name(),
             getResetIntervalInMs());
+    LatencyMetric globalBaselineCalcBuildLatencyGauge =
+        new RebalanceLatencyGauge(
+            WagedRebalancerMetricNames.GlobalBaselineCalcBuildLatencyGauge.name(),
+            getResetIntervalInMs());
+    LatencyMetric globalBaselineCalcSolveLatencyGauge =
+        new RebalanceLatencyGauge(
+            WagedRebalancerMetricNames.GlobalBaselineCalcSolveLatencyGauge.name(),
+            getResetIntervalInMs());
+    LatencyMetric partialRebalanceBuildLatencyGauge =
+        new RebalanceLatencyGauge(WagedRebalancerMetricNames.PartialRebalanceBuildLatencyGauge.name(),
+            getResetIntervalInMs());
+    LatencyMetric partialRebalanceSolveLatencyGauge =
+        new RebalanceLatencyGauge(WagedRebalancerMetricNames.PartialRebalanceSolveLatencyGauge.name(),
+            getResetIntervalInMs());
     LatencyMetric stateReadLatencyGauge =
         new RebalanceLatencyGauge(WagedRebalancerMetricNames.StateReadLatencyGauge.name(),
             getResetIntervalInMs());
@@ -177,12 +206,19 @@ private void createMetrics() {
         new RebalanceCounter(WagedRebalancerMetricNames.EmergencyRebalanceCounter.name());
     CountMetric rebalanceOverwriteCounter =
         new RebalanceCounter(WagedRebalancerMetricNames.RebalanceOverwriteCounter.name());
+    CountMetric partialRebalanceReplicaMovementCounter =
+        new RebalanceCounter(
+            WagedRebalancerMetricNames.PartialRebalanceReplicaMovementCounter.name());
 
     // Add metrics to WagedRebalancerMetricCollector
     addMetric(globalBaselineCalcLatencyGauge);
     addMetric(partialRebalanceLatencyGauge);
     addMetric(emergencyRebalanceLatencyGauge);
     addMetric(rebalanceOverwriteLatencyGauge);
+    addMetric(globalBaselineCalcBuildLatencyGauge);
+    addMetric(globalBaselineCalcSolveLatencyGauge);
+    addMetric(partialRebalanceBuildLatencyGauge);
+    addMetric(partialRebalanceSolveLatencyGauge);
     addMetric(stateReadLatencyGauge);
     addMetric(stateWriteLatencyGauge);
     addMetric(baselineDivergenceGauge);
@@ -206,5 +242,6 @@ private void createMetrics() {
     addMetric(partialRebalanceCounter);
     addMetric(emergencyRebalanceCounter);
     addMetric(rebalanceOverwriteCounter);
+    addMetric(partialRebalanceReplicaMovementCounter);
   }
 }

helix-core/src/test/java/org/apache/helix/controller/rebalancer/util/TestResourceUsageCalculator.java

@@ -60,6 +60,40 @@ public void testMeasureBaselineDivergence(Map<String, Map<String, Map<String, St
         0.0d);
   }
 
+  @Test
+  public void testCountReplicaMovements() {
+    Map<String, ResourceAssignment> previous = buildResourceAssignment(ImmutableMap.of("DB",
+        ImmutableMap.of(
+            "DB_0", ImmutableMap.of("host_A", "MASTER", "host_B", "SLAVE", "host_C", "SLAVE"),
+            "DB_1", ImmutableMap.of("host_A", "MASTER", "host_B", "SLAVE", "host_C", "SLAVE"))));
+
+    // Identical assignment => no movements.
+    Assert.assertEquals(ResourceUsageCalculator.countReplicaMovements(previous, previous), 0L);
+
+    // Empty new assignment => nothing was placed, so nothing moved.
+    Assert.assertEquals(
+        ResourceUsageCalculator.countReplicaMovements(previous, Collections.emptyMap()), 0L);
+
+    // Empty previous assignment => every replica in the new assignment is a fresh placement (2 x 3).
+    Assert.assertEquals(
+        ResourceUsageCalculator.countReplicaMovements(Collections.emptyMap(), previous), 6L);
+
+    // Mixed changes:
+    //  - DB_0: host_C's SLAVE relocates to host_D (1 movement). The host_C drop is NOT counted.
+    //  - DB_1: host_A MASTER->SLAVE and host_B SLAVE->MASTER (2 state changes); host_C unchanged.
+    // Total = 3.
+    Map<String, ResourceAssignment> updated = buildResourceAssignment(ImmutableMap.of("DB",
+        ImmutableMap.of(
+            "DB_0", ImmutableMap.of("host_A", "MASTER", "host_B", "SLAVE", "host_D", "SLAVE"),
+            "DB_1", ImmutableMap.of("host_A", "SLAVE", "host_B", "MASTER", "host_C", "SLAVE"))));
+    Assert.assertEquals(ResourceUsageCalculator.countReplicaMovements(previous, updated), 3L);
+
+    // A resource absent from the previous assignment => all its replicas count as movements.
+    Map<String, ResourceAssignment> newResource = buildResourceAssignment(ImmutableMap.of("DB2",
+        ImmutableMap.of("DB2_0", ImmutableMap.of("host_A", "MASTER", "host_B", "SLAVE"))));
+    Assert.assertEquals(ResourceUsageCalculator.countReplicaMovements(previous, newResource), 2L);
+  }
+
   @Test
   public void testCalculateAveragePartitionWeight() {
     Map<String, Map<String, Integer>> partitionCapacityMap = ImmutableMap.of(

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

@@ -59,6 +59,7 @@
 import org.apache.helix.monitoring.metrics.MetricCollector;
 import org.apache.helix.monitoring.metrics.WagedRebalancerMetricCollector;
 import org.apache.helix.monitoring.metrics.model.CountMetric;
+import org.apache.helix.monitoring.metrics.model.LatencyMetric;
 import org.apache.helix.monitoring.metrics.model.RatioMetric;
 import org.mockito.stubbing.Answer;
 import org.testng.Assert;
@@ -205,6 +206,26 @@ public void testWagedRebalanceMetrics()
         WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceCounter.name(),
         CountMetric.class).getLastEmittedMetricValue(), 1L);
 
+    // O1: both the build and solve sub-phase latencies were measured for the baseline and the
+    // partial rebalance (a measured gauge reports >= 0; an unmeasured one would still be -1).
+    Assert.assertTrue((long) metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.GlobalBaselineCalcBuildLatencyGauge
+            .name(), LatencyMetric.class).getLastEmittedMetricValue() >= 0L);
+    Assert.assertTrue((long) metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.GlobalBaselineCalcSolveLatencyGauge
+            .name(), LatencyMetric.class).getLastEmittedMetricValue() >= 0L);
+    Assert.assertTrue((long) metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceBuildLatencyGauge
+            .name(), LatencyMetric.class).getLastEmittedMetricValue() >= 0L);
+    Assert.assertTrue((long) metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceSolveLatencyGauge
+            .name(), LatencyMetric.class).getLastEmittedMetricValue() >= 0L);
+
+    // O2: the from-scratch assignment moved replicas, so the partial movement counter is positive.
+    Assert.assertTrue((long) metricCollector.getMetric(
+        WagedRebalancerMetricCollector.WagedRebalancerMetricNames.PartialRebalanceReplicaMovementCounter
+            .name(), CountMetric.class).getLastEmittedMetricValue() > 0L);
+
     // Wait for asyncReportBaselineDivergenceGauge to complete and verify.
     Assert.assertTrue(TestHelper.verify(() -> (double) metricCollector.getMetric(
         WagedRebalancerMetricCollector.WagedRebalancerMetricNames.BaselineDivergenceGauge.name(),