From 3c1b98fbe1e725161e3f226e48db0b304cee6222 Mon Sep 17 00:00:00 2001
From: Briana Augenreich <baugenreich@hubspot.com>
Date: Tue, 15 Mar 2022 15:31:40 -0400
Subject: [PATCH 1/4] Hubspot Backport: HBASE-26798 improve balancer visibility

---
 .../balancer/StochasticLoadBalancer.java      | 45 ++++++++++++++-----
 1 file changed, 35 insertions(+), 10 deletions(-)

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 4a0dc6fe7d9b..a19e8b389858 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
@@ -132,6 +132,7 @@ public class StochasticLoadBalancer extends BaseLoadBalancer {
   protected static final String COST_FUNCTIONS_COST_FUNCTIONS_KEY =
           "hbase.master.balancer.stochastic.additionalCostFunctions";
   public static final String OVERALL_COST_FUNCTION_NAME = "Overall";
+  public static final String WIGHTED_IMBALANCE_COST_NAME = "Imbalance";
 
   protected static final Random RANDOM = new Random(System.currentTimeMillis());
   private static final Logger LOG = LoggerFactory.getLogger(StochasticLoadBalancer.class);
@@ -323,7 +324,7 @@ private void updateBalancerTableLoadInfo(TableName tableName,
     initCosts(cluster);
     curOverallCost = computeCost(cluster, Double.MAX_VALUE);
     System.arraycopy(tempFunctionCosts, 0, curFunctionCosts, 0, curFunctionCosts.length);
-    updateStochasticCosts(tableName, curOverallCost, curFunctionCosts);
+    updateStochasticCosts(tableName, curOverallCost, curOverallCost/sumMultiplier, curFunctionCosts);
   }
 
   @Override
@@ -424,7 +425,7 @@ protected boolean needsBalance(TableName tableName, Cluster cluster) {
   @InterfaceAudience.Private
   Cluster.Action nextAction(Cluster cluster) {
     return candidateGenerators.get(RANDOM.nextInt(candidateGenerators.size()))
-            .generate(cluster);
+          .generate(cluster);
   }
 
   /**
@@ -480,7 +481,7 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
     double currentCost = computeCost(cluster, Double.MAX_VALUE);
     curOverallCost = currentCost;
     System.arraycopy(tempFunctionCosts, 0, curFunctionCosts, 0, curFunctionCosts.length);
-    updateStochasticCosts(tableName, curOverallCost, curFunctionCosts);
+    updateStochasticCosts(tableName, curOverallCost, currentCost / sumMultiplier, curFunctionCosts);
     double initCost = currentCost;
     double newCost;
 
@@ -505,8 +506,8 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
       }
     }
     LOG.info("Start StochasticLoadBalancer.balancer, initial weighted average imbalance={}," +
-        " functionCost={} computedMaxSteps={}",
-      currentCost / sumMultiplier, functionCost(), computedMaxSteps);
+        " sumMultiplier={} sumCost={} functionCost={} computedMaxSteps={}",
+      currentCost / sumMultiplier, sumMultiplier, currentCost, functionCost(), computedMaxSteps);
 
     final String initFunctionTotalCosts = totalCostsPerFunc();
     // Perform a stochastic walk to see if we can get a good fit.
@@ -532,6 +533,9 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
         curOverallCost = currentCost;
         System.arraycopy(tempFunctionCosts, 0, curFunctionCosts, 0, curFunctionCosts.length);
       } else {
+        //LOG hidden move cost if greatly impacting ability to find better plan
+        logMoveCosts(newCost, tempFunctionCosts);
+
         // Put things back the way they were before.
         // TODO: undo by remembering old values
         Action undoAction = action.undoAction();
@@ -548,7 +552,7 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
     metricsBalancer.balanceCluster(endTime - startTime);
 
     if (initCost > currentCost) {
-      updateStochasticCosts(tableName, curOverallCost, curFunctionCosts);
+      updateStochasticCosts(tableName, curOverallCost, currentCost / sumMultiplier, curFunctionCosts);
       plans = createRegionPlans(cluster);
       LOG.info("Finished computing new moving plan. Computation took {} ms" +
           " to try {} different iterations.  Found a solution that moves " +
@@ -559,9 +563,11 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
       sendRegionPlansToRingBuffer(plans, currentCost, initCost, initFunctionTotalCosts, step);
       return plans;
     }
+
     LOG.info("Could not find a better moving plan.  Tried {} different configurations in "
-        + "{} ms, and did not find anything with an imbalance score less than {}", step,
+        + "{} ms, and did not find anything with an imbalance score less than {}.", step,
       endTime - startTime, initCost / sumMultiplier);
+
     return null;
   }
 
@@ -603,10 +609,24 @@ private void sendRegionPlansToRingBuffer(List<RegionPlan> plans, double currentC
     }
   }
 
+  private void logMoveCosts(double currentCost, double [] currentFunctionCosts){
+    for (int i = 0; i < costFunctions.size(); i++) {
+      CostFunction costFunction = costFunctions.get(i);
+      if (costFunction instanceof MoveCostFunction){
+        String costFunctionName = costFunction.getClass().getSimpleName();
+        double moveCost = currentFunctionCosts[i];
+        double costPercent = (currentCost == 0) ? 0 : ( moveCost/ currentCost);
+        if (costPercent > .50 ){
+          LOG.info("Move cost greatly impacting overall cost of improvement plan. currentCost={} percentageOfCost={}. Consider lowering moveCost multiplier.",moveCost, costPercent );
+        }
+      }
+
+    }
+  }
   /**
    * update costs to JMX
    */
-  private void updateStochasticCosts(TableName tableName, double overall, double[] subCosts) {
+  private void updateStochasticCosts(TableName tableName, double overall, double weightedImbalance, double[] subCosts) {
     if (tableName == null) {
       return;
     }
@@ -618,13 +638,17 @@ private void updateStochasticCosts(TableName tableName, double overall, double[]
       balancer.updateStochasticCost(tableName.getNameAsString(),
         OVERALL_COST_FUNCTION_NAME, "Overall cost", overall);
 
-      // each cost function
+      // weighted imbalance
+      balancer.updateStochasticCost(tableName.getNameAsString(),
+        WIGHTED_IMBALANCE_COST_NAME, "Weighted imbalance", weightedImbalance);
+
+      // each cost function as a percent of overall cost
       for (int i = 0; i < costFunctions.size(); i++) {
         CostFunction costFunction = costFunctions.get(i);
         String costFunctionName = costFunction.getClass().getSimpleName();
         double costPercent = (overall == 0) ? 0 : (subCosts[i] / overall);
         // TODO: cost function may need a specific description
-        balancer.updateStochasticCost(tableName.getNameAsString(), costFunctionName,
+        balancer.updateStochasticCost(tableName.getNameAsString(), costFunctionName + "Percentage",
           "The percent of " + costFunctionName, costPercent);
       }
     }
@@ -967,6 +991,7 @@ protected double cost() {
       // Don't let this single balance move more than the max moves.
       // This allows better scaling to accurately represent the actual cost of a move.
       if (moveCost > maxMoves) {
+        LOG.info("Calculated moves exceed maxMoves {}. Greatly increasing move cost. ", maxMoves);
         return 1000000;   // return a number much greater than any of the other cost
       }
 

From bbc43db69e8f717c13941dd699cc95f4a1b6d988 Mon Sep 17 00:00:00 2001
From: Briana Augenreich <baugenreich@hubspot.com>
Date: Thu, 17 Mar 2022 13:40:11 -0400
Subject: [PATCH 2/4] PR feedback + fix tests

---
 .../balancer/StochasticLoadBalancer.java      | 13 ++++++------
 .../TestStochasticBalancerJmxMetrics.java     | 21 +++++++++++++------
 2 files changed, 22 insertions(+), 12 deletions(-)

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 a19e8b389858..c834e3288707 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
@@ -481,7 +481,7 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
     double currentCost = computeCost(cluster, Double.MAX_VALUE);
     curOverallCost = currentCost;
     System.arraycopy(tempFunctionCosts, 0, curFunctionCosts, 0, curFunctionCosts.length);
-    updateStochasticCosts(tableName, curOverallCost, currentCost / sumMultiplier, curFunctionCosts);
+    updateStochasticCosts(tableName, curOverallCost, curOverallCost / sumMultiplier, curFunctionCosts);
     double initCost = currentCost;
     double newCost;
 
@@ -552,7 +552,7 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
     metricsBalancer.balanceCluster(endTime - startTime);
 
     if (initCost > currentCost) {
-      updateStochasticCosts(tableName, curOverallCost, currentCost / sumMultiplier, curFunctionCosts);
+      updateStochasticCosts(tableName, curOverallCost, curOverallCost / sumMultiplier, curFunctionCosts);
       plans = createRegionPlans(cluster);
       LOG.info("Finished computing new moving plan. Computation took {} ms" +
           " to try {} different iterations.  Found a solution that moves " +
@@ -615,9 +615,10 @@ private void logMoveCosts(double currentCost, double [] currentFunctionCosts){
       if (costFunction instanceof MoveCostFunction){
         String costFunctionName = costFunction.getClass().getSimpleName();
         double moveCost = currentFunctionCosts[i];
-        double costPercent = (currentCost == 0) ? 0 : ( moveCost/ currentCost);
+        double costPercent = (currentCost == 0) ? 0 : (moveCost/ currentCost);
         if (costPercent > .50 ){
-          LOG.info("Move cost greatly impacting overall cost of improvement plan. currentCost={} percentageOfCost={}. Consider lowering moveCost multiplier.",moveCost, costPercent );
+          LOG.info("{} greatly impacting overall cost of improvement plan. currentCost={} percentageOfCost={}. Consider lowering moveCost multiplier.", costFunctionName, moveCost, costPercent );
+          break;
         }
       }
 
@@ -627,6 +628,7 @@ private void logMoveCosts(double currentCost, double [] currentFunctionCosts){
    * update costs to JMX
    */
   private void updateStochasticCosts(TableName tableName, double overall, double weightedImbalance, double[] subCosts) {
+
     if (tableName == null) {
       return;
     }
@@ -648,7 +650,7 @@ private void updateStochasticCosts(TableName tableName, double overall, double w
         String costFunctionName = costFunction.getClass().getSimpleName();
         double costPercent = (overall == 0) ? 0 : (subCosts[i] / overall);
         // TODO: cost function may need a specific description
-        balancer.updateStochasticCost(tableName.getNameAsString(), costFunctionName + "Percentage",
+        balancer.updateStochasticCost(tableName.getNameAsString(), costFunctionName,
           "The percent of " + costFunctionName, costPercent);
       }
     }
@@ -991,7 +993,6 @@ protected double cost() {
       // Don't let this single balance move more than the max moves.
       // This allows better scaling to accurately represent the actual cost of a move.
       if (moveCost > maxMoves) {
-        LOG.info("Calculated moves exceed maxMoves {}. Greatly increasing move cost. ", maxMoves);
         return 1000000;   // return a number much greater than any of the other cost
       }
 
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
index ab1c76c2e9dd..54d7a91b6da0 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.HashSet;
 import java.util.Hashtable;
 import java.util.Iterator;
@@ -60,7 +61,6 @@
 
 @Category({ MiscTests.class, MediumTests.class })
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
-@Ignore
 public class TestStochasticBalancerJmxMetrics extends BalancerTestBase {
 
   @ClassRule
@@ -95,6 +95,7 @@ public static void setupBeforeClass() throws Exception {
 
     conf.setClass("hbase.util.ip.to.rack.determiner", MockMapping.class, DNSToSwitchMapping.class);
     conf.setFloat("hbase.master.balancer.stochastic.maxMovePercent", 0.75f);
+//    conf.setFloat("hbase.master.balancer.stochastic.moveCost", 0);
     conf.setFloat("hbase.regions.slop", 0.0f);
     conf.set(CoprocessorHost.REGIONSERVER_COPROCESSOR_CONF_KEY, JMXListener.class.getName());
     Random rand = new Random();
@@ -130,6 +131,7 @@ public static void tearDownAfterClass() throws Exception {
    */
   @Test
   public void testJmxMetrics_EnsembleMode() throws Exception {
+
     loadBalancer = new StochasticLoadBalancer();
 
     conf.setBoolean(HConstants.HBASE_MASTER_LOADBALANCE_BYTABLE, false);
@@ -138,14 +140,16 @@ public void testJmxMetrics_EnsembleMode() throws Exception {
     TableName tableName = HConstants.ENSEMBLE_TABLE_NAME;
     Map<ServerName, List<RegionInfo>> clusterState = mockClusterServers(mockCluster_ensemble);
     loadBalancer.balanceTable(tableName, clusterState);
+    Thread.sleep(Duration.ofSeconds(5).toMillis());
 
     String[] tableNames = new String[] { tableName.getNameAsString() };
     String[] functionNames = loadBalancer.getCostFunctionNames();
+
     Set<String> jmxMetrics = readJmxMetricsWithRetry();
     Set<String> expectedMetrics = getExpectedJmxMetrics(tableNames, functionNames);
 
-    // printMetrics(jmxMetrics, "existing metrics in ensemble mode");
-    // printMetrics(expectedMetrics, "expected metrics in ensemble mode");
+     printMetrics(jmxMetrics, "existing metrics in ensemble mode");
+     printMetrics(expectedMetrics, "expected metrics in ensemble mode");
 
     // assert that every expected is in the JMX
     for (String expected : expectedMetrics) {
@@ -159,6 +163,7 @@ public void testJmxMetrics_EnsembleMode() throws Exception {
    */
   @Test
   public void testJmxMetrics_PerTableMode() throws Exception {
+
     loadBalancer = new StochasticLoadBalancer();
 
     conf.setBoolean(HConstants.HBASE_MASTER_LOADBALANCE_BYTABLE, true);
@@ -174,23 +179,25 @@ public void testJmxMetrics_PerTableMode() throws Exception {
     TableName tableName = TableName.valueOf(TABLE_NAME_1);
     Map<ServerName, List<RegionInfo>> clusterState = mockClusterServers(mockCluster_pertable_1);
     loadBalancer.balanceTable(tableName, clusterState);
-
+    Thread.sleep(Duration.ofSeconds(3).toMillis());
     // table 2
     tableName = TableName.valueOf(TABLE_NAME_2);
     clusterState = mockClusterServers(mockCluster_pertable_2);
     loadBalancer.balanceTable(tableName, clusterState);
+    Thread.sleep(Duration.ofSeconds(3).toMillis());
 
     // table hbase:namespace
     tableName = TableName.valueOf(TABLE_NAME_NAMESPACE);
     clusterState = mockClusterServers(mockCluster_pertable_namespace);
     loadBalancer.balanceTable(tableName, clusterState);
 
+    Thread.sleep(Duration.ofSeconds(5).toMillis());
     String[] tableNames = new String[] { TABLE_NAME_1, TABLE_NAME_2, TABLE_NAME_NAMESPACE };
     Set<String> jmxMetrics = readJmxMetricsWithRetry();
     Set<String> expectedMetrics = getExpectedJmxMetrics(tableNames, functionNames);
 
-    // printMetrics(jmxMetrics, "existing metrics in per-table mode");
-    // printMetrics(expectedMetrics, "expected metrics in per-table mode");
+     printMetrics(jmxMetrics, "existing metrics in per-table mode");
+     printMetrics(expectedMetrics, "expected metrics in per-table mode");
 
     // assert that every expected is in the JMX
     for (String expected : expectedMetrics) {
@@ -202,6 +209,7 @@ public void testJmxMetrics_PerTableMode() throws Exception {
   private Set<String> readJmxMetricsWithRetry() throws IOException {
     final int count = 0;
     for (int i = 0; i < 10; i++) {
+
       Set<String> metrics = readJmxMetrics();
       if (metrics != null) {
         return metrics;
@@ -268,6 +276,7 @@ private Set<String> getExpectedJmxMetrics(String[] tableNames, String[] function
 
     for (String tableName : tableNames) {
       ret.add(StochasticLoadBalancer.composeAttributeName(tableName, "Overall"));
+      ret.add(StochasticLoadBalancer.composeAttributeName(tableName, "Imbalance"));
       for (String functionName : functionNames) {
         String metricsName = StochasticLoadBalancer.composeAttributeName(tableName, functionName);
         ret.add(metricsName);

From 498517c65869ed676150eeb3c7007a912b4a44de Mon Sep 17 00:00:00 2001
From: Briana Augenreich <baugenreich@hubspot.com>
Date: Mon, 21 Mar 2022 15:22:04 -0400
Subject: [PATCH 3/4] PR feedback improvements

---
 .../master/balancer/StochasticLoadBalancer.java    | 14 ++++++++++----
 .../balancer/TestStochasticBalancerJmxMetrics.java | 14 +++++++-------
 2 files changed, 17 insertions(+), 11 deletions(-)

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 c834e3288707..4417e960d0cb 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
@@ -512,6 +512,7 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
     final String initFunctionTotalCosts = totalCostsPerFunc();
     // Perform a stochastic walk to see if we can get a good fit.
     long step;
+    int moveCostLogCounter = 0;
 
     for (step = 0; step < computedMaxSteps; step++) {
       Cluster.Action action = nextAction(cluster);
@@ -534,7 +535,8 @@ public synchronized List<RegionPlan> balanceTable(TableName tableName, Map<Serve
         System.arraycopy(tempFunctionCosts, 0, curFunctionCosts, 0, curFunctionCosts.length);
       } else {
         //LOG hidden move cost if greatly impacting ability to find better plan
-        logMoveCosts(newCost, tempFunctionCosts);
+        logMoveCosts(newCost, tempFunctionCosts, moveCostLogCounter);
+        moveCostLogCounter++;
 
         // Put things back the way they were before.
         // TODO: undo by remembering old values
@@ -609,15 +611,19 @@ private void sendRegionPlansToRingBuffer(List<RegionPlan> plans, double currentC
     }
   }
 
-  private void logMoveCosts(double currentCost, double [] currentFunctionCosts){
+  private void logMoveCosts(double currentCost, double [] currentFunctionCosts, int logCount){
+    if (logCount > 5){
+      // no op - so as to not flood the logs
+      return;
+    }
     for (int i = 0; i < costFunctions.size(); i++) {
       CostFunction costFunction = costFunctions.get(i);
       if (costFunction instanceof MoveCostFunction){
         String costFunctionName = costFunction.getClass().getSimpleName();
         double moveCost = currentFunctionCosts[i];
-        double costPercent = (currentCost == 0) ? 0 : (moveCost/ currentCost);
+        double costPercent = (currentCost == 0) ? 0 : (moveCost/ currentCost) * 100;
         if (costPercent > .50 ){
-          LOG.info("{} greatly impacting overall cost of improvement plan. currentCost={} percentageOfCost={}. Consider lowering moveCost multiplier.", costFunctionName, moveCost, costPercent );
+          LOG.info("{} may be impacting the ability to find an improved plan. calculatedMoveCost={} percentageOfPlanCost={}%. Consider lowering moveCost multiplier.", costFunctionName, moveCost, costPercent );
           break;
         }
       }
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
index 54d7a91b6da0..7536ace6a060 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
@@ -20,7 +20,6 @@
 import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
-import java.time.Duration;
 import java.util.HashSet;
 import java.util.Hashtable;
 import java.util.Iterator;
@@ -28,6 +27,7 @@
 import java.util.Map;
 import java.util.Random;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import javax.management.MBeanAttributeInfo;
 import javax.management.MBeanInfo;
 import javax.management.MBeanServerConnection;
@@ -52,7 +52,6 @@
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.FixMethodOrder;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runners.MethodSorters;
@@ -140,7 +139,6 @@ public void testJmxMetrics_EnsembleMode() throws Exception {
     TableName tableName = HConstants.ENSEMBLE_TABLE_NAME;
     Map<ServerName, List<RegionInfo>> clusterState = mockClusterServers(mockCluster_ensemble);
     loadBalancer.balanceTable(tableName, clusterState);
-    Thread.sleep(Duration.ofSeconds(5).toMillis());
 
     String[] tableNames = new String[] { tableName.getNameAsString() };
     String[] functionNames = loadBalancer.getCostFunctionNames();
@@ -179,19 +177,19 @@ public void testJmxMetrics_PerTableMode() throws Exception {
     TableName tableName = TableName.valueOf(TABLE_NAME_1);
     Map<ServerName, List<RegionInfo>> clusterState = mockClusterServers(mockCluster_pertable_1);
     loadBalancer.balanceTable(tableName, clusterState);
-    Thread.sleep(Duration.ofSeconds(3).toMillis());
+
     // table 2
     tableName = TableName.valueOf(TABLE_NAME_2);
     clusterState = mockClusterServers(mockCluster_pertable_2);
     loadBalancer.balanceTable(tableName, clusterState);
-    Thread.sleep(Duration.ofSeconds(3).toMillis());
+
 
     // table hbase:namespace
     tableName = TableName.valueOf(TABLE_NAME_NAMESPACE);
     clusterState = mockClusterServers(mockCluster_pertable_namespace);
     loadBalancer.balanceTable(tableName, clusterState);
 
-    Thread.sleep(Duration.ofSeconds(5).toMillis());
+
     String[] tableNames = new String[] { TABLE_NAME_1, TABLE_NAME_2, TABLE_NAME_NAMESPACE };
     Set<String> jmxMetrics = readJmxMetricsWithRetry();
     Set<String> expectedMetrics = getExpectedJmxMetrics(tableNames, functionNames);
@@ -211,7 +209,7 @@ private Set<String> readJmxMetricsWithRetry() throws IOException {
     for (int i = 0; i < 10; i++) {
 
       Set<String> metrics = readJmxMetrics();
-      if (metrics != null) {
+      if (metrics != null && metrics.size() > 2) {
         return metrics;
       }
       LOG.warn("Failed to get jmxmetrics... sleeping, retrying; " + i + " of " + count + " times");
@@ -227,6 +225,8 @@ private Set<String> readJmxMetrics() throws IOException {
     JMXConnector connector = null;
     ObjectName target = null;
     MBeanServerConnection mb = null;
+    CountDownLatch latch = new CountDownLatch(1);
+
     try {
       connector =
           JMXConnectorFactory.connect(JMXListener.buildJMXServiceURL(connectorPort, connectorPort));

From c127d84016244ddba31641d59c05fc291ced5a65 Mon Sep 17 00:00:00 2001
From: Briana Augenreich <baugenreich@hubspot.com>
Date: Wed, 23 Mar 2022 15:49:16 -0400
Subject: [PATCH 4/4] Final testing tweaks

---
 .../TestStochasticBalancerJmxMetrics.java     | 25 +++++++++----------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
index 7536ace6a060..be1256c8220a 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticBalancerJmxMetrics.java
@@ -18,7 +18,6 @@
 package org.apache.hadoop.hbase.master.balancer;
 
 import static org.junit.Assert.assertTrue;
-
 import java.io.IOException;
 import java.util.HashSet;
 import java.util.Hashtable;
@@ -27,7 +26,6 @@
 import java.util.Map;
 import java.util.Random;
 import java.util.Set;
-import java.util.concurrent.CountDownLatch;
 import javax.management.MBeanAttributeInfo;
 import javax.management.MBeanInfo;
 import javax.management.MBeanServerConnection;
@@ -94,7 +92,6 @@ public static void setupBeforeClass() throws Exception {
 
     conf.setClass("hbase.util.ip.to.rack.determiner", MockMapping.class, DNSToSwitchMapping.class);
     conf.setFloat("hbase.master.balancer.stochastic.maxMovePercent", 0.75f);
-//    conf.setFloat("hbase.master.balancer.stochastic.moveCost", 0);
     conf.setFloat("hbase.regions.slop", 0.0f);
     conf.set(CoprocessorHost.REGIONSERVER_COPROCESSOR_CONF_KEY, JMXListener.class.getName());
     Random rand = new Random();
@@ -143,7 +140,7 @@ public void testJmxMetrics_EnsembleMode() throws Exception {
     String[] tableNames = new String[] { tableName.getNameAsString() };
     String[] functionNames = loadBalancer.getCostFunctionNames();
 
-    Set<String> jmxMetrics = readJmxMetricsWithRetry();
+    Set<String> jmxMetrics = readJmxMetricsWithRetry(functionNames.length + 1);
     Set<String> expectedMetrics = getExpectedJmxMetrics(tableNames, functionNames);
 
      printMetrics(jmxMetrics, "existing metrics in ensemble mode");
@@ -169,9 +166,10 @@ public void testJmxMetrics_PerTableMode() throws Exception {
 
     // NOTE the size is normally set in setClusterMetrics, for test purpose, we set it manually
     // Tables: hbase:namespace, table1, table2
-    // Functions: costFunctions, overall
+    // Functions: costFunctions, overall, imbalance
     String[] functionNames = loadBalancer.getCostFunctionNames();
-    loadBalancer.updateMetricsSize(3 * (functionNames.length + 1));
+
+    loadBalancer.updateMetricsSize(3 * (functionNames.length + 2));
 
     // table 1
     TableName tableName = TableName.valueOf(TABLE_NAME_1);
@@ -191,11 +189,11 @@ public void testJmxMetrics_PerTableMode() throws Exception {
 
 
     String[] tableNames = new String[] { TABLE_NAME_1, TABLE_NAME_2, TABLE_NAME_NAMESPACE };
-    Set<String> jmxMetrics = readJmxMetricsWithRetry();
+    Set<String> jmxMetrics = readJmxMetricsWithRetry(3 * (functionNames.length + 2));
     Set<String> expectedMetrics = getExpectedJmxMetrics(tableNames, functionNames);
 
-     printMetrics(jmxMetrics, "existing metrics in per-table mode");
-     printMetrics(expectedMetrics, "expected metrics in per-table mode");
+    printMetrics(jmxMetrics, "existing metrics in per-table mode");
+    printMetrics(expectedMetrics, "expected metrics in per-table mode");
 
     // assert that every expected is in the JMX
     for (String expected : expectedMetrics) {
@@ -204,12 +202,13 @@ public void testJmxMetrics_PerTableMode() throws Exception {
     }
   }
 
-  private Set<String> readJmxMetricsWithRetry() throws IOException {
+  private Set<String> readJmxMetricsWithRetry(int expectedMetricsCount) throws IOException {
     final int count = 0;
-    for (int i = 0; i < 10; i++) {
+    for (int i = 0; i < 20; i++) {
 
       Set<String> metrics = readJmxMetrics();
-      if (metrics != null && metrics.size() > 2) {
+
+      if (metrics != null && metrics.size() >= expectedMetricsCount){
         return metrics;
       }
       LOG.warn("Failed to get jmxmetrics... sleeping, retrying; " + i + " of " + count + " times");
@@ -225,7 +224,6 @@ private Set<String> readJmxMetrics() throws IOException {
     JMXConnector connector = null;
     ObjectName target = null;
     MBeanServerConnection mb = null;
-    CountDownLatch latch = new CountDownLatch(1);
 
     try {
       connector =
@@ -295,5 +293,6 @@ private static void printMetrics(Set<String> metrics, String info) {
     for (String str : metrics) {
       LOG.info(" ++++ " + str);
     }
+
   }
 }