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

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

commit 11d2cf9750d5b45f0e3b526ce4e61eaae97bae81
Author: Ray Mattingly <[email protected]>
AuthorDate: Tue Jan 21 12:54:24 2025 -0500

    HBASE-29071 StochasticLoadBalancer candidate generators should use a Map, 
rather than ordinal based indexing (#6598)
    
    Co-authored-by: Ray Mattingly <[email protected]>
---
 .../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 +-
 ...ochasticLoadBalancerRegionReplicaWithRacks.java |   2 +-
 11 files changed, 143 insertions(+), 73 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 bc31a317aeb..385ed564c7d 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
@@ -68,12 +68,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;
   }
 
@@ -409,8 +410,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);
     }
   }
 
@@ -478,8 +480,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/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;

Reply via email to