alluxio.underfs.tos.TOSLowLevelOutputStream Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of alluxio-underfs-tos Show documentation
Tinder Object Storage Under File System implementation
The newest version!
/*
 * 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.underfs.tos;

import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.underfs.ObjectLowLevelOutputStream;

import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.volcengine.tos.TOSV2;
import com.volcengine.tos.TosException;
import com.volcengine.tos.comm.io.TosRepeatableBoundedFileInputStream;
import com.volcengine.tos.model.object.AbortMultipartUploadInput;
import com.volcengine.tos.model.object.CompleteMultipartUploadV2Input;
import com.volcengine.tos.model.object.CompleteMultipartUploadV2Output;
import com.volcengine.tos.model.object.CreateMultipartUploadInput;
import com.volcengine.tos.model.object.CreateMultipartUploadOutput;
import com.volcengine.tos.model.object.PutObjectInput;
import com.volcengine.tos.model.object.UploadPartV2Input;
import com.volcengine.tos.model.object.UploadPartV2Output;
import com.volcengine.tos.model.object.UploadedPartV2;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import javax.annotation.Nullable;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * Object storage low output stream for TOS.
 */
@NotThreadSafe
public class TOSLowLevelOutputStream extends ObjectLowLevelOutputStream {
  private static final Logger LOG = LoggerFactory.getLogger(TOSLowLevelOutputStream.class);

  /**
   * The TOS client to interact with TOS.
   */
  protected TOSV2 mClient;
  /**
   * Tags for the uploaded part, provided by TOS after uploading.
   */
  private final List mTags = Collections.synchronizedList(new ArrayList<>());

  /**
   * The upload id of this multipart upload.
   */
  protected volatile String mUploadId;

  private String mContentHash;

  /**
   * Constructs a new stream for writing a file.
   *
   * @param bucketName the name of the bucket
   * @param key        the key of the file
   * @param tosClient  the TOS client to upload the file with
   * @param executor   a thread pool executor
   * @param ufsConf    the object store under file system configuration
   */
  public TOSLowLevelOutputStream(String bucketName, String key, TOSV2 tosClient,
                                 ListeningExecutorService executor,
                                 AlluxioConfiguration ufsConf) {
    super(bucketName, key, executor,
        ufsConf.getBytes(PropertyKey.UNDERFS_TOS_STREAMING_UPLOAD_PARTITION_SIZE), ufsConf);
    mClient = Preconditions.checkNotNull(tosClient);
  }

  @Override
  protected void abortMultiPartUploadInternal() throws IOException {
    try {
      AbortMultipartUploadInput input = new AbortMultipartUploadInput().setBucket(mBucketName)
          .setKey(mKey).setUploadID(mUploadId);
      getClient().abortMultipartUpload(input);
    } catch (TosException e) {
      LOG.debug("failed to abort multi part upload. upload id: {}", mUploadId, e);
      throw AlluxioTosException.from(String.format(
          "failed to upload part. key: %s uploadId: %s",
          mKey, mUploadId), e);
    }
  }

  @Override
  protected void uploadPartInternal(
      File file,
      int partNumber,
      boolean isLastPart,
      @Nullable String md5)
      throws IOException {
    long fileSize = file.length();
    long partSize = mPartitionSize;
    try (FileInputStream content = new FileInputStream(file)) {
      InputStream wrappedContent = new TosRepeatableBoundedFileInputStream(content, mPartitionSize);
      if (isLastPart) {
        partSize = fileSize;
      }
      final UploadPartV2Input input = new UploadPartV2Input()
          .setBucket(mBucketName)
          .setKey(mKey)
          .setUploadID(mUploadId)
          .setPartNumber(partNumber)
          .setContentLength(partSize)
          .setContent(wrappedContent);
      UploadPartV2Output output = getClient().uploadPart(input);
      mTags.add(new UploadedPartV2().setPartNumber(partNumber).setEtag(output.getEtag()));
    } catch (IOException e) {
      LOG.debug("failed to upload part.", e);
      throw new IOException(String.format(
          "failed to upload part. key: %s part number: %s uploadId: %s",
          mKey, partNumber, mUploadId), e);
    }
  }

  @Override
  protected void initMultiPartUploadInternal() throws IOException {
    try {
      CreateMultipartUploadInput create =
          new CreateMultipartUploadInput().setBucket(mBucketName).setKey(mKey);
      CreateMultipartUploadOutput output = getClient().createMultipartUpload(create);
      mUploadId = output.getUploadID();
    } catch (TosException e) {
      LOG.debug("failed to init multi part upload", e);
      throw AlluxioTosException.from("failed to init multi part upload", e);
    }
  }

  @Override
  protected void completeMultiPartUploadInternal() {
    try {
      LOG.debug("complete multi part {}", mUploadId);
      CompleteMultipartUploadV2Input complete = new CompleteMultipartUploadV2Input()
          .setBucket(mBucketName)
          .setKey(mKey)
          .setUploadID(mUploadId);
      complete.setUploadedParts(mTags);
      CompleteMultipartUploadV2Output completedOutput =
          getClient().completeMultipartUpload(complete);
      mContentHash = completedOutput.getEtag();
    } catch (TosException e) {
      LOG.debug("failed to complete multi part upload", e);
      throw AlluxioTosException.from(
          String.format("failed to complete multi part upload, key: %s, upload id: %s",
              mKey, mUploadId), e);
    }
  }

  @Override
  protected void createEmptyObject(String key) {
    try {
      PutObjectInput putObjectInput = new PutObjectInput()
          .setBucket(mBucketName)
          .setKey(key)
          .setContent(new ByteArrayInputStream(new byte[0]))
          .setContentLength(0);
      mContentHash = getClient().putObject(putObjectInput).getEtag();
    } catch (TosException e) {
      LOG.debug("failed to create empty object", e);
      throw AlluxioTosException.from(e);
    }
  }

  @Override
  protected void putObject(String key, File file, @Nullable String md5) throws IOException {
    try (InputStream content = Files.newInputStream(file.toPath())) {
      PutObjectInput putObjectInput = new PutObjectInput()
          .setBucket(mBucketName)
          .setKey(key)
          .setContent(content)
          .setContentLength(file.length()); // Set the correct content length
      mContentHash = getClient().putObject(putObjectInput).getEtag();
    } catch (TosException e) {
      LOG.debug("failed to put object", e);
      throw AlluxioTosException.from(e);
    }
  }

  @Override
  public Optional getContentHash() {
    return Optional.ofNullable(mContentHash);
  }

  protected TOSV2 getClient() {
    return mClient;
  }
}