alluxio-master-2.x/core/server/proxy/src/main/java/alluxio/proxy/s3/MultipartUploadCleaner.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.proxy.s3;

import alluxio.AlluxioURI;
import alluxio.client.file.FileSystem;
import alluxio.client.file.URIStatus;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.AlluxioException;
import alluxio.exception.FileDoesNotExistException;
import alluxio.grpc.DeletePOptions;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * A lazy method (not scan the whole fileSystem to find tmp directory) to
 * support automatic abortion s3 multipartUpload after a timeout.
 */
public class MultipartUploadCleaner {
  private static final Logger LOG = LoggerFactory.getLogger(MultipartUploadCleaner.class);

  private final long mTimeout;
  private final int mRetry;
  private final long mRetryDelay;
  private final ScheduledThreadPoolExecutor mExecutor;
  private final ConcurrentHashMap<AbortTask, ScheduledFuture<?>> mTasks;

  private static volatile MultipartUploadCleaner sInstance = null;

  /**
   * Creates a new instance of {@link MultipartUploadCleaner}.
   */
  private MultipartUploadCleaner() {
    mTimeout =
        Configuration.getMs(PropertyKey.PROXY_S3_MULTIPART_UPLOAD_CLEANER_TIMEOUT);
    mRetry =
        Configuration.getInt(PropertyKey.PROXY_S3_MULTIPART_UPLOAD_CLEANER_RETRY_COUNT);
    mRetryDelay =
        Configuration.getMs(PropertyKey.PROXY_S3_MULTIPART_UPLOAD_CLEANER_RETRY_DELAY);
    mExecutor = new ScheduledThreadPoolExecutor(
        Configuration.getInt(PropertyKey.PROXY_S3_MULTIPART_UPLOAD_CLEANER_POOL_SIZE));
    mTasks = new ConcurrentHashMap<>();
  }

  /**
   *@return instance of {@link MultipartUploadCleaner}
   */
  public static MultipartUploadCleaner getInstance() {
    if (sInstance == null) {
      synchronized (MultipartUploadCleaner.class) {
        if (sInstance == null) {
          sInstance = new MultipartUploadCleaner();
        }
      }
    }
    return sInstance;
  }

  /**
   * Shutdown cleaner.
   */
  public static void shutdown() {
    if (sInstance != null) {
      synchronized (MultipartUploadCleaner.class) {
        if (sInstance != null) {
          sInstance.mExecutor.shutdownNow();
          sInstance = null;
        }
      }
    }
  }

  /**
   * Schedule a task to clean multipart upload.
   *
   * @param metaFs instance of {@link FileSystem} - used for metadata operations
   * @param userFs instance of {@link FileSystem} - under the scope of a user agent
   * @param bucket bucket name
   * @param object object name
   * @param uploadId multipart upload tmp directory fileId
   * @return true if add a abort task
   */
  public static boolean apply(final FileSystem metaFs, final FileSystem userFs,
                              final String bucket, final String object, String uploadId)
      throws IOException, AlluxioException {
    final MultipartUploadCleaner cleaner = getInstance();
    // Use schedule pool do everything.
    return cleaner.apply(new AbortTask(metaFs, userFs, bucket, object, uploadId), 0);
  }

  /**
   * Schedule a task to clean multipart upload.
   *
   * @param task abort multipart upload task
   * @param delay delay time
   * @return true if add a abort task
   */
  private boolean apply(AbortTask task, long delay) {
    final ScheduledFuture<?> f = mExecutor.schedule(task, delay, TimeUnit.MILLISECONDS);
    mTasks.put(task, f);
    return true;
  }

  /**
   * Cancel schedule task.
   *
   * @param metaFs instance of {@link FileSystem} - used for metadata operations
   * @param userFs instance of {@link FileSystem} - under the scope of a user agent
   * @param bucket bucket name
   * @param object object name
   * @param uploadId multipart upload tmp directory fileId
   */
  public static void cancelAbort(final FileSystem metaFs, final FileSystem userFs,
                                 final String bucket, final String object, final String uploadId) {
    final MultipartUploadCleaner cleaner = getInstance();
    AbortTask task = new AbortTask(metaFs, userFs, bucket, object, uploadId);
    if (cleaner.containsTaskRecord(task)) {
      ScheduledFuture<?> f = cleaner.removeTaskRecord(task);
      if (f != null) {
        f.cancel(false);
      }
    }
  }

  /**
   * Remove finished task.
   */
  private ScheduledFuture<?> removeTaskRecord(AbortTask task) {
    return mTasks.remove(task);
  }

  /**
   * Check task if already exist.
   */
  private boolean containsTaskRecord(AbortTask task) {
    return mTasks.containsKey(task);
  }

  /**
   * Detect if abort task can retry.
   *
   * @param task abort task
   * @return if task can retry
   */
  private boolean canRetry(AbortTask task) {
    return task.mRetryCount < mRetry;
  }

  /**
   * Get retry delay.
   *
   * @return delay time
   */
  public long getRetryDelay() {
    return mRetryDelay;
  }

  /**
   * Try to abort a multipart upload if it was timeout.
   *
   * @param metaFs instance of {@link FileSystem} - used for metadata operations
   * @param userFs instance of {@link FileSystem} - under the scope of a user agent
   * @param bucket the bucket name
   * @param object the object name
   * @param uploadId multipart upload tmp directory fileId
   * @return delay time, non-positive values indicate to not retry this method
   */
  public long tryAbortMultipartUpload(FileSystem metaFs, FileSystem userFs,
                                      String bucket, String object, String uploadId)
      throws IOException, AlluxioException {
    final String bucketPath = S3RestUtils.parsePath(AlluxioURI.SEPARATOR + bucket);
    final AlluxioURI multipartTempDirUri = new AlluxioURI(
        S3RestUtils.getMultipartTemporaryDirForObject(bucketPath, object, uploadId));
    try {
      URIStatus status = S3RestUtils.checkStatusesForUploadId(metaFs, userFs,
          multipartTempDirUri, uploadId).get(0);
      // Check if multipart upload has exceeded its timeout
      final long curTime = System.currentTimeMillis();
      long delay = status.getLastModificationTimeMs() + mTimeout - curTime;
      if (delay > 0) { return delay; }
      // Abort the multipart upload
      userFs.delete(multipartTempDirUri, DeletePOptions.newBuilder().setRecursive(true).build());
      metaFs.delete(new AlluxioURI(S3RestUtils.getMultipartMetaFilepathForUploadId(uploadId)),
          DeletePOptions.newBuilder().build());
      LOG.info("Timeout exceeded, aborting multipart upload (bucket {}: object: {}, uploadId: {}).",
          object, bucket, uploadId);
    } catch (FileDoesNotExistException e) {
      return 0; // do not retry, multipart upload has been completed/aborted already
    }
    return 0; // do not retry, multipart upload has been completed/aborted by this method
  }

  /**
   * Abort Multipart upload task.
   */
  public static class AbortTask implements Runnable {
    private final FileSystem mMetaFs;
    private final FileSystem mUserFs;
    private final String mBucket;
    private final String mObject;
    private final String mUploadId;
    private int mRetryCount;
    private final MultipartUploadCleaner mCleaner;

    /**
     * Creates a new instance of {@link AbortTask}.
     *
     * @param metaFs instance of {@link FileSystem} - used for metadata operations
     * @param userFs instance of {@link FileSystem} - under the scope of a user agent
     * @param bucket the bucket name
     * @param object the object name
     * @param uploadId multipart upload tmp directory fileId
     */
    public AbortTask(final FileSystem metaFs, final FileSystem userFs, final String bucket,
                     final String object, final String uploadId) {
      mMetaFs = metaFs;
      mUserFs = userFs;
      mBucket = bucket;
      mObject = object;
      mUploadId = uploadId;
      mRetryCount = 0;
      mCleaner = getInstance();
    }

    @Override
    public void run() {
      try {
        long delay = mCleaner.tryAbortMultipartUpload(mMetaFs, mUserFs,
            mBucket, mObject, mUploadId);
        if (delay > 0) {
          mCleaner.apply(this, delay);
        } else {
          mCleaner.removeTaskRecord(this);
        }
      } catch (IOException | AlluxioException e) {
        mRetryCount++;
        LOG.error("Failed to abort multipart upload (bucket: {}, object: {}, uploadId: {}) "
                + "after {} retries with error {}.", mObject, mBucket, mUploadId, mRetryCount, e);
        e.printStackTrace();
        if (mCleaner.canRetry(this)) {
          mCleaner.apply(this, mCleaner.getRetryDelay());
        }
      }
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }
      AbortTask abortTask = (AbortTask) o;
      return mBucket.equals(abortTask.mBucket)
          && mObject.equals(abortTask.mObject)
          && mUploadId.equals(abortTask.mUploadId);
    }

    @Override
    public int hashCode() {
      return Objects.hash(mBucket, mObject, mUploadId);
    }

    @Override
    public String toString() {
      return "AbortTask{"
          + "mBucket='" + mBucket + '\''
          + ", mObject='" + mObject + '\''
          + ", mUploadId=" + mUploadId
          + ", mRetryCount=" + mRetryCount
          + '}';
    }
  }
}
