HBASE-25093 the RSGroupBasedLoadBalancer#retainAssignment throws NPE (#2534)

Signed-off-by: Duo Zhang <zhangduo@apache.org>

hbase-rsgroup/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupBasedLoadBalancer.java

@@ -18,6 +18,7 @@
 
 package org.apache.hadoop.hbase.rsgroup;
 
+import edu.umd.cs.findbugs.annotations.NonNull;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -171,58 +172,37 @@ public class RSGroupBasedLoadBalancer implements RSGroupableBalancer {
   }
 
   @Override
+  @NonNull
   public Map<ServerName, List<RegionInfo>> roundRobinAssignment(List<RegionInfo> regions,
-    List<ServerName> servers) throws HBaseIOException {
+      List<ServerName> servers) throws HBaseIOException {
     Map<ServerName, List<RegionInfo>> assignments = Maps.newHashMap();
     List<Pair<List<RegionInfo>, List<ServerName>>> pairs =
       generateGroupAssignments(regions, servers);
     for (Pair<List<RegionInfo>, List<ServerName>> pair : pairs) {
-      Map<ServerName, List<RegionInfo>> result = this.internalBalancer
+      Map<ServerName, List<RegionInfo>> result =
-        .roundRobinAssignment(pair.getFirst(), pair.getSecond());
+          this.internalBalancer.roundRobinAssignment(pair.getFirst(), pair.getSecond());
-      if (result != null) {
+      result.forEach((server, regionInfos) -> assignments
-        result.forEach((server, regionInfos) ->
+          .computeIfAbsent(server, s -> Lists.newArrayList()).addAll(regionInfos));
-          assignments.computeIfAbsent(server, s -> Lists.newArrayList()).addAll(regionInfos));
-      }
     }
     return assignments;
   }
 
   @Override
+  @NonNull
   public Map<ServerName, List<RegionInfo>> retainAssignment(Map<RegionInfo, ServerName> regions,
-    List<ServerName> servers) throws HBaseIOException {
+      List<ServerName> servers) throws HBaseIOException {
     try {
       Map<ServerName, List<RegionInfo>> assignments = new TreeMap<>();
-      ListMultimap<String, RegionInfo> groupToRegion = ArrayListMultimap.create();
+      List<Pair<List<RegionInfo>, List<ServerName>>> pairs =
-      for (RegionInfo region : regions.keySet()) {
+          generateGroupAssignments(Lists.newArrayList(regions.keySet()), servers);
-        String groupName = rsGroupInfoManager.getRSGroupOfTable(region.getTable());
+      for (Pair<List<RegionInfo>, List<ServerName>> pair : pairs) {
-        if (groupName == null) {
+        List<RegionInfo> regionList = pair.getFirst();
-          LOG.debug("Group not found for table " + region.getTable() + ", using default");
+        Map<RegionInfo, ServerName> currentAssignmentMap = Maps.newTreeMap();
-          groupName = RSGroupInfo.DEFAULT_GROUP;
+        regionList.forEach(r -> currentAssignmentMap.put(r, regions.get(r)));
-        }
+        Map<ServerName, List<RegionInfo>> pairResult =
-        groupToRegion.put(groupName, region);
+            this.internalBalancer.retainAssignment(currentAssignmentMap, pair.getSecond());
-      }
+        pairResult.forEach((server, rs) -> assignments
-      for (String group : groupToRegion.keySet()) {
+            .computeIfAbsent(server, s -> Lists.newArrayList()).addAll(rs));
-        Map<RegionInfo, ServerName> currentAssignmentMap = new TreeMap<RegionInfo, ServerName>();
-        List<RegionInfo> regionList = groupToRegion.get(group);
-        RSGroupInfo info = rsGroupInfoManager.getRSGroup(group);
-        List<ServerName> candidateList = filterOfflineServers(info, servers);
-        if (fallbackEnabled && candidateList.isEmpty()) {
-          candidateList = getFallBackCandidates(servers);
-        }
-        for (RegionInfo region : regionList) {
-          currentAssignmentMap.put(region, regions.get(region));
-        }
-        if (candidateList.size() > 0) {
-          assignments
-            .putAll(this.internalBalancer.retainAssignment(currentAssignmentMap, candidateList));
-        } else {
-          if (LOG.isDebugEnabled()) {
-            LOG.debug("No available servers for group {} to assign regions: {}", group,
-              RegionInfo.getShortNameToLog(regionList));
-          }
-          assignments.computeIfAbsent(LoadBalancer.BOGUS_SERVER_NAME, s -> new ArrayList<>())
-            .addAll(regionList);
-        }
       }
       return assignments;
     } catch (IOException e) {

hbase-server/src/main/java/org/apache/hadoop/hbase/favored/FavoredNodeLoadBalancer.java

@@ -22,12 +22,14 @@ import static org.apache.hadoop.hbase.favored.FavoredNodesPlan.Position.PRIMARY;
 import static org.apache.hadoop.hbase.favored.FavoredNodesPlan.Position.SECONDARY;
 import static org.apache.hadoop.hbase.favored.FavoredNodesPlan.Position.TERTIARY;
 
+import edu.umd.cs.findbugs.annotations.NonNull;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.HBaseIOException;
 import org.apache.hadoop.hbase.HBaseInterfaceAudience;
@@ -161,6 +163,7 @@ public class FavoredNodeLoadBalancer extends BaseLoadBalancer implements Favored
   }
 
   @Override
+  @NonNull
   public Map<ServerName, List<RegionInfo>> roundRobinAssignment(List<RegionInfo> regions,
       List<ServerName> servers) throws HBaseIOException {
     Map<ServerName, List<RegionInfo>> assignmentMap;

hbase-server/src/main/java/org/apache/hadoop/hbase/master/LoadBalancer.java

@@ -18,7 +18,7 @@
  */
 package org.apache.hadoop.hbase.master;
 
-import edu.umd.cs.findbugs.annotations.Nullable;
+import edu.umd.cs.findbugs.annotations.NonNull;
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
@@ -103,10 +103,9 @@ public interface LoadBalancer extends Configurable, Stoppable, ConfigurationObse
    * @param servers
    * @return Map of servername to regioninfos
    */
-  Map<ServerName, List<RegionInfo>> roundRobinAssignment(
+  @NonNull
-    List<RegionInfo> regions,
+  Map<ServerName, List<RegionInfo>> roundRobinAssignment(List<RegionInfo> regions,
-    List<ServerName> servers
+      List<ServerName> servers) throws HBaseIOException;
-  ) throws HBaseIOException;
 
   /**
    * Assign regions to the previously hosting region server
@@ -114,11 +113,9 @@ public interface LoadBalancer extends Configurable, Stoppable, ConfigurationObse
    * @param servers
    * @return List of plans
    */
-  @Nullable
+  @NonNull
-  Map<ServerName, List<RegionInfo>> retainAssignment(
+  Map<ServerName, List<RegionInfo>> retainAssignment(Map<RegionInfo, ServerName> regions,
-    Map<RegionInfo, ServerName> regions,
+      List<ServerName> servers) throws HBaseIOException;
-    List<ServerName> servers
-  ) throws HBaseIOException;
 
   /**
    * Get a random region server from the list

hbase-server/src/main/java/org/apache/hadoop/hbase/master/assignment/AssignmentManager.java

@@ -2159,12 +2159,8 @@ public class AssignmentManager {
     final ProcedureEvent<?>[] events = new ProcedureEvent[regions.size()];
     final long st = System.currentTimeMillis();
 
-    if (plan == null) {
-      throw new HBaseIOException("unable to compute plans for regions=" + regions.size());
-    }
-
     if (plan.isEmpty()) {
-      return;
+      throw new HBaseIOException("unable to compute plans for regions=" + regions.size());
     }
 
     int evcount = 0;

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

@@ -18,6 +18,7 @@
  */
 package org.apache.hadoop.hbase.master.balancer;
 
+import edu.umd.cs.findbugs.annotations.NonNull;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -1131,11 +1132,9 @@ public abstract class BaseLoadBalancer implements LoadBalancer {
    * If master is configured to carry system tables only, in here is
    * where we figure what to assign it.
    */
+  @NonNull
   protected Map<ServerName, List<RegionInfo>> assignMasterSystemRegions(
       Collection<RegionInfo> regions, List<ServerName> servers) {
-    if (servers == null || regions == null || regions.isEmpty()) {
-      return null;
-    }
     Map<ServerName, List<RegionInfo>> assignments = new TreeMap<>();
     if (this.maintenanceMode || this.onlySystemTablesOnMaster) {
       if (masterServerName != null && servers.contains(masterServerName)) {
@@ -1266,15 +1265,16 @@ public abstract class BaseLoadBalancer implements LoadBalancer {
    *
    * @param regions all regions
    * @param servers all servers
-   * @return map of server to the regions it should take, or null if no
+   * @return map of server to the regions it should take, or emptyMap if no
-   *         assignment is possible (ie. no regions or no servers)
+   *         assignment is possible (ie. no servers)
    */
   @Override
+  @NonNull
   public Map<ServerName, List<RegionInfo>> roundRobinAssignment(List<RegionInfo> regions,
       List<ServerName> servers) throws HBaseIOException {
     metricsBalancer.incrMiscInvocations();
     Map<ServerName, List<RegionInfo>> assignments = assignMasterSystemRegions(regions, servers);
-    if (assignments != null && !assignments.isEmpty()) {
+    if (!assignments.isEmpty()) {
       servers = new ArrayList<>(servers);
       // Guarantee not to put other regions on master
       servers.remove(masterServerName);
@@ -1284,14 +1284,17 @@ public abstract class BaseLoadBalancer implements LoadBalancer {
         regions.removeAll(masterRegions);
       }
     }
-    if (this.maintenanceMode || regions == null || regions.isEmpty()) {
+    /**
+     * only need assign system table
+     */
+    if (this.maintenanceMode || regions.isEmpty()) {
       return assignments;
     }
 
     int numServers = servers == null ? 0 : servers.size();
     if (numServers == 0) {
       LOG.warn("Wanted to do round robin assignment but no servers to assign to");
-      return null;
+      return Collections.emptyMap();
     }
 
     // TODO: instead of retainAssignment() and roundRobinAssignment(), we should just run the
@@ -1406,15 +1409,17 @@ public abstract class BaseLoadBalancer implements LoadBalancer {
    *
    * @param regions regions and existing assignment from meta
    * @param servers available servers
-   * @return map of servers and regions to be assigned to them
+   * @return map of servers and regions to be assigned to them, or emptyMap if no
+   *           assignment is possible (ie. no servers)
    */
   @Override
+  @NonNull
   public Map<ServerName, List<RegionInfo>> retainAssignment(Map<RegionInfo, ServerName> regions,
       List<ServerName> servers) throws HBaseIOException {
     // Update metrics
     metricsBalancer.incrMiscInvocations();
     Map<ServerName, List<RegionInfo>> assignments = assignMasterSystemRegions(regions.keySet(), servers);
-    if (assignments != null && !assignments.isEmpty()) {
+    if (!assignments.isEmpty()) {
       servers = new ArrayList<>(servers);
       // Guarantee not to put other regions on master
       servers.remove(masterServerName);
@@ -1429,7 +1434,7 @@ public abstract class BaseLoadBalancer implements LoadBalancer {
     int numServers = servers == null ? 0 : servers.size();
     if (numServers == 0) {
       LOG.warn("Wanted to do retain assignment but no servers to assign to");
-      return null;
+      return Collections.emptyMap();
     }
     if (numServers == 1) { // Only one server, nothing fancy we can do here
       ServerName server = servers.get(0);

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

@@ -23,6 +23,7 @@ import static org.apache.hadoop.hbase.favored.FavoredNodesPlan.Position.PRIMARY;
 import static org.apache.hadoop.hbase.favored.FavoredNodesPlan.Position.SECONDARY;
 import static org.apache.hadoop.hbase.favored.FavoredNodesPlan.Position.TERTIARY;
 
+import edu.umd.cs.findbugs.annotations.NonNull;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -31,6 +32,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+
 import org.apache.hadoop.hbase.HBaseIOException;
 import org.apache.hadoop.hbase.ServerMetrics;
 import org.apache.hadoop.hbase.ServerName;
@@ -109,6 +111,7 @@ public class FavoredStochasticBalancer extends StochasticLoadBalancer implements
    * secondary and tertiary as per favored nodes constraints.
    */
   @Override
+  @NonNull
   public Map<ServerName, List<RegionInfo>> roundRobinAssignment(List<RegionInfo> regions,
       List<ServerName> servers) throws HBaseIOException {
 
@@ -116,7 +119,7 @@ public class FavoredStochasticBalancer extends StochasticLoadBalancer implements
 
     Set<RegionInfo> regionSet = Sets.newHashSet(regions);
     Map<ServerName, List<RegionInfo>> assignmentMap = assignMasterSystemRegions(regions, servers);
-    if (assignmentMap != null && !assignmentMap.isEmpty()) {
+    if (!assignmentMap.isEmpty()) {
       servers = new ArrayList<>(servers);
       // Guarantee not to put other regions on master
       servers.remove(masterServerName);
@@ -367,14 +370,15 @@ public class FavoredStochasticBalancer extends StochasticLoadBalancer implements
    * Reuse BaseLoadBalancer's retainAssignment, but generate favored nodes when its missing.
    */
   @Override
+  @NonNull
   public Map<ServerName, List<RegionInfo>> retainAssignment(Map<RegionInfo, ServerName> regions,
       List<ServerName> servers) throws HBaseIOException {
 
     Map<ServerName, List<RegionInfo>> assignmentMap = Maps.newHashMap();
     Map<ServerName, List<RegionInfo>> result = super.retainAssignment(regions, servers);
-    if (result == null || result.isEmpty()) {
+    if (result.isEmpty()) {
       LOG.warn("Nothing to assign to, probably no servers or no regions");
-      return null;
+      return result;
     }
 
     // Guarantee not to put other regions on master

hbase-server/src/test/java/org/apache/hadoop/hbase/TestZooKeeper.java

@@ -21,8 +21,10 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import edu.umd.cs.findbugs.annotations.NonNull;
 import java.util.List;
 import java.util.Map;
+
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.client.Admin;
 import org.apache.hadoop.hbase.client.ColumnFamilyDescriptorBuilder;
@@ -282,6 +284,7 @@ public class TestZooKeeper {
     static boolean retainAssignCalled = false;
 
     @Override
+    @NonNull
     public Map<ServerName, List<RegionInfo>> retainAssignment(
         Map<RegionInfo, ServerName> regions, List<ServerName> servers) throws HBaseIOException {
       retainAssignCalled = true;