content: /*
   * 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.conf.Configuration;
  import alluxio.conf.PropertyKey;
  import alluxio.exception.ExceptionMessage;
  import alluxio.exception.JobDoesNotExistException;
  import alluxio.grpc.OperationType;
  import alluxio.job.CmdConfig;
  import alluxio.job.cmd.load.LoadCliConfig;
  import alluxio.job.cmd.migrate.MigrateCliConfig;
  import alluxio.job.cmd.persist.PersistCmdConfig;
  import alluxio.job.wire.CmdStatusBlock;
  import alluxio.job.wire.SimpleJobStatusBlock;
  import alluxio.job.wire.Status;
  import alluxio.master.job.JobMaster;
  import alluxio.master.job.common.CmdInfo;
  import alluxio.master.job.plan.PlanTracker;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.Executors;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.stream.Collectors;
  import javax.annotation.concurrent.ThreadSafe;
  
  /**
   * CmdJobTracker to schedule a Cmd job to run.
   */
  @ThreadSafe
  public class CmdJobTracker implements AutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(CmdJobTracker.class);
    private final Map<Long, CmdInfo> mInfoMap = new ConcurrentHashMap<>(0, 0.95f,
        Math.max(8, 2 * Runtime.getRuntime().availableProcessors()));
    private final DistLoadCliRunner mDistLoadCliRunner;
    private final MigrateCliRunner mMigrateCliRunner;
    private final PersistRunner mPersistRunner;
    protected FileSystemContext mFsContext;
    public static final String DELIMITER = ",";
    private final ScheduledExecutorService mScheduleCleanExecutor;
    private final Long mTraceRetentionTime;
  
    private final PlanTracker mPlanTracker;
  
    /**
     * Create a new instance of {@link CmdJobTracker}.
     * @param fsContext filesystem context
     * @param jobMaster the job master
     * @param planTracker the planTracker
     */
    public CmdJobTracker(FileSystemContext fsContext,
                     JobMaster jobMaster, PlanTracker planTracker) {
      mFsContext = fsContext;
      mDistLoadCliRunner = new DistLoadCliRunner(mFsContext, jobMaster);
      mMigrateCliRunner = new MigrateCliRunner(mFsContext, jobMaster);
      mPersistRunner = new PersistRunner(mFsContext, jobMaster);
      mScheduleCleanExecutor = Executors.newSingleThreadScheduledExecutor();
      mScheduleCleanExecutor.scheduleAtFixedRate(this::
              cleanExpiredJobInfos, 60, 600, TimeUnit.SECONDS);
      mTraceRetentionTime = Configuration.getMs(
          PropertyKey.JOB_MASTER_JOB_TRACE_RETENTION_TIME);
      mPlanTracker = planTracker;
    }
  
    /**
     * Constructor with runner providers.
     * @param fsContext Filesystem context
     * @param distLoadCliRunner DistributedLoad runner
     * @param migrateCliRunner DistributedCopy runner
     * @param persistRunner Persist runner
     * @param retentionTime job retention time
     * @param planTracker the planTracker
     */
    public CmdJobTracker(FileSystemContext fsContext,
                         DistLoadCliRunner distLoadCliRunner,
                         MigrateCliRunner migrateCliRunner,
                         PersistRunner persistRunner,
                         Long retentionTime,
                         PlanTracker planTracker
    ) {
      mFsContext = fsContext;
      mDistLoadCliRunner = distLoadCliRunner;
      mMigrateCliRunner = migrateCliRunner;
      mPersistRunner = persistRunner;
      mScheduleCleanExecutor = Executors.newSingleThreadScheduledExecutor();
      mScheduleCleanExecutor.scheduleAtFixedRate(this::
              cleanExpiredJobInfos, 60, 600, TimeUnit.SECONDS);
      mTraceRetentionTime = retentionTime;
      mPlanTracker = planTracker;
    }
  
    /**
     * Run the PlanTracker to trigger a Cmd job based on CmdConfig.
     * @param cmdConfig the distributed command job config
     * @param jobControlId job control id for the command
     */
    public synchronized void run(CmdConfig cmdConfig, long jobControlId) throws
            JobDoesNotExistException, IOException {
      runDistributedCommand(cmdConfig, jobControlId);
    }
  
    private void runDistributedCommand(CmdConfig cmdConfig, long jobControlId)
            throws JobDoesNotExistException, IOException {
      CmdInfo cmdInfo;
      switch (cmdConfig.getOperationType()) {
        case DIST_LOAD:
          LoadCliConfig loadCliConfig = (LoadCliConfig) cmdConfig;
          int batchSize = loadCliConfig.getBatchSize();
          AlluxioURI filePath = new AlluxioURI(loadCliConfig.getFilePath());
          int replication = loadCliConfig.getReplication();
          Set<String> workerSet = loadCliConfig.getWorkerSet();
          Set<String> excludedWorkerSet = loadCliConfig.getExcludedWorkerSet();
          Set<String> localityIds =  loadCliConfig.getLocalityIds();
          Set<String> excludedLocalityIds = loadCliConfig.getExcludedLocalityIds();
          boolean directCache = loadCliConfig.getDirectCache();
          cmdInfo = mDistLoadCliRunner.runDistLoad(batchSize, filePath, replication, workerSet,
                  excludedWorkerSet, localityIds, excludedLocalityIds, directCache, jobControlId);
          break;
        case DIST_CP:
          MigrateCliConfig migrateCliConfig = (MigrateCliConfig) cmdConfig;
          AlluxioURI srcPath = new AlluxioURI(migrateCliConfig.getSource());
          AlluxioURI dstPath = new AlluxioURI(migrateCliConfig.getDestination());
          LOG.info("run a dist cp command, job control id: {}, cmd config: {}",
              jobControlId, cmdConfig);
          cmdInfo = mMigrateCliRunner.runDistCp(srcPath, dstPath,
              migrateCliConfig.getOverWrite(), migrateCliConfig.getBatchSize(),
              jobControlId);
          break;
        case PERSIST:
          PersistCmdConfig persistCmdConfig = (PersistCmdConfig) cmdConfig;
          cmdInfo = mPersistRunner.runPersistJob(persistCmdConfig, jobControlId);
          break;
        default:
          throw new JobDoesNotExistException(
                  ExceptionMessage.JOB_DEFINITION_DOES_NOT_EXIST.getMessage(cmdConfig.getName()));
      }
  
      mInfoMap.put(cmdInfo.getJobControlId(), cmdInfo);
    }
  
    /**
     * Get status information for a CMD.
     * @param jobControlId jobControlId to trace a CMD
     * @return the Command level status
     */
    public Status getCmdStatus(long jobControlId) throws JobDoesNotExistException {
      if (!mInfoMap.containsKey(jobControlId)) {
        throw new JobDoesNotExistException(
                ExceptionMessage.JOB_DEFINITION_DOES_NOT_EXIST.getMessage(jobControlId));
      }
  
      CmdInfo cmdInfo = mInfoMap.get(jobControlId);
  
      if (cmdInfo.getCmdRunAttempt().isEmpty()) { // If no attempts created,
        // that means the files are loaded already, set status to complete
        return Status.COMPLETED;
      }
  
      int completed = 0;
      boolean finished = true;
      boolean failed = false;
      boolean canceled = false;
      for (CmdRunAttempt attempt : cmdInfo.getCmdRunAttempt()) {
        Status s = attempt.checkJobStatus();
        if (!s.isFinished()) {
          finished = false;
          break;
        }
        if (!failed && s == Status.FAILED) {
          failed = true;
          attempt.printFailed();
        }
        if (!canceled && s == Status.CANCELED) {
          canceled = true;
        }
        if (s == Status.COMPLETED) {
          completed++;
        }
      }
  
      if (finished) {
        if (failed) {
          return Status.FAILED; // FAILED has higher priority than CANCELED
        }
        if (canceled) {
          return Status.CANCELED;
        }
        if (completed == cmdInfo.getCmdRunAttempt().size()) {
          return Status.COMPLETED;
        }
      }
  
      return Status.RUNNING;
    }
  
    /**
     * @param statusList status list filter
     * @return cmd ids matching conditions
     */
    public Set<Long> findCmdIds(List<Status> statusList) throws JobDoesNotExistException {
      Set<Long> set = new HashSet<>();
      for (Map.Entry<Long, CmdInfo> x : mInfoMap.entrySet()) {
        if (statusList.isEmpty()
                || statusList.contains(getCmdStatus(
                        x.getValue().getJobControlId()))) {
          Long key = x.getKey();
          set.add(key);
        }
      }
      return set;
    }
  
    /**
     * @return all failed file paths
     */
    public Set<String> findAllFailedPaths() {
      Set<String> set = new HashSet<>();
      for (Map.Entry<Long, CmdInfo> x : mInfoMap.entrySet()) {
        long jobControlId = x.getKey();
        try {
          set.addAll(findFailedPaths(jobControlId));
        } catch (JobDoesNotExistException e) {
          LOG.info("skip because of no such a command id.");
        }
      }
      return set;
    }
  
    /**
     * @param jobControlId jobControlId
     * @return failed file paths
     */
    public Set<String> findFailedPaths(long jobControlId) throws JobDoesNotExistException {
      if (!mInfoMap.containsKey(jobControlId)) {
        throw new JobDoesNotExistException(
                ExceptionMessage.JOB_DEFINITION_DOES_NOT_EXIST.getMessage(jobControlId));
      }
  
      CmdInfo cmdInfo = mInfoMap.get(jobControlId);
  
      if (cmdInfo.getCmdRunAttempt().isEmpty()) { // If no attempts are created, throws an Exception
        throw new JobDoesNotExistException(
                ExceptionMessage.JOB_DEFINITION_DOES_NOT_EXIST.getMessage(jobControlId));
      }
  
      return cmdInfo.getCmdRunAttempt().stream()
              .map(CmdRunAttempt::getFailedFiles).flatMap(Collection::stream)
              .collect(Collectors.toSet());
    }
  
    /**
     * Get a cmdStatusBlock information.
     * @param jobControlId command id
     * @return CmdStatusBlock
     */
    public CmdStatusBlock getCmdStatusBlock(long jobControlId)
            throws JobDoesNotExistException {
      if (!mInfoMap.containsKey(jobControlId)) {
        throw new JobDoesNotExistException(
                ExceptionMessage.JOB_DEFINITION_DOES_NOT_EXIST.getMessage(jobControlId));
      }
  
      CmdInfo cmdInfo = mInfoMap.get(jobControlId);
  
      if (cmdInfo.getCmdRunAttempt().isEmpty()) { // If no attempts are created,
        // that means the files are loaded already, set status to complete
        return new CmdStatusBlock(cmdInfo.getJobControlId(), Collections.EMPTY_LIST,
                cmdInfo.getOperationType());
      }
  
      List<SimpleJobStatusBlock> blockList = cmdInfo.getCmdRunAttempt().stream()
              .map(attempt -> new SimpleJobStatusBlock(attempt.getJobId(),
                      attempt.checkJobStatus(),
                      attempt.getFilePath(),
                      String.join(DELIMITER, attempt.getFailedFiles())))
              .collect(Collectors.toList());
      return new CmdStatusBlock(cmdInfo.getJobControlId(), blockList, cmdInfo.getOperationType());
    }
  
    private void cleanExpiredJobInfos() {
      long currentTime = System.currentTimeMillis();
      for (Map.Entry<Long, CmdInfo> x : mInfoMap.entrySet()) {
        CmdInfo cmdInfo = x.getValue();
        List<Long> cleanedJobsId = new ArrayList<>();
        if (OperationType.DIST_LOAD.equals(cmdInfo.getOperationType())
                && currentTime - cmdInfo.getJobSubmissionTime() > mTraceRetentionTime) {
          try {
            Status jobStatus = getCmdStatus(cmdInfo.getJobControlId());
            if (jobStatus.isFinished()) {
              for (CmdRunAttempt runAttempt : cmdInfo.getCmdRunAttempt()) {
                cleanedJobsId.add(runAttempt.getJobId());
              }
              mPlanTracker.removeJobs(cleanedJobsId);
              mInfoMap.remove(cmdInfo.getJobControlId());
              LOG.info("JobControlId:{} has been cleaned in CmdJobTracker,"
                  + " client will not trace the job anymore.The filePaths in CmdInfo are:{}",
                      cmdInfo.getJobControlId(), String.join(", ", cmdInfo.getFilePath()));
            }
          } catch (JobDoesNotExistException e) {
            LOG.warn("JobControlId:{} can not find in CmdJobTracker when clean expired Job"
                + "with unexpected exception.The filePaths in CmdInfo are:{}",
                    cmdInfo.getJobControlId(), String.join(", ", cmdInfo.getFilePath()));
          }
        }
      }
    }
  
    @Override
    public void close() throws Exception {
      mScheduleCleanExecutor.shutdown();
    }
  }
  