001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase.master.balancer;
019
020import static org.junit.Assert.assertEquals;
021import static org.junit.Assert.assertFalse;
022import static org.junit.Assert.assertNull;
023import static org.junit.Assert.assertTrue;
024import static org.junit.Assert.fail;
025import static org.mockito.Mockito.mock;
026import static org.mockito.Mockito.when;
027
028import java.util.ArrayList;
029import java.util.HashMap;
030import java.util.HashSet;
031import java.util.List;
032import java.util.Map;
033import java.util.Random;
034import java.util.Set;
035import java.util.TreeMap;
036import java.util.concurrent.ThreadLocalRandom;
037import org.apache.hadoop.conf.Configuration;
038import org.apache.hadoop.hbase.ClusterMetrics;
039import org.apache.hadoop.hbase.HBaseClassTestRule;
040import org.apache.hadoop.hbase.HBaseConfiguration;
041import org.apache.hadoop.hbase.HConstants;
042import org.apache.hadoop.hbase.RegionMetrics;
043import org.apache.hadoop.hbase.ServerMetrics;
044import org.apache.hadoop.hbase.ServerName;
045import org.apache.hadoop.hbase.Size;
046import org.apache.hadoop.hbase.TableName;
047import org.apache.hadoop.hbase.client.RegionInfo;
048import org.apache.hadoop.hbase.client.TableDescriptor;
049import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
050import org.apache.hadoop.hbase.master.RegionPlan;
051import org.apache.hadoop.hbase.testclassification.LargeTests;
052import org.apache.hadoop.hbase.util.Bytes;
053import org.apache.hadoop.hbase.util.EnvironmentEdgeManager;
054import org.apache.hadoop.hbase.util.Pair;
055import org.junit.BeforeClass;
056import org.junit.ClassRule;
057import org.junit.Test;
058import org.junit.experimental.categories.Category;
059import org.slf4j.Logger;
060import org.slf4j.LoggerFactory;
061
062import org.apache.hbase.thirdparty.com.google.common.collect.Lists;
063
064@Category({ LargeTests.class })
065public class TestCacheAwareLoadBalancer extends BalancerTestBase {
066  @ClassRule
067  public static final HBaseClassTestRule CLASS_RULE =
068    HBaseClassTestRule.forClass(TestCacheAwareLoadBalancer.class);
069
070  private static final Logger LOG = LoggerFactory.getLogger(TestCacheAwareLoadBalancer.class);
071
072  private static CacheAwareLoadBalancer loadBalancer;
073
074  static List<ServerName> servers;
075
076  static List<TableDescriptor> tableDescs;
077
078  static Map<TableName, String> tableMap = new HashMap<>();
079
080  static TableName[] tables = new TableName[] { TableName.valueOf("dt1"), TableName.valueOf("dt2"),
081    TableName.valueOf("dt3"), TableName.valueOf("dt4") };
082
083  private static List<ServerName> generateServers(int numServers) {
084    List<ServerName> servers = new ArrayList<>(numServers);
085    Random rand = ThreadLocalRandom.current();
086    for (int i = 0; i < numServers; i++) {
087      String host = "server" + rand.nextInt(100000);
088      int port = rand.nextInt(60000);
089      servers.add(ServerName.valueOf(host, port, -1));
090    }
091    return servers;
092  }
093
094  private static List<TableDescriptor> constructTableDesc(boolean hasBogusTable) {
095    List<TableDescriptor> tds = Lists.newArrayList();
096    for (int i = 0; i < tables.length; i++) {
097      TableDescriptor htd = TableDescriptorBuilder.newBuilder(tables[i]).build();
098      tds.add(htd);
099    }
100    return tds;
101  }
102
103  private ServerMetrics mockServerMetricsWithRegionCacheInfo(ServerName server,
104    List<RegionInfo> regionsOnServer, float currentCacheRatio, List<RegionInfo> oldRegionCacheInfo,
105    int oldRegionCachedSize, int regionSize) {
106    ServerMetrics serverMetrics = mock(ServerMetrics.class);
107    Map<byte[], RegionMetrics> regionLoadMap = new TreeMap<>(Bytes.BYTES_COMPARATOR);
108    for (RegionInfo info : regionsOnServer) {
109      RegionMetrics rl = mock(RegionMetrics.class);
110      when(rl.getReadRequestCount()).thenReturn(0L);
111      when(rl.getWriteRequestCount()).thenReturn(0L);
112      when(rl.getMemStoreSize()).thenReturn(Size.ZERO);
113      when(rl.getStoreFileSize()).thenReturn(Size.ZERO);
114      when(rl.getCurrentRegionCachedRatio()).thenReturn(currentCacheRatio);
115      when(rl.getRegionSizeMB()).thenReturn(new Size(regionSize, Size.Unit.MEGABYTE));
116      regionLoadMap.put(info.getRegionName(), rl);
117    }
118    when(serverMetrics.getRegionMetrics()).thenReturn(regionLoadMap);
119    Map<String, Integer> oldCacheRatioMap = new HashMap<>();
120    for (RegionInfo info : oldRegionCacheInfo) {
121      oldCacheRatioMap.put(info.getEncodedName(), oldRegionCachedSize);
122    }
123    when(serverMetrics.getRegionCachedInfo()).thenReturn(oldCacheRatioMap);
124    return serverMetrics;
125  }
126
127  @BeforeClass
128  public static void beforeAllTests() throws Exception {
129    servers = generateServers(3);
130    tableDescs = constructTableDesc(false);
131    Configuration conf = HBaseConfiguration.create();
132    conf.set(HConstants.BUCKET_CACHE_PERSISTENT_PATH_KEY, "prefetch_file_list");
133    loadBalancer = new CacheAwareLoadBalancer();
134    loadBalancer.setClusterInfoProvider(new DummyClusterInfoProvider(conf));
135    loadBalancer.loadConf(conf);
136  }
137
138  @Test
139  public void testRegionsNotCachedOnOldServerAndCurrentServer() throws Exception {
140    // The regions are not cached on old server as well as the current server. This causes
141    // skewness in the region allocation which should be fixed by the balancer
142
143    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
144    ServerName server0 = servers.get(0);
145    ServerName server1 = servers.get(1);
146    ServerName server2 = servers.get(2);
147
148    // Simulate that the regions previously hosted by server1 are now hosted on server0
149    List<RegionInfo> regionsOnServer0 = randomRegions(10);
150    List<RegionInfo> regionsOnServer1 = randomRegions(0);
151    List<RegionInfo> regionsOnServer2 = randomRegions(5);
152
153    clusterState.put(server0, regionsOnServer0);
154    clusterState.put(server1, regionsOnServer1);
155    clusterState.put(server2, regionsOnServer2);
156
157    // Mock cluster metrics
158    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
159    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
160      0.0f, new ArrayList<>(), 0, 10));
161    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
162      0.0f, new ArrayList<>(), 0, 10));
163    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
164      0.0f, new ArrayList<>(), 0, 10));
165    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
166    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
167    loadBalancer.updateClusterMetrics(clusterMetrics);
168
169    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
170      (Map) mockClusterServersWithTables(clusterState);
171    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
172    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
173    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
174    for (RegionPlan plan : plans) {
175      if (plan.getSource().equals(server0)) {
176        regionsMovedFromServer0.add(plan.getRegionInfo());
177        if (!targetServers.containsKey(plan.getDestination())) {
178          targetServers.put(plan.getDestination(), new ArrayList<>());
179        }
180        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
181      }
182    }
183    // should move 5 regions from server0 to server 1
184    assertEquals(5, regionsMovedFromServer0.size());
185    assertEquals(5, targetServers.get(server1).size());
186  }
187
188  @Test
189  public void testRegionsPartiallyCachedOnOldServerAndNotCachedOnCurrentServer() throws Exception {
190    // The regions are partially cached on old server but not cached on the current server
191
192    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
193    ServerName server0 = servers.get(0);
194    ServerName server1 = servers.get(1);
195    ServerName server2 = servers.get(2);
196
197    // Simulate that the regions previously hosted by server1 are now hosted on server0
198    List<RegionInfo> regionsOnServer0 = randomRegions(10);
199    List<RegionInfo> regionsOnServer1 = randomRegions(0);
200    List<RegionInfo> regionsOnServer2 = randomRegions(5);
201
202    clusterState.put(server0, regionsOnServer0);
203    clusterState.put(server1, regionsOnServer1);
204    clusterState.put(server2, regionsOnServer2);
205
206    // Mock cluster metrics
207
208    // Mock 5 regions from server0 were previously hosted on server1
209    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);
210
211    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
212    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
213      0.0f, new ArrayList<>(), 0, 10));
214    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
215      0.0f, oldCachedRegions, 6, 10));
216    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
217      0.0f, new ArrayList<>(), 0, 10));
218    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
219    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
220    loadBalancer.updateClusterMetrics(clusterMetrics);
221
222    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
223      (Map) mockClusterServersWithTables(clusterState);
224    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
225    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
226    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
227    for (RegionPlan plan : plans) {
228      if (plan.getSource().equals(server0)) {
229        regionsMovedFromServer0.add(plan.getRegionInfo());
230        if (!targetServers.containsKey(plan.getDestination())) {
231          targetServers.put(plan.getDestination(), new ArrayList<>());
232        }
233        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
234      }
235    }
236    // should move 5 regions from server0 to server1
237    assertEquals(5, regionsMovedFromServer0.size());
238    assertEquals(5, targetServers.get(server1).size());
239    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
240  }
241
242  @Test
243  public void testThrottlingRegionBeyondThreshold() throws Exception {
244    Configuration conf = HBaseConfiguration.create();
245    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
246    balancer.setClusterInfoProvider(new DummyClusterInfoProvider(conf));
247    balancer.loadConf(conf);
248    balancer.initialize();
249    ServerName server0 = servers.get(0);
250    ServerName server1 = servers.get(1);
251    Pair<ServerName, Float> regionRatio = new Pair<>();
252    regionRatio.setFirst(server0);
253    regionRatio.setSecond(1.0f);
254    balancer.regionCacheRatioOnOldServerMap.put("region1", regionRatio);
255    RegionInfo mockedInfo = mock(RegionInfo.class);
256    when(mockedInfo.getEncodedName()).thenReturn("region1");
257    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
258    long startTime = EnvironmentEdgeManager.currentTime();
259    balancer.throttle(plan);
260    long endTime = EnvironmentEdgeManager.currentTime();
261    assertTrue((endTime - startTime) < 10);
262  }
263
264  @Test
265  public void testThrottlingRegionBelowThreshold() throws Exception {
266    Configuration conf = HBaseConfiguration.create();
267    conf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, 100);
268    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
269    balancer.setClusterInfoProvider(new DummyClusterInfoProvider(conf));
270    balancer.loadConf(conf);
271    balancer.initialize();
272    ServerName server0 = servers.get(0);
273    ServerName server1 = servers.get(1);
274    Pair<ServerName, Float> regionRatio = new Pair<>();
275    regionRatio.setFirst(server0);
276    regionRatio.setSecond(0.1f);
277    balancer.regionCacheRatioOnOldServerMap.put("region1", regionRatio);
278    RegionInfo mockedInfo = mock(RegionInfo.class);
279    when(mockedInfo.getEncodedName()).thenReturn("region1");
280    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
281    long startTime = EnvironmentEdgeManager.currentTime();
282    balancer.throttle(plan);
283    long endTime = EnvironmentEdgeManager.currentTime();
284    assertTrue((endTime - startTime) >= 100);
285  }
286
287  @Test
288  public void testThrottlingCacheRatioUnknownOnTarget() throws Exception {
289    Configuration conf = HBaseConfiguration.create();
290    conf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, 100);
291    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
292    balancer.setClusterInfoProvider(new DummyClusterInfoProvider(conf));
293    balancer.loadConf(conf);
294    balancer.initialize();
295    ServerName server0 = servers.get(0);
296    ServerName server1 = servers.get(1);
297    ServerName server3 = servers.get(2);
298    // setting region cache ratio 100% on server 3, though this is not the target in the region plan
299    Pair<ServerName, Float> regionRatio = new Pair<>();
300    regionRatio.setFirst(server3);
301    regionRatio.setSecond(1.0f);
302    balancer.regionCacheRatioOnOldServerMap.put("region1", regionRatio);
303    RegionInfo mockedInfo = mock(RegionInfo.class);
304    when(mockedInfo.getEncodedName()).thenReturn("region1");
305    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
306    long startTime = EnvironmentEdgeManager.currentTime();
307    balancer.throttle(plan);
308    long endTime = EnvironmentEdgeManager.currentTime();
309    assertTrue((endTime - startTime) >= 100);
310  }
311
312  @Test
313  public void testThrottlingCacheRatioUnknownForRegion() throws Exception {
314    Configuration conf = HBaseConfiguration.create();
315    conf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, 100);
316    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
317    balancer.setClusterInfoProvider(new DummyClusterInfoProvider(conf));
318    balancer.loadConf(conf);
319    balancer.initialize();
320    ServerName server0 = servers.get(0);
321    ServerName server1 = servers.get(1);
322    ServerName server3 = servers.get(2);
323    // No cache ratio available for region1
324    RegionInfo mockedInfo = mock(RegionInfo.class);
325    when(mockedInfo.getEncodedName()).thenReturn("region1");
326    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
327    long startTime = EnvironmentEdgeManager.currentTime();
328    balancer.throttle(plan);
329    long endTime = EnvironmentEdgeManager.currentTime();
330    assertTrue((endTime - startTime) >= 100);
331  }
332
333  @Test
334  public void testRegionPlansSortedByCacheRatioOnTarget() throws Exception {
335    // The regions are fully cached on old server
336
337    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
338    ServerName server0 = servers.get(0);
339    ServerName server1 = servers.get(1);
340    ServerName server2 = servers.get(2);
341
342    // Simulate on RS with all regions, and two RSes with no regions
343    List<RegionInfo> regionsOnServer0 = randomRegions(15);
344    List<RegionInfo> regionsOnServer1 = randomRegions(0);
345    List<RegionInfo> regionsOnServer2 = randomRegions(0);
346
347    clusterState.put(server0, regionsOnServer0);
348    clusterState.put(server1, regionsOnServer1);
349    clusterState.put(server2, regionsOnServer2);
350
351    // Mock cluster metrics
352    // Mock 5 regions from server0 were previously hosted on server1
353    List<RegionInfo> oldCachedRegions1 = regionsOnServer0.subList(5, 10);
354    List<RegionInfo> oldCachedRegions2 = regionsOnServer0.subList(10, regionsOnServer0.size());
355    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
356    // mock server metrics to set cache ratio as 0 in the RS 0
357    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
358      0.0f, new ArrayList<>(), 0, 10));
359    // mock server metrics to set cache ratio as 1 in the RS 1
360    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
361      0.0f, oldCachedRegions1, 10, 10));
362    // mock server metrics to set cache ratio as .8 in the RS 2
363    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
364      0.0f, oldCachedRegions2, 8, 10));
365    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
366    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
367    loadBalancer.updateClusterMetrics(clusterMetrics);
368
369    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
370      (Map) mockClusterServersWithTables(clusterState);
371    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
372    LOG.debug("plans size: {}", plans.size());
373    LOG.debug("plans: {}", plans);
374    LOG.debug("server1 name: {}", server1.getServerName());
375    // assert the plans are in descending order from the most cached to the least cached
376    int highCacheCount = 0;
377    for (RegionPlan plan : plans) {
378      LOG.debug("plan region: {}, target server: {}", plan.getRegionInfo().getEncodedName(),
379        plan.getDestination().getServerName());
380      if (highCacheCount < 5) {
381        LOG.debug("Count: {}", highCacheCount);
382        assertTrue(oldCachedRegions1.contains(plan.getRegionInfo()));
383        assertFalse(oldCachedRegions2.contains(plan.getRegionInfo()));
384        highCacheCount++;
385      } else {
386        assertTrue(oldCachedRegions2.contains(plan.getRegionInfo()));
387        assertFalse(oldCachedRegions1.contains(plan.getRegionInfo()));
388      }
389    }
390
391  }
392
393  @Test
394  public void testRegionsFullyCachedOnOldServerAndNotCachedOnCurrentServers() throws Exception {
395    // The regions are fully cached on old server
396
397    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
398    ServerName server0 = servers.get(0);
399    ServerName server1 = servers.get(1);
400    ServerName server2 = servers.get(2);
401
402    // Simulate that the regions previously hosted by server1 are now hosted on server0
403    List<RegionInfo> regionsOnServer0 = randomRegions(10);
404    List<RegionInfo> regionsOnServer1 = randomRegions(0);
405    List<RegionInfo> regionsOnServer2 = randomRegions(5);
406
407    clusterState.put(server0, regionsOnServer0);
408    clusterState.put(server1, regionsOnServer1);
409    clusterState.put(server2, regionsOnServer2);
410
411    // Mock cluster metrics
412
413    // Mock 5 regions from server0 were previously hosted on server1
414    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);
415
416    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
417    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
418      0.0f, new ArrayList<>(), 0, 10));
419    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
420      0.0f, oldCachedRegions, 10, 10));
421    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
422      0.0f, new ArrayList<>(), 0, 10));
423    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
424    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
425    loadBalancer.updateClusterMetrics(clusterMetrics);
426
427    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
428      (Map) mockClusterServersWithTables(clusterState);
429    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
430    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
431    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
432    for (RegionPlan plan : plans) {
433      if (plan.getSource().equals(server0)) {
434        regionsMovedFromServer0.add(plan.getRegionInfo());
435        if (!targetServers.containsKey(plan.getDestination())) {
436          targetServers.put(plan.getDestination(), new ArrayList<>());
437        }
438        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
439      }
440    }
441    // should move 5 regions from server0 to server1
442    assertEquals(5, regionsMovedFromServer0.size());
443    assertEquals(5, targetServers.get(server1).size());
444    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
445  }
446
447  @Test
448  public void testRegionsFullyCachedOnOldAndCurrentServers() throws Exception {
449    // The regions are fully cached on old server
450
451    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
452    ServerName server0 = servers.get(0);
453    ServerName server1 = servers.get(1);
454    ServerName server2 = servers.get(2);
455
456    // Simulate that the regions previously hosted by server1 are now hosted on server0
457    List<RegionInfo> regionsOnServer0 = randomRegions(10);
458    List<RegionInfo> regionsOnServer1 = randomRegions(0);
459    List<RegionInfo> regionsOnServer2 = randomRegions(5);
460
461    clusterState.put(server0, regionsOnServer0);
462    clusterState.put(server1, regionsOnServer1);
463    clusterState.put(server2, regionsOnServer2);
464
465    // Mock cluster metrics
466
467    // Mock 5 regions from server0 were previously hosted on server1
468    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);
469
470    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
471    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
472      1.0f, new ArrayList<>(), 0, 10));
473    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
474      1.0f, oldCachedRegions, 10, 10));
475    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
476      1.0f, new ArrayList<>(), 0, 10));
477    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
478    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
479    loadBalancer.updateClusterMetrics(clusterMetrics);
480
481    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
482      (Map) mockClusterServersWithTables(clusterState);
483    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
484    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
485    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
486    for (RegionPlan plan : plans) {
487      if (plan.getSource().equals(server0)) {
488        regionsMovedFromServer0.add(plan.getRegionInfo());
489        if (!targetServers.containsKey(plan.getDestination())) {
490          targetServers.put(plan.getDestination(), new ArrayList<>());
491        }
492        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
493      }
494    }
495    // should move 5 regions from server0 to server1
496    assertEquals(5, regionsMovedFromServer0.size());
497    assertEquals(5, targetServers.get(server1).size());
498    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
499  }
500
501  @Test
502  public void testRegionsPartiallyCachedOnOldServerAndCurrentServer() throws Exception {
503    // The regions are partially cached on old server
504
505    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
506    ServerName server0 = servers.get(0);
507    ServerName server1 = servers.get(1);
508    ServerName server2 = servers.get(2);
509
510    // Simulate that the regions previously hosted by server1 are now hosted on server0
511    List<RegionInfo> regionsOnServer0 = randomRegions(10);
512    List<RegionInfo> regionsOnServer1 = randomRegions(0);
513    List<RegionInfo> regionsOnServer2 = randomRegions(5);
514
515    clusterState.put(server0, regionsOnServer0);
516    clusterState.put(server1, regionsOnServer1);
517    clusterState.put(server2, regionsOnServer2);
518
519    // Mock cluster metrics
520
521    // Mock 5 regions from server0 were previously hosted on server1
522    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);
523
524    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
525    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
526      0.2f, new ArrayList<>(), 0, 10));
527    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
528      0.0f, oldCachedRegions, 6, 10));
529    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
530      1.0f, new ArrayList<>(), 0, 10));
531    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
532    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
533    loadBalancer.updateClusterMetrics(clusterMetrics);
534
535    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
536      (Map) mockClusterServersWithTables(clusterState);
537    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
538    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
539    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
540    for (RegionPlan plan : plans) {
541      if (plan.getSource().equals(server0)) {
542        regionsMovedFromServer0.add(plan.getRegionInfo());
543        if (!targetServers.containsKey(plan.getDestination())) {
544          targetServers.put(plan.getDestination(), new ArrayList<>());
545        }
546        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
547      }
548    }
549    assertEquals(5, regionsMovedFromServer0.size());
550    assertEquals(5, targetServers.get(server1).size());
551    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
552  }
553
554  @Test
555  public void testBalancerNotThrowNPEWhenBalancerPlansIsNull() throws Exception {
556    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
557    ServerName server0 = servers.get(0);
558    ServerName server1 = servers.get(1);
559    ServerName server2 = servers.get(2);
560
561    List<RegionInfo> regionsOnServer0 = randomRegions(5);
562    List<RegionInfo> regionsOnServer1 = randomRegions(5);
563    List<RegionInfo> regionsOnServer2 = randomRegions(5);
564
565    clusterState.put(server0, regionsOnServer0);
566    clusterState.put(server1, regionsOnServer1);
567    clusterState.put(server2, regionsOnServer2);
568
569    // Mock cluster metrics
570    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
571    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
572      0.0f, new ArrayList<>(), 0, 10));
573    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
574      0.0f, new ArrayList<>(), 0, 10));
575    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
576      0.0f, new ArrayList<>(), 0, 10));
577
578    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
579    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
580    loadBalancer.updateClusterMetrics(clusterMetrics);
581
582    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
583      (Map) mockClusterServersWithTables(clusterState);
584    try {
585      List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
586      assertNull(plans);
587    } catch (NullPointerException npe) {
588      fail("NPE should not be thrown");
589    }
590  }
591}