From 7c24ed4f4573d15f194a5ced7130fae4dcfe20c2 Mon Sep 17 00:00:00 2001
From: Duo Zhang <zhangduo@apache.org>
Date: Thu, 20 May 2021 20:58:53 +0800
Subject: [PATCH] HBASE-25897 TestRetainAssignmentOnRestart is flaky after
 HBASE-25032 (#3281)

Signed-off-by: Xiaolin Ha <haxiaolin@apache.org>
---
 .../apache/hadoop/hbase/master/HMaster.java   |  3 +-
 .../master/TestRetainAssignmentOnRestart.java | 34 +++++++++++++++++--
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/HMaster.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/HMaster.java
index 0702e3d839e..c8239c38640 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/HMaster.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/HMaster.java
@@ -1498,7 +1498,8 @@ public class HMaster extends HRegionServer implements MasterServices {
     }
   }
 
-  private void startProcedureExecutor() throws IOException {
+  // will be override in UT
+  protected void startProcedureExecutor() throws IOException {
     procedureExecutor.startWorkers();
   }
 
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/TestRetainAssignmentOnRestart.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/TestRetainAssignmentOnRestart.java
index ee4da743d21..4a5491ed8f8 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/TestRetainAssignmentOnRestart.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/TestRetainAssignmentOnRestart.java
@@ -22,13 +22,15 @@ import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.util.List;
 import java.util.Map;
-
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.HBaseClassTestRule;
 import org.apache.hadoop.hbase.HConstants;
 import org.apache.hadoop.hbase.MiniHBaseCluster;
 import org.apache.hadoop.hbase.ServerName;
+import org.apache.hadoop.hbase.StartMiniClusterOption;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.RegionInfo;
 import org.apache.hadoop.hbase.master.procedure.ServerCrashProcedure;
@@ -52,6 +54,33 @@ public class TestRetainAssignmentOnRestart extends AbstractTestRestartCluster {
 
   private static int NUM_OF_RS = 3;
 
+  public static final class HMasterForTest extends HMaster {
+
+    public HMasterForTest(Configuration conf) throws IOException {
+      super(conf);
+    }
+
+    @Override
+    protected void startProcedureExecutor() throws IOException {
+      // only start procedure executor when we have all the regionservers ready to take regions
+      new Thread(() -> {
+        for (;;) {
+          if (getServerManager().createDestinationServersList().size() == NUM_OF_RS) {
+            try {
+              HMasterForTest.super.startProcedureExecutor();
+            } catch (IOException e) {
+              throw new UncheckedIOException(e);
+            }
+          }
+          try {
+            Thread.sleep(1000);
+          } catch (InterruptedException e) {
+          }
+        }
+      }).start();
+    }
+  }
+
   @Override
   protected boolean splitWALCoordinatedByZk() {
     return true;
@@ -205,7 +234,8 @@ public class TestRetainAssignmentOnRestart extends AbstractTestRestartCluster {
       HConstants.ZK_CONNECTION_REGISTRY_CLASS);
     // Enable retain assignment during ServerCrashProcedure
     UTIL.getConfiguration().setBoolean(ServerCrashProcedure.MASTER_SCP_RETAIN_ASSIGNMENT, true);
-    UTIL.startMiniCluster(NUM_OF_RS);
+    UTIL.startMiniCluster(StartMiniClusterOption.builder().masterClass(HMasterForTest.class)
+      .numRegionServers(NUM_OF_RS).build());
 
     // Turn off balancer
     UTIL.getMiniHBaseCluster().getMaster().getMasterRpcServices().synchronousBalanceSwitch(false);