alluxio-master-2.x/core/client/fs/src/main/java/alluxio/client/block/policy/CapacityBasedDeterministicHashPolicy.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.client.block.policy;

import alluxio.client.block.BlockWorkerInfo;
import alluxio.client.block.policy.options.GetWorkerOptions;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.wire.WorkerNetAddress;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.Streams;
import org.apache.commons.codec.digest.MurmurHash3;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import java.util.TreeMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;

/**
 * A policy that pseudo-randomly distributes blocks between workers according to their capacity,
 * so that the probability a worker is chosen is equal to the ratio of its capacity over total
 * capacity of all workers, provided that the blocks requested follow a uniform distribution.
 * If sharding is disabled, the same block is always assigned to the same worker. If sharding
 * is enabled, the block is assigned to a fixed set of workers.
 *
 * The target worker is determined by the following algorithm:
 * 1. build a cumulative distribution function by adding up all workers and their capacities.
 *    workers are sorted by their host name alphabetically.
 *    if worker A has 90 GB, B has 10 GB and C has 900 GB, the CDF looks like
 *    | 0 ... 90 | 90 ... 100 | 100 ... 1000 |
 *    | worker A |  worker B  |   worker C   |
 * 2. find a fixed starting point in [0, totalCapacity) determined by the hashed block id.
 *    | 0 ... 90 | 90 ... 100 | 100 ... 1000 |
 *    | worker A |  worker B  |   worker C   |
 *                     ^ start = 95
 * 3. find the corresponding worker in the CDF.
 *    which is worker B in this example
 * 4. if #shards = 1, this worker is selected. otherwise, find a set of candidates:
 *    4.1 hashed_block_id(0) = block id
 *    4.2 for i in [1, #shards], hashed_block_id(i) = hash(hashed_block_id(i-1))
 *    4.3 find the worker whose position corresponds to hashed_block_id(i) in the CDF,
 *        and add it to the candidates set
 *    4.4 repeat 4.2 - 4.4
 * 5. select a random worker in the candidate set
 *
 * The difference between this policy and {@link CapacityBaseRandomPolicy} is that this policy
 * uses the hashed block ID as the index to choose the target worker, so that the same block is
 * always routed to the same set of workers.
 *
 * Both this policy and {@link DeterministicHashPolicy} choose workers based the hashed block ID.
 * The difference is that {@link DeterministicHashPolicy} uniformly distributes the blocks among
 * the configured number of shards, while this policy chooses workers based on a distribution of
 * their normalized capacity.
 *
 * @see CapacityBaseRandomPolicy
 * @see DeterministicHashPolicy
 */
public class CapacityBasedDeterministicHashPolicy implements BlockLocationPolicy {
  private final int mShards;

  /**
   * Constructor required by
   * {@link BlockLocationPolicy.Factory#create(Class, AlluxioConfiguration)}.
   * @param conf Alluxio configuration
   */
  public CapacityBasedDeterministicHashPolicy(AlluxioConfiguration conf) {
    int numShards =
        conf.getInt(PropertyKey.USER_UFS_BLOCK_READ_LOCATION_POLICY_DETERMINISTIC_HASH_SHARDS);
    Preconditions.checkArgument(numShards >= 1, "number of shards must be no less than 1");
    mShards = numShards;
  }

  @Override
  public Optional<WorkerNetAddress> getWorker(GetWorkerOptions options) {
    TreeMap<Long, BlockWorkerInfo> capacityCdf = new TreeMap<>();
    AtomicLong totalCapacity = new AtomicLong(0);
    Streams.stream(options.getBlockWorkerInfos())
        .filter(workerInfo -> workerInfo.getCapacityBytes() >= options.getBlockInfo().getLength())
        // sort by hostname to guarantee two workers with the same capacity has a defined order
        .sorted(Comparator.comparing(w -> w.getNetAddress().getHost()))
        .forEach(workerInfo -> {
          capacityCdf.put(totalCapacity.get(), workerInfo);
          totalCapacity.getAndAdd(workerInfo.getCapacityBytes());
        });
    if (totalCapacity.get() == 0 || capacityCdf.isEmpty()) {
      return Optional.empty();
    }
    long blockId = options.getBlockInfo().getBlockId();
    BlockWorkerInfo chosenWorker = pickWorker(capacityCdf, blockId, totalCapacity.get());
    return Optional.of(chosenWorker.getNetAddress());
  }

  private BlockWorkerInfo pickWorker(TreeMap<Long, BlockWorkerInfo> capacityCdf,
      long blockId, long totalCapacity) {
    if (mShards == 1) {
      // if no sharding, simply return the worker corresponding to the start point
      long startPoint = Math.abs(hashBlockId(blockId)) % totalCapacity;
      return capacityCdf.floorEntry(startPoint).getValue();
    }
    long hashedBlockId = blockId;
    List<BlockWorkerInfo> candidates = new ArrayList<>();
    for (int i = 1; i <= Math.min(mShards, capacityCdf.size()); i++) {
      hashedBlockId = hashBlockId(hashedBlockId);
      BlockWorkerInfo candidate = capacityCdf
          .floorEntry(Math.abs(hashedBlockId) % totalCapacity) // non-null as capacities >= 0
          .getValue();
      candidates.add(candidate);
    }
    return getRandomCandidate(candidates);
  }

  @VisibleForTesting
  protected long hashBlockId(long blockId) {
    return MurmurHash3.hash64(blockId);
  }

  @VisibleForTesting
  protected BlockWorkerInfo getRandomCandidate(List<BlockWorkerInfo> candidates) {
    int randomIndex = ThreadLocalRandom.current().nextInt(candidates.size());
    return candidates.get(randomIndex);
  }
}
