From 5d449d563feaf79116f83c0b38ddc1932b568341 Mon Sep 17 00:00:00 2001
From: Tang Tianhang <tianhang.tang@Tangs-MacBook-Pro.local>
Date: Fri, 23 Sep 2022 17:07:40 +0800
Subject: [PATCH] HBASE-27388 make balance byTable config take effect when use
 CLI

---
 .../hbase/rsgroup/RSGroupInfoManagerImpl.java | 60 ++++++++++++-------
 1 file changed, 40 insertions(+), 20 deletions(-)

diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupInfoManagerImpl.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupInfoManagerImpl.java
index 5332c6e84506..20e67203e395 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupInfoManagerImpl.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupInfoManagerImpl.java
@@ -1143,29 +1143,49 @@ Map<TableName, Map<ServerName, List<RegionInfo>>> getRSGroupAssignmentsByTable(
     TableStateManager tableStateManager, String groupName) throws IOException {
     Map<TableName, Map<ServerName, List<RegionInfo>>> result = Maps.newHashMap();
     Set<TableName> tablesInGroupCache = new HashSet<>();
-    for (Map.Entry<RegionInfo, ServerName> entry : masterServices.getAssignmentManager()
-      .getRegionStates().getRegionAssignments().entrySet()) {
-      RegionInfo region = entry.getKey();
-      TableName tn = region.getTable();
-      ServerName server = entry.getValue();
-      if (isTableInGroup(tn, groupName, tablesInGroupCache)) {
-        if (
-          tableStateManager.isTableState(tn, TableState.State.DISABLED, TableState.State.DISABLING)
-        ) {
-          continue;
+    RSGroupInfo rsGroupInfo = getRSGroupInfo(groupName);
+    if (!masterServices.getConfiguration().getBoolean(
+      HConstants.HBASE_MASTER_LOADBALANCE_BYTABLE, false)) {
+      Map<ServerName, List<RegionInfo>> serverMap =
+        new HashMap<>();
+      for (ServerName serverName : masterServices.getServerManager().getOnlineServers().keySet()) {
+        if (rsGroupInfo.containsServer(serverName.getAddress())) {
+          serverMap.put(serverName, new ArrayList<>());
         }
-        if (region.isSplitParent()) {
-          continue;
+      }
+      for (Map.Entry<RegionInfo, ServerName> entry :
+        masterServices.getAssignmentManager().getRegionStates().getRegionAssignments().entrySet()) {
+        TableName currTable = entry.getKey().getTable();
+        ServerName currServer = entry.getValue();
+        RegionInfo currRegion = entry.getKey();
+        if (isTableInGroup(currTable, groupName, tablesInGroupCache)) {
+          if (!serverMap.containsKey(currServer)) {
+            serverMap.put(currServer, new ArrayList<>());
+          }
+          serverMap.get(currServer).add(currRegion);
         }
-        result.computeIfAbsent(tn, k -> new HashMap<>())
-          .computeIfAbsent(server, k -> new ArrayList<>()).add(region);
       }
-    }
-    RSGroupInfo rsGroupInfo = getRSGroupInfo(groupName);
-    for (ServerName serverName : masterServices.getServerManager().getOnlineServers().keySet()) {
-      if (rsGroupInfo.containsServer(serverName.getAddress())) {
-        for (Map<ServerName, List<RegionInfo>> map : result.values()) {
-          map.computeIfAbsent(serverName, k -> Collections.emptyList());
+      result.put(HConstants.ENSEMBLE_TABLE_NAME, serverMap);
+    } else {
+      for (Map.Entry<RegionInfo, ServerName> entry : masterServices.getAssignmentManager().getRegionStates().getRegionAssignments().entrySet()) {
+        RegionInfo region = entry.getKey();
+        TableName tn = region.getTable();
+        ServerName server = entry.getValue();
+        if (isTableInGroup(tn, groupName, tablesInGroupCache)) {
+          if (tableStateManager.isTableState(tn, TableState.State.DISABLED, TableState.State.DISABLING)) {
+            continue;
+          }
+          if (region.isSplitParent()) {
+            continue;
+          }
+          result.computeIfAbsent(tn, k -> new HashMap<>()).computeIfAbsent(server, k -> new ArrayList<>()).add(region);
+        }
+      }
+      for (ServerName serverName : masterServices.getServerManager().getOnlineServers().keySet()) {
+        if (rsGroupInfo.containsServer(serverName.getAddress())) {
+          for (Map<ServerName, List<RegionInfo>> map : result.values()) {
+            map.computeIfAbsent(serverName, k -> Collections.emptyList());
+          }
         }
       }
     }