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

Co-authored-by: Ray Mattingly <rmattingly@hubspot.com>
Signed-off-by: Nick Dimiduk <ndimiduk@apache.org>

Author: rmdmattingly

hbase-server/src/main/java/org/apache/hadoop/hbase/master/balancer/CacheAwareLoadBalancer.java

@@ -67,12 +67,13 @@ public synchronized void loadConf(Configuration configuration) {
   }
 
   @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 @@ protected void regionMoved(int region, int oldServer, int newServer) {
       });
     }
 
-    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 @@ private int getServerWithBestCacheRatioForRegion(int region) {
     }
 
     @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);
     }
   }
 }

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 @@ protected void regionMoved(int region, int oldServer, int newServer) {
 
   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 @@ protected static double scale(double min, double max, double value) {
 
     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();
-  }
 }

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

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 @@ private double getWeightedLocality(int region, int entity) {
   }
 
   @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);
   }
 }

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 @@ protected void regionMoved(int region, int oldServer, int newServer) {
   }
 
   @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);
   }
 }

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 @@ protected double cost() {
     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 @@ protected final long costPerGroup(Int2IntCounterMap colocatedReplicaCounts) {
     });
     return cost.longValue();
   }
-
-  @Override
-  public final void updateWeight(double[] weights) {
-    weights[StochasticLoadBalancer.GeneratorType.RACK.ordinal()] += cost();
-  }
 }