alluxio-master-2.x/stress/shell/src/main/java/alluxio/stress/cli/RpcBenchPreparationUtils.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.stress.cli;

import static alluxio.stress.rpc.TierAlias.MEM;

import alluxio.ClientContext;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.ConfigProperty;
import alluxio.master.MasterClientContext;
import alluxio.stress.rpc.TierAlias;
import alluxio.util.executor.ExecutorServiceFactories;
import alluxio.util.network.NetworkAddressUtils;
import alluxio.wire.WorkerNetAddress;
import alluxio.worker.block.BlockMasterClient;
import alluxio.worker.block.BlockStoreLocation;

import com.google.common.base.Supplier;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Utilities for the preparation step in RPC benchmark testing.
 */
public class RpcBenchPreparationUtils {
  private static final Logger LOG = LoggerFactory.getLogger(RpcBenchPreparationUtils.class);

  public static final long CAPACITY = 20L * 1024 * 1024 * 1024; // 20GB
  public static final Map<String, Long> CAPACITY_MEM = ImmutableMap.of(MEM.toString(), CAPACITY);
  public static final Map<String, Long> USED_MEM_EMPTY = ImmutableMap.of(MEM.toString(), 0L);
  public static final BlockStoreLocation BLOCK_LOCATION_MEM =
      new BlockStoreLocation(MEM.toString(), 0, MEM.toString());
  public static final Map<String, List<String>> LOST_STORAGE =
      ImmutableMap.of(MEM.toString(), ImmutableList.of());
  public static final List<ConfigProperty> EMPTY_CONFIG = ImmutableList.of();

  private RpcBenchPreparationUtils() {}

  /**
   * Prepare all relevant block IDs on the master side concurrently.
   *
   * @param locToBlocks a map from block location to block IDs
   */
  public static void prepareBlocksInMaster(Map<BlockStoreLocation, List<Long>> locToBlocks)
      throws InterruptedException {
    // Since the task is I/O bound, set concurrency set to 4x CPUs
    int concurrency = Runtime.getRuntime().availableProcessors() * 4;
    // Partition the wanted block IDs to smaller jobs in order to utilize concurrency
    List<List<Long>> jobs = new ArrayList<>();
    long totalBlocks = 0;
    for (Map.Entry<BlockStoreLocation, List<Long>> e : locToBlocks.entrySet()) {
      List<Long> v = e.getValue();
      totalBlocks += v.size();
      jobs.addAll(Lists.partition(v, Math.min(v.size() / concurrency, 1_000)));
    }
    final long totalBlocksFinal = totalBlocks;

    LOG.info("Split block ID generation into {} jobs", jobs.size());
    for (List<Long> job : jobs) {
      LOG.debug("Block ids: [{},{}]", job.get(0), job.get(job.size() - 1));
    }
    ExecutorService pool =
        ExecutorServiceFactories.fixedThreadPool("rpc-bench-prepare", concurrency).create();

    long blockSize = Configuration.getBytes(PropertyKey.USER_BLOCK_SIZE_BYTES_DEFAULT);
    CompletableFuture[] futures = new CompletableFuture[jobs.size()];
    AtomicInteger progress = new AtomicInteger(0);
    for (int i = 0; i < jobs.size(); i++) {
      List<Long> job = jobs.get(i);
      final int batchIndex = i;
      final int batchSize = job.size();
      CompletableFuture<Void> future = CompletableFuture.supplyAsync((Supplier<Void>) () -> {
        BlockMasterClient client =
            new BlockMasterClient(MasterClientContext
                .newBuilder(ClientContext.create())
                .build());
        for (Long blockId : job) {
          try {
            client.commitBlockInUfs(blockId, blockSize);
          } catch (IOException e) {
            LOG.error("Failed to commitBlockInUfs in batch {}, blockId={} total={}",
                batchIndex, blockId, totalBlocksFinal, e);
          }
        }
        long finishedCount = progress.addAndGet(batchSize);
        LOG.info("Generated {}th batch of {} blocks, {}% completed",
            batchIndex, batchSize, String.format("%.2f", 100.0 * finishedCount / totalBlocksFinal));
        return null;
      }, pool);
      futures[i] = (future);
    }

    LOG.info("Collect all results");
    try {
      CompletableFuture.allOf(futures).join();
    } finally {
      pool.shutdownNow();
      pool.awaitTermination(30, TimeUnit.SECONDS);
    }
  }

  /**
   * Reserves a list of worker IDs from the master.
   */
  static Deque<Long> prepareWorkerIds(BlockMasterClient client, int numWorkers) throws IOException {
    Deque<Long> workerPool = new ArrayDeque<>();
    int freePort = 40000;
    for (int i = 0; i < numWorkers; i++) {
      LOG.info("Preparing worker {}", i);

      // Prepare a worker ID
      // The addresses are different in order to get different worker IDs
      long workerId;
      String hostname = NetworkAddressUtils.getLocalHostName(500);
      LOG.info("Detected local hostname {}", hostname);
      WorkerNetAddress address = new WorkerNetAddress().setHost(hostname)
              .setDataPort(freePort++)
              .setRpcPort(freePort++)
              .setWebPort(freePort++);
      workerId = client.getId(address);
      LOG.info("Created worker ID {} on {}", workerId, address);
      workerPool.offer(workerId);
    }
    return workerPool;
  }

  static void registerWorkers(BlockMasterClient client, Deque<Long> workerIds) throws IOException {
    for (long w : workerIds) {
      LOG.info("Worker {} registering", w);
      client.register(w,
          ImmutableList.of(MEM.toString()),
          CAPACITY_MEM,
          USED_MEM_EMPTY,
          ImmutableMap.of(BLOCK_LOCATION_MEM, ImmutableList.of()),
          LOST_STORAGE, // lost storage
          EMPTY_CONFIG); // extra config
    }
    LOG.info("All workers registered");
  }

  /**
   * Generates block IDs according to the storage tier/dir setup.
   * In order to avoid block ID colliding with existing blocks, this will generate IDs
   * decreasingly from the {@link Long#MAX_VALUE}.
   *
   * @param tiersConfig the tier/dir block counts
   * @return a map of location to generated block lists
   */
  public static Map<BlockStoreLocation, List<Long>> generateBlockIdOnTiers(
          Map<TierAlias, List<Integer>> tiersConfig) {
    Map<BlockStoreLocation, List<Long>> blockMap = new HashMap<>();

    long blockIdStart = Long.MAX_VALUE;
    for (Map.Entry<TierAlias, List<Integer>> tierConfig : tiersConfig.entrySet()) {
      List<Integer> dirConfigs = tierConfig.getValue();
      for (int i = 0; i < dirConfigs.size(); i++) {
        int dirNumBlocks = dirConfigs.get(i);
        LOG.info("Found dir on tier {} with {} blocks", tierConfig.getKey(), dirNumBlocks);
        BlockStoreLocation loc = new BlockStoreLocation(
            tierConfig.getKey().toString(), i, tierConfig.getKey().toString());
        List<Long> blockIds = generateDecreasingNumbers(blockIdStart, dirNumBlocks);
        blockMap.put(loc, blockIds);
        blockIdStart -= dirNumBlocks;
      }
    }
    return blockMap;
  }

  private static List<Long> generateDecreasingNumbers(long start, int count) {
    LOG.info("Generating block Ids ({}, {}]", start - count, start);
    List<Long> list = new ArrayList<>(count);
    for (long i = 0; i < count; i++) {
      list.add(start - i);
    }
    return list;
  }
}
