alluxio-master-2.x/core/server/worker/src/main/java/alluxio/worker/grpc/UfsFallbackBlockWriteHandler.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.worker.grpc;

import alluxio.conf.Configuration;
import alluxio.exception.runtime.ResourceExhaustedRuntimeException;
import alluxio.grpc.WriteRequestCommand;
import alluxio.grpc.WriteResponse;
import alluxio.metrics.MetricInfo;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.proto.dataserver.Protocol;
import alluxio.security.authentication.AuthenticatedUserInfo;
import alluxio.underfs.UfsManager;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.options.CreateOptions;
import alluxio.worker.BlockUtils;
import alluxio.worker.block.CreateBlockOptions;
import alluxio.worker.block.DefaultBlockWorker;
import alluxio.worker.block.meta.TempBlockMeta;

import com.google.common.base.Preconditions;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Optional;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * This handler handles UFS block write request. Instead of writing a block to tiered storage, this
 * handler writes the block into UFS .
 */
@alluxio.annotation.SuppressFBWarnings(value = "BC_UNCONFIRMED_CAST_OF_RETURN_VALUE",
    justification = "false positive with superclass generics, "
        + "see more description in https://sourceforge.net/p/findbugs/bugs/1242/")
@NotThreadSafe
public final class UfsFallbackBlockWriteHandler
    extends AbstractWriteHandler<BlockWriteRequestContext> {
  private static final Logger LOG = LoggerFactory.getLogger(UfsFallbackBlockWriteHandler.class);

  /** The Block Worker which handles blocks stored in the Alluxio storage of the worker. */
  private final DefaultBlockWorker mWorker;
  private final UfsManager mUfsManager;
  private final BlockWriteHandler mBlockWriteHandler;
  private final boolean mDomainSocketEnabled;

  /**
   * Creates an instance of {@link UfsFallbackBlockWriteHandler}.
   *
   * @param blockWorker the block worker
   * @param userInfo the authenticated user info
   * @param domainSocketEnabled whether using a domain socket
   */
  UfsFallbackBlockWriteHandler(DefaultBlockWorker blockWorker, UfsManager ufsManager,
      StreamObserver<WriteResponse> responseObserver, AuthenticatedUserInfo userInfo,
      boolean domainSocketEnabled) {
    super(responseObserver, userInfo);
    mWorker = blockWorker;
    mUfsManager = ufsManager;
    mBlockWriteHandler =
        new BlockWriteHandler(blockWorker, responseObserver, userInfo, domainSocketEnabled);
    mDomainSocketEnabled = domainSocketEnabled;
  }

  @Override
  protected BlockWriteRequestContext createRequestContext(alluxio.grpc.WriteRequest msg)
      throws Exception {
    BlockWriteRequestContext context = new BlockWriteRequestContext(msg, FILE_BUFFER_SIZE);
    if (mDomainSocketEnabled) {
      context.setCounter(MetricsSystem.counter(MetricKey.WORKER_BYTES_WRITTEN_DOMAIN.getName()));
      context.setMeter(MetricsSystem
          .meter(MetricKey.WORKER_BYTES_WRITTEN_DOMAIN_THROUGHPUT.getName()));
    } else {
      context.setCounter(MetricsSystem.counter(MetricKey.WORKER_BYTES_WRITTEN_REMOTE.getName()));
      context.setMeter(MetricsSystem
          .meter(MetricKey.WORKER_BYTES_WRITTEN_REMOTE_THROUGHPUT.getName()));
    }
    BlockWriteRequest request = context.getRequest();
    Preconditions.checkState(request.hasCreateUfsBlockOptions());
    // if it is already a UFS fallback from short-circuit write, avoid writing to local again
    context.setWritingToLocal(!request.getCreateUfsBlockOptions().getFallback());
    if (context.isWritingToLocal()) {
      mWorker.createBlock(request.getSessionId(), request.getId(), request.getTier(),
          new CreateBlockOptions(null, request.getMediumType(), FILE_BUFFER_SIZE));
    }
    return context;
  }

  @Override
  protected void completeRequest(BlockWriteRequestContext context)
      throws Exception {
    if (context.isWritingToLocal()) {
      mBlockWriteHandler.completeRequest(context);
    } else {
      mWorker.commitBlockInUfs(context.getRequest().getId(), context.getPos());
      if (context.getOutputStream() != null) {
        context.getOutputStream().close();
        context.setOutputStream(null);
      }
    }
    if (context.getUfsResource() != null) {
      context.getUfsResource().close();
    }
  }

  @Override
  protected void cancelRequest(BlockWriteRequestContext context) throws Exception {
    if (context.isWritingToLocal()) {
      mBlockWriteHandler.cancelRequest(context);
    } else {
      if (context.getOutputStream() != null) {
        context.getOutputStream().close();
        context.setOutputStream(null);
      }
      if (context.getUfsResource() != null) {
        context.getUfsResource().get().deleteExistingFile(context.getUfsPath());
      }
    }
    if (context.getUfsResource() != null) {
      context.getUfsResource().close();
    }
  }

  @Override
  protected void cleanupRequest(BlockWriteRequestContext context) throws Exception {
    if (context.isWritingToLocal()) {
      mBlockWriteHandler.cleanupRequest(context);
    } else {
      cancelRequest(context);
    }
  }

  @Override
  protected void flushRequest(BlockWriteRequestContext context) throws Exception {
    if (context.isWritingToLocal()) {
      mBlockWriteHandler.flushRequest(context);
    } else if (context.getOutputStream() != null) {
      context.getOutputStream().flush();
    }
  }

  @Override
  protected void writeBuf(BlockWriteRequestContext context,
      StreamObserver<WriteResponse> responseObserver, DataBuffer buf, long pos) throws Exception {
    if (context.isWritingToLocal()) {
      // TODO(binfan): change signature of writeBuf to pass current offset and length of buffer.
      // Currently pos is the calculated offset after writeBuf succeeds.
      long posBeforeWrite = pos - buf.readableBytes();
      try {
        mBlockWriteHandler.writeBuf(context, responseObserver, buf, pos);
        return;
      } catch (ResourceExhaustedRuntimeException e) {
        LOG.warn("Not enough space to write block {} to local worker, fallback to UFS. "
            + " {} bytes have been written.",
            context.getRequest().getId(), posBeforeWrite);
        context.setWritingToLocal(false);
      }
      // close the block writer first
      if (context.getBlockWriter() != null) {
        context.getBlockWriter().close();
      }
      // prepare the UFS block and transfer data from the temp block to UFS
      createUfsBlock(context);
      if (posBeforeWrite > 0) {
        transferToUfsBlock(context, posBeforeWrite);
      }
      // close the original block writer and remove the temp file
      mBlockWriteHandler.cancelRequest(context);
    }
    if (context.getOutputStream() == null) {
      createUfsBlock(context);
    }
    buf.readBytes(context.getOutputStream(), buf.readableBytes());
  }

  @Override
  protected void handleCommand(WriteRequestCommand command, BlockWriteRequestContext context)
      throws Exception {
    if (command.hasCreateUfsBlockOptions()
        && command.getOffset() == 0
        && command.getCreateUfsBlockOptions().hasBytesInBlockStore()) {
      long ufsFallbackInitBytes = command.getCreateUfsBlockOptions().getBytesInBlockStore();
      context.setPos(context.getPos() + ufsFallbackInitBytes);
      initUfsFallback(context);
    }
  }

  @Override
  protected String getLocationInternal(BlockWriteRequestContext context) {
    Protocol.CreateUfsBlockOptions createUfsBlockOptions =
        context.getRequest().getCreateUfsBlockOptions();

    if (createUfsBlockOptions == null) {
      return String.format("blockId-%d", context.getRequest().getId());
    }

    UfsManager.UfsClient ufsClient;
    try {
      ufsClient = mUfsManager.get(createUfsBlockOptions.getMountId());
    } catch (Throwable e) {
      return String.format("blockId-%d", context.getRequest().getId());
    }

    return BlockUtils.getUfsBlockPath(ufsClient, context.getRequest().getId());
  }

  private void initUfsFallback(BlockWriteRequestContext context) throws Exception {
    Preconditions.checkState(!context.isWritingToLocal());
    if (context.getOutputStream() == null) {
      createUfsBlock(context);
    }
    // transfer data from the temp block to UFS
    transferToUfsBlock(context, context.getPos());
  }

  /**
   * Creates a UFS block and initialize it with bytes read from block store.
   *
   * @param context context of this request
   */
  private void createUfsBlock(BlockWriteRequestContext context)
      throws Exception {
    BlockWriteRequest request = context.getRequest();
    Protocol.CreateUfsBlockOptions createUfsBlockOptions = request.getCreateUfsBlockOptions();
    UfsManager.UfsClient ufsClient = mUfsManager.get(createUfsBlockOptions.getMountId());
    alluxio.resource.CloseableResource<UnderFileSystem> ufsResource =
        ufsClient.acquireUfsResource();
    context.setUfsResource(ufsResource);
    String ufsString = MetricsSystem.escape(ufsClient.getUfsMountPointUri());
    String ufsPath = BlockUtils.getUfsBlockPath(ufsClient, request.getId());
    UnderFileSystem ufs = ufsResource.get();
    // Set the atomic flag to be true to ensure only the creation of this file is atomic on close.
    OutputStream ufsOutputStream =
        ufs.createNonexistingFile(ufsPath,
            CreateOptions.defaults(Configuration.global()).setEnsureAtomic(true)
                .setCreateParent(true));
    context.setOutputStream(ufsOutputStream);
    context.setUfsPath(ufsPath);

    MetricKey counterKey = MetricKey.WORKER_BYTES_WRITTEN_UFS;
    MetricKey meterKey = MetricKey.WORKER_BYTES_WRITTEN_UFS_THROUGHPUT;
    context.setCounter(MetricsSystem.counterWithTags(counterKey.getName(),
        counterKey.isClusterAggregated(), MetricInfo.TAG_UFS, ufsString));
    context.setMeter(MetricsSystem.meterWithTags(meterKey.getName(),
        meterKey.isClusterAggregated(), MetricInfo.TAG_UFS, ufsString));
  }

  /**
   * Transfers data from block store to UFS.
   *
   * @param context context of this request
   * @param pos number of bytes in block store to write in the UFS block
   */
  private void transferToUfsBlock(BlockWriteRequestContext context, long pos) throws IOException {
    OutputStream ufsOutputStream = context.getOutputStream();
    long blockId = context.getRequest().getId();
    Optional<TempBlockMeta> block = mWorker.getBlockStore().getTempBlockMeta(blockId);
    Preconditions.checkState(block.isPresent()
        && Files.copy(Paths.get(block.get().getPath()), ufsOutputStream) == pos);
  }
}
