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.worker.grpc;
  
  import alluxio.grpc.WriteResponse;
  import alluxio.metrics.MetricKey;
  import alluxio.metrics.MetricsSystem;
  import alluxio.network.protocol.databuffer.DataBuffer;
  import alluxio.security.authentication.AuthenticatedUserInfo;
  import alluxio.worker.block.BlockWorker;
  import alluxio.worker.block.CreateBlockOptions;
  
  import com.codahale.metrics.Counter;
  import com.google.common.base.Preconditions;
  import io.grpc.stub.StreamObserver;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import javax.annotation.concurrent.NotThreadSafe;
  
  /**
   * This handler handles block write request. Check more information in
   * {@link AbstractWriteHandler}.
   */
  @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 BlockWriteHandler extends AbstractWriteHandler<BlockWriteRequestContext> {
    private static final Logger LOG = LoggerFactory.getLogger(BlockWriteHandler.class);
    /** Metrics. */
    private static final Counter RPC_WRITE_COUNT =
        MetricsSystem.counterWithTags(MetricKey.WORKER_ACTIVE_RPC_WRITE_COUNT.getName(),
              MetricKey.WORKER_ACTIVE_RPC_WRITE_COUNT.isClusterAggregated());
  
    /** The Block Worker which handles blocks stored in the Alluxio storage of the worker. */
    private final BlockWorker mWorker;
  
    private final boolean mDomainSocketEnabled;
  
    /**
     * Creates an instance of {@link BlockWriteHandler}.
     *
     * @param blockWorker the block worker
     * @param responseObserver the stream observer for the write response
     * @param userInfo the authenticated user info
     * @param domainSocketEnabled whether reading block over domain socket
     */
    BlockWriteHandler(BlockWorker blockWorker, StreamObserver<WriteResponse> responseObserver,
        AuthenticatedUserInfo userInfo, boolean domainSocketEnabled) {
      super(responseObserver, userInfo);
      mWorker = blockWorker;
      mDomainSocketEnabled = domainSocketEnabled;
    }
  
    @Override
    protected BlockWriteRequestContext createRequestContext(alluxio.grpc.WriteRequest msg) {
      long bytesToReserve = FILE_BUFFER_SIZE;
      if (msg.getCommand().hasSpaceToReserve()) {
        bytesToReserve = msg.getCommand().getSpaceToReserve();
      }
      BlockWriteRequestContext context = new BlockWriteRequestContext(msg, bytesToReserve);
      BlockWriteRequest request = context.getRequest();
      mWorker.createBlock(request.getSessionId(), request.getId(), request.getTier(),
          new CreateBlockOptions(null, request.getMediumType(), bytesToReserve));
      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()));
      }
      RPC_WRITE_COUNT.inc();
      return context;
    }
  
    @Override
    protected void completeRequest(BlockWriteRequestContext context) throws Exception {
      WriteRequest request = context.getRequest();
      if (context.getBlockWriter() != null) {
        context.getBlockWriter().close();
      }
      mWorker.commitBlock(request.getSessionId(), request.getId(), request.getPinOnCreate());
      RPC_WRITE_COUNT.dec();
    }
  
    @Override
    protected void cancelRequest(BlockWriteRequestContext context) throws Exception {
      WriteRequest request = context.getRequest();
      if (context.getBlockWriter() != null) {
        context.getBlockWriter().close();
      }
      mWorker.abortBlock(request.getSessionId(), request.getId());
      RPC_WRITE_COUNT.dec();
    }
  
    @Override
    protected void cleanupRequest(BlockWriteRequestContext context) throws Exception {
      if (context.getBlockWriter() != null) {
        context.getBlockWriter().close();
      }
      mWorker.cleanupSession(context.getRequest().getSessionId());
  
      // Decrement RPC counter only if the request wasn't completed/canceled already
      if (!context.isDoneUnsafe()) {
        RPC_WRITE_COUNT.dec();
      }
    }
  
    @Override
    protected void flushRequest(BlockWriteRequestContext context) {
      // This is a no-op because block worker does not support flush currently.
    }
  
    @Override
    protected void writeBuf(BlockWriteRequestContext context,
        StreamObserver<WriteResponse> observer, DataBuffer buf, long pos) throws Exception {
      Preconditions.checkState(context != null);
      WriteRequest request = context.getRequest();
      long bytesReserved = context.getBytesReserved();
      if (bytesReserved < pos) {
        long bytesToReserve = Math.max(FILE_BUFFER_SIZE, pos - bytesReserved);
        // Allocate enough space in the existing temporary block for the write.
        mWorker.requestSpace(request.getSessionId(), request.getId(), bytesToReserve);
        context.setBytesReserved(bytesReserved + bytesToReserve);
      }
      if (context.getBlockWriter() == null) {
        context.setBlockWriter(
            mWorker.createBlockWriter(request.getSessionId(), request.getId()));
      }
      Preconditions.checkState(context.getBlockWriter() != null);
      int sz = buf.readableBytes();
      Preconditions.checkState(context.getBlockWriter().append(buf)  == sz);
    }
  
    @Override
    protected String getLocationInternal(BlockWriteRequestContext context) {
      return String.format("temp-block-session-%d-id-%d", context.getRequest().getSessionId(),
          context.getRequest().getId());
    }
  }
  