Source code

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