alluxio-master-2.x/shell/src/main/java/alluxio/cli/fs/command/LoadCommand.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.annotation.PublicApi;
import alluxio.cli.CommandUtils;
import alluxio.client.block.BlockStoreClient;
import alluxio.client.block.policy.BlockLocationPolicy;
import alluxio.client.block.stream.BlockInStream;
import alluxio.client.block.stream.BlockWorkerClient;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.URIStatus;
import alluxio.client.file.options.InStreamOptions;
import alluxio.collections.Pair;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.exception.AlluxioException;
import alluxio.exception.status.InvalidArgumentException;
import alluxio.grpc.CacheRequest;
import alluxio.grpc.JobProgressReportFormat;
import alluxio.grpc.LoadJobPOptions;
import alluxio.grpc.OpenFilePOptions;
import alluxio.job.JobDescription;
import alluxio.job.LoadJobRequest;
import alluxio.proto.dataserver.Protocol;
import alluxio.resource.CloseableResource;
import alluxio.util.FileSystemOptionsUtils;
import alluxio.util.FormatUtils;
import alluxio.wire.BlockInfo;
import alluxio.wire.WorkerNetAddress;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;

import java.io.IOException;
import java.util.List;
import java.util.Optional;
import java.util.OptionalLong;
import javax.annotation.concurrent.ThreadSafe;

/**
 * Loads a file or directory in Alluxio space, making it resident in Alluxio.
 */
@ThreadSafe
@PublicApi
public final class LoadCommand extends AbstractFileSystemCommand {
  private static final JobProgressReportFormat DEFAULT_FORMAT = JobProgressReportFormat.TEXT;
  private static final String JOB_TYPE = "load";
  private static final Option LOCAL_OPTION =
      Option.builder()
          .longOpt("local")
          .required(false)
          .hasArg(false)
          .desc("load the file to local worker.")
          .build();
  private static final Option SUBMIT_OPTION = Option.builder()
      .longOpt("submit")
      .required(false)
      .hasArg(false)
      .desc("Submit load job to Alluxio master, update job options if already exists.")
      .build();

  private static final Option STOP_OPTION = Option.builder()
      .longOpt("stop")
      .required(false)
      .hasArg(false)
      .desc("Stop a load job if it's still running.")
      .build();

  private static final Option PROGRESS_OPTION = Option.builder()
      .longOpt("progress")
      .required(false)
      .hasArg(false)
      .desc("Get progress report of a load job.")
      .build();

  private static final Option PARTIAL_LISTING_OPTION = Option.builder()
      .longOpt("partial-listing")
      .required(false)
      .hasArg(false)
      .desc("Use partial directory listing. This limits the memory usage "
          + "and starts load sooner for larger directory. But progress "
          + "report cannot report on the total number of files because the "
          + "whole directory is not listed yet.")
      .build();

  private static final Option VERIFY_OPTION = Option.builder()
      .longOpt("verify")
      .required(false)
      .hasArg(false)
      .desc("Run verification when load finish and load new files if any.")
      .build();

  private static final Option BANDWIDTH_OPTION = Option.builder()
      .longOpt("bandwidth")
      .required(false)
      .hasArg(true)
      .desc("Single worker read bandwidth limit.")
      .build();

  private static final Option PROGRESS_FORMAT = Option.builder()
      .longOpt("format")
      .required(false)
      .hasArg(true)
      .desc("Format of the progress report, supports TEXT and JSON. If not "
          + "set, TEXT is used.")
      .build();

  private static final Option PROGRESS_VERBOSE = Option.builder()
      .longOpt("verbose")
      .required(false)
      .hasArg(false)
      .desc("Whether to return a verbose progress report with detailed errors")
      .build();

  /**
   * Constructs a new instance to load a file or directory in Alluxio space.
   *
   * @param fsContext the filesystem of Alluxio
   */
  public LoadCommand(FileSystemContext fsContext) {
    super(fsContext);
  }

  @Override
  public String getCommandName() {
    return "load";
  }

  @Override
  public Options getOptions() {
    return new Options()
        .addOption(BANDWIDTH_OPTION)
        .addOption(PARTIAL_LISTING_OPTION)
        .addOption(VERIFY_OPTION)
        .addOption(SUBMIT_OPTION)
        .addOption(STOP_OPTION)
        .addOption(PROGRESS_OPTION)
        .addOption(PROGRESS_FORMAT)
        .addOption(PROGRESS_VERBOSE)
        .addOption(LOCAL_OPTION);
  }

  @Override
  public int run(CommandLine cl) throws AlluxioException, IOException {
    String[] args = cl.getArgs();
    AlluxioURI path = new AlluxioURI(args[0]);
    if (isOldFormat(cl)) {
      runWildCardCmd(path, cl);
      return 0;
    }

    if (path.containsWildcard()) {
      throw new UnsupportedOperationException("Load does not support wildcard path");
    }

    if (cl.hasOption(SUBMIT_OPTION.getLongOpt())) {
      OptionalLong bandwidth = OptionalLong.empty();
      if (cl.hasOption(BANDWIDTH_OPTION.getLongOpt())) {
        bandwidth = OptionalLong.of(FormatUtils.parseSpaceSize(
            cl.getOptionValue(BANDWIDTH_OPTION.getLongOpt())));
      }
      return submitLoad(
          path,
          bandwidth,
          cl.hasOption(PARTIAL_LISTING_OPTION.getLongOpt()),
          cl.hasOption(VERIFY_OPTION.getLongOpt()));
    }

    if (cl.hasOption(STOP_OPTION.getLongOpt())) {
      return stopLoad(path);
    }
    JobProgressReportFormat format = DEFAULT_FORMAT;
    if (cl.hasOption(PROGRESS_OPTION.getLongOpt())) {
      if (cl.hasOption(PROGRESS_FORMAT.getLongOpt())) {
        format = JobProgressReportFormat.valueOf(cl.getOptionValue(PROGRESS_FORMAT.getLongOpt()));
      }
      return getProgress(path, format, cl.hasOption(PROGRESS_VERBOSE.getLongOpt()));
    }

    return 0;
  }

  @Override
  public String getUsage() {
    return "For backward compatibility: load [--local] <path>\n"
        + "For distributed load:\n"
        + "\tload <path> --submit [--bandwidth N] [--verify] [--partial-listing]\n"
        + "\tload <path> --stop\n"
        + "\tload <path> --progress [--format TEXT|JSON] [--verbose]\n";
  }

  @Override
  public String getDescription() {
    return "Loads a file or directory in Alluxio space, makes it resident in Alluxio.";
  }

  @Override
  public void validateArgs(CommandLine cl) throws InvalidArgumentException {
    CommandUtils.checkNumOfArgsNoLessThan(this, cl, 1);
    if (!isOldFormat(cl)) {
      int commands = 0;
      if (cl.hasOption(SUBMIT_OPTION.getLongOpt())) {
        commands++;
      }
      if (cl.hasOption(STOP_OPTION.getLongOpt())) {
        commands++;
      }
      if (cl.hasOption(PROGRESS_OPTION.getLongOpt())) {
        commands++;
      }
      if (commands != 1) {
        throw new InvalidArgumentException("Must have one of submit / stop / progress");
      }
    }
  }

  private int submitLoad(AlluxioURI path, OptionalLong bandwidth,
      boolean usePartialListing, boolean verify) {
    LoadJobPOptions.Builder options = alluxio.grpc.LoadJobPOptions
        .newBuilder().setPartialListing(usePartialListing).setVerify(verify);
    if (bandwidth.isPresent()) {
      options.setBandwidth(bandwidth.getAsLong());
    }
    LoadJobRequest job = new LoadJobRequest(path.getPath(), options.build());
    try {
      Optional<String> jobId = mFileSystem.submitJob(job);
      if (jobId.isPresent()) {
        System.out.printf("Load '%s' is successfully submitted. JobId: %s%n", path, jobId.get());
      } else {
        System.out.printf("Load already running for path '%s', updated the job with "
                + "new bandwidth: %s, verify: %s%n",
            path,
            bandwidth.isPresent() ? String.valueOf(bandwidth.getAsLong()) : "unlimited",
            verify);
      }
      return 0;
    } catch (StatusRuntimeException e) {
      System.out.println("Failed to submit load job " + path + ": " + e.getMessage());
      return -1;
    }
  }

  private int stopLoad(AlluxioURI path) {
    try {
      if (mFileSystem.stopJob(JobDescription
          .newBuilder()
          .setPath(path.getPath())
          .setType(JOB_TYPE)
          .build())) {
        System.out.printf("Load '%s' is successfully stopped.%n", path);
      }
      else {
        System.out.printf("Cannot find load job for path %s, it might have already been "
            + "stopped or finished%n", path);
      }
      return 0;
    } catch (StatusRuntimeException e) {
      System.out.println("Failed to stop load job " + path + ": " + e.getMessage());
      return -1;
    }
  }

  private int getProgress(AlluxioURI path, JobProgressReportFormat format,
      boolean verbose) {
    try {
      System.out.println("Progress for loading path '" + path + "':");
      System.out.println(mFileSystem.getJobProgress(JobDescription
          .newBuilder()
          .setPath(path.getPath())
          .setType(JOB_TYPE)
          .build(), format, verbose));
      return 0;
    } catch (StatusRuntimeException e) {
      if (e.getStatus().getCode() == Status.Code.NOT_FOUND) {
        System.out.println("Load for path '" + path + "' cannot be found.");
        return -2;
      }
      System.out.println("Failed to get progress for load job " + path + ": " + e.getMessage());
      return -1;
    }
  }

  private boolean isOldFormat(CommandLine cl) {
    return cl.getOptions().length == 0
        || (cl.getOptions().length == 1 && cl.hasOption(LOCAL_OPTION.getLongOpt()));
  }

  @Override
  protected void runPlainPath(AlluxioURI plainPath, CommandLine cl)
      throws AlluxioException, IOException {
    Preconditions.checkState(
        isOldFormat(cl),
        "The new load command should not hit this code path");
    oldLoad(plainPath, cl.hasOption(LOCAL_OPTION.getLongOpt()));
  }

  /**
   * Loads a file or directory in Alluxio space, makes it resident in Alluxio.
   *
   * @param filePath The {@link AlluxioURI} path to load into Alluxio
   * @param local whether to load data to local worker even when the data is already loaded remotely
   */
  private void oldLoad(AlluxioURI filePath, boolean local)
      throws AlluxioException, IOException {
    URIStatus status = mFileSystem.getStatus(filePath);
    if (status.isFolder()) {
      List<URIStatus> statuses = mFileSystem.listStatus(filePath);
      for (URIStatus uriStatus : statuses) {
        AlluxioURI newPath = new AlluxioURI(uriStatus.getPath());
        oldLoad(newPath, local);
      }
    } else {
      if (local) {
        if (!mFsContext.hasNodeLocalWorker()) {
          System.out.println(
              "When local option is specified, there must be a local worker available");
          return;
        }
      } else if (status.getInAlluxioPercentage() == 100) {
        // The file has already been fully loaded into Alluxio.
        System.out.println(filePath + " already in Alluxio fully");
        return;
      }
      runLoadTask(filePath, status, local);
    }
    System.out.println(filePath + " loaded");
  }

  private void runLoadTask(AlluxioURI filePath, URIStatus status, boolean local)
      throws IOException {
    AlluxioConfiguration conf = mFsContext.getPathConf(filePath);
    OpenFilePOptions options = FileSystemOptionsUtils.openFileDefaults(conf);
    BlockLocationPolicy policy = Preconditions.checkNotNull(
        BlockLocationPolicy.Factory
            .create(conf.getClass(PropertyKey.USER_UFS_BLOCK_READ_LOCATION_POLICY), conf),
        "UFS read location policy Required when loading files");
    WorkerNetAddress dataSource;
    List<Long> blockIds = status.getBlockIds();
    for (long blockId : blockIds) {
      if (local) {
        dataSource = mFsContext.getNodeLocalWorker();
      } else { // send request to data source
        BlockStoreClient blockStore = BlockStoreClient.create(mFsContext);
        Pair<WorkerNetAddress, BlockInStream.BlockInStreamSource> dataSourceAndType = blockStore
            .getDataSourceAndType(status.getBlockInfo(blockId), status, policy, ImmutableMap.of());
        dataSource = dataSourceAndType.getFirst();
      }
      Protocol.OpenUfsBlockOptions openUfsBlockOptions =
          new InStreamOptions(status, options, conf, mFsContext).getOpenUfsBlockOptions(blockId);
      BlockInfo info = status.getBlockInfo(blockId);
      long blockLength = info.getLength();
      String host = dataSource.getHost();
      // issues#11172: If the worker is in a container, use the container hostname
      // to establish the connection.
      if (!dataSource.getContainerHost().equals("")) {
        host = dataSource.getContainerHost();
      }
      CacheRequest request = CacheRequest.newBuilder().setBlockId(blockId).setLength(blockLength)
          .setOpenUfsBlockOptions(openUfsBlockOptions).setSourceHost(host)
          .setSourcePort(dataSource.getDataPort()).build();
      try (CloseableResource<BlockWorkerClient> blockWorker =
          mFsContext.acquireBlockWorkerClient(dataSource)) {
        blockWorker.get().cache(request);
      } catch (Exception e) {
        throw new RuntimeException(String.format("Failed to complete cache request from %s for "
            + "block %d of file %s: %s", dataSource, blockId, status.getPath(), e), e);
      }
    }
  }
}
