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.job.plan.load;
  
  import alluxio.AlluxioURI;
  import alluxio.Constants;
  import alluxio.client.block.BlockWorkerInfo;
  import alluxio.client.file.URIStatus;
  import alluxio.collections.Pair;
  import alluxio.exception.status.FailedPreconditionException;
  import alluxio.job.RunTaskContext;
  import alluxio.job.SelectExecutorsContext;
  import alluxio.job.plan.AbstractVoidPlanDefinition;
  import alluxio.job.plan.load.LoadDefinition.LoadTask;
  import alluxio.job.util.JobUtils;
  import alluxio.job.util.SerializableVoid;
  import alluxio.util.CommonUtils;
  import alluxio.wire.BlockLocation;
  import alluxio.wire.FileBlockInfo;
  import alluxio.wire.TieredIdentity.LocalityTier;
  import alluxio.wire.WorkerInfo;
  import alluxio.wire.WorkerNetAddress;
  
  import com.google.common.base.MoreObjects;
  import com.google.common.collect.LinkedListMultimap;
  import com.google.common.collect.Lists;
  import com.google.common.collect.Multimap;
  import com.google.common.collect.Sets;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.Serializable;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.stream.Collectors;
  import javax.annotation.concurrent.NotThreadSafe;
  
  /**
   * A simple loading job that loads the blocks of a file in a distributed and round-robin fashion.
   */
  @NotThreadSafe
  public final class LoadDefinition
      extends AbstractVoidPlanDefinition<LoadConfig, ArrayList<LoadTask>> {
    private static final Logger LOG = LoggerFactory.getLogger(LoadDefinition.class);
    private static final int MAX_BUFFER_SIZE = 500 * Constants.MB;
    private static final int JOBS_PER_WORKER = 10;
  
    /**
     * Constructs a new {@link LoadDefinition}.
     */
    public LoadDefinition() {
    }
  
    @Override
    public Set<Pair<WorkerInfo, ArrayList<LoadTask>>> selectExecutors(LoadConfig config,
        List<WorkerInfo> jobWorkerInfoList, SelectExecutorsContext context)
        throws Exception {
      Map<String, WorkerInfo> jobWorkersByAddress = jobWorkerInfoList.stream()
          .collect(Collectors.toMap(info -> info.getAddress().getHost(), info -> info));
      // Filter out workers which have no local job worker available.
      List<String> missingJobWorkerHosts = new ArrayList<>();
      List<BlockWorkerInfo> workers = new ArrayList<>();
      for (BlockWorkerInfo worker : context.getFsContext().getCachedWorkers()) {
        if (jobWorkersByAddress.containsKey(worker.getNetAddress().getHost())) {
          String workerHost = worker.getNetAddress().getHost().toUpperCase();
          if (!isEmptySet(config.getExcludedWorkerSet())
              && config.getExcludedWorkerSet().contains(workerHost)) {
            continue;
          }
          // If specified the locality id, the candidate worker must match one at least
          boolean match = false;
          if (worker.getNetAddress().getTieredIdentity().getTiers() != null) {
            if (!(isEmptySet(config.getLocalityIds())
                && isEmptySet(config.getExcludedLocalityIds()))) {
              boolean exclude = false;
              for (LocalityTier tier : worker.getNetAddress().getTieredIdentity().getTiers()) {
                if (!isEmptySet(config.getExcludedLocalityIds())
                    && config.getExcludedLocalityIds().contains(tier.getValue().toUpperCase())) {
                  exclude = true;
                  break;
                }
                if (!isEmptySet(config.getLocalityIds())
                    && config.getLocalityIds().contains(tier.getValue().toUpperCase())) {
                  match = true;
                  break;
                }
              }
              if (exclude) {
                continue;
              }
            }
          }
          // Add current worker as candidate worker
          // if it matches the given locality id or contained in the given worker set
          // Or user specified neither worker-set nor locality id
          if ((isEmptySet(config.getWorkerSet()) && isEmptySet(config.getLocalityIds()))
              || match
              || (!isEmptySet(config.getWorkerSet())
                  && config.getWorkerSet().contains(workerHost))) {
            workers.add(worker);
          }
        } else {
          LOG.warn("Worker on host {} has no local job worker", worker.getNetAddress().getHost());
          missingJobWorkerHosts.add(worker.getNetAddress().getHost());
        }
      }
      // Mapping from worker to block ids which that worker is supposed to load.
      Multimap<WorkerInfo, LoadTask> assignments = LinkedListMultimap.create();
      AlluxioURI uri = new AlluxioURI(config.getFilePath());
      for (FileBlockInfo blockInfo : context.getFileSystem().getStatus(uri).getFileBlockInfos()) {
        List<BlockWorkerInfo> workersWithoutBlock = getWorkersWithoutBlock(workers, blockInfo);
        int neededReplicas = config.getReplication() - blockInfo.getBlockInfo().getLocations().size();
        if (workersWithoutBlock.size() < neededReplicas) {
          String missingJobWorkersMessage = "";
          if (!missingJobWorkerHosts.isEmpty()) {
            missingJobWorkersMessage = ". The following workers could not be used because they have "
                + "no local job workers: " + missingJobWorkerHosts;
          }
          throw new FailedPreconditionException(String.format(
              "Failed to find enough block workers to replicate to. Needed %s but only found %s. "
              + "Available workers without the block: %s" + missingJobWorkersMessage,
              neededReplicas, workersWithoutBlock.size(), workersWithoutBlock));
        }
        Collections.shuffle(workersWithoutBlock);
        for (int i = 0; i < neededReplicas; i++) {
          String address = workersWithoutBlock.get(i).getNetAddress().getHost();
          WorkerInfo jobWorker = jobWorkersByAddress.get(address);
          assignments.put(jobWorker, new LoadTask(blockInfo.getBlockInfo().getBlockId(),
              workersWithoutBlock.get(i).getNetAddress(), config.getFilePath()));
        }
      }
  
      Set<Pair<WorkerInfo, ArrayList<LoadTask>>> result = Sets.newHashSet();
      for (Map.Entry<WorkerInfo, Collection<LoadTask>> assignment : assignments.asMap().entrySet()) {
        Collection<LoadTask> loadTasks = assignment.getValue();
        List<List<LoadTask>> partitionedTasks =
            CommonUtils.partition(Lists.newArrayList(loadTasks), JOBS_PER_WORKER);
  
        for (List<LoadTask> tasks : partitionedTasks) {
          if (!tasks.isEmpty()) {
            result.add(new Pair<>(assignment.getKey(), Lists.newArrayList(tasks)));
          }
        }
      }
  
      return result;
    }
  
    /*
     * @param blockWorkers a list of block workers
     * @param blockInfo information about a block
     * @return the block worker hosts which are not storing the specified block
     */
    private List<BlockWorkerInfo> getWorkersWithoutBlock(List<BlockWorkerInfo> blockWorkers,
        FileBlockInfo blockInfo) {
      Set<WorkerNetAddress> blockLocations = blockInfo.getBlockInfo().getLocations().stream()
          .map(BlockLocation::getWorkerAddress)
          .collect(Collectors.toSet());
  
      return blockWorkers.stream()
          .filter(worker -> !blockLocations.contains(worker.getNetAddress()))
          .collect(Collectors.toList());
    }
  
    @Override
    public SerializableVoid runTask(LoadConfig config, ArrayList<LoadTask> tasks,
        RunTaskContext context) throws Exception {
  
      // TODO(jiri): Replace with internal client that uses file ID once the internal client is
      // factored out of the core server module. The reason to prefer using file ID for this job is
      // to avoid the the race between "replicate" and "rename", so that even a file to replicate is
      // renamed, the job is still working on the correct file.
      URIStatus status = context.getFileSystem().getStatus(new AlluxioURI(config.getFilePath()));
  
      for (LoadTask task : tasks) {
        JobUtils.loadBlock(status, context.getFsContext(), task.getBlockId(),
            task.getWorkerNetAddress(), config.isDirectCache());
        LOG.info("Loaded file " + config.getFilePath() + " block " + task.getBlockId());
      }
      return null;
    }
  
    private boolean isEmptySet(Set s) {
      return s == null || s.isEmpty();
    }
  
    /**
     * A task representing loading a block into the memory of a worker.
     */
    public static class LoadTask implements Serializable {
      private static final long serialVersionUID = 2028545900913354425L;
      final long mBlockId;
      final String mFilePath;
      final WorkerNetAddress mWorkerNetAddress;
  
      /**
       * @param blockId the id of the block to load
       * @param workerNetAddress worker net address
       * @param filePath file path for task
       */
      public LoadTask(long blockId, WorkerNetAddress workerNetAddress, String filePath) {
        mBlockId = blockId;
        mWorkerNetAddress = workerNetAddress;
        mFilePath = filePath;
      }
  
      /**
       * @return the block id
       */
      public long getBlockId() {
        return mBlockId;
      }
  
      /**
       * @return the file path
       */
      public String getFilePath() {
        return mFilePath;
      }
  
      /**
       * @return worker net address
       */
      public WorkerNetAddress getWorkerNetAddress() {
        return mWorkerNetAddress;
      }
  
      @Override
      public String toString() {
        return MoreObjects.toStringHelper(this)
            .add("FilePath", mFilePath) // first letter has to be Upper case for parsing
            .add("blockId", mBlockId)
            .add("workerNetAddress", mWorkerNetAddress)
            .toString();
      }
    }
  
    @Override
    public Class<LoadConfig> getJobConfigClass() {
      return LoadConfig.class;
    }
  }
  