From d6bdd6de87962bb5ba77431a39c5a18c7c4a8bbc Mon Sep 17 00:00:00 2001
From: Huginn <63332600+Huginn-kio@users.noreply.github.com>
Date: Wed, 16 Sep 2026 21:23:00 +0800
Subject: [PATCH] HBASE-27878 Fix NPE in balancing rsgroup

---
 .../hbase/rsgroup/RSGroupInfoManagerImpl.java | 38 ++++++++-----------
 .../hbase/rsgroup/TestRSGroupsBalance.java    | 32 ++++++++++++++--
 2 files changed, 43 insertions(+), 27 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 783d7c7dd335..1bb4ab2c68b1 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
@@ -1142,31 +1142,23 @@ private Map<String, RegionState> rsGroupGetRegionsInTransition(String groupName)
   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;
-        }
-        if (region.isSplitParent()) {
-          continue;
-        }
-        result.computeIfAbsent(tn, k -> new HashMap<>())
-          .computeIfAbsent(server, k -> new ArrayList<>()).add(region);
-      }
-    }
     RSGroupInfo rsGroupInfo = getRSGroupInfo(groupName);
-    for (ServerName serverName : masterServices.getServerManager().getOnlineServers().keySet()) {
+    List<ServerName> onlineServersInGroup = new ArrayList<>();
+    for (ServerName serverName : masterServices.getServerManager().getOnlineServersList()) {
       if (rsGroupInfo.containsServer(serverName.getAddress())) {
-        for (Map<ServerName, List<RegionInfo>> map : result.values()) {
-          map.computeIfAbsent(serverName, k -> Collections.emptyList());
-        }
+        onlineServersInGroup.add(serverName);
+      }
+    }
+    Map<TableName, Map<ServerName, List<RegionInfo>>> assignments = masterServices
+      .getAssignmentManager().getRegionStates()
+      .getAssignmentsForBalancer(tableStateManager, onlineServersInGroup);
+
+    Set<TableName> tablesInGroupCache = new HashSet<>();
+    for (Map.Entry<TableName, Map<ServerName, List<RegionInfo>>> entry : assignments.entrySet()) {
+      TableName tableName = entry.getKey();
+      if (isTableInGroup(tableName, groupName, tablesInGroupCache)) {
+        result.put(tableName, entry.getValue());
+        LOG.debug("Adding assignments for {}: {}", tableName, entry.getValue());
       }
     }
     return result;
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/rsgroup/TestRSGroupsBalance.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/rsgroup/TestRSGroupsBalance.java
index e596145218a9..fc6b28b5b4d4 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/rsgroup/TestRSGroupsBalance.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/rsgroup/TestRSGroupsBalance.java
@@ -36,6 +36,7 @@
 import org.apache.hadoop.hbase.client.TableDescriptor;
 import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
 import org.apache.hadoop.hbase.master.HMaster;
+import org.apache.hadoop.hbase.master.assignment.RegionStateNode;
 import org.apache.hadoop.hbase.testclassification.MediumTests;
 import org.apache.hadoop.hbase.testclassification.RSGroupTests;
 import org.apache.hadoop.hbase.util.Bytes;
@@ -222,9 +223,32 @@ public void testGetRSGroupAssignmentsByTable() throws Exception {
 
     HMaster master = TEST_UTIL.getMiniHBaseCluster().getMaster();
     RSGroupInfoManagerImpl gm = (RSGroupInfoManagerImpl) master.getRSGroupInfoManager();
-    Map<TableName, Map<ServerName, List<RegionInfo>>> assignments =
-      gm.getRSGroupAssignmentsByTable(master.getTableStateManager(), RSGroupInfo.DEFAULT_GROUP);
-    assertFalse(assignments.containsKey(disableTableName));
-    assertTrue(assignments.containsKey(tableName));
+    RegionInfo regionWithNullServer = ADMIN.getRegions(tableName).get(0);
+    RegionStateNode regionNode = master.getAssignmentManager()
+      .getRegionStates()
+      .getRegionStateNode(regionWithNullServer);
+    ServerName originalServer = setRegionLocation(regionNode, null);
+    try {
+      Map<TableName, Map<ServerName, List<RegionInfo>>> assignments =
+        gm.getRSGroupAssignmentsByTable(master.getTableStateManager(), RSGroupInfo.DEFAULT_GROUP);
+      assertFalse(assignments.containsKey(disableTableName));
+      assertTrue(assignments.containsKey(tableName));
+      assertFalse(assignments.get(tableName).containsKey(null));
+      assertFalse(assignments.get(tableName)
+        .values()
+        .stream()
+        .anyMatch(regions -> regions.contains(regionWithNullServer)));
+    } finally {
+      setRegionLocation(regionNode, originalServer);
+    }
+  }
+
+  private static ServerName setRegionLocation(RegionStateNode regionNode, ServerName serverName) {
+    regionNode.lock();
+    try {
+      return regionNode.setRegionLocation(serverName);
+    } finally {
+      regionNode.unlock();
+    }
   }
 }