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

import alluxio.AlluxioURI;
import alluxio.Constants;
import alluxio.client.WriteType;
import alluxio.client.block.BlockWorkerInfo;
import alluxio.client.file.FileInStream;
import alluxio.client.file.FileOutStream;
import alluxio.client.file.FileSystem;
import alluxio.client.file.URIStatus;
import alluxio.collections.Pair;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.FileAlreadyExistsException;
import alluxio.grpc.CreateFilePOptions;
import alluxio.grpc.OpenFilePOptions;
import alluxio.grpc.ReadPType;
import alluxio.grpc.WritePType;
import alluxio.job.RunTaskContext;
import alluxio.job.SelectExecutorsContext;
import alluxio.job.plan.AbstractVoidPlanDefinition;
import alluxio.job.util.JobUtils;
import alluxio.job.util.SerializableVoid;
import alluxio.util.io.PathUtils;
import alluxio.wire.WorkerInfo;

import com.google.common.base.Preconditions;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;

/**
 * A job that migrates a single file from source path to a destination path.
 * The source and the destination can be across mount points.
 */
public final class MigrateDefinition
    extends AbstractVoidPlanDefinition<MigrateConfig, MigrateCommand> {
  private static final Logger LOG = LoggerFactory.getLogger(MigrateDefinition.class);

  private static final int JOBS_PER_WORKER = 10;

  private final Random mRandom = new Random();

  /**
   * Constructs a new {@link MigrateDefinition}.
   */
  public MigrateDefinition() {
  }

  private void checkMigrateValid(MigrateConfig config, FileSystem fs) throws Exception {
    AlluxioURI source = new AlluxioURI(config.getSource());
    AlluxioURI destination = new AlluxioURI(config.getDestination());
    // The source cannot be a prefix of the destination -
    // that would be migrating a path inside itself.
    if (PathUtils.hasPrefix(destination.toString(), source.toString())) {
      throw new RuntimeException(ExceptionMessage.MIGRATE_CANNOT_BE_TO_SUBDIRECTORY.getMessage(
          source, config.getDestination()));
    }
  }

  /**
   * {@inheritDoc}
   *
   * Assigns each worker to migrate whichever files it has the most blocks for.
   * If no worker has blocks for a file, a random worker is chosen.
   * @return
   */
  @Override
  public Set<Pair<WorkerInfo, MigrateCommand>> selectExecutors(MigrateConfig config,
      List<WorkerInfo> jobWorkerInfoList, SelectExecutorsContext context) throws Exception {
    AlluxioURI source = new AlluxioURI(config.getSource());
    AlluxioURI destination = new AlluxioURI(config.getDestination());
    if (source.equals(destination)) {
      return Sets.newHashSet();
    }
    checkMigrateValid(config, context.getFileSystem());
    Preconditions.checkState(!jobWorkerInfoList.isEmpty(), "No workers are available");

    URIStatus status = context.getFileSystem().getStatus(source);
    ConcurrentMap<String, WorkerInfo> hostnameToWorker = Maps.newConcurrentMap();
    for (WorkerInfo workerInfo : jobWorkerInfoList) {
      hostnameToWorker.put(workerInfo.getAddress().getHost(), workerInfo);
    }
    List<BlockWorkerInfo> alluxioWorkerInfoList = context.getFsContext().getCachedWorkers();
    if (status.isFolder()) {
      throw new RuntimeException("Cannot migrate directory");
    } else {
      WorkerInfo bestJobWorker =
          getBestJobWorker(status, alluxioWorkerInfoList, jobWorkerInfoList, hostnameToWorker);
      Set<Pair<WorkerInfo, MigrateCommand>> result = Sets.newHashSet();
      result.add(new Pair<>(bestJobWorker,
          new MigrateCommand(status.getPath(), destination.getPath())));
      return result;
    }
  }

  private WorkerInfo getBestJobWorker(URIStatus status, List<BlockWorkerInfo> alluxioWorkerInfoList,
      List<WorkerInfo> jobWorkerInfoList, Map<String, WorkerInfo> hostnameToJobWorker) {
    BlockWorkerInfo bestWorker =
        JobUtils.getWorkerWithMostBlocks(alluxioWorkerInfoList, status.getFileBlockInfos());
    if (bestWorker == null) {
      // Nobody has blocks, choose a random worker.
      bestWorker = alluxioWorkerInfoList.get(mRandom.nextInt(jobWorkerInfoList.size()));
    }
    // Map the best Alluxio worker to a job worker.
    WorkerInfo worker = hostnameToJobWorker.get(bestWorker.getNetAddress().getHost());
    if (worker == null) {
      return jobWorkerInfoList.get(new Random().nextInt(jobWorkerInfoList.size()));
    }
    return worker;
  }

  /**
   * {@inheritDoc}
   *
   * Migrates the file specified in the config to the configured path. If the destination path is a
   * directory, the file is migrated inside that directory.
   */
  @Override
  public SerializableVoid runTask(MigrateConfig config, MigrateCommand command,
      RunTaskContext context) throws Exception {
    WriteType writeType = config.getWriteType() == null
        ? Configuration.getEnum(PropertyKey.USER_FILE_WRITE_TYPE_DEFAULT, WriteType.class)
        : config.getWriteType();
    migrate(command, writeType.toProto(), context.getFileSystem(), config.isOverwrite());
    return null;
  }

  /*
   * @param command the migrate command to execute
   * @param writeType the write type to use for the moved file
   * @param fileSystem the Alluxio file system
   * @param overwrite whether to overwrite destination
   */
  private static void migrate(MigrateCommand command,
                              WritePType writeType,
                              FileSystem fileSystem,
                              boolean overwrite) throws Exception {
    String source = command.getSource();
    String destination = command.getDestination();
    LOG.debug("Migrating {} to {}", source, destination);

    CreateFilePOptions createOptions =
        CreateFilePOptions.newBuilder().setWriteType(writeType).build();
    OpenFilePOptions openFileOptions =
        OpenFilePOptions.newBuilder().setReadType(ReadPType.NO_CACHE).build();
    AlluxioURI destinationUri = new AlluxioURI(destination);
    String tmpPath =
        PathUtils.temporaryFileName(System.currentTimeMillis(), destinationUri.toString());
    AlluxioURI tmpUri = new AlluxioURI(tmpPath);
    boolean retry;
    boolean processOverwrite = false;
    do {
      retry = false;
      try (FileInStream in = fileSystem.openFile(new AlluxioURI(source), openFileOptions);
           FileOutStream out = fileSystem.createFile(destinationUri, createOptions)) {
        try {
          IOUtils.copyLarge(in, out, new byte[8 * Constants.MB]);
          if (processOverwrite) {
            fileSystem.delete(tmpUri);
            fileSystem.rename(destinationUri, tmpUri);
          }
        } catch (Throwable t) {
          // Since we only catch CancelledException here if we get interrupted, so we want to keep
          // the following clean up process not interrupted. If we have other exception, the flag
          // doesn't matter here.
          Thread.interrupted();
          try {
            out.cancel();
            fileSystem.delete(destinationUri);
          } catch (Throwable t2) {
            t.addSuppressed(t2);
          }
          throw t;
        }
      } catch (FileAlreadyExistsException e) {
        if (overwrite) {
          AlluxioURI switchUri = destinationUri;
          destinationUri = tmpUri;
          tmpUri = switchUri;
          retry = true;
          processOverwrite = true;
        } else {
          throw e;
        }
      }
    } while (retry);
  }

  @Override
  public Class<MigrateConfig> getJobConfigClass() {
    return MigrateConfig.class;
  }
}
