/*
 * 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;

/** An implementation of the {@link org.apache.hadoop.hbase.master.LoadBalancer} that assigns regions
 * based on the amount they are cached on a given server. A region can move across the region
 * servers whenever a region server shuts down or crashes. The region server preserves the cache
 * periodically and restores the cache when it is restarted. This balancer implements a mechanism
 * where it maintains the amount by which a region is cached on a region server. During balancer
 * run, a region plan is generated that takes into account this cache information and tries to
 * move the regions so that the cache is minimally impacted.
 */

import static org.apache.hadoop.hbase.HConstants.BUCKET_CACHE_PERSISTENT_PATH_KEY;

import java.math.BigDecimal;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Deque;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ThreadLocalRandom;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.ClusterMetrics;
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.master.RackManager;
import org.apache.hadoop.hbase.master.RegionPlan;
import org.apache.hadoop.hbase.util.Pair;
import org.apache.yetus.audience.InterfaceAudience;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@InterfaceAudience.Private
public class CacheAwareLoadBalancer extends StochasticLoadBalancer {
  private static final Logger LOG = LoggerFactory.getLogger(CacheAwareLoadBalancer.class);

  public static final String CACHE_RATIO_THRESHOLD =
    "hbase.master.balancer.stochastic.throttling.cacheRatio";
  public static final float CACHE_RATIO_THRESHOLD_DEFAULT = 0.8f;

  /**
   * Below this cache ratio on the current host, a move may be considered for the free-space
   * heuristic.
   */
  public static final String LOW_CACHE_RATIO_FOR_RELOCATION_KEY =
    "hbase.master.balancer.cacheaware.lowCacheRatioThreshold";
  public static final float LOW_CACHE_RATIO_FOR_RELOCATION_DEFAULT = 0.35f;

  /**
   * Optimistic region cache ratio assumed for cost purposes when a better host has free cache space
   * (actual warmup is not modeled).
   */
  public static final String POTENTIAL_CACHE_RATIO_AFTER_MOVE_KEY =
    "hbase.master.balancer.cacheaware.potentialCacheRatioAfterMove";
  public static final float POTENTIAL_CACHE_RATIO_AFTER_MOVE_DEFAULT = 0.95f;

  /**
   * Minimum free block cache on a target server, as a multiple of the region's on-disk size in
   * bytes, required to count that server as a relocation opportunity.
   */
  public static final String MIN_FREE_CACHE_SPACE_FACTOR_KEY =
    "hbase.master.balancer.cacheaware.minFreeCacheSpaceFactor";
  public static final float MIN_FREE_CACHE_SPACE_FACTOR_DEFAULT = 1.0f;

  public Float ratioThreshold;

  private Long sleepTime;
  private Configuration configuration;

  private float lowCacheRatioThreshold;
  private float potentialCacheRatioAfterMove;
  private float minFreeCacheSpaceFactor;

  private BigDecimal simulatedRatio = BigDecimal.ZERO;

  @Override
  public void loadConf(Configuration configuration) {
    this.configuration = configuration;
    this.costFunctions = new ArrayList<>();
    super.loadConf(configuration);
    ratioThreshold =
      this.configuration.getFloat(CACHE_RATIO_THRESHOLD, CACHE_RATIO_THRESHOLD_DEFAULT);
    sleepTime = configuration.getLong(MOVE_THROTTLING, MOVE_THROTTLING_DEFAULT.toMillis());
    lowCacheRatioThreshold = configuration.getFloat(LOW_CACHE_RATIO_FOR_RELOCATION_KEY,
      LOW_CACHE_RATIO_FOR_RELOCATION_DEFAULT);
    potentialCacheRatioAfterMove = configuration.getFloat(POTENTIAL_CACHE_RATIO_AFTER_MOVE_KEY,
      POTENTIAL_CACHE_RATIO_AFTER_MOVE_DEFAULT);
    minFreeCacheSpaceFactor =
      configuration.getFloat(MIN_FREE_CACHE_SPACE_FACTOR_KEY, MIN_FREE_CACHE_SPACE_FACTOR_DEFAULT);
  }

  @Override
  protected Map<Class<? extends CandidateGenerator>, CandidateGenerator>
    createCandidateGenerators(Configuration conf) {
    Map<Class<? extends CandidateGenerator>, CandidateGenerator> candidateGenerators =
      new HashMap<>(2);
    candidateGenerators.put(CacheAwareSkewnessCandidateGenerator.class,
      new CacheAwareSkewnessCandidateGenerator());
    candidateGenerators.put(CacheAwareCandidateGenerator.class, new CacheAwareCandidateGenerator());
    return candidateGenerators;
  }

  @Override
  protected List<CostFunction> createCostFunctions(Configuration configuration) {
    List<CostFunction> costFunctions = new ArrayList<>();
    addCostFunction(costFunctions, new CacheAwareRegionSkewnessCostFunction(configuration));
    addCostFunction(costFunctions, new CacheAwareCostFunction(configuration));
    return costFunctions;
  }

  private void addCostFunction(List<CostFunction> costFunctions, CostFunction costFunction) {
    if (costFunction.getMultiplier() > 0) {
      costFunctions.add(costFunction);
    }
  }

  @Override
  public void updateClusterMetrics(ClusterMetrics clusterMetrics) {
    this.clusterStatus = clusterMetrics;
    updateRegionLoad();
  }

  protected Map<ServerName, Long> getServerBlockCacheFreeBytes() {
    if (clusterStatus == null) {
      return null;
    }
    Map<ServerName, Long> map = new HashMap<>();
    clusterStatus.getLiveServerMetrics().forEach((sn, sm) -> map.put(sn, sm.getCacheFreeSize()));
    return map;
  }

  @Override
  protected BalancerClusterState createState(Map<ServerName, List<RegionInfo>> clusterState,
    Map<String, Deque<BalancerRegionLoad>> loads, RegionHDFSBlockLocationFinder finder,
    RackManager rackManager) {
    return new BalancerClusterState(clusterState, loads, finder, rackManager,
      regionCacheRatioOnOldServerMap, getServerBlockCacheFreeBytes());
  }

  /**
   * Collect the amount of region cached for all the regions from all the active region servers.
   */
  private void updateRegionLoad() {
    loads = new HashMap<>();
    regionCacheRatioOnOldServerMap = new HashMap<>();
    Map<String, Pair<ServerName, Integer>> regionCacheRatioOnCurrentServerMap = new HashMap<>();

    // Build current region cache statistics
    clusterStatus.getLiveServerMetrics().forEach((ServerName sn, ServerMetrics sm) -> {
      // Create a map of region and the server where it is currently hosted
      sm.getRegionMetrics().forEach((byte[] regionName, RegionMetrics rm) -> {
        String regionEncodedName = RegionInfo.encodeRegionName(regionName);

        Deque<BalancerRegionLoad> rload = new ArrayDeque<>();

        // Get the total size of the hFiles in this region
        int regionSizeMB = (int) rm.getRegionSizeMB().get(Size.Unit.MEGABYTE);

        rload.add(new BalancerRegionLoad(rm));
        // Maintain a map of region and its total size. This is needed to calculate the cache
        // ratios for the regions cached on old region servers
        regionCacheRatioOnCurrentServerMap.put(regionEncodedName, new Pair<>(sn, regionSizeMB));
        loads.put(regionEncodedName, rload);
      });
    });

    // Build cache statistics for the regions hosted previously on old region servers
    clusterStatus.getLiveServerMetrics().forEach((ServerName sn, ServerMetrics sm) -> {
      // Find if a region was previously hosted on a server other than the one it is currently
      // hosted on.
      sm.getRegionCachedInfo().forEach((String regionEncodedName, Integer regionSizeInCache) -> {
        // If the region is found in regionCacheRatioOnCurrentServerMap, it is currently hosted on
        // this server
        if (regionCacheRatioOnCurrentServerMap.containsKey(regionEncodedName)) {
          ServerName currentServer =
            regionCacheRatioOnCurrentServerMap.get(regionEncodedName).getFirst();
          if (!ServerName.isSameAddress(currentServer, sn)) {
            int regionSizeMB =
              regionCacheRatioOnCurrentServerMap.get(regionEncodedName).getSecond();
            // The coldDataSize accounts for data size classified as "cold" by DataTieringManager,
            // which should be kept out of cache. We calculate cache ratio on old server based
            // only on the hot data size for the region (regionSizeMB - coldDataSize), as we
            // don't want to move regions with low cache ratio due to data classified as cold.
            int coldDataSize = sm.getRegionColdDataSize().getOrDefault(regionEncodedName, 0);
            float regionCacheRatioOnOldServer = (regionSizeMB - coldDataSize) <= 0
              ? 0.0f
              : (float) regionSizeInCache / (regionSizeMB - coldDataSize);
            regionCacheRatioOnOldServerMap.put(regionEncodedName,
              new Pair<>(sn, regionCacheRatioOnOldServer));
          }
        }
      });
    });
  }

  private RegionInfo getRegionInfoByEncodedName(BalancerClusterState cluster, String regionName) {
    Optional<RegionInfo> regionInfoOptional =
      Arrays.stream(cluster.regions).filter((RegionInfo ri) -> {
        return regionName.equals(ri.getEncodedName());
      }).findFirst();

    if (regionInfoOptional.isPresent()) {
      return regionInfoOptional.get();
    }
    return null;
  }

  @Override
  public long getThrottleDurationMs(RegionPlan plan) {
    Pair<ServerName, Float> rsRatio = this.regionCacheRatioOnOldServerMap.get(plan.getRegionName());
    if (
      rsRatio != null && plan.getDestination().equals(rsRatio.getFirst())
        && rsRatio.getSecond() >= ratioThreshold
    ) {
      LOG.debug("Moving region {} to server {} with cache ratio {}. No throttling needed.",
        plan.getRegionInfo().getEncodedName(), plan.getDestination(), rsRatio.getSecond());
      return 0L;
    } else {
      if (rsRatio != null) {
        LOG.debug("Moving region {} to server {} with cache ratio: {}. Throttling move for {}ms.",
          plan.getRegionInfo().getEncodedName(), plan.getDestination(),
          plan.getDestination().equals(rsRatio.getFirst()) ? rsRatio.getSecond() : "unknown",
          sleepTime);
      } else {
        LOG.debug(
          "Moving region {} to server {} with no cache ratio info for the region. "
            + "Throttling move for {}ms.",
          plan.getRegionInfo().getEncodedName(), plan.getDestination(), sleepTime);
      }
      return sleepTime;
    }
  }

  @Override
  protected List<RegionPlan> balanceTable(TableName tableName,
    Map<ServerName, List<RegionInfo>> loadOfOneTable) {
    final Map<String, Pair<ServerName, Float>> snapshot = new HashMap<>();
    snapshot.putAll(this.regionCacheRatioOnOldServerMap);
    List<RegionPlan> plans = super.balanceTable(tableName, loadOfOneTable);
    if (plans == null) {
      return plans;
    }
    plans.sort((p1, p2) -> {
      Pair<ServerName, Float> pair1 = snapshot.get(p1.getRegionName());
      Float ratio1 =
        pair1 == null ? 0 : pair1.getFirst().equals(p1.getDestination()) ? pair1.getSecond() : 0f;
      Pair<ServerName, Float> pair2 = snapshot.get(p2.getRegionName());
      Float ratio2 =
        pair2 == null ? 0 : pair2.getFirst().equals(p2.getDestination()) ? pair2.getSecond() : 0f;
      return ratio1.compareTo(ratio2) * (-1);
    });
    return plans;
  }

  private class CacheAwareCandidateGenerator extends CandidateGenerator {
    @Override
    protected BalanceAction generate(BalancerClusterState cluster) {
      simulatedRatio = BigDecimal.ZERO;
      // Move the regions to the servers they were previously hosted on based on the cache ratio
      if (
        !regionCacheRatioOnOldServerMap.isEmpty()
          && regionCacheRatioOnOldServerMap.entrySet().iterator().hasNext()
      ) {
        Map.Entry<String, Pair<ServerName, Float>> regionCacheRatioServerMap =
          regionCacheRatioOnOldServerMap.entrySet().iterator().next();
        // Get the server where this region was previously hosted
        String regionEncodedName = regionCacheRatioServerMap.getKey();
        RegionInfo regionInfo = getRegionInfoByEncodedName(cluster, regionEncodedName);
        if (regionInfo == null) {
          LOG.warn("Region {} not found", regionEncodedName);
          regionCacheRatioOnOldServerMap.remove(regionEncodedName);
          return BalanceAction.NULL_ACTION;
        }
        if (regionInfo.isMetaRegion() || regionInfo.getTable().isSystemTable()) {
          regionCacheRatioOnOldServerMap.remove(regionEncodedName);
          return BalanceAction.NULL_ACTION;
        }
        int regionIndex = cluster.regionsToIndex.get(regionInfo);
        int oldServerIndex = cluster.serversToIndex
          .get(regionCacheRatioOnOldServerMap.get(regionEncodedName).getFirst().getAddress());
        if (oldServerIndex < 0) {
          LOG.warn("Server previously hosting region {} not found", regionEncodedName);
          regionCacheRatioOnOldServerMap.remove(regionEncodedName);
          return BalanceAction.NULL_ACTION;
        }

        float oldRegionCacheRatio =
          cluster.getOrComputeRegionCacheRatio(regionIndex, oldServerIndex);
        int currentServerIndex = cluster.regionIndexToServerIndex[regionIndex];
        float currentRegionCacheRatio =
          cluster.getOrComputeRegionCacheRatio(regionIndex, currentServerIndex);

        BalanceAction action = generatePlan(cluster, regionIndex, currentServerIndex,
          currentRegionCacheRatio, oldServerIndex, oldRegionCacheRatio);
        regionCacheRatioOnOldServerMap.remove(regionEncodedName);
        return action;
      }
      return generatePlanForFreeCacheSpace(cluster);
    }

    private BalanceAction generatePlanForFreeCacheSpace(BalancerClusterState cluster) {
      if (cluster.serverBlockCacheFreeSize == null) {
        return BalanceAction.NULL_ACTION;
      }
      List<BalanceAction> possibleActions = new ArrayList<>();
      Map<Integer, Long> serverFreeCacheAfterAction = new HashMap<>();
      for (int region = 0; region < cluster.numRegions; region++) {
        RegionInfo regionInfo = cluster.regions[region];
        if (regionInfo.isMetaRegion() || regionInfo.getTable().isSystemTable()) {
          continue;
        }
        int currentServer = cluster.regionIndexToServerIndex[region];
        float ratio = cluster.getSumRegionCacheAndColdDataRatio(region);
        if (ratio >= lowCacheRatioThreshold) {
          continue;
        }
        int regionSizeMb = cluster.getRegionSizeMinusColdDataMB(region);
        if (regionSizeMb <= 0) {
          continue;
        }
        long bytesNeeded = (long) (regionSizeMb * 1024L * 1024L * minFreeCacheSpaceFactor);
        for (int server = 0; server < cluster.numServers; server++) {
          // Skips current server for region, as we can't generate a move to same server
          if (server == currentServer) {
            continue;
          }
          serverFreeCacheAfterAction.putIfAbsent(server, cluster.serverBlockCacheFreeSize[server]);
          if (serverFreeCacheAfterAction.get(server) >= bytesNeeded) {
            serverFreeCacheAfterAction.compute(server, (s, freeCache) -> freeCache - bytesNeeded);
            possibleActions.add(getAction(currentServer, region, server, -1));
          }
        }
      }
      if (!possibleActions.isEmpty()) {
        BalanceAction action =
          possibleActions.get(ThreadLocalRandom.current().nextInt(possibleActions.size()));
        LOG.debug("region {} had sum ratio {}",
          cluster.regions[((MoveRegionAction) action).getRegion()].getEncodedName(),
          cluster.getSumRegionCacheAndColdDataRatio(((MoveRegionAction) action).getRegion()));
        return action;
      }
      return BalanceAction.NULL_ACTION;
    }

    private BalanceAction generatePlan(BalancerClusterState cluster, int regionIndex,
      int currentServerIndex, float cacheRatioOnCurrentServer, int oldServerIndex,
      float cacheRatioOnOldServer) {
      return moveRegionToOldServer(cluster, regionIndex, currentServerIndex,
        cacheRatioOnCurrentServer, oldServerIndex, cacheRatioOnOldServer)
          ? getAction(currentServerIndex, regionIndex, oldServerIndex, -1)
          : generatePlanForFreeCacheSpace(cluster);
    }

    private boolean moveRegionToOldServer(BalancerClusterState cluster, int regionIndex,
      int currentServerIndex, float cacheRatioOnCurrentServer, int oldServerIndex,
      float cacheRatioOnOldServer) {
      // Find if the region has already moved by comparing the current server index with the
      // current server index. This can happen when other candidate generator has moved the region
      if (currentServerIndex < 0 || oldServerIndex < 0) {
        return false;
      }

      float cacheRatioDiffThreshold = 0.6f;

      // Conditions for moving the region

      // If the region is fully cached on the old server, move the region back
      if (cacheRatioOnOldServer == 1.0f) {
        if (LOG.isDebugEnabled()) {
          LOG.debug("Region {} moved to the old server {} as it is fully cached there",
            cluster.regions[regionIndex].getEncodedName(), cluster.servers[oldServerIndex]);
        }
        return true;
      }

      // Move the region back to the old server if it is cached equally on both the servers
      if (cacheRatioOnCurrentServer == cacheRatioOnOldServer) {
        if (LOG.isDebugEnabled()) {
          LOG.debug(
            "Region {} moved from {} to {} as the region is cached {} equally on both servers",
            cluster.regions[regionIndex].getEncodedName(), cluster.servers[currentServerIndex],
            cluster.servers[oldServerIndex], cacheRatioOnCurrentServer);
        }
        return true;
      }

      // If the region is not fully cached on either of the servers, move the region back to the
      // old server if the region cache ratio on the current server is still much less than the old
      // server
      if (
        cacheRatioOnOldServer > 0.0f
          && cacheRatioOnCurrentServer / cacheRatioOnOldServer < cacheRatioDiffThreshold
      ) {
        if (LOG.isDebugEnabled()) {
          LOG.debug(
            "Region {} moved from {} to {} as region cache ratio {} is better than the current "
              + "cache ratio {}",
            cluster.regions[regionIndex].getEncodedName(), cluster.servers[currentServerIndex],
            cluster.servers[oldServerIndex], cacheRatioOnOldServer, cacheRatioOnCurrentServer);
        }
        return true;
      }

      if (LOG.isDebugEnabled()) {
        LOG.debug(
          "Region {} not moved from {} to {} with current cache ratio {} and old cache ratio {}",
          cluster.regions[regionIndex], cluster.servers[currentServerIndex],
          cluster.servers[oldServerIndex], cacheRatioOnCurrentServer, cacheRatioOnOldServer);
      }
      return false;
    }
  }

  private class CacheAwareSkewnessCandidateGenerator extends LoadCandidateGenerator {
    @Override
    BalanceAction pickRandomRegions(BalancerClusterState cluster, int thisServer, int otherServer) {
      simulatedRatio = BigDecimal.ZERO;
      // First move all the regions which were hosted previously on some other server back to their
      // old servers
      if (
        !regionCacheRatioOnOldServerMap.isEmpty()
          && regionCacheRatioOnOldServerMap.entrySet().iterator().hasNext()
      ) {
        // Get the first region index in the historical cache ratio list
        Map.Entry<String, Pair<ServerName, Float>> regionEntry =
          regionCacheRatioOnOldServerMap.entrySet().iterator().next();
        String regionEncodedName = regionEntry.getKey();

        RegionInfo regionInfo = getRegionInfoByEncodedName(cluster, regionEncodedName);
        if (regionInfo == null) {
          LOG.warn("Region {} does not exist", regionEncodedName);
          regionCacheRatioOnOldServerMap.remove(regionEncodedName);
          return BalanceAction.NULL_ACTION;
        }
        if (regionInfo.isMetaRegion() || regionInfo.getTable().isSystemTable()) {
          regionCacheRatioOnOldServerMap.remove(regionEncodedName);
          return BalanceAction.NULL_ACTION;
        }

        int regionIndex = cluster.regionsToIndex.get(regionInfo);

        // Get the current host name for this region
        thisServer = cluster.regionIndexToServerIndex[regionIndex];

        // Get the old server index
        otherServer = cluster.serversToIndex.get(regionEntry.getValue().getFirst().getAddress());

        regionCacheRatioOnOldServerMap.remove(regionEncodedName);

        if (otherServer < 0) {
          // The old server has been moved to other host and hence, the region cannot be moved back
          // to the old server
          if (LOG.isDebugEnabled()) {
            LOG.debug(
              "CacheAwareSkewnessCandidateGenerator: Region {} not moved to the old "
                + "server {} as the server does not exist",
              regionEncodedName, regionEntry.getValue().getFirst().getHostname());
          }
          return BalanceAction.NULL_ACTION;
        }

        if (LOG.isDebugEnabled()) {
          LOG.debug(
            "CacheAwareSkewnessCandidateGenerator: Region {} moved from {} to {} as it "
              + "was hosted there earlier",
            regionEncodedName, cluster.servers[thisServer].getHostname(),
            cluster.servers[otherServer].getHostname());
        }

        return getAction(thisServer, regionIndex, otherServer, -1);
      }

      if (thisServer < 0 || otherServer < 0) {
        return BalanceAction.NULL_ACTION;
      }

      int regionIndexToMove = pickLeastCachedRegion(cluster, thisServer);
      if (regionIndexToMove < 0) {
        if (LOG.isDebugEnabled()) {
          LOG.debug("CacheAwareSkewnessCandidateGenerator: No region found for movement");
        }
        return BalanceAction.NULL_ACTION;
      }
      if (LOG.isDebugEnabled()) {
        LOG.debug(
          "CacheAwareSkewnessCandidateGenerator: Region {} moved from {} to {} as it is "
            + "least cached on current server",
          cluster.regions[regionIndexToMove].getEncodedName(),
          cluster.servers[thisServer].getHostname(), cluster.servers[otherServer].getHostname());
      }
      return getAction(thisServer, regionIndexToMove, otherServer, -1);
    }

    private int pickLeastCachedRegion(BalancerClusterState cluster, int thisServer) {
      float minCacheRatio = Float.MAX_VALUE;
      int leastCachedRegion = -1;
      for (int i = 0; i < cluster.regionsPerServer[thisServer].length; i++) {
        int regionIndex = cluster.regionsPerServer[thisServer][i];

        float cacheRatioOnCurrentServer =
          cluster.getOrComputeRegionCacheRatio(regionIndex, thisServer);
        if (cacheRatioOnCurrentServer < minCacheRatio) {
          minCacheRatio = cacheRatioOnCurrentServer;
          leastCachedRegion = regionIndex;
        }
      }
      return leastCachedRegion;
    }
  }

  static class CacheAwareRegionSkewnessCostFunction extends CostFunction {
    static final String REGION_COUNT_SKEW_COST_KEY =
      "hbase.master.balancer.stochastic.regionCountCost";
    static final float DEFAULT_REGION_COUNT_SKEW_COST = 20;
    private final DoubleArrayCost cost = new DoubleArrayCost();

    CacheAwareRegionSkewnessCostFunction(Configuration conf) {
      // Load multiplier should be the greatest as it is the most general way to balance data.
      this.setMultiplier(conf.getFloat(REGION_COUNT_SKEW_COST_KEY, DEFAULT_REGION_COUNT_SKEW_COST));
    }

    @Override
    void prepare(BalancerClusterState cluster) {
      super.prepare(cluster);
      cost.prepare(cluster.numServers);
      cost.applyCostsChange(costs -> {
        for (int i = 0; i < cluster.numServers; i++) {
          costs[i] = cluster.regionsPerServer[i].length;
        }
      });
    }

    @Override
    protected double cost() {
      return cost.cost();
    }

    @Override
    protected void regionMoved(int region, int oldServer, int newServer) {
      cost.applyCostsChange(costs -> {
        costs[oldServer] = cluster.regionsPerServer[oldServer].length;
        costs[newServer] = cluster.regionsPerServer[newServer].length;
      });
    }

    @Override
    public final void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
      weights.merge(LoadCandidateGenerator.class, cost(), Double::sum);
    }
  }

  class CacheAwareCostFunction extends CostFunction {
    private static final String CACHE_COST_KEY = "hbase.master.balancer.stochastic.cacheCost";
    private double cacheRatio;
    private double bestCacheRatio;
    private final float lowCacheRatioThreshold;
    private final float potentialCacheRatioAfterMove;
    private final float minFreeCacheSpaceFactor;

    private static final float DEFAULT_CACHE_COST = 20;

    CacheAwareCostFunction(Configuration conf) {
      boolean isPersistentCache = conf.get(BUCKET_CACHE_PERSISTENT_PATH_KEY) != null;
      // Disable the CacheAwareCostFunction if the cached file list persistence is not enabled
      this.setMultiplier(
        !isPersistentCache ? 0.0f : conf.getFloat(CACHE_COST_KEY, DEFAULT_CACHE_COST));
      bestCacheRatio = 0.0;
      cacheRatio = 0.0;
      lowCacheRatioThreshold =
        conf.getFloat(LOW_CACHE_RATIO_FOR_RELOCATION_KEY, LOW_CACHE_RATIO_FOR_RELOCATION_DEFAULT);
      potentialCacheRatioAfterMove = Math.min(1.0f, conf
        .getFloat(POTENTIAL_CACHE_RATIO_AFTER_MOVE_KEY, POTENTIAL_CACHE_RATIO_AFTER_MOVE_DEFAULT));
      minFreeCacheSpaceFactor =
        conf.getFloat(MIN_FREE_CACHE_SPACE_FACTOR_KEY, MIN_FREE_CACHE_SPACE_FACTOR_DEFAULT);
    }

    @Override
    void prepare(BalancerClusterState cluster) {
      super.prepare(cluster);
      recomputeCacheRatio(cluster);
      if (LOG.isDebugEnabled()) {
        LOG.debug("CacheAwareCostFunction: Cost: {}", 1 - cacheRatio);
      }
    }

    private void recomputeCacheRatio(BalancerClusterState cluster) {
      double[] currentWeighted = computeCurrentWeightedContributions(cluster);
      double currentSum = 0.0;
      double bestCacheSum = 0.0;
      for (int region = 0; region < cluster.numRegions; region++) {
        currentSum += currentWeighted[region];
        // here we only get the server index where this region cache ratio is the highest
        int serverIndexBestCache = cluster.getOrComputeServerWithBestRegionCachedRatio()[region];
        // get the highest cacheRatio for this region on the current state of allocations
        double currentHighestCache =
          cluster.getOrComputeWeightedRegionCacheRatio(region, serverIndexBestCache);
        // Get a hypothetical best cache ratio for this region if any server has enough free cache
        // to host it.
        double potentialHighestCache = potentialBestWeightedFromFreeCache(cluster, region);
        bestCacheSum += Math.max(currentHighestCache, potentialHighestCache);
      }
      bestCacheRatio = bestCacheSum;
      if (bestCacheSum <= 0.0) {
        cacheRatio = cluster.numRegions == 0 ? 1.0 : 0.0;
      } else {
        cacheRatio = Math.min(1.0, currentSum / bestCacheSum);
      }
    }

    private double[] computeCurrentWeightedContributions(BalancerClusterState cluster) {
      int totalRegions = cluster.numRegions;
      double[] contrib = new double[totalRegions];
      for (int r = 0; r < totalRegions; r++) {
        int s = cluster.regionIndexToServerIndex[r];
        int sizeMb = cluster.getRegionSizeMinusColdDataMB(r);
        if (sizeMb <= 0) {
          contrib[r] = 0.0;
          continue;
        }
        boolean movedInSimulation = cluster.initialRegionIndexToServerIndex[r] != s;
        if (
          cluster.serverBlockCacheFreeSize != null && movedInSimulation
            && cluster.getSumRegionCacheAndColdDataRatio(r) < lowCacheRatioThreshold
        ) {
          LOG.debug("Region {} is simulated moved to new server {}",
            cluster.regions[r].getEncodedName(), cluster.servers[s].getHostname());
          long bytesNeeded = (long) (sizeMb * 1024L * 1024L * minFreeCacheSpaceFactor);
          if (cluster.serverBlockCacheFreeSize[s] >= bytesNeeded) {
            contrib[r] = sizeMb * potentialCacheRatioAfterMove;
            continue;
          }
        }
        contrib[r] = cluster.getOrComputeWeightedRegionCacheRatio(r, s);
      }
      return contrib;
    }

    /*
     * If this region is cold in metrics and at least one RS (including its current host) reports
     * enough free block cache to hold it, return an optimistic weighted cache score ({@link
     * #potentialCacheRatioAfterMove} * region MB) so placement is not considered optimal solely
     * from low ratios when capacity exists somewhere in the cluster.
     */
    private double potentialBestWeightedFromFreeCache(BalancerClusterState cluster, int region) {
      float observedRatio = cluster.getSumRegionCacheAndColdDataRatio(region);
      if (observedRatio >= lowCacheRatioThreshold) {
        return 0.0;
      }
      int regionSizeMb = cluster.getRegionSizeMinusColdDataMB(region);
      if (regionSizeMb <= 0) {
        return 0.0;
      }
      long regionSizeBytes = (long) regionSizeMb * 1024L * 1024L;
      long requiredFree = (long) (regionSizeBytes * minFreeCacheSpaceFactor);
      for (int s = 0; s < cluster.numServers; s++) {
        if (cluster.serverBlockCacheFreeSize[s] >= requiredFree) {
          return regionSizeMb * potentialCacheRatioAfterMove;
        }
      }
      return 0.0;
    }

    @Override
    protected double cost() {
      return scale(0, 1, 1 - cacheRatio);
    }

    @Override
    protected void regionMoved(int region, int oldServer, int newServer) {
      double regionCacheRatioOnOldServer =
        cluster.getOrComputeWeightedRegionCacheRatio(region, oldServer);
      if (simulatedRatio.equals(BigDecimal.ZERO)) {
        double potentialCachedSizeOnNewServer =
          cluster.getRegionSizeMinusColdDataMB(region) * potentialCacheRatioAfterMove;
        boolean simulateCacheBasedOnFreeSpace =
          cluster.getOrComputeRegionCacheRatio(region, oldServer) < lowCacheRatioThreshold
            && cluster.serverBlockCacheFreeSize[newServer] >= potentialCachedSizeOnNewServer;
        double regionCacheRatioOnNewServer = simulateCacheBasedOnFreeSpace
          ? potentialCachedSizeOnNewServer
          : cluster.getOrComputeWeightedRegionCacheRatio(region, newServer);
        double cacheRatioDiff = regionCacheRatioOnNewServer - regionCacheRatioOnOldServer;
        double normalizedDelta = bestCacheRatio == 0.0 ? 0.0 : cacheRatioDiff / bestCacheRatio;
        LOG.debug(
          "simulating moving region {} using simulateCacheBasedOnFreeSpace={} "
            + "got a normalized delta of {} to be added to cacheRatio: {}",
          cluster.regions[region].getEncodedName(), simulateCacheBasedOnFreeSpace, normalizedDelta,
          cacheRatio);
        simulatedRatio = BigDecimal.valueOf(normalizedDelta);
        cacheRatio += normalizedDelta;
        if (cacheRatio < 0.0 || cacheRatio > 1.0) {
          LOG.info(
            "Recomputing cacheRatio after calculating impact of region move: \n "
              + "CacheAwareCostFunction:regionMoved:region:{}:from:{}:to:{}:"
              + "regionCacheRatioOnOldServer:{}:regionCacheRatioOnNewServer:{}:"
              + "bestRegionCacheRatio:{}:cacheRatio:{}",
            cluster.regions[region].getEncodedName(), cluster.servers[oldServer].getHostname(),
            cluster.servers[newServer].getHostname(), regionCacheRatioOnOldServer,
            regionCacheRatioOnNewServer, bestCacheRatio, cacheRatio);
          recomputeCacheRatio(cluster);
        }
      } else {
        // This means we are in an undoAction call and need to reverse the cache delta applied in
        // the region move simulation
        cacheRatio -= simulatedRatio.doubleValue();
      }
    }

    private int getServerWithBestCacheRatioForRegion(int region) {
      return cluster.getOrComputeServerWithBestRegionCachedRatio()[region];
    }

    @Override
    public void updateWeight(Map<Class<? extends CandidateGenerator>, Double> weights) {
      weights.merge(LoadCandidateGenerator.class, cost(), Double::sum);
    }
  }
}