You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hbase.apache.org by rm...@apache.org on 2025/01/23 18:36:38 UTC

(hbase) branch HBASE-29071-branch-2.6 created (now 510ca313555)

This is an automated email from the ASF dual-hosted git repository.

rmattingly pushed a change to branch HBASE-29071-branch-2.6
in repository https://gitbox.apache.org/repos/asf/hbase.git


      at 510ca313555 HBASE-29071 StochasticLoadBalancer candidate generators should use a Map, rather than ordinal based indexing (#6598) (#6621)

This branch includes the following new commits:

     new 510ca313555 HBASE-29071 StochasticLoadBalancer candidate generators should use a Map, rather than ordinal based indexing (#6598) (#6621)

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.



(hbase) 01/01: HBASE-29071 StochasticLoadBalancer candidate generators should use a Map, rather than ordinal based indexing (#6598) (#6621)

Posted by rm...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

rmattingly pushed a commit to branch HBASE-29071-branch-2.6
in repository https://gitbox.apache.org/repos/asf/hbase.git

commit 510ca313555fd753753122eb254053bf06a85666
Author: Ray Mattingly <rm...@apache.org>
AuthorDate: Thu Jan 23 13:35:03 2025 -0500

    HBASE-29071 StochasticLoadBalancer candidate generators should use a Map, rather than ordinal based indexing (#6598) (#6621)
    
    Co-authored-by: Ray Mattingly <rm...@hubspot.com>
    Signed-off-by: Nick Dimiduk <nd...@apache.org>
---
 .../master/balancer/CacheAwareLoadBalancer.java    |  20 ++--
 .../hadoop/hbase/master/balancer/CostFunction.java |  21 ++--
 .../master/balancer/FavoredStochasticBalancer.java |  17 ++-
 .../master/balancer/LocalityBasedCostFunction.java |   5 +-
 .../balancer/RegionCountSkewCostFunction.java      |   5 +-
 .../RegionReplicaGroupingCostFunction.java         |  11 +-
 .../master/balancer/StochasticLoadBalancer.java    | 120 +++++++++++++++------
 .../LoadOnlyFavoredStochasticBalancer.java         |  11 +-
 .../balancer/TestStochasticLoadBalancer.java       |   2 +-
 .../TestStochasticLoadBalancerBalanceCluster.java  |   2 +-
 ...estStochasticLoadBalancerHeterogeneousCost.java |   3 +-
 ...ochasticLoadBalancerRegionReplicaWithRacks.java |   2 +-
 12 files changed, 145 insertions(+), 74 deletions(-)

diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CacheAwareLoadBalancer.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CacheAwareLoadBalancer.java
index d73769a3971..a954ef47ad9 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CacheAwareLoadBalancer.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CacheAwareLoadBalancer.java
@@ -67,12 +67,13 @@ public class CacheAwareLoadBalancer extends StochasticLoadBalancer {
   }
 
   @Override
-  protected List<CandidateGenerator> createCandidateGenerators() {
-    List<CandidateGenerator> candidateGenerators = new ArrayList<>(2);
-    candidateGenerators.add(GeneratorFunctionType.LOAD.ordinal(),
+  protected Map<Class<? extends CandidateGenerator>, CandidateGenerator>
+    createCandidateGenerators() {
+    Map<Class<? extends CandidateGenerator>, CandidateGenerator> candidateGenerators =
+      new HashMap<>(2);
+    candidateGenerators.put(CacheAwareSkewnessCandidateGenerator.class,
       new CacheAwareSkewnessCandidateGenerator());
-    candidateGenerators.add(GeneratorFunctionType.CACHE_RATIO.ordinal(),
-      new CacheAwareCandidateGenerator());
+    candidateGenerators.put(CacheAwareCandidateGenerator.class, new CacheAwareCandidateGenerator());
     return candidateGenerators;
   }
 
@@ -403,8 +404,9 @@ public class CacheAwareLoadBalancer extends StochasticLoadBalancer {
       });
     }
 
-    public final void updateWeight(double[] weights) {
-      weights[GeneratorFunctionType.LOAD.ordinal()] += cost();
+    @Override
+    public final void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
+      weights.merge(LoadCandidateGenerator.class, cost(), Double::sum);
     }
   }
 
@@ -472,8 +474,8 @@ public class CacheAwareLoadBalancer extends StochasticLoadBalancer {
     }
 
     @Override
-    public final void updateWeight(double[] weights) {
-      weights[GeneratorFunctionType.CACHE_RATIO.ordinal()] += cost();
+    public void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
+      weights.merge(LoadCandidateGenerator.class, cost(), Double::sum);
     }
   }
 }
diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CostFunction.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CostFunction.java
index 372cc54c8b2..1dcd4580b1a 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CostFunction.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CostFunction.java
@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.master.balancer;
 
+import java.util.Map;
 import org.apache.yetus.audience.InterfaceAudience;
 
 /**
@@ -85,6 +86,16 @@ abstract class CostFunction {
 
   protected abstract double cost();
 
+  /**
+   * Add the cost of this cost function to the weight of the candidate generator that is optimized
+   * for this cost function. By default it is the RandomCandiateGenerator for a cost function.
+   * Called once per init or after postAction.
+   * @param weights the weights for every generator.
+   */
+  public void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
+    weights.merge(RandomCandidateGenerator.class, cost(), Double::sum);
+  }
+
   /**
    * Scale the value between 0 and 1.
    * @param min   Min value
@@ -106,14 +117,4 @@ abstract class CostFunction {
 
     return Math.max(0d, Math.min(1d, (value - min) / (max - min)));
   }
-
-  /**
-   * Add the cost of this cost function to the weight of the candidate generator that is optimized
-   * for this cost function. By default it is the RandomCandiateGenerator for a cost function.
-   * Called once per init or after postAction.
-   * @param weights the weights for every generator.
-   */
-  public void updateWeight(double[] weights) {
-    weights[StochasticLoadBalancer.GeneratorType.RANDOM.ordinal()] += cost();
-  }
 }
diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/FavoredStochasticBalancer.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/FavoredStochasticBalancer.java
index 7fc867bf1b9..db4c7c95b65 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/FavoredStochasticBalancer.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/FavoredStochasticBalancer.java
@@ -75,18 +75,25 @@ public class FavoredStochasticBalancer extends StochasticLoadBalancer
   private static final Logger LOG = LoggerFactory.getLogger(FavoredStochasticBalancer.class);
   private FavoredNodesManager fnm;
 
+  public void setFavoredNodesManager(FavoredNodesManager fnm) {
+    this.fnm = fnm;
+  }
+
   @Override
-  protected List<CandidateGenerator> createCandidateGenerators() {
-    List<CandidateGenerator> fnPickers = new ArrayList<>(2);
-    fnPickers.add(new FavoredNodeLoadPicker());
-    fnPickers.add(new FavoredNodeLocalityPicker());
+  protected Map<Class<? extends CandidateGenerator>, CandidateGenerator>
+    createCandidateGenerators() {
+    Map<Class<? extends CandidateGenerator>, CandidateGenerator> fnPickers = new HashMap<>(2);
+    fnPickers.put(FavoredNodeLoadPicker.class, new FavoredNodeLoadPicker());
+    fnPickers.put(FavoredNodeLocalityPicker.class, new FavoredNodeLocalityPicker());
     return fnPickers;
   }
 
   /** Returns any candidate generator in random */
   @Override
   protected CandidateGenerator getRandomGenerator() {
-    return candidateGenerators.get(ThreadLocalRandom.current().nextInt(candidateGenerators.size()));
+    Class<? extends CandidateGenerator> clazz = shuffledGeneratorClasses.get()
+      .get(ThreadLocalRandom.current().nextInt(candidateGenerators.size()));
+    return candidateGenerators.get(clazz);
   }
 
   @Override
diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/LocalityBasedCostFunction.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/LocalityBasedCostFunction.java
index 678c9a3e9ad..fc2d6cc12fb 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/LocalityBasedCostFunction.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/LocalityBasedCostFunction.java
@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.master.balancer;
 
+import java.util.Map;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.master.balancer.BalancerClusterState.LocalityType;
 import org.apache.yetus.audience.InterfaceAudience;
@@ -89,7 +90,7 @@ abstract class LocalityBasedCostFunction extends CostFunction {
   }
 
   @Override
-  public final void updateWeight(double[] weights) {
-    weights[StochasticLoadBalancer.GeneratorType.LOCALITY.ordinal()] += cost();
+  public final void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
+    weights.merge(LocalityBasedCandidateGenerator.class, cost(), Double::sum);
   }
 }
diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionCountSkewCostFunction.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionCountSkewCostFunction.java
index 9f542640b8d..7843ce163b0 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionCountSkewCostFunction.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionCountSkewCostFunction.java
@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.master.balancer;
 
+import java.util.Map;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.yetus.audience.InterfaceAudience;
 
@@ -62,7 +63,7 @@ class RegionCountSkewCostFunction extends CostFunction {
   }
 
   @Override
-  public final void updateWeight(double[] weights) {
-    weights[StochasticLoadBalancer.GeneratorType.LOAD.ordinal()] += cost();
+  public final void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
+    weights.merge(LoadCandidateGenerator.class, cost(), Double::sum);
   }
 }
diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionReplicaGroupingCostFunction.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionReplicaGroupingCostFunction.java
index 025f540b76f..43eaceb4f77 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionReplicaGroupingCostFunction.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/RegionReplicaGroupingCostFunction.java
@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.master.balancer;
 
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
 import org.agrona.collections.Hashing;
 import org.agrona.collections.Int2IntCounterMap;
@@ -72,6 +73,11 @@ abstract class RegionReplicaGroupingCostFunction extends CostFunction {
     return scale(0, maxCost, totalCost);
   }
 
+  @Override
+  public final void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
+    weights.merge(RegionReplicaRackCandidateGenerator.class, cost(), Double::sum);
+  }
+
   /**
    * For each primary region, it computes the total number of replicas in the array (numReplicas)
    * and returns a sum of numReplicas-1 squared. For example, if the server hosts regions a, b, c,
@@ -91,9 +97,4 @@ abstract class RegionReplicaGroupingCostFunction extends CostFunction {
     });
     return cost.longValue();
   }
-
-  @Override
-  public final void updateWeight(double[] weights) {
-    weights[StochasticLoadBalancer.GeneratorType.RACK.ordinal()] += cost();
-  }
 }
diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/StochasticLoadBalancer.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/StochasticLoadBalancer.java
index 507d1a8e369..46efb4d499d 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/StochasticLoadBalancer.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/StochasticLoadBalancer.java
@@ -22,11 +22,15 @@ import java.lang.reflect.Constructor;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.StringJoiner;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.ClusterMetrics;
 import org.apache.hadoop.hbase.HBaseInterfaceAudience;
@@ -50,6 +54,9 @@ import org.apache.yetus.audience.InterfaceAudience;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.hbase.thirdparty.com.google.common.base.Preconditions;
+import org.apache.hbase.thirdparty.com.google.common.base.Suppliers;
+
 /**
  * <p>
  * This is a best effort load balancer. Given a Cost function F(C) =&gt; x It will randomly try and
@@ -142,7 +149,6 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
   private boolean isBalancerDecisionRecording = false;
   private boolean isBalancerRejectionRecording = false;
 
-  protected List<CandidateGenerator> candidateGenerators;
   Map<String, Pair<ServerName, Float>> regionCacheRatioOnOldServerMap = new HashMap<>();
 
   protected List<CostFunction> costFunctions; // FindBugs: Wants this protected;
@@ -155,7 +161,6 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
     RACK
   }
 
-  private double[] weightsOfGenerators;
   // To save currently configed sum of multiplier. Defaulted at 1 for cases that carry high cost
   private float sumMultiplier;
   // to save and report costs to JMX
@@ -176,6 +181,17 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
    */
   NamedQueueRecorder namedQueueRecorder;
 
+  private final Map<Class<? extends CandidateGenerator>, Double> weightsOfGenerators =
+    new HashMap<>();
+  protected Map<Class<? extends CandidateGenerator>, CandidateGenerator> candidateGenerators;
+  protected final Supplier<List<Class<? extends CandidateGenerator>>> shuffledGeneratorClasses =
+    Suppliers.memoizeWithExpiration(() -> {
+      List<Class<? extends CandidateGenerator>> shuffled =
+        new ArrayList<>(candidateGenerators.keySet());
+      Collections.shuffle(shuffled);
+      return shuffled;
+    }, 5, TimeUnit.SECONDS);
+
   /**
    * The constructor that pass a MetricsStochasticBalancer to BaseLoadBalancer to replace its
    * default MetricsBalancer
@@ -223,16 +239,20 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
 
   @RestrictedApi(explanation = "Should only be called in tests", link = "",
       allowedOnPath = ".*/src/test/.*")
-  List<CandidateGenerator> getCandidateGenerators() {
+  Map<Class<? extends CandidateGenerator>, CandidateGenerator> getCandidateGenerators() {
     return this.candidateGenerators;
   }
 
-  protected List<CandidateGenerator> createCandidateGenerators() {
-    List<CandidateGenerator> candidateGenerators = new ArrayList<CandidateGenerator>(4);
-    candidateGenerators.add(GeneratorType.RANDOM.ordinal(), new RandomCandidateGenerator());
-    candidateGenerators.add(GeneratorType.LOAD.ordinal(), new LoadCandidateGenerator());
-    candidateGenerators.add(GeneratorType.LOCALITY.ordinal(), localityCandidateGenerator);
-    candidateGenerators.add(GeneratorType.RACK.ordinal(),
+  protected Map<Class<? extends CandidateGenerator>, CandidateGenerator>
+    createCandidateGenerators() {
+    Map<Class<? extends CandidateGenerator>, CandidateGenerator> candidateGenerators =
+      new HashMap<>(5);
+    candidateGenerators.put(RandomCandidateGenerator.class, new RandomCandidateGenerator());
+    candidateGenerators.put(LoadCandidateGenerator.class, new LoadCandidateGenerator());
+    candidateGenerators.put(LocalityBasedCandidateGenerator.class, localityCandidateGenerator);
+    candidateGenerators.put(RegionReplicaCandidateGenerator.class,
+      new RegionReplicaCandidateGenerator());
+    candidateGenerators.put(RegionReplicaRackCandidateGenerator.class,
       new RegionReplicaRackCandidateGenerator());
     return candidateGenerators;
   }
@@ -432,34 +452,54 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
 
   @RestrictedApi(explanation = "Should only be called in tests", link = "",
       allowedOnPath = ".*(/src/test/.*|StochasticLoadBalancer).java")
-  BalanceAction nextAction(BalancerClusterState cluster) {
-    return getRandomGenerator().generate(cluster);
+  Pair<CandidateGenerator, BalanceAction> nextAction(BalancerClusterState cluster) {
+    CandidateGenerator generator = getRandomGenerator();
+    return Pair.newPair(generator, generator.generate(cluster));
   }
 
   /**
    * Select the candidate generator to use based on the cost of cost functions. The chance of
-   * selecting a candidate generator is propotional to the share of cost of all cost functions among
-   * all cost functions that benefit from it.
+   * selecting a candidate generator is proportional to the share of cost of all cost functions
+   * among all cost functions that benefit from it.
    */
   protected CandidateGenerator getRandomGenerator() {
-    double sum = 0;
-    for (int i = 0; i < weightsOfGenerators.length; i++) {
-      sum += weightsOfGenerators[i];
-      weightsOfGenerators[i] = sum;
+    Preconditions.checkState(!candidateGenerators.isEmpty(), "No candidate generators available.");
+    List<Class<? extends CandidateGenerator>> generatorClasses = shuffledGeneratorClasses.get();
+    List<Double> partialSums = new ArrayList<>(generatorClasses.size());
+    double sum = 0.0;
+    for (Class<? extends CandidateGenerator> clazz : generatorClasses) {
+      double weight = weightsOfGenerators.getOrDefault(clazz, 0.0);
+      sum += weight;
+      partialSums.add(sum);
     }
-    if (sum == 0) {
-      return candidateGenerators.get(0);
-    }
-    for (int i = 0; i < weightsOfGenerators.length; i++) {
-      weightsOfGenerators[i] /= sum;
+
+    // If the sum of all weights is zero, fall back to any generator
+    if (sum == 0.0) {
+      return pickAnyGenerator(generatorClasses);
     }
+
     double rand = ThreadLocalRandom.current().nextDouble();
-    for (int i = 0; i < weightsOfGenerators.length; i++) {
-      if (rand <= weightsOfGenerators[i]) {
-        return candidateGenerators.get(i);
+    // Normalize partial sums so that the last one should be exactly 1.0
+    for (int i = 0; i < partialSums.size(); i++) {
+      partialSums.set(i, partialSums.get(i) / sum);
+    }
+
+    // Generate a random number and pick the first generator whose partial sum is >= rand
+    for (int i = 0; i < partialSums.size(); i++) {
+      if (rand <= partialSums.get(i)) {
+        return candidateGenerators.get(generatorClasses.get(i));
       }
     }
-    return candidateGenerators.get(candidateGenerators.size() - 1);
+
+    // Fallback: if for some reason we didn't return above, return any generator
+    return pickAnyGenerator(generatorClasses);
+  }
+
+  private CandidateGenerator
+    pickAnyGenerator(List<Class<? extends CandidateGenerator>> generatorClasses) {
+    Class<? extends CandidateGenerator> randomClass =
+      generatorClasses.get(ThreadLocalRandom.current().nextInt(candidateGenerators.size()));
+    return candidateGenerators.get(randomClass);
   }
 
   @RestrictedApi(explanation = "Should only be called in tests", link = "",
@@ -559,9 +599,12 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
     final String initFunctionTotalCosts = totalCostsPerFunc();
     // Perform a stochastic walk to see if we can get a good fit.
     long step;
-
+    Map<Class<? extends CandidateGenerator>, Long> generatorToStepCount = new HashMap<>();
+    Map<Class<? extends CandidateGenerator>, Long> generatorToApprovedActionCount = new HashMap<>();
     for (step = 0; step < computedMaxSteps; step++) {
-      BalanceAction action = nextAction(cluster);
+      Pair<CandidateGenerator, BalanceAction> nextAction = nextAction(cluster);
+      CandidateGenerator generator = nextAction.getFirst();
+      BalanceAction action = nextAction.getSecond();
 
       if (action.getType() == BalanceAction.Type.NULL) {
         continue;
@@ -569,12 +612,14 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
 
       cluster.doAction(action);
       updateCostsAndWeightsWithAction(cluster, action);
+      generatorToStepCount.merge(generator.getClass(), 1L, Long::sum);
 
       newCost = computeCost(cluster, currentCost);
 
       // Should this be kept?
       if (newCost < currentCost) {
         currentCost = newCost;
+        generatorToApprovedActionCount.merge(generator.getClass(), 1L, Long::sum);
 
         // save for JMX
         curOverallCost = currentCost;
@@ -593,6 +638,15 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
     }
     long endTime = EnvironmentEdgeManager.currentTime();
 
+    StringJoiner joiner = new StringJoiner("\n");
+    joiner.add("CandidateGenerator activity summary:");
+    generatorToStepCount.forEach((generator, count) -> {
+      long approvals = generatorToApprovedActionCount.getOrDefault(generator, 0L);
+      joiner.add(String.format(" - %s: %d steps, %d approvals", generator.getSimpleName(), count,
+        approvals));
+    });
+    LOG.debug(joiner.toString());
+
     metricsBalancer.balanceCluster(endTime - startTime);
 
     if (initCost > currentCost) {
@@ -785,8 +839,10 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
   @RestrictedApi(explanation = "Should only be called in tests", link = "",
       allowedOnPath = ".*(/src/test/.*|StochasticLoadBalancer).java")
   void initCosts(BalancerClusterState cluster) {
-    // Initialize the weights of generator every time
-    weightsOfGenerators = new double[this.candidateGenerators.size()];
+    weightsOfGenerators.clear();
+    for (Class<? extends CandidateGenerator> clazz : candidateGenerators.keySet()) {
+      weightsOfGenerators.put(clazz, 0.0);
+    }
     for (CostFunction c : costFunctions) {
       c.prepare(cluster);
       c.updateWeight(weightsOfGenerators);
@@ -797,8 +853,8 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
       allowedOnPath = ".*(/src/test/.*|StochasticLoadBalancer).java")
   void updateCostsAndWeightsWithAction(BalancerClusterState cluster, BalanceAction action) {
     // Reset all the weights to 0
-    for (int i = 0; i < weightsOfGenerators.length; i++) {
-      weightsOfGenerators[i] = 0;
+    for (Class<? extends CandidateGenerator> clazz : candidateGenerators.keySet()) {
+      weightsOfGenerators.put(clazz, 0.0);
     }
     for (CostFunction c : costFunctions) {
       if (c.isNeeded()) {
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/LoadOnlyFavoredStochasticBalancer.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/LoadOnlyFavoredStochasticBalancer.java
index ba6ee4af843..d658f7cfa16 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/LoadOnlyFavoredStochasticBalancer.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/LoadOnlyFavoredStochasticBalancer.java
@@ -17,8 +17,8 @@
  */
 package org.apache.hadoop.hbase.master.balancer;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * Used for FavoredNode unit tests
@@ -26,9 +26,10 @@ import java.util.List;
 public class LoadOnlyFavoredStochasticBalancer extends FavoredStochasticBalancer {
 
   @Override
-  protected List<CandidateGenerator> createCandidateGenerators() {
-    List<CandidateGenerator> fnPickers = new ArrayList<>(1);
-    fnPickers.add(new FavoredNodeLoadPicker());
+  protected Map<Class<? extends CandidateGenerator>, CandidateGenerator>
+    createCandidateGenerators() {
+    Map<Class<? extends CandidateGenerator>, CandidateGenerator> fnPickers = new HashMap<>(1);
+    fnPickers.put(FavoredNodeLoadPicker.class, new FavoredNodeLoadPicker());
     return fnPickers;
   }
 }
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancer.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancer.java
index e59f4dd85a7..9dc7dab6562 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancer.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancer.java
@@ -420,7 +420,7 @@ public class TestStochasticLoadBalancer extends BalancerTestBase {
       loadBalancer.initCosts(cluster);
       for (int i = 0; i != runs; ++i) {
         final double expectedCost = loadBalancer.computeCost(cluster, Double.MAX_VALUE);
-        BalanceAction action = loadBalancer.nextAction(cluster);
+        BalanceAction action = loadBalancer.nextAction(cluster).getSecond();
         cluster.doAction(action);
         loadBalancer.updateCostsAndWeightsWithAction(cluster, action);
         BalanceAction undoAction = action.undoAction();
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerBalanceCluster.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerBalanceCluster.java
index f9255fb1bfd..c18f29e770a 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerBalanceCluster.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerBalanceCluster.java
@@ -53,7 +53,7 @@ public class TestStochasticLoadBalancerBalanceCluster extends BalancerTestBase {
   @Test
   public void testBalanceCluster() throws Exception {
     conf.setLong(StochasticLoadBalancer.MAX_STEPS_KEY, 20000000L);
-    setMaxRunTime(Duration.ofMillis(1500));
+    setMaxRunTime(Duration.ofMillis(2500));
     loadBalancer.onConfigurationChange(conf);
 
     for (int[] mockCluster : clusterStateMocks) {
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerHeterogeneousCost.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerHeterogeneousCost.java
index 41bcb87d4b2..960783a8467 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerHeterogeneousCost.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerHeterogeneousCost.java
@@ -78,7 +78,8 @@ public class TestStochasticLoadBalancerHeterogeneousCost extends BalancerTestBas
     when(services.getConfiguration()).thenReturn(conf);
     BalancerTestBase.loadBalancer.setMasterServices(services);
     loadBalancer.initialize();
-    loadBalancer.getCandidateGenerators().add(new FairRandomCandidateGenerator());
+    loadBalancer.getCandidateGenerators().put(FairRandomCandidateGenerator.class,
+      new FairRandomCandidateGenerator());
   }
 
   @Test
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerRegionReplicaWithRacks.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerRegionReplicaWithRacks.java
index 6ecf73314ad..89bebaa5402 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerRegionReplicaWithRacks.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancerRegionReplicaWithRacks.java
@@ -80,7 +80,7 @@ public class TestStochasticLoadBalancerRegionReplicaWithRacks extends BalancerTe
   public void testRegionReplicationOnLargeClusterWithRacks() {
     conf.setBoolean("hbase.master.balancer.stochastic.runMaxSteps", true);
     conf.setLong(StochasticLoadBalancer.MAX_STEPS_KEY, 100000000L);
-    setMaxRunTime(Duration.ofSeconds(5));
+    setMaxRunTime(Duration.ofSeconds(10));
     loadBalancer.onConfigurationChange(conf);
     int numNodes = 100;
     int numRegions = numNodes * 30;