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

import alluxio.AlluxioURI;
import alluxio.cli.fs.FileSystemShellUtils;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.exception.AlluxioException;
import alluxio.job.CmdConfig;
import alluxio.job.cmd.load.LoadCliConfig;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

/**
 * Utilities Loads a file or directory in Alluxio space, makes it resident in memory.
 */
public final class DistributedLoadUtils {
  private static final Logger LOG = LoggerFactory.getLogger(DistributedLoadUtils.class);
  private static final int DEFAULT_REPLICATION = 1;
  private static final String DEFAULT_FAILURE_FILE_PATH =
      "./logs/user/distributedLoad_%s_failures.csv";
  private static final Option REPLICATION_OPTION =
      Option.builder()
          .longOpt("replication")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .type(Number.class)
          .argName("replicas")
          .desc("Number of block replicas of each loaded file, default: " + DEFAULT_REPLICATION)
          .build();
  private static final Option ACTIVE_JOB_COUNT_OPTION =
      Option.builder()
          .longOpt("active-jobs")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .type(Number.class)
          .argName("active job count")
          .desc("Number of active jobs that can run at the same time. Later jobs must wait. "
              + "The default upper limit is "
              + AbstractDistributedJobCommand.DEFAULT_ACTIVE_JOBS)
          .build();
  private static final Option INDEX_FILE =
      Option.builder()
          .longOpt("index")
          .required(false)
          .hasArg(true)
          .numberOfArgs(0)
          .type(String.class)
          .argName("index file")
          .desc("Name of the index file that lists all files to be loaded")
          .build();
  private static final Option HOSTS_OPTION =
      Option.builder()
          .longOpt("hosts")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("hosts")
          .desc("A list of worker hosts separated by comma."
              + " When host and locality options are not set,"
              + " all hosts will be selected unless explicitly excluded by setting excluded option"
              + "('excluded-hosts', 'excluded-host-file', 'excluded-locality'"
              + " and 'excluded-locality-file')."
              + " Only one of the 'hosts' and 'host-file' should be set,"
              + " and it should not be set with excluded option together.")
          .build();
  private static final Option HOST_FILE_OPTION =
      Option.builder()
          .longOpt("host-file")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("host-file")
          .desc("Host File contains worker hosts, each line has a worker host."
              + " When host and locality options are not set,"
              + " all hosts will be selected unless explicitly excluded by setting excluded option"
              + "('excluded-hosts', 'excluded-host-file', 'excluded-locality'"
              + " and 'excluded-locality-file')."
              + " Only one of the 'hosts' and 'host-file' should be set,"
              + " and it should not be set with excluded option together.")
          .build();
  private static final Option EXCLUDED_HOSTS_OPTION =
      Option.builder()
          .longOpt("excluded-hosts")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("excluded-hosts")
          .desc("A list of excluded worker hosts separated by comma."
              + " Only one of the 'excluded-hosts' and 'excluded-host-file' should be set,"
              + " and it should not be set with 'hosts', 'host-file', 'locality'"
              + " and 'locality-file' together.")
          .build();
  private static final Option EXCLUDED_HOST_FILE_OPTION =
      Option.builder()
          .longOpt("excluded-host-file")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("excluded-host-file")
          .desc("Host File contains excluded worker hosts, each line has a worker host."
              + " Only one of the 'excluded-hosts' and 'excluded-host-file' should be set,"
              + " and it should not be set with 'hosts', 'host-file', 'locality'"
              + " and 'locality-file' together.")
          .build();
  private static final Option LOCALITY_OPTION =
      Option.builder()
          .longOpt("locality")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("locality")
          .desc("A list of worker locality separated by comma."
              + " When host and locality options are not set,"
              + " all hosts will be selected unless explicitly excluded by setting excluded option"
              + "('excluded-hosts', 'excluded-host-file', 'excluded-locality'"
              + " and 'excluded-locality-file')."
              + " Only one of the 'locality' and 'locality-file' should be set,"
              + " and it should not be set with excluded option together.")
          .build();
  private static final Option LOCALITY_FILE_OPTION =
      Option.builder()
          .longOpt("locality-file")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("locality-file")
          .argName("locality-file")
          .desc("Locality File contains worker localities, each line has a worker locality."
              + " When host and locality options are not set,"
              + " all hosts will be selected unless explicitly excluded by setting excluded option"
              + "('excluded-hosts', 'excluded-host-file', 'excluded-locality'"
              + " and 'excluded-locality-file')."
              + " Only one of the 'locality' and 'locality-file' should be set,"
              + " and it should not be set with excluded option together.")
          .build();
  private static final Option EXCLUDED_LOCALITY_OPTION =
      Option.builder()
          .longOpt("excluded-locality")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("excluded-locality")
          .desc("A list of excluded worker locality separated by comma."
              + " Only one of the 'excluded-locality' and 'excluded-locality-file' should be set,"
              + " and it should not be set with 'hosts', 'host-file', 'locality'"
              + " and 'locality-file' together.")
          .build();
  private static final Option EXCLUDED_LOCALITY_FILE_OPTION =
      Option.builder()
          .longOpt("excluded-locality-file")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .argName("excluded-locality-file")
          .desc("Locality File contains excluded worker localities,"
              + " each line has a worker locality."
              + " Only one of the 'excluded-locality' and 'excluded-locality-file' should be set,"
              + " and it should not be set with 'hosts', 'host-file', 'locality'"
              + " and 'locality-file' together.")
          .build();
  private static final Option BATCH_SIZE_OPTION =
      Option.builder()
          .longOpt("batch-size")
          .required(false)
          .hasArg(true)
          .numberOfArgs(1)
          .type(Number.class)
          .argName("batch-size")
          .desc("Number of files per request")
          .build();
  private static final Option PASSIVE_CACHE_OPTION =
      Option.builder()
          .longOpt("passive-cache")
          .required(false)
          .hasArg(false)
          .desc("Use passive-cache as the cache implementation,"
              + " turn on to use the old cache through read implementation. "
              + "Passive-cache is default when there's no option set or "
              + "both options are set for cache implementation."
              + "Notice that this flag is temporary, "
              + "and it would retire after direct cache graduate from experimental stage")
          .build();
  private static final Option DIRECT_CACHE_OPTION =
      Option.builder()
          .longOpt("direct-cache")
          .required(false)
          .hasArg(false)
          .desc("Use direct cache request as the cache implementation,"
              + " turn on to use the new cache through cache manager implementation. "
              + "Notice that this flag is temporary, "
              + "and it would retire after direct cache graduate from experimental stage")
          .build();

  private DistributedLoadUtils() {} // prevent instantiation

  /**
   * Get the general options that distributed load needs to use.
   * Suitable for file semantics and structured data semantics.
   * @return general options
   */
  public static Options getDistLoadCommonOptions() {
    return new Options()
        .addOption(REPLICATION_OPTION)
        .addOption(ACTIVE_JOB_COUNT_OPTION)
        .addOption(HOSTS_OPTION)
        .addOption(HOST_FILE_OPTION)
        .addOption(EXCLUDED_HOSTS_OPTION)
        .addOption(EXCLUDED_HOST_FILE_OPTION)
        .addOption(LOCALITY_OPTION)
        .addOption(LOCALITY_FILE_OPTION)
        .addOption(EXCLUDED_LOCALITY_OPTION)
        .addOption(EXCLUDED_LOCALITY_FILE_OPTION)
        .addOption(PASSIVE_CACHE_OPTION)
        .addOption(DIRECT_CACHE_OPTION)
        .addOption(BATCH_SIZE_OPTION)
        .addOption(AbstractDistributedJobCommand.ASYNC_OPTION);
  }

  /**
   * Get the options of distributed load needs to use in file semantics.
   * @return options in file semantics
   */
  public static Options getDistLoadFileSemanticsOptions() {
    return getDistLoadCommonOptions()
        .addOption(INDEX_FILE);
  }

  /**
   * Get the general usage of distributed load.
   * Suitable for file semantics and structured data semantics.
   * @return usage message
   */
  public static String getDistLoadCommonUsage() {
    return "[--replication <num>] [--active-jobs <num>] [--batch-size <num>] "
        + "[--host-file <hostFilePath>] "
        + "[--excluded-hosts <host1>,<host2>,...,<hostN>] [--excluded-host-file <hostFilePath>] "
        + "[--locality <locality1>,<locality2>,...,<localityN>] "
        + "[--locality-file <localityFilePath>] "
        + "[--excluded-locality <locality1>,<locality2>,...,<localityN>] "
        + "[--excluded-locality-file <localityFilePath>] "
        + "[--passive-cache] "
        + "[--direct-cache] ";
  }

  /**
   * Get usage of distributed load in file semantics.
   * @return usage message
   */
  public static String getDistLoadFileSemanticsUsage() {
    return getDistLoadCommonUsage()
        + "[--index] [--hosts <host1>,<host2>,...,<hostN>] ";
  }

  /**
   * Run the actual distributedLoad command.
   * @param command Distributed Load command instance
   * @param cl CommandLine instance of command
   * @param pathStr path to be loaded
   * @return the result of running the command
   */
  public static int distributedLoad(AbstractDistributedJobCommand command, CommandLine cl,
                                     String pathStr)
      throws AlluxioException, IOException  {
    AlluxioConfiguration conf = command.mFsContext.getClusterConf();
    int defaultBatchSize = conf.getInt(PropertyKey.JOB_REQUEST_BATCH_SIZE);
    int replication = FileSystemShellUtils.getIntArg(cl, REPLICATION_OPTION, DEFAULT_REPLICATION);
    int batchSize = FileSystemShellUtils.getIntArg(cl, BATCH_SIZE_OPTION, defaultBatchSize);
    boolean directCache = !cl.hasOption(PASSIVE_CACHE_OPTION.getLongOpt()) && cl.hasOption(
        DIRECT_CACHE_OPTION.getLongOpt());
    boolean async = cl.hasOption(AbstractDistributedJobCommand.ASYNC_OPTION.getLongOpt());
    if (async) {
      System.out.println("Entering async submission mode. ");
    }

    Set<String> workerSet = new HashSet<>();
    Set<String> excludedWorkerSet = new HashSet<>();
    Set<String> localityIds = new HashSet<>();
    Set<String> excludedLocalityIds = new HashSet<>();
    if (cl.hasOption(HOST_FILE_OPTION.getLongOpt())) {
      String hostFile = cl.getOptionValue(HOST_FILE_OPTION.getLongOpt()).trim();
      readLinesToSet(workerSet, hostFile);
    } else if (cl.hasOption(HOSTS_OPTION.getLongOpt())) {
      String argOption = cl.getOptionValue(HOSTS_OPTION.getLongOpt()).trim();
      readItemsFromOptionString(workerSet, argOption);
    }
    if (cl.hasOption(EXCLUDED_HOST_FILE_OPTION.getLongOpt())) {
      String hostFile = cl.getOptionValue(EXCLUDED_HOST_FILE_OPTION.getLongOpt()).trim();
      readLinesToSet(excludedWorkerSet, hostFile);
    } else if (cl.hasOption(EXCLUDED_HOSTS_OPTION.getLongOpt())) {
      String argOption = cl.getOptionValue(EXCLUDED_HOSTS_OPTION.getLongOpt()).trim();
      readItemsFromOptionString(excludedWorkerSet, argOption);
    }
    if (cl.hasOption(LOCALITY_FILE_OPTION.getLongOpt())) {
      String localityFile = cl.getOptionValue(LOCALITY_FILE_OPTION.getLongOpt()).trim();
      readLinesToSet(localityIds, localityFile);
    } else if (cl.hasOption(LOCALITY_OPTION.getLongOpt())) {
      String argOption = cl.getOptionValue(LOCALITY_OPTION.getLongOpt()).trim();
      readItemsFromOptionString(localityIds, argOption);
    }
    if (cl.hasOption(EXCLUDED_LOCALITY_FILE_OPTION.getLongOpt())) {
      String localityFile = cl.getOptionValue(EXCLUDED_LOCALITY_FILE_OPTION.getLongOpt()).trim();
      readLinesToSet(excludedLocalityIds, localityFile);
    } else if (cl.hasOption(EXCLUDED_LOCALITY_OPTION.getLongOpt())) {
      String argOption = cl.getOptionValue(EXCLUDED_LOCALITY_OPTION.getLongOpt()).trim();
      readItemsFromOptionString(excludedLocalityIds, argOption);
    }

    System.out.println("Please wait for command submission to finish..");

    long jobControlId;
    if (!cl.hasOption(INDEX_FILE.getLongOpt())) {
      AlluxioURI path = new AlluxioURI(pathStr);
      jobControlId = DistributedLoadUtils.runDistLoad(command, path, replication, batchSize,
          workerSet, excludedWorkerSet, localityIds, excludedLocalityIds, directCache);
      if (!async) {
        System.out.format("Submitted successfully, jobControlId = %s%n"
            + "Waiting for the command to finish ...%n", jobControlId);
        command.waitForCmd(jobControlId);
        command.postProcessing(jobControlId);
      } else {
        System.out.format("Submitted distLoad job successfully, jobControlId = %s%n",
            jobControlId);
      }
    } else {
      try (BufferedReader reader = new BufferedReader(new FileReader(pathStr))) {
        for (String filename; (filename = reader.readLine()) != null;) {
          AlluxioURI path = new AlluxioURI(filename);
          jobControlId = DistributedLoadUtils.runDistLoad(command, path, replication, batchSize,
              workerSet, excludedWorkerSet, localityIds, excludedLocalityIds, directCache);
          if (!async) {
            System.out.format("Submitted successfully, jobControlId = %s%n"
                + "Waiting for the command to finish ...%n", jobControlId);
            command.waitForCmd(jobControlId);
            command.postProcessing(jobControlId);
          } else {
            System.out.format("Submitted distLoad job successfully, jobControlId = %s%n",
                jobControlId);
          }
        }
      }
    }

    Set<String> failures = command.getFailedFiles();
    if (failures.size() > 0) {
      command.processFailures(pathStr, failures, DEFAULT_FAILURE_FILE_PATH);
    }
    return 0;
  }

  /**
   * Run the actual distributedLoad command.
   * @param command Distributed Load command instance
   * @param filePath file path to load
   * @param replication Number of block replicas of each loaded file
   * @param batchSize Batch size for loading
   * @param workerSet A set of worker hosts to load data
   * @param excludedWorkerSet A set of worker hosts can not to load data
   * @param localityIds The locality identify set
   * @param excludedLocalityIds A set of worker locality identify can not to load data
   * @param directCache use direct cache request or cache through read
   * @return job Control ID
   */
  public static long runDistLoad(AbstractDistributedJobCommand command, AlluxioURI filePath,
                                     int replication, int batchSize, Set<String> workerSet,
                                     Set<String> excludedWorkerSet, Set<String> localityIds,
                                     Set<String> excludedLocalityIds, boolean directCache) {
    CmdConfig cmdConfig = new LoadCliConfig(filePath.getPath(), batchSize, replication, workerSet,
        excludedWorkerSet, localityIds, excludedLocalityIds, directCache);
    return command.submit(cmdConfig);
  }

  private static void readItemsFromOptionString(Set<String> localityIds, String argOption) {
    for (String locality : StringUtils.split(argOption, ",")) {
      locality = locality.trim().toUpperCase();
      if (!locality.isEmpty()) {
        localityIds.add(locality);
      }
    }
  }

  private static void readLinesToSet(Set<String> workerSet, String hostFile) throws IOException {
    try (BufferedReader reader = new BufferedReader(new FileReader(hostFile))) {
      for (String worker; (worker = reader.readLine()) != null;) {
        worker = worker.trim().toUpperCase();
        if (!worker.isEmpty()) {
          workerSet.add(worker);
        }
      }
    }
  }
}
