alluxio-master-2.x/job/server/src/main/java/alluxio/master/job/tracker/DistLoadCliRunner.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.master.job.tracker;

import alluxio.AlluxioURI;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.URIStatus;
import alluxio.exception.AlluxioException;
import alluxio.grpc.ListStatusPOptions;
import alluxio.grpc.OperationType;
import alluxio.job.JobConfig;
import alluxio.job.plan.BatchedJobConfig;
import alluxio.job.plan.load.LoadConfig;
import alluxio.job.wire.JobSource;
import alluxio.master.job.JobMaster;
import alluxio.master.job.common.CmdInfo;
import alluxio.master.job.metrics.DistributedCmdMetrics;
import alluxio.retry.CountingRetry;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.LongAdder;
import java.util.stream.Collectors;

/**
 * A config runner for a DistLoadCli job.
 */
public class DistLoadCliRunner extends AbstractCmdRunner {
  private static final Logger LOG = LoggerFactory.getLogger(DistLoadCliRunner.class);

  /**
   * constructor.
   * @param fsContext
   * @param jobMaster
   */
  public DistLoadCliRunner(FileSystemContext fsContext, JobMaster jobMaster) {
    super(fsContext, jobMaster);
  }

  /**
   * Run a DistLoad command.
   * @param batchSize size for batched jobs
   * @param filePath The {@link AlluxioURI} path to load into Alluxio memory
   * @param replication Number of block replicas of each loaded file
   * @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
   * @param jobControlId THe parent level job control ID
   * @return CmdInfo
   */
  public CmdInfo runDistLoad(
          int batchSize, AlluxioURI filePath,
          int replication, Set<String> workerSet,
          Set<String> excludedWorkerSet, Set<String> localityIds,
          Set<String> excludedLocalityIds, boolean directCache,
          long jobControlId) throws IOException {
    long submissionTime = System.currentTimeMillis();
    List<URIStatus> filePool = new ArrayList<>(batchSize);

    //only use the source path as the file parameter.
    List<String> path = Lists.newArrayList(filePath.getPath());

    CmdInfo cmdInfo = new CmdInfo(jobControlId, OperationType.DIST_LOAD, JobSource.CLI,
            submissionTime, path);

    try {
      load(filePath, batchSize, replication, workerSet, excludedWorkerSet, localityIds,
              excludedLocalityIds, directCache, filePool, cmdInfo);
    } catch (IOException | AlluxioException e) {
      LOG.warn(String.format("DistributedLoad job is failing for path = %s!", filePath.getPath()));
      LOG.error(e.getMessage());
      throw new IOException(e.getMessage());
    }

    // add all the jobs left in the pool.
    if (filePool.size() > 0) {
      submitDistLoad(filePool, replication, workerSet, excludedWorkerSet, localityIds,
              excludedLocalityIds, directCache, cmdInfo);
      filePool.clear();
    }
    return cmdInfo;
  }

  /**
   * Loads a file or directory in Alluxio space, makes it resident in memory.
   *
   * @param filePath The {@link AlluxioURI} path to load into Alluxio memory
   * @param replication Number of block replicas of each loaded file
   * @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 pool The pool for batched jobs
   * @param cmdInfo  Command information
   * @throws AlluxioException when Alluxio exception occurs
   * @throws IOException when non-Alluxio exception occurs
   */
  private void load(AlluxioURI filePath, int batchSize, int replication, Set<String> workerSet,
                    Set<String> excludedWorkerSet, Set<String> localityIds,
                    Set<String> excludedLocalityIds, boolean directCache,
                    List<URIStatus> pool, CmdInfo cmdInfo)
          throws IOException, AlluxioException {
    ListStatusPOptions options = ListStatusPOptions.newBuilder().setRecursive(true).build();
    LongAdder incompleteCount = new LongAdder();
    mFileSystem.iterateStatus(filePath, options, uriStatus -> {
      if (!uriStatus.isFolder()) {
        if (!uriStatus.isCompleted()) {
          incompleteCount.increment();
          System.out.printf("Ignored load because: %s is in incomplete status",
                  uriStatus.getPath());
          return;
        }
        AlluxioURI fileURI = new AlluxioURI(uriStatus.getPath());
        if (uriStatus.getInAlluxioPercentage() == 100 && replication == 1) {
          // The file has already been fully loaded into Alluxio.
          System.out.println(fileURI + " is already fully loaded in Alluxio");
          return;
        }
        pool.add(uriStatus);
        if (pool.size() == batchSize) {
          submitDistLoad(pool, replication, workerSet, excludedWorkerSet, localityIds,
                  excludedLocalityIds, directCache, cmdInfo);
          pool.clear();
        }
      }
    });
    if (incompleteCount.longValue() > 0) {
      System.out.printf("Ignore load %d paths because they are in incomplete status",
              incompleteCount.longValue());
    }
  }

  // Submit a child job within a distributed command job.
  private void submitDistLoad(List<URIStatus> pool, int replication,
                              Set<String> workerSet, Set<String> excludedWorkerSet,
                              Set<String> localityIds, Set<String> excludedLocalityIds,
                              boolean directCache, CmdInfo cmdInfo) {
    if (mSubmitted.size() >= mActiveJobs) {
      waitForCmdJob();
    }

    CmdRunAttempt attempt = new CmdRunAttempt(new CountingRetry(3), mJobMaster);
    setJobConfigAndFileMetrics(pool, replication, workerSet, excludedWorkerSet,
            localityIds, excludedLocalityIds, directCache, attempt);
    mSubmitted.add(attempt);
    cmdInfo.addCmdRunAttempt(attempt);
    attempt.run();
  }

  // Create a JobConfig and set file count and size for the DistLoad job.
  protected void setJobConfigAndFileMetrics(
          List<URIStatus> filePath, int replication,
          Set<String> workerSet, Set<String> excludedWorkerSet,
          Set<String> localityIds, Set<String> excludedLocalityIds,
          boolean directCache, CmdRunAttempt attempt) {
    JobConfig jobConfig;
    long fileCount = 0;
    long fileSize = 0;
    String filePathString = filePath.stream().map(URIStatus::getPath)
            .collect(Collectors.joining(CmdJobTracker.DELIMITER));

    if (filePath.size() == 1) {
      URIStatus status = filePath.iterator().next();
      String source = status.getPath();
      jobConfig = new LoadConfig(source, replication,
              workerSet, excludedWorkerSet, localityIds, excludedLocalityIds, directCache);
      fileCount = DEFAULT_FILE_COUNT;
      fileSize = DistributedCmdMetrics.getFileSize(source, mFileSystem, new CountingRetry(3));
    } else {
      HashSet<Map<String, String>> configs = Sets.newHashSet();
      ObjectMapper oMapper = new ObjectMapper();
      for (URIStatus status : filePath) {
        LoadConfig loadConfig = new LoadConfig(status.getPath(), replication, workerSet,
                excludedWorkerSet, localityIds, excludedLocalityIds, directCache);
        Map<String, String> map = oMapper.convertValue(loadConfig, Map.class);
        configs.add(map);
        fileSize += DistributedCmdMetrics.getFileSize(status.getPath(),
                mFileSystem, new CountingRetry(3));
      }
      fileCount = filePath.size();
      jobConfig = new BatchedJobConfig(LoadConfig.NAME, configs);
    }
    attempt.setFileCount(fileCount);
    attempt.setFileSize(fileSize);
    attempt.setConfig(jobConfig);
    attempt.setFilePath(filePathString);
  }
}

