alluxio-master-2.x/core/server/worker/src/main/java/alluxio/worker/grpc/AbstractWriteHandler.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.Constants;
import alluxio.RpcSensitiveConfigMask;
import alluxio.client.block.stream.GrpcDataWriter;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.status.AlluxioStatusException;
import alluxio.exception.status.InvalidArgumentException;
import alluxio.grpc.WriteRequest;
import alluxio.grpc.WriteRequestCommand;
import alluxio.grpc.WriteResponse;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.network.protocol.databuffer.NioDataBuffer;
import alluxio.security.authentication.AuthenticatedUserInfo;
import alluxio.util.LogUtils;
import alluxio.util.logging.SamplingLogger;

import com.codahale.metrics.Counter;
import com.codahale.metrics.Meter;
import com.google.common.base.Preconditions;
import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.internal.SerializingExecutor;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.Semaphore;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * This class handles {@link WriteRequest}s.
 *
 * Protocol: Check {@link GrpcDataWriter} for more information.
 * 1. The write handler streams data from the client and writes them. The flow control is handled
 *    by gRPC server.
 * 3. A complete or cancel event signifies the completion of this request.
 * 4. When an error occurs, the response stream is closed.
 *
 * Threading model:
 * The write handler runs the gRPC event threads. gRPC serializes the events so there will not be
 * events running concurrently for same request.
 *
 * @param <T> type of write request
 */
@NotThreadSafe
abstract class AbstractWriteHandler<T extends WriteRequestContext<?>> {
  private static final Logger LOG = LoggerFactory.getLogger(AbstractWriteHandler.class);
  private static final Logger SLOW_WRITE_LOG = new SamplingLogger(LOG, 5 * Constants.MINUTE_MS);
  private static final long SLOW_WRITE_MS =
      Configuration.getMs(PropertyKey.WORKER_REMOTE_IO_SLOW_THRESHOLD);
  public static final long FILE_BUFFER_SIZE = Constants.MB;

  /** The observer for sending response messages. */
  private final StreamObserver<WriteResponse> mResponseObserver;
  /** The executor for running write tasks asynchronously in the submission order. */
  private final SerializingExecutor mSerializingExecutor;
  /** The semaphore to control the number of write tasks queued up in the executor.*/
  private final Semaphore mSemaphore = new Semaphore(
      Configuration.getInt(PropertyKey.WORKER_NETWORK_WRITER_BUFFER_SIZE_MESSAGES), true);

  /**
   * This is initialized only once for a whole file or block in
   * {@link AbstractWriteHandler#write(WriteRequest)}.
   * It is safe to read those final primitive fields (e.g. mId, mSessionId) if mError is not set.
   *
   * Using "volatile" because we want any value change of this variable to be
   * visible across all gRPC threads.
   */
  private volatile T mContext;

  protected AuthenticatedUserInfo mUserInfo;

  /**
   * Creates an instance of {@link AbstractWriteHandler}.
   *
   * @param responseObserver the response observer for the write request
   * @param userInfo the authenticated user info
   */
  AbstractWriteHandler(StreamObserver<WriteResponse> responseObserver,
      AuthenticatedUserInfo userInfo) {
    mResponseObserver = responseObserver;
    mUserInfo = userInfo;
    mSerializingExecutor = new SerializingExecutor(GrpcExecutors.BLOCK_WRITER_EXECUTOR);
  }

  /**
   * Handles write request.
   *
   * @param writeRequest the request from the client
   */
  public void write(WriteRequest writeRequest) {
    if (!tryAcquireSemaphore()) {
      return;
    }
    mSerializingExecutor.execute(() -> {
      try {
        if (mContext == null) {
          LOG.debug("Received write request {}.",
              RpcSensitiveConfigMask.CREDENTIAL_FIELD_MASKER.maskObjects(LOG, writeRequest));
          try {
            mContext = createRequestContext(writeRequest);
          } catch (Exception e) {
            // abort() assumes context is initialized.
            // Reply with the error in order to prevent clients getting stuck.
            replyError(new Error(AlluxioStatusException.fromThrowable(e), true));
            throw e;
          }
        } else {
          Preconditions.checkState(!mContext.isDoneUnsafe(),
              "invalid request after write request is completed.");
        }
        if (mContext.isDoneUnsafe() || mContext.getError() != null) {
          return;
        }
        validateWriteRequest(writeRequest);
        if (writeRequest.hasCommand()) {
          WriteRequestCommand command = writeRequest.getCommand();
          if (command.getFlush()) {
            flush();
          } else {
            handleCommand(command, mContext);
          }
        } else {
          Preconditions.checkState(writeRequest.hasChunk(),
              "write request is missing data chunk in non-command message");
          ByteString data = writeRequest.getChunk().getData();
          Preconditions.checkState(data != null && data.size() > 0,
              "invalid data size from write request message");
          writeData(new NioDataBuffer(data.asReadOnlyByteBuffer(), data.size()));
        }
      } catch (Exception e) {
        LogUtils.warnWithException(LOG, "Exception occurred while processing write request {}.",
            writeRequest, e);
        abort(new Error(AlluxioStatusException.fromThrowable(e), true));
      } finally {
        mSemaphore.release();
      }
    });
  }

  /**
   * Handles write request with data message.
   *
   * @param request the request from the client
   * @param buffer the data associated with the request
   */
  public void writeDataMessage(WriteRequest request, DataBuffer buffer) {
    if (buffer == null) {
      write(request);
      return;
    }
    boolean releaseBuf = true;
    try {
      Preconditions.checkState(!request.hasCommand(),
          "write request command should not come with data buffer");
      Preconditions.checkState(buffer.readableBytes() > 0,
          "invalid data size from write request message");
      if (!tryAcquireSemaphore()) {
        return;
      }
      releaseBuf = false;
      mSerializingExecutor.execute(() -> {
        try {
          writeData(buffer);
        } finally {
          mSemaphore.release();
        }
      });
    } finally {
      if (releaseBuf) {
        buffer.release();
      }
    }
  }

  /**
   * Handles request complete event.
   */
  public void onCompleted() {
    mSerializingExecutor.execute(() -> {
      Preconditions.checkState(mContext != null);
      try {
        completeRequest(mContext);
        replySuccess();
      } catch (Exception e) {
        LogUtils.warnWithException(LOG, "Exception occurred while completing write request {}.",
            mContext.getRequest(), e);
        abort(new Error(AlluxioStatusException.fromThrowable(e), true));
      }
    });
  }

  /**
   * Handles request cancellation event.
   */
  public void onCancel() {
    mSerializingExecutor.execute(() -> {
      try {
        cancelRequest(mContext);
        replyCancel();
      } catch (Exception e) {
        LogUtils.warnWithException(LOG, "Exception occurred while cancelling write request {}.",
            mContext.getRequest(), e);
        abort(new Error(AlluxioStatusException.fromThrowable(e), true));
      }
    });
  }

  /**
   * Handles errors from the request.
   *
   * @param cause the exception
   */
  public void onError(Throwable cause) {
    if (cause instanceof StatusRuntimeException
        && ((StatusRuntimeException) cause).getStatus().getCode() == Status.Code.CANCELLED) {
      // Cancellation is already handled.
      return;
    }
    mSerializingExecutor.execute(() -> {
      LogUtils.warnWithException(LOG, "Exception thrown while handling write request {}",
          mContext == null ? "unknown" : mContext.getRequest(), cause);
      abort(new Error(AlluxioStatusException.fromThrowable(cause), false));
    });
  }

  private boolean tryAcquireSemaphore() {
    try {
      mSemaphore.acquire();
    } catch (InterruptedException e) {
      LOG.warn("write data request {} is interrupted: {}",
          mContext == null ? "unknown" : mContext.getRequest(), e.getMessage());
      abort(new Error(AlluxioStatusException.fromThrowable(e), true));
      Thread.currentThread().interrupt();
      return false;
    }
    return true;
  }

  /**
   * Validates a block write request.
   *
   * @param request the block write request
   * @throws InvalidArgumentException if the write request is invalid
   */
  @GuardedBy("mLock")
  private void validateWriteRequest(alluxio.grpc.WriteRequest request)
      throws InvalidArgumentException {
    if (request.hasCommand() && request.getCommand().hasOffset()
        && request.getCommand().getOffset() != mContext.getPos()) {
      throw new InvalidArgumentException(String.format(
          "Offsets do not match [received: %d, expected: %d].",
          request.getCommand().getOffset(), mContext.getPos()));
    }
  }

  private void writeData(DataBuffer buf) {
    try {
      if (mContext.isDoneUnsafe() || mContext.getError() != null) {
        return;
      }
      int readableBytes = buf.readableBytes();
      mContext.setPos(mContext.getPos() + readableBytes);

      long writeStartMs = System.currentTimeMillis();
      writeBuf(mContext, mResponseObserver, buf, mContext.getPos());
      long writeMs = System.currentTimeMillis() - writeStartMs;

      if (writeMs >= SLOW_WRITE_MS) {
        // A single write call took much longer than expected.

        String prefix = String
            .format("Writing buffer for remote write took longer than %s ms. handler: %s",
                SLOW_WRITE_MS, this.getClass().getName());

        // Do not template the handler class, so the sampling log can distinguish between
        // different handler types
        SLOW_WRITE_LOG.warn(prefix + " id: {} location: {} bytes: {} durationMs: {}",
            mContext.getRequest().getId(), getLocation(), readableBytes, writeMs);
      }

      incrementMetrics(readableBytes);
    } catch (Exception e) {
      LOG.error("Failed to write data for request {}", mContext.getRequest(), e);
      abort(new Error(AlluxioStatusException.fromThrowable(e), true));
    } finally {
      buf.release();
    }
  }

  private void flush() {
    try {
      flushRequest(mContext);
      replyFlush();
    } catch (Exception e) {
      LOG.error("Failed to flush for write request {}", mContext.getRequest(), e);
      abort(new Error(AlluxioStatusException.fromThrowable(e), true));
    }
  }

  /**
   * Abort the write process due to error.
   *
   * @param error the error
   */
  private void abort(Error error) {
    try {
      if (mContext == null || mContext.getError() != null || mContext.isDoneUnsafe()) {
        // Note, we may reach here via events due to network errors bubbling up before
        // mContext is initialized, or stream error after the request is finished.
        return;
      }
      mContext.setError(error);
      cleanupRequest(mContext);
    } catch (Exception e) {
      LOG.warn("Failed to cleanup states with error {}.", e.toString());
    } finally {
      replyError();
    }
  }

  /**
   * Creates a new request context.
   *
   * @param msg the block write request
   */
  protected abstract T createRequestContext(WriteRequest msg) throws Exception;

  /**
   * Completes this write. This is called when the write completes.
   *  @param context context of the request to complete
   */
  protected abstract void completeRequest(T context) throws Exception;

  /**
   * Cancels this write. This is called when the client issues a cancel request.
   *
   * @param context context of the request to complete
   */
  protected abstract void cancelRequest(T context) throws Exception;

  /**
   * Cleans up this write. This is called when the write request is aborted due to any exception
   * or session timeout.
   *
   * @param context context of the request to complete
   */
  protected abstract void cleanupRequest(T context) throws Exception;

  /**
   * Flushes the buffered data. Flush only happens after write.
   *
   * @param context context of the request to complete
   */
  protected abstract void flushRequest(T context) throws Exception;

  /**
   * Writes the buffer.
   * @param context context of the request to complete
   * @param responseObserver the response observer
   * @param buf the buffer
   * @param pos the pos
   */
  protected abstract void writeBuf(T context, StreamObserver<WriteResponse> responseObserver,
      DataBuffer buf, long pos) throws Exception;

  /**
   * @param context the context of the request
   * @return an informational string of the location the writer is writing to
   */
  protected abstract String getLocationInternal(T context);

  /**
   * @return an informational string of the location the writer is writing to
   */
  public String getLocation() {
    if (mContext == null) {
      return "null";
    }
    return getLocationInternal(mContext);
  }

  /**
   * Handles a command in the write request.
   *
   * @param command the command to be handled
   */
  protected void handleCommand(WriteRequestCommand command, T context) throws Exception {
    // no additional command is handled by default
  }

  /**
   * Writes a response to signify the success of the write request.
   */
  private void replySuccess() {
    mContext.setDoneUnsafe(true);
    mContext.getContentHash().ifPresent(contentHash -> mResponseObserver.onNext(
        WriteResponse.newBuilder().setContentHash(contentHash).setOffset(
            mContext.getPos()).build()));
    mResponseObserver.onCompleted();
  }

  /**
   * Writes a response to signify the successful cancellation of the write request.
   */
  private void replyCancel() {
    mContext.setDoneUnsafe(true);
    mResponseObserver.onCompleted();
  }

  /**
   * Writes an error response.
   */
  private void replyError() {
    replyError(Preconditions.checkNotNull(mContext.getError()));
  }

  /**
   * Writes an error response.
   */
  private void replyError(Error error) {
    if (error.isNotifyClient()) {
      mResponseObserver.onError(error.getCause().toGrpcStatusException());
    }
  }

  /**
   * Writes a response to signify the successful flush.
   */
  private void replyFlush() {
    mResponseObserver.onNext(
        WriteResponse.newBuilder().setOffset(mContext.getPos()).build());
  }

  /**
   * @param bytesWritten bytes written
   */
  private void incrementMetrics(long bytesWritten) {
    Counter counter = mContext.getCounter();
    Meter meter = mContext.getMeter();
    Preconditions.checkState(counter != null, "counter");
    Preconditions.checkState(meter != null, "meter");
    counter.inc(bytesWritten);
    meter.mark(bytesWritten);
  }
}
