HBASE-30093 Start stochastic search timeout after initialization

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

@@ -585,7 +585,6 @@ protected List<RegionPlan> balanceTable(TableName tableName,
     BalancerClusterState cluster = createState(loadOfOneTable, loads, finder, rackManager);
 
     long startTime = EnvironmentEdgeManager.currentTime();
-    cluster.setStopRequestedAt(startTime + maxRunningTime);
 
     initCosts(cluster);
     balancerConditionals.loadClusterState(cluster);
@@ -636,6 +635,7 @@ protected List<RegionPlan> balanceTable(TableName tableName,
       currentCost / sumMultiplier, functionCost(), computedMaxSteps);
 
     final String initFunctionTotalCosts = totalCostsPerFunc();
+    cluster.setStopRequestedAt(EnvironmentEdgeManager.currentTime() + maxRunningTime);
     // Perform a stochastic walk to see if we can get a good fit.
     long step;
     boolean planImprovedConditionals = false;

hbase-balancer/src/test/java/org/apache/hadoop/hbase/master/balancer/TestStochasticLoadBalancer.java

@@ -27,6 +27,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -44,6 +45,7 @@
 import org.apache.hadoop.hbase.Size;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.master.RackManager;
 import org.apache.hadoop.hbase.master.RegionPlan;
 import org.apache.hadoop.hbase.testclassification.MasterTests;
 import org.apache.hadoop.hbase.testclassification.MediumTests;
@@ -496,6 +498,42 @@ public void testCostAfterUndoAction() {
     }
   }
 
+  @Test
+  public void testSearchDeadlineIsSetAfterInitialization() throws Exception {
+    Configuration testConf = new Configuration(conf);
+    testConf.setInt(StochasticLoadBalancer.MAX_STEPS_KEY, 0);
+    boolean[] needsBalanceCalled = { false };
+    boolean[] deadlineSet = { false };
+    StochasticLoadBalancer balancer =
+      new StochasticLoadBalancer(new DummyMetricsStochasticBalancer()) {
+        @Override
+        protected BalancerClusterState createState(Map<ServerName, List<RegionInfo>> clusterState,
+          Map<String, Deque<BalancerRegionLoad>> loads, RegionHDFSBlockLocationFinder finder,
+          RackManager rackManager) {
+          return new BalancerClusterState(clusterState, loads, finder, rackManager) {
+            @Override
+            void setStopRequestedAt(long stopRequestedAt) {
+              deadlineSet[0] = true;
+              assertTrue(needsBalanceCalled[0]);
+              super.setStopRequestedAt(stopRequestedAt);
+            }
+          };
+        }
+
+        @Override
+        boolean needsBalance(TableName tableName, BalancerClusterState cluster) {
+          needsBalanceCalled[0] = true;
+          return true;
+        }
+      };
+    balancer.setClusterInfoProvider(new DummyClusterInfoProvider(testConf));
+    balancer.initialize();
+
+    balancer.balanceTable(HConstants.ENSEMBLE_TABLE_NAME, createServerMap(2, 2, 1, 1, 1));
+
+    assertTrue(deadlineSet[0]);
+  }
+
   @Test
   public void testTableSkewCost() {
     Configuration conf = HBaseConfiguration.create();