alluxio-master-2.x/core/client/fs/src/main/java/alluxio/client/block/stream/BufferCachingGrpcDataReader.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.client.file.FileSystemContext;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.ReadRequest;
import alluxio.grpc.ReadResponse;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.network.protocol.databuffer.NioDataBuffer;
import alluxio.resource.CloseableResource;
import alluxio.resource.LockResource;
import alluxio.wire.WorkerNetAddress;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * A gRPC data reader that responsible for reading from a specific block.
 *
 * The current implementation cached the block data from the beginning to
 * the largest index being read.
 */
@NotThreadSafe
public class BufferCachingGrpcDataReader {
  private static final Logger LOG = LoggerFactory.getLogger(BufferCachingGrpcDataReader.class);

  private final WorkerNetAddress mAddress;
  private final CloseableResource<BlockWorkerClient> mClient;
  private final long mDataTimeoutMs;
  private final ReadRequest mReadRequest;
  private final GrpcBlockingStream<ReadRequest, ReadResponse> mStream;
  /**
   * Count the number of threads that are accessing the same block together.
   * When no thread is accessing this block, the cached data will be GCed.
   */
  private final AtomicInteger mRefCount = new AtomicInteger(0);
  private final AtomicInteger mBufferCount = new AtomicInteger(0);
  @GuardedBy("mBufferLocks")
  private final DataBuffer[] mDataBuffers;
  private final ReentrantReadWriteLock mBufferLocks = new ReentrantReadWriteLock();

  /** The next pos to read. */
  @VisibleForTesting
  protected long mPosToRead;

  /**
   * Creates an instance of {@link BufferCachingGrpcDataReader}.
   *
   * @param address the data server address
   * @param client the block worker client to read data from
   * @param dataTimeoutMs the maximum time to wait for a data response
   * @param readRequest the read request
   * @param stream the underlying gRPC stream to read data
   */
  @VisibleForTesting
  protected BufferCachingGrpcDataReader(WorkerNetAddress address,
      CloseableResource<BlockWorkerClient> client, long dataTimeoutMs,
      ReadRequest readRequest, GrpcBlockingStream<ReadRequest, ReadResponse> stream) {
    mAddress = address;
    Objects.requireNonNull(client);
    mClient = client;
    mDataTimeoutMs = dataTimeoutMs;
    mPosToRead = readRequest.getOffset();
    mReadRequest = readRequest;
    Objects.requireNonNull(stream);
    mStream = stream;
    long blockSize = mReadRequest.getLength() + mReadRequest.getOffset();
    long chunkSize = mReadRequest.getChunkSize();
    int buffCount = (int) (blockSize / chunkSize);
    if ((blockSize % chunkSize) != 0) {
      buffCount += 1;
    }
    mDataBuffers = new DataBuffer[buffCount];
  }

  /**
   * Reads a specific chunk from the block.
   *
   * @param index the chunk index
   * @return the chunk data if exists
   */
  @Nullable
  public DataBuffer readChunk(int index) throws IOException {
    if (index >= mDataBuffers.length) {
      return null;
    }

    if (index >= mBufferCount.get()) {
      try (LockResource ignored = new LockResource(mBufferLocks.writeLock())) {
        while (index >= mBufferCount.get()) {
          DataBuffer buffer = readChunk();
          mDataBuffers[mBufferCount.get()] = buffer;
          mBufferCount.incrementAndGet();
        }
      }
    }

    return mDataBuffers[index];
  }

  /**
   * Reads a chunk of data.
   *
   * @return a chunk of data
   */
  @Nullable
  @VisibleForTesting
  protected DataBuffer readChunk() throws IOException {
    Preconditions.checkState(!mClient.get().isShutdown(),
        "Data reader is closed while reading data chunks.");
    ReadResponse response = mStream.receive(mDataTimeoutMs);
    if (response == null) {
      return null;
    }
    Preconditions.checkState(response.hasChunk() && response.getChunk().hasData(),
        "response should always contain chunk");

    ByteBuffer byteBuffer = response.getChunk().getData().asReadOnlyByteBuffer();
    DataBuffer buffer = new NioDataBuffer(byteBuffer, byteBuffer.remaining());
    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);
    }
    Preconditions.checkState(mPosToRead - mReadRequest.getOffset() <= mReadRequest.getLength());
    return buffer;
  }

  /**
   * Closes the {@link BufferCachingGrpcDataReader}.
   */
  public void close() throws IOException {
    try {
      if (mClient.get().isShutdown()) {
        return;
      }
      mStream.close();
      mStream.waitForComplete(mDataTimeoutMs);
    } finally {
      mClient.close();
    }
  }

  /**
   * Increases the reference count and return the current count.
   *
   * @return the incremented count
   */
  public int ref() {
    return mRefCount.incrementAndGet();
  }

  /**
   * Decreases the reference count and return the current count.
   *
   * @return the decremented count
   */
  public int deRef() {
    return mRefCount.decrementAndGet();
  }

  /**
   * @return the current count
   */
  public int getRefCount() {
    return mRefCount.get();
  }

  /**
   * Creates an instance of {@link BufferCachingGrpcDataReader} for block reads.
   *
   * @param context the file system context
   * @param address the worker address
   * @param readRequest the read request
   * @return a new {@link BufferCachingGrpcDataReader}
   */
  public static BufferCachingGrpcDataReader create(FileSystemContext context,
      WorkerNetAddress address, ReadRequest readRequest) throws IOException {
    AlluxioConfiguration alluxioConf = context.getClusterConf();
    int readerBufferSizeMessages = alluxioConf
        .getInt(PropertyKey.USER_STREAMING_READER_BUFFER_SIZE_MESSAGES);
    long dataTimeoutMs = alluxioConf.getMs(PropertyKey.USER_STREAMING_DATA_READ_TIMEOUT);

    CloseableResource<BlockWorkerClient> client = context.acquireBlockWorkerClient(address);

    String desc = "BufferCachingGrpcDataReader";
    if (LOG.isDebugEnabled()) { // More detailed description when debug logging is enabled
      desc = String.format("BufferCachingGrpcDataReader(request=%s,address=%s)",
          readRequest, address);
    }
    GrpcBlockingStream<ReadRequest, ReadResponse> stream = null;
    try {
      // Stream here cannot be GrpcDataMessagingBlockingStream
      // DataBuffer.getReadOnlyByteBuffer is used to clone a copy in SharedDataReader.readChunk.
      // getReadOnlyByteBuffer is not implemented in DataBuffer
      // returned from GrpcDataMessagingBlockingStream.
      stream = new GrpcBlockingStream<>(client.get()::readBlock, readerBufferSizeMessages,
          desc);
      stream.send(readRequest, dataTimeoutMs);
    } catch (Exception e) {
      if (stream != null) {
        stream.close();
      }
      client.close();
      throw e;
    }
    return new BufferCachingGrpcDataReader(address, client, dataTimeoutMs, readRequest, stream);
  }
}

