/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.hadoop.hbase.master.balancer;

import static org.junit.Assert.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadLocalRandom;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.ClusterMetrics;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.HConstants;
import org.apache.hadoop.hbase.RegionMetrics;
import org.apache.hadoop.hbase.ServerMetrics;
import org.apache.hadoop.hbase.ServerName;
import org.apache.hadoop.hbase.Size;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.RegionInfo;
import org.apache.hadoop.hbase.client.TableDescriptor;
import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
import org.apache.hadoop.hbase.master.MasterServices;
import org.apache.hadoop.hbase.master.RegionPlan;
import org.apache.hadoop.hbase.testclassification.LargeTests;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.hbase.util.EnvironmentEdgeManager;
import org.apache.hadoop.hbase.util.Pair;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.hbase.thirdparty.com.google.common.collect.Lists;

@Tag(LargeTests.TAG)
public class TestCacheAwareLoadBalancer extends BalancerTestBase {

  private static final Logger LOG = LoggerFactory.getLogger(TestCacheAwareLoadBalancer.class);

  private static CacheAwareLoadBalancer loadBalancer;

  static List<ServerName> servers;

  static List<TableDescriptor> tableDescs;

  static Map<TableName, String> tableMap = new HashMap<>();

  static TableName[] tables = new TableName[] { TableName.valueOf("dt1"), TableName.valueOf("dt2"),
    TableName.valueOf("dt3"), TableName.valueOf("dt4") };

  private static List<ServerName> generateServers(int numServers) {
    List<ServerName> servers = new ArrayList<>(numServers);
    Random rand = ThreadLocalRandom.current();
    for (int i = 0; i < numServers; i++) {
      String host = "server" + rand.nextInt(100000);
      int port = rand.nextInt(60000);
      servers.add(ServerName.valueOf(host, port, -1));
    }
    return servers;
  }

  private static List<TableDescriptor> constructTableDesc(boolean hasBogusTable) {
    List<TableDescriptor> tds = Lists.newArrayList();
    for (int i = 0; i < tables.length; i++) {
      TableDescriptor htd = TableDescriptorBuilder.newBuilder(tables[i]).build();
      tds.add(htd);
    }
    return tds;
  }

  private ServerMetrics mockServerMetricsWithRegionCacheInfo(ServerName server,
    List<RegionInfo> regionsOnServer, float currentCacheRatio, List<RegionInfo> oldRegionCacheInfo,
    int oldRegionCachedSize, int regionSize) {
    ServerMetrics serverMetrics = mock(ServerMetrics.class);
    Map<byte[], RegionMetrics> regionLoadMap = new TreeMap<>(Bytes.BYTES_COMPARATOR);
    for (RegionInfo info : regionsOnServer) {
      RegionMetrics rl = mock(RegionMetrics.class);
      when(rl.getReadRequestCount()).thenReturn(0L);
      when(rl.getWriteRequestCount()).thenReturn(0L);
      when(rl.getMemStoreSize()).thenReturn(Size.ZERO);
      when(rl.getStoreFileSize()).thenReturn(Size.ZERO);
      when(rl.getCurrentRegionCachedRatio()).thenReturn(currentCacheRatio);
      when(rl.getRegionSizeMB()).thenReturn(new Size(regionSize, Size.Unit.MEGABYTE));
      regionLoadMap.put(info.getRegionName(), rl);
    }
    when(serverMetrics.getRegionMetrics()).thenReturn(regionLoadMap);
    Map<String, Integer> oldCacheRatioMap = new HashMap<>();
    for (RegionInfo info : oldRegionCacheInfo) {
      oldCacheRatioMap.put(info.getEncodedName(), oldRegionCachedSize);
    }
    when(serverMetrics.getRegionCachedInfo()).thenReturn(oldCacheRatioMap);
    return serverMetrics;
  }

  @BeforeAll
  public static void beforeAllTests() throws Exception {
    servers = generateServers(3);
    tableDescs = constructTableDesc(false);
    Configuration conf = HBaseConfiguration.create();
    conf.set(HConstants.BUCKET_CACHE_PERSISTENT_PATH_KEY, "prefetch_file_list");
    loadBalancer = new CacheAwareLoadBalancer();
    MasterServices services = mock(MasterServices.class);
    when(services.getConfiguration()).thenReturn(conf);
    loadBalancer.setMasterServices(services);
    loadBalancer.loadConf(conf);
  }

  @Test
  public void testRegionsNotCachedOnOldServerAndCurrentServer() throws Exception {
    // The regions are not cached on old server as well as the current server. This causes
    // skewness in the region allocation which should be fixed by the balancer

    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    // Simulate that the regions previously hosted by server1 are now hosted on server0
    List<RegionInfo> regionsOnServer0 = randomRegions(10);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(5);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock cluster metrics
    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      0.0f, new ArrayList<>(), 0, 10));
    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
    for (RegionPlan plan : plans) {
      if (plan.getSource().equals(server0)) {
        regionsMovedFromServer0.add(plan.getRegionInfo());
        if (!targetServers.containsKey(plan.getDestination())) {
          targetServers.put(plan.getDestination(), new ArrayList<>());
        }
        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
      }
    }
    // should move 5 regions from server0 to server 1
    assertEquals(5, regionsMovedFromServer0.size());
    assertEquals(5, targetServers.get(server1).size());
  }

  /**
   * Regions on the overloaded RS report low block-cache ratio; no RS reports prefetch/historical
   * cache for those regions (so {@link CacheAwareLoadBalancer.CacheAwareCandidateGenerator} has no
   * "old server" to prefer). Another RS has ample free block cache. The balancer should still emit
   * plans that shed load from the hot RS onto the idle RS with spare cache capacity.
   */
  @Test
  public void testLowCacheRatioNoHistoricalCacheRelocatesWhenTargetHasFreeBlockCache()
    throws Exception {
    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    List<RegionInfo> regionsOnServer0 = randomRegions(10);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(5);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Below LOW_CACHE_RATIO_FOR_RELOCATION_DEFAULT (0.35);
    ServerMetrics sm0 = mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0, 0.1f,
      new ArrayList<>(), 0, 10);
    when(sm0.getCacheFreeSize()).thenReturn(0L);
    ServerMetrics sm1 = mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1, 0.0f,
      new ArrayList<>(), 0, 10);
    // Simulates 1GB free cache space on server1
    when(sm1.getCacheFreeSize()).thenReturn(1024L * 1024 * 1024);
    ServerMetrics sm2 = mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2, 1.0f,
      new ArrayList<>(), 0, 10);
    when(sm2.getCacheFreeSize()).thenReturn(0L);

    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, sm0);
    serverMetricsMap.put(server1, sm1);
    serverMetricsMap.put(server2, sm2);
    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    assertTrue(loadBalancer.regionCacheRatioOnOldServerMap.isEmpty());

    Map<TableName, Map<ServerName, List<RegionInfo>>> loadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    List<RegionPlan> plans = loadBalancer.balanceCluster(loadOfAllTable);
    assertNotNull(plans);

    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
    for (RegionPlan plan : plans) {
      if (plan.getSource().equals(server0)) {
        regionsMovedFromServer0.add(plan.getRegionInfo());
        if (!targetServers.containsKey(plan.getDestination())) {
          targetServers.put(plan.getDestination(), new ArrayList<>());
        }
        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
      }
    }
    assertEquals(5, regionsMovedFromServer0.size());
    assertNotNull(targetServers.get(server1));
    assertEquals(5, targetServers.get(server1).size());
  }

  @Test
  public void testRegionsPartiallyCachedOnOldServerAndNotCachedOnCurrentServer() throws Exception {
    // The regions are partially cached on old server but not cached on the current server

    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    // Simulate that the regions previously hosted by server1 are now hosted on server0
    List<RegionInfo> regionsOnServer0 = randomRegions(10);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(5);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock cluster metrics

    // Mock 5 regions from server0 were previously hosted on server1
    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);

    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      0.0f, oldCachedRegions, 6, 10));
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      0.0f, new ArrayList<>(), 0, 10));
    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
    for (RegionPlan plan : plans) {
      if (plan.getSource().equals(server0)) {
        regionsMovedFromServer0.add(plan.getRegionInfo());
        if (!targetServers.containsKey(plan.getDestination())) {
          targetServers.put(plan.getDestination(), new ArrayList<>());
        }
        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
      }
    }
    // should move 5 regions from server0 to server1
    assertEquals(5, regionsMovedFromServer0.size());
    assertEquals(5, targetServers.get(server1).size());
    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
  }

  @Test
  public void testThrottlingRegionBeyondThreshold() throws Exception {
    Configuration conf = HBaseConfiguration.create();
    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
    MasterServices services = mock(MasterServices.class);
    when(services.getConfiguration()).thenReturn(conf);
    balancer.setMasterServices(services);
    balancer.loadConf(conf);
    balancer.initialize();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    Pair<ServerName, Float> regionRatio = new Pair<>();
    regionRatio.setFirst(server0);
    regionRatio.setSecond(1.0f);
    balancer.regionCacheRatioOnOldServerMap.put("region1", regionRatio);
    RegionInfo mockedInfo = mock(RegionInfo.class);
    when(mockedInfo.getEncodedName()).thenReturn("region1");
    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
    long startTime = EnvironmentEdgeManager.currentTime();
    balancer.throttle(plan);
    long endTime = EnvironmentEdgeManager.currentTime();
    assertTrue((endTime - startTime) < 10);
  }

  @Test
  public void testThrottlingRegionBelowThreshold() throws Exception {
    Configuration conf = HBaseConfiguration.create();
    conf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, 100);
    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
    MasterServices services = mock(MasterServices.class);
    when(services.getConfiguration()).thenReturn(conf);
    balancer.setMasterServices(services);
    balancer.loadConf(conf);
    balancer.initialize();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    Pair<ServerName, Float> regionRatio = new Pair<>();
    regionRatio.setFirst(server0);
    regionRatio.setSecond(0.1f);
    balancer.regionCacheRatioOnOldServerMap.put("region1", regionRatio);
    RegionInfo mockedInfo = mock(RegionInfo.class);
    when(mockedInfo.getEncodedName()).thenReturn("region1");
    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
    long startTime = EnvironmentEdgeManager.currentTime();
    balancer.throttle(plan);
    long endTime = EnvironmentEdgeManager.currentTime();
    assertTrue((endTime - startTime) >= 100);
  }

  @Test
  public void testThrottlingCacheRatioUnknownOnTarget() throws Exception {
    Configuration conf = HBaseConfiguration.create();
    conf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, 100);
    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
    MasterServices services = mock(MasterServices.class);
    when(services.getConfiguration()).thenReturn(conf);
    balancer.setMasterServices(services);
    balancer.loadConf(conf);
    balancer.initialize();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server3 = servers.get(2);
    // setting region cache ratio 100% on server 3, though this is not the target in the region plan
    Pair<ServerName, Float> regionRatio = new Pair<>();
    regionRatio.setFirst(server3);
    regionRatio.setSecond(1.0f);
    balancer.regionCacheRatioOnOldServerMap.put("region1", regionRatio);
    RegionInfo mockedInfo = mock(RegionInfo.class);
    when(mockedInfo.getEncodedName()).thenReturn("region1");
    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
    long startTime = EnvironmentEdgeManager.currentTime();
    balancer.throttle(plan);
    long endTime = EnvironmentEdgeManager.currentTime();
    assertTrue((endTime - startTime) >= 100);
  }

  @Test
  public void testThrottlingCacheRatioUnknownForRegion() throws Exception {
    Configuration conf = HBaseConfiguration.create();
    conf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, 100);
    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
    MasterServices services = mock(MasterServices.class);
    when(services.getConfiguration()).thenReturn(conf);
    balancer.setMasterServices(services);
    balancer.loadConf(conf);
    balancer.initialize();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server3 = servers.get(2);
    // No cache ratio available for region1
    RegionInfo mockedInfo = mock(RegionInfo.class);
    when(mockedInfo.getEncodedName()).thenReturn("region1");
    RegionPlan plan = new RegionPlan(mockedInfo, server1, server0);
    long startTime = EnvironmentEdgeManager.currentTime();
    balancer.throttle(plan);
    long endTime = EnvironmentEdgeManager.currentTime();
    assertTrue((endTime - startTime) >= 100);
  }

  @Test
  public void testRegionPlansSortedByCacheRatioOnTarget() throws Exception {
    // The regions are fully cached on old server

    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    // Simulate on RS with all regions, and two RSes with no regions
    List<RegionInfo> regionsOnServer0 = randomRegions(15);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(0);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock cluster metrics
    // Mock 5 regions from server0 were previously hosted on server1
    List<RegionInfo> oldCachedRegions1 = regionsOnServer0.subList(5, 10);
    List<RegionInfo> oldCachedRegions2 = regionsOnServer0.subList(10, regionsOnServer0.size());
    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    // mock server metrics to set cache ratio as 0 in the RS 0
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      0.0f, new ArrayList<>(), 0, 10));
    // mock server metrics to set cache ratio as 1 in the RS 1
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      0.0f, oldCachedRegions1, 10, 10));
    // mock server metrics to set cache ratio as .8 in the RS 2
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      0.0f, oldCachedRegions2, 8, 10));
    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
    LOG.debug("plans size: {}", plans.size());
    LOG.debug("plans: {}", plans);
    LOG.debug("server1 name: {}", server1.getServerName());
    // assert the plans are in descending order from the most cached to the least cached
    int highCacheCount = 0;
    for (RegionPlan plan : plans) {
      LOG.debug("plan region: {}, target server: {}", plan.getRegionInfo().getEncodedName(),
        plan.getDestination().getServerName());
      if (highCacheCount < 5) {
        LOG.debug("Count: {}", highCacheCount);
        assertTrue(oldCachedRegions1.contains(plan.getRegionInfo()));
        assertFalse(oldCachedRegions2.contains(plan.getRegionInfo()));
        highCacheCount++;
      } else {
        assertTrue(oldCachedRegions2.contains(plan.getRegionInfo()));
        assertFalse(oldCachedRegions1.contains(plan.getRegionInfo()));
      }
    }

  }

  @Test
  public void testRegionsFullyCachedOnOldServerAndNotCachedOnCurrentServers() throws Exception {
    // The regions are fully cached on old server

    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    // Simulate that the regions previously hosted by server1 are now hosted on server0
    List<RegionInfo> regionsOnServer0 = randomRegions(10);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(5);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock cluster metrics

    // Mock 5 regions from server0 were previously hosted on server1
    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);

    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      0.0f, oldCachedRegions, 10, 10));
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      0.0f, new ArrayList<>(), 0, 10));
    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
    for (RegionPlan plan : plans) {
      if (plan.getSource().equals(server0)) {
        regionsMovedFromServer0.add(plan.getRegionInfo());
        if (!targetServers.containsKey(plan.getDestination())) {
          targetServers.put(plan.getDestination(), new ArrayList<>());
        }
        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
      }
    }
    // should move 5 regions from server0 to server1
    assertEquals(5, regionsMovedFromServer0.size());
    assertEquals(5, targetServers.get(server1).size());
    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
  }

  @Test
  public void testRegionsFullyCachedOnOldAndCurrentServers() throws Exception {
    // The regions are fully cached on old server

    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    // Simulate that the regions previously hosted by server1 are now hosted on server0
    List<RegionInfo> regionsOnServer0 = randomRegions(10);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(5);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock cluster metrics

    // Mock 5 regions from server0 were previously hosted on server1
    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);

    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      1.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      1.0f, oldCachedRegions, 10, 10));
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      1.0f, new ArrayList<>(), 0, 10));
    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
    for (RegionPlan plan : plans) {
      if (plan.getSource().equals(server0)) {
        regionsMovedFromServer0.add(plan.getRegionInfo());
        if (!targetServers.containsKey(plan.getDestination())) {
          targetServers.put(plan.getDestination(), new ArrayList<>());
        }
        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
      }
    }
    // should move 5 regions from server0 to server1
    assertEquals(5, regionsMovedFromServer0.size());
    assertEquals(5, targetServers.get(server1).size());
    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
  }

  @Test
  public void testRegionsPartiallyCachedOnOldServerAndCurrentServer() throws Exception {
    // The regions are partially cached on old server

    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    // Simulate that the regions previously hosted by server1 are now hosted on server0
    List<RegionInfo> regionsOnServer0 = randomRegions(10);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(5);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock cluster metrics

    // Mock 5 regions from server0 were previously hosted on server1
    List<RegionInfo> oldCachedRegions = regionsOnServer0.subList(5, regionsOnServer0.size() - 1);

    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      0.2f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      0.0f, oldCachedRegions, 6, 10));
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      1.0f, new ArrayList<>(), 0, 10));
    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
    Set<RegionInfo> regionsMovedFromServer0 = new HashSet<>();
    Map<ServerName, List<RegionInfo>> targetServers = new HashMap<>();
    for (RegionPlan plan : plans) {
      if (plan.getSource().equals(server0)) {
        regionsMovedFromServer0.add(plan.getRegionInfo());
        if (!targetServers.containsKey(plan.getDestination())) {
          targetServers.put(plan.getDestination(), new ArrayList<>());
        }
        targetServers.get(plan.getDestination()).add(plan.getRegionInfo());
      }
    }
    assertEquals(5, regionsMovedFromServer0.size());
    assertEquals(5, targetServers.get(server1).size());
    assertTrue(targetServers.get(server1).containsAll(oldCachedRegions));
  }

  @Test
  public void testBalancerNotThrowNPEWhenBalancerPlansIsNull() throws Exception {
    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    List<RegionInfo> regionsOnServer0 = randomRegions(5);
    List<RegionInfo> regionsOnServer1 = randomRegions(5);
    List<RegionInfo> regionsOnServer2 = randomRegions(5);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock cluster metrics
    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      0.0f, new ArrayList<>(), 0, 10));

    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    loadBalancer.updateClusterMetrics(clusterMetrics);

    Map<TableName, Map<ServerName, List<RegionInfo>>> LoadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);
    try {
      List<RegionPlan> plans = loadBalancer.balanceCluster(LoadOfAllTable);
      assertNull(plans);
    } catch (NullPointerException npe) {
      fail("NPE should not be thrown");
    }
  }

  /**
   * This test verifies that when loadConf/onConfigurationChange is called on a
   * CacheAwareLoadBalancer while a balance run is in progress, the configuration update: 1. Does
   * not block (returns quickly without waiting for balancing to finish) 2. Does not affect the
   * ongoing balance run (the configuration used during balancing remains the old one) 3. Is applied
   * correctly after the balance run completes
   */
  @Test
  @Timeout(60)
  public void testConfigUpdateDuringBalance() throws Exception {
    Float expectedOldRatioThreshold = 0.8f;
    Float expectedNewRatioThreshold = 0.95f;
    long throttleTimeMs = 10000;

    // Actual old ratio threshold used during balance
    float[] actualOldRatioThresholdDuringBalance = new float[1];

    Configuration conf = HBaseConfiguration.create();
    conf.set(HConstants.BUCKET_CACHE_PERSISTENT_PATH_KEY, "prefetch_file_list");
    conf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, throttleTimeMs);
    conf.setFloat(CacheAwareLoadBalancer.CACHE_RATIO_THRESHOLD, expectedOldRatioThreshold);

    CacheAwareLoadBalancer balancer = new CacheAwareLoadBalancer();
    MasterServices services = mock(MasterServices.class);
    when(services.getConfiguration()).thenReturn(conf);
    balancer.setMasterServices(services);
    balancer.loadConf(conf);
    balancer.initialize();

    Map<ServerName, List<RegionInfo>> clusterState = new HashMap<>();
    ServerName server0 = servers.get(0);
    ServerName server1 = servers.get(1);
    ServerName server2 = servers.get(2);

    // Setup cluster: all 3 regions on server0 (unbalanced)
    List<RegionInfo> regionsOnServer0 = randomRegions(3);
    List<RegionInfo> regionsOnServer1 = randomRegions(0);
    List<RegionInfo> regionsOnServer2 = randomRegions(0);

    clusterState.put(server0, regionsOnServer0);
    clusterState.put(server1, regionsOnServer1);
    clusterState.put(server2, regionsOnServer2);

    // Mock metrics: NO cache info for any region = all will be throttled
    Map<ServerName, ServerMetrics> serverMetricsMap = new TreeMap<>();
    serverMetricsMap.put(server0, mockServerMetricsWithRegionCacheInfo(server0, regionsOnServer0,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server1, mockServerMetricsWithRegionCacheInfo(server1, regionsOnServer1,
      0.0f, new ArrayList<>(), 0, 10));
    serverMetricsMap.put(server2, mockServerMetricsWithRegionCacheInfo(server2, regionsOnServer2,
      0.0f, new ArrayList<>(), 0, 10));

    ClusterMetrics clusterMetrics = mock(ClusterMetrics.class);
    when(clusterMetrics.getLiveServerMetrics()).thenReturn(serverMetricsMap);
    balancer.updateClusterMetrics(clusterMetrics);

    final Map<TableName, Map<ServerName, List<RegionInfo>>> loadOfAllTable =
      (Map) mockClusterServersWithTables(clusterState);

    // Verify initial configuration
    assertEquals(expectedOldRatioThreshold, balancer.ratioThreshold, 0.001f);

    CountDownLatch balanceStarted = new CountDownLatch(1);
    CountDownLatch updateConfigInitiated = new CountDownLatch(1);

    long[] configUpdateDuration = new long[1];
    long[] balanceDuration = new long[1];

    ExecutorService executor = Executors.newFixedThreadPool(2);

    try {
      // Thread 1 Simulate similar flow to HMaster.balance() which holds synchronized(balancer) for
      // the duration of balance
      Future<Long> balanceFuture = executor.submit(() -> {
        try {
          long start = EnvironmentEdgeManager.currentTime();
          synchronized (balancer) {
            try {
              // Simulate beginning of HMaster.balance() mark balancing window open
              balancer.onBalancingStart();
              balanceStarted.countDown();
              List<RegionPlan> plans = balancer.balanceCluster(loadOfAllTable);

              LOG.info("Balance generated {} plans, executing with throttling",
                plans != null ? plans.size() : 0);

              if (plans != null) {
                for (int i = 0; i < plans.size(); i++) {
                  RegionPlan plan = plans.get(i);
                  balancer.throttle(plan);
                }
              }
              // Wait until config update is initiated while balance is still in progress
              updateConfigInitiated.await();

              // Old config should still be visible during current balance run
              actualOldRatioThresholdDuringBalance[0] = balancer.ratioThreshold;
            } finally {
              balancer.onBalancingComplete();
            }
          }
          return EnvironmentEdgeManager.currentTime() - start;
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      });

      // Thread 2: Simulate update_all_config / onConfigurationChange
      Future<Long> configUpdateFuture = executor.submit(() -> {
        try {
          // Wait for balance to start
          balanceStarted.await();
          long startTime = EnvironmentEdgeManager.currentTime();

          // Call onConfigurationChange - should NOT hang
          Configuration newConf = HBaseConfiguration.create();
          newConf.set(HConstants.BUCKET_CACHE_PERSISTENT_PATH_KEY, "prefetch_file_list");
          newConf.setLong(CacheAwareLoadBalancer.MOVE_THROTTLING, 10000);
          newConf.setFloat(CacheAwareLoadBalancer.CACHE_RATIO_THRESHOLD, expectedNewRatioThreshold);
          balancer.onConfigurationChange(newConf);
          updateConfigInitiated.countDown();

          return EnvironmentEdgeManager.currentTime() - startTime;
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      });

      // Wait for both threads to complete
      configUpdateDuration[0] = configUpdateFuture.get();
      balanceDuration[0] = balanceFuture.get();
      System.out.println("Balance duration (ms): " + balanceDuration[0]);
      System.out.println("Config update duration (ms): " + configUpdateDuration[0]);

      // Verify that ratio threshold used during balance is stll the old
      assertEquals(expectedOldRatioThreshold, actualOldRatioThresholdDuringBalance[0], 0.001f);

      // Verify that config updated successfully after balance completed
      assertEquals(expectedNewRatioThreshold, balancer.ratioThreshold, 0.001f);

      // Verify that config update didn't hang/timeout waiting for balance
      assertTrue(configUpdateDuration[0] < balanceDuration[0]);

    } finally {
      executor.shutdownNow();
    }
  }
}
