diff --git a/src/main/java/org/apache/hadoop/hbase/master/AssignmentManager.java b/src/main/java/org/apache/hadoop/hbase/master/AssignmentManager.java index 8755317..31d8014 100644 --- a/src/main/java/org/apache/hadoop/hbase/master/AssignmentManager.java +++ b/src/main/java/org/apache/hadoop/hbase/master/AssignmentManager.java @@ -3431,6 +3431,7 @@ public class AssignmentManager extends ZooKeeperListener { } } Map onlineSvrs = this.serverManager.getOnlineServers(); + List drainingServers = this.serverManager.getDrainingServersList(); // Take care of servers w/o assignments. for (Map> map : result.values()) { for (Map.Entry svrEntry: onlineSvrs.entrySet()) { @@ -3438,6 +3439,7 @@ public class AssignmentManager extends ZooKeeperListener { map.put(svrEntry.getKey(), new ArrayList()); } } + map.keySet().removeAll(drainingServers); } return result; } diff --git a/src/test/java/org/apache/hadoop/hbase/master/TestAssignmentManager.java b/src/test/java/org/apache/hadoop/hbase/master/TestAssignmentManager.java index 2f9779f..0a9748a 100644 --- a/src/test/java/org/apache/hadoop/hbase/master/TestAssignmentManager.java +++ b/src/test/java/org/apache/hadoop/hbase/master/TestAssignmentManager.java @@ -24,12 +24,10 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import java.io.IOException; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.util.*; import java.util.concurrent.atomic.AtomicBoolean; +import com.google.common.collect.Lists; import org.apache.hadoop.hbase.HBaseConfiguration; import org.apache.hadoop.hbase.HBaseTestingUtility; import org.apache.hadoop.hbase.HConstants; @@ -92,6 +90,8 @@ public class TestAssignmentManager { new ServerName("example.org", 1234, 5678); private static final ServerName SERVERNAME_B = new ServerName("example.org", 0, 5678); + private static final ServerName SERVERNAME_C = + new ServerName("example.org", 123, 5678); private static final HRegionInfo REGIONINFO = new HRegionInfo(Bytes.toBytes("t"), HConstants.EMPTY_START_ROW, HConstants.EMPTY_START_ROW); @@ -134,24 +134,23 @@ public class TestAssignmentManager { // Mock a ServerManager. Say server SERVERNAME_{A,B} are online. Also // make it so if close or open, we return 'success'. - this.serverManager = Mockito.mock(ServerManager.class); - Mockito.when(this.serverManager.isServerOnline(SERVERNAME_A)).thenReturn(true); - Mockito.when(this.serverManager.isServerOnline(SERVERNAME_B)).thenReturn(true); + this.serverManager = mockManager(SERVERNAME_A, SERVERNAME_B); + } + + private ServerManager mockManager(ServerName... servers) throws IOException { + ServerManager serverManager = Mockito.mock(ServerManager.class); final Map onlineServers = new HashMap(); - onlineServers.put(SERVERNAME_B, new HServerLoad()); - onlineServers.put(SERVERNAME_A, new HServerLoad()); - Mockito.when(this.serverManager.getOnlineServersList()).thenReturn( + for (ServerName server : servers) { + Mockito.when(serverManager.isServerOnline(server)).thenReturn(true); + onlineServers.put(server, new HServerLoad()); + Mockito.when(serverManager.sendRegionClose(server, REGIONINFO, -1)).thenReturn(true); + Mockito.when(serverManager.sendRegionOpen(server, REGIONINFO, -1)). + thenReturn(RegionOpeningState.OPENED); + } + Mockito.when(serverManager.getOnlineServersList()).thenReturn( new ArrayList(onlineServers.keySet())); - Mockito.when(this.serverManager.getOnlineServers()).thenReturn(onlineServers); - Mockito.when(this.serverManager.sendRegionClose(SERVERNAME_A, REGIONINFO, -1)). - thenReturn(true); - Mockito.when(this.serverManager.sendRegionClose(SERVERNAME_B, REGIONINFO, -1)). - thenReturn(true); - // Ditto on open. - Mockito.when(this.serverManager.sendRegionOpen(SERVERNAME_A, REGIONINFO, -1)). - thenReturn(RegionOpeningState.OPENED); - Mockito.when(this.serverManager.sendRegionOpen(SERVERNAME_B, REGIONINFO, -1)). - thenReturn(RegionOpeningState.OPENED); + Mockito.when(serverManager.getOnlineServers()).thenReturn(onlineServers); + return serverManager; } @After @@ -397,6 +396,35 @@ public class TestAssignmentManager { } } + @Test + public void testGettingAssignmentsExcludesDrainingServers() throws Exception { + List availableServers = + Lists.newArrayList(SERVERNAME_A, SERVERNAME_B, SERVERNAME_C); + ServerManager serverManager = mockManager(availableServers.toArray(new ServerName[0])); + + + ExecutorService executor = startupMasterExecutor("testAssignmentsWithRSInDraining"); + CatalogTracker ct = Mockito.mock(CatalogTracker.class); + + LoadBalancer balancer = LoadBalancerFactory.getLoadBalancer(server.getConfiguration()); + + Mockito.when(serverManager.getDrainingServersList()).thenReturn( + Lists.newArrayList(SERVERNAME_C)); + AssignmentManager am = new AssignmentManager(this.server, serverManager, ct, balancer, executor); + + for (ServerName availableServer : availableServers) { + HRegionInfo info = Mockito.mock(HRegionInfo.class); + Mockito.when(info.getEncodedName()).thenReturn(UUID.randomUUID().toString()); + am.regionOnline(info, availableServer); + } + + Map>> result = am.getAssignmentsByTable(); + for (Map> map : result.values()) { + System.out.println(map.keySet()); + assertFalse(map.containsKey(SERVERNAME_C)); + } + } + /** * Run a simple server shutdown handler. * @throws KeeperException