alluxio-master-2.x/core/client/fs/src/main/java/alluxio/client/block/stream/GrpcDataReader.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.client.block.stream;

import alluxio.Constants;
import alluxio.client.file.FileSystemContext;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.DataMessage;
import alluxio.grpc.ReadRequest;
import alluxio.grpc.ReadResponse;
import alluxio.grpc.ReadResponseMarshaller;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.network.protocol.databuffer.NioDataBuffer;
import alluxio.resource.CloseableResource;
import alluxio.util.logging.SamplingLogger;
import alluxio.wire.WorkerNetAddress;

import com.codahale.metrics.Timer;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.ByteBuffer;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * A gRPC data reader that streams a region from gRPC data server.
 *
 * Protocol:
 * 1. The client sends a read request (id, offset, length).
 * 2. Once the server receives the request, it streams chunks to the client. The streaming pauses
 *    if the server's buffer is full and resumes if the buffer is not full.
 * 3. The client reads chunks from the stream using an iterator.
 * 4. The client can cancel the read request at anytime. The cancel request is ignored by the
 *    server if everything has been sent to channel.
 * 5. To make it simple to handle errors, the channel is closed if any error occurs.
 */
@NotThreadSafe
public final class GrpcDataReader implements DataReader {
  private static final Logger LOG = LoggerFactory.getLogger(GrpcDataReader.class);
  private static final Logger SLOW_CLOSE_LOG = new SamplingLogger(LOG, Constants.MINUTE_MS);

  private final long mDataTimeoutMs;
  private final boolean mDetailedMetricsEnabled;
  private final CloseableResource<BlockWorkerClient> mClient;
  private final ReadRequest mReadRequest;
  private final WorkerNetAddress mAddress;

  private final GrpcBlockingStream<ReadRequest, ReadResponse> mStream;
  private final ReadResponseMarshaller mMarshaller;
  private final long mCloseWaitMs;

  /** The next pos to read. */
  private long mPosToRead;

  /**
   * Creates an instance of {@link GrpcDataReader}.
   *
   * @param context the file system context
   * @param address the data server address
   * @param readRequest the read request
   */
  private GrpcDataReader(FileSystemContext context, WorkerNetAddress address,
      ReadRequest readRequest) throws IOException {
    mAddress = address;
    mPosToRead = readRequest.getOffset();
    mReadRequest = readRequest;
    AlluxioConfiguration alluxioConf = context.getClusterConf();
    mDataTimeoutMs = alluxioConf.getMs(PropertyKey.USER_STREAMING_DATA_READ_TIMEOUT);
    mDetailedMetricsEnabled = alluxioConf.getBoolean(PropertyKey.USER_BLOCK_READ_METRICS_ENABLED);
    mMarshaller = new ReadResponseMarshaller();
    mClient = context.acquireBlockWorkerClient(address);
    mCloseWaitMs = alluxioConf.getMs(PropertyKey.USER_STREAMING_READER_CLOSE_TIMEOUT);
    int readerBufferSizeMessages = alluxioConf
        .getInt(PropertyKey.USER_STREAMING_READER_BUFFER_SIZE_MESSAGES);

    try {
      if (alluxioConf.getBoolean(PropertyKey.USER_STREAMING_ZEROCOPY_ENABLED)) {
        String desc = "Zero Copy GrpcDataReader";
        if (LOG.isDebugEnabled()) { // More detailed description when debug logging is enabled
          desc = MoreObjects.toStringHelper(this)
              .add("request", mReadRequest)
              .add("address", address)
              .toString();
        }
        mStream = new GrpcDataMessageBlockingStream<>(mClient.get()::readBlock,
            readerBufferSizeMessages,
            desc, null, mMarshaller);
      } else {
        String desc = "GrpcDataReader";
        if (LOG.isDebugEnabled()) { // More detailed description when debug logging is enabled
          desc = MoreObjects.toStringHelper(this)
              .add("request", mReadRequest)
              .add("address", address)
              .toString();
        }
        mStream = new GrpcBlockingStream<>(mClient.get()::readBlock, readerBufferSizeMessages,
            desc);
      }
      mStream.send(mReadRequest, mDataTimeoutMs);
    } catch (Exception e) {
      mClient.close();
      throw e;
    }
  }

  @Override
  public long pos() {
    return mPosToRead;
  }

  @Override
  public DataBuffer readChunk() throws IOException {
    if (mDetailedMetricsEnabled) {
      try (Timer.Context ignored = MetricsSystem
          .timer(MetricKey.CLIENT_BLOCK_READ_CHUNK_REMOTE.getName()).time()) {
        return readChunkInternal();
      }
    }
    return readChunkInternal();
  }

  private DataBuffer readChunkInternal() throws IOException {
    Preconditions.checkState(!mClient.get().isShutdown(),
        "Data reader is closed while reading data chunks.");
    DataBuffer buffer = null;
    ReadResponse response = null;
    if (mStream instanceof GrpcDataMessageBlockingStream) {
      DataMessage<ReadResponse, DataBuffer> message =
          ((GrpcDataMessageBlockingStream<ReadRequest, ReadResponse>) mStream)
              .receiveDataMessage(mDataTimeoutMs);
      if (message != null) {
        response = message.getMessage();
        buffer = message.getBuffer();
        if (buffer == null && response.hasChunk() && response.getChunk().hasData()) {
          // falls back to use chunk message for compatibility
          ByteBuffer byteBuffer = response.getChunk().getData().asReadOnlyByteBuffer();
          buffer = new NioDataBuffer(byteBuffer, byteBuffer.remaining());
        }
        Preconditions.checkState(buffer != null, "response should always contain chunk");
      }
    } else {
      response = mStream.receive(mDataTimeoutMs);
      if (response != null) {
        Preconditions.checkState(response.hasChunk() && response.getChunk().hasData(),
            "response should always contain chunk");
        ByteBuffer byteBuffer = response.getChunk().getData().asReadOnlyByteBuffer();
        buffer = new NioDataBuffer(byteBuffer, byteBuffer.remaining());
      }
    }
    if (response == null) {
      return null;
    }
    mPosToRead += buffer.readableBytes();
    try {
      mStream.send(mReadRequest.toBuilder().setOffsetReceived(mPosToRead).build());
    } catch (Exception e) {
      // nothing is done as the receipt is sent at best effort
      LOG.debug("Failed to send receipt of data to worker {} for request {}: {}.", mAddress,
          mReadRequest, e.getMessage());
    }
    Preconditions.checkState(mPosToRead - mReadRequest.getOffset() <= mReadRequest.getLength());
    return buffer;
  }

  @Override
  public void close() throws IOException {
    try {
      if (mClient.get().isShutdown()) {
        return;
      }
      mStream.close();

      // When a reader is closed, there is technically nothing the client requires from the server.
      // However, the server does need to clean up resources for a client close(), including closing
      // or canceling any temp blocks. Therefore, we should wait for some amount of time for the
      // server to finish cleanup, but it should not be very long (since the client is finished
      // with the read). Also, if there is any error when waiting for the complete, it should be
      // ignored since again, the client is completely finished with the read.
      try {
        // Wait a short time for the server to finish the close, and then let the client continue.
        if (mCloseWaitMs > 0) {
          mStream.waitForComplete(mCloseWaitMs);
        }
      } catch (IOException | RuntimeException e) {
        // ignore any errors
        SLOW_CLOSE_LOG.warn(
            "Closing gRPC read stream took longer than {}ms, moving on. blockId: {}, address: {}",
            mCloseWaitMs, mReadRequest.getBlockId(), mAddress);
      }
    } finally {
      mMarshaller.close();
      mClient.close();
    }
  }

  /**
   * Factory class to create {@link GrpcDataReader}s.
   */
  public static class Factory implements DataReader.Factory {
    private final FileSystemContext mContext;
    private final WorkerNetAddress mAddress;
    private final ReadRequest.Builder mReadRequestBuilder;

    /**
     * Creates an instance of {@link GrpcDataReader.Factory} for block reads.
     *
     * @param context the file system context
     * @param address the worker address
     * @param readRequestBuilder the builder of read request
     */
    public Factory(FileSystemContext context, WorkerNetAddress address,
        ReadRequest.Builder readRequestBuilder) {
      mContext = context;
      mAddress = address;
      mReadRequestBuilder = readRequestBuilder;
    }

    @Override
    public DataReader create(long offset, long len) throws IOException {
      return new GrpcDataReader(mContext, mAddress,
          mReadRequestBuilder.setOffset(offset).setLength(len).build());
    }

    @Override
    public void close() throws IOException {}
  }
}

