alluxio-master-2.x/core/client/fs/src/main/java/alluxio/client/block/stream/SharedGrpcDataReader.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.grpc.ReadRequest;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.network.protocol.databuffer.NioDataBuffer;
import alluxio.resource.LockResource;
import alluxio.wire.WorkerNetAddress;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.hash.HashFunction;
import com.google.common.hash.Hashing;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import javax.annotation.Nullable;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * A shared gRPC data reader that cache blocks data for multi-thread accessing.
 *
 * It follows GrpcDataReader protocol and takes strong assumption:
 * Parallel read to the same file happens on the same time, so that read request is
 * serialized by kernel
 */
@NotThreadSafe
public class SharedGrpcDataReader implements DataReader {
  private static final int BLOCK_LOCK_NUM = 32;
  //
  // BLOCK_LOCKS is used to ensure thread-safety when referencing or updating the shared data
  // reader for block i in different FileInStream instances.
  // BLOCK_READERS is a ConcurrentHashMap from block id to its shared data reader, as different
  // DataReader may be needed to handle different blocks at the same time.
  //
  /** An array of locks to guard cached data readers based on block id. */
  private static final ReentrantReadWriteLock[] BLOCK_LOCKS =
      new ReentrantReadWriteLock[BLOCK_LOCK_NUM];
  /** A map from block id to the block's cached data reader. */
  private static final ConcurrentHashMap<Long, BufferCachingGrpcDataReader> BLOCK_READERS =
      new ConcurrentHashMap<>();
  /** A hashing function to map block id to one of the locks. */
  private static final HashFunction HASH_FUNC = Hashing.murmur3_32_fixed();

  static {
    for (int i = 0; i < BLOCK_LOCK_NUM; i++) {
      BLOCK_LOCKS[i] = new ReentrantReadWriteLock();
    }
  }

  private static ReentrantReadWriteLock getLock(long blockId) {
    return BLOCK_LOCKS[Math.floorMod(HASH_FUNC.hashLong(blockId).asInt(), BLOCK_LOCKS.length)];
  }

  private final long mBlockId;
  private final BufferCachingGrpcDataReader mCachedDataReader;
  private final long mChunkSize;

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

  /**
   * Creates an instance of {@link SharedGrpcDataReader}.
   *
   * @param readRequest the read request
   * @param reader the cached Grpc data reader for the given block
   */
  @VisibleForTesting
  protected SharedGrpcDataReader(ReadRequest readRequest, BufferCachingGrpcDataReader reader) {
    mChunkSize = readRequest.getChunkSize();
    mPosToRead = readRequest.getOffset();
    mBlockId = readRequest.getBlockId();
    mCachedDataReader = reader;
  }

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

  /**
   * Seeks to a specific position.
   *
   * @param pos the position to seek to
   */
  public void seek(long pos) {
    mPosToRead = pos;
  }

  @Override
  @Nullable
  public DataBuffer readChunk() throws IOException {
    int index = (int) (mPosToRead / mChunkSize);
    DataBuffer chunk = mCachedDataReader.readChunk(index);
    if (chunk == null) {
      return null;
    }
    ByteBuffer bb = chunk.getReadOnlyByteBuffer();
    // Force the buffer to align to chunk size
    bb.position((int) (mPosToRead % mChunkSize));
    mPosToRead += mChunkSize - mPosToRead % mChunkSize;

    return new NioDataBuffer(bb, bb.remaining());
  }

  @Override
  public void close() throws IOException {
    if (mCachedDataReader.deRef() > 0) {
      return;
    }
    try (LockResource ignored = new LockResource(getLock(mBlockId).writeLock())) {
      if (mCachedDataReader.getRefCount() == 0) {
        BLOCK_READERS.remove(mBlockId);
        mCachedDataReader.close();
      }
    }
  }

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

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

    @alluxio.annotation.SuppressFBWarnings(
        value = "AT_OPERATION_SEQUENCE_ON_CONCURRENT_ABSTRACTION",
        justification = "operation is still atomic guarded by block lock")
    @Override
    public DataReader create(long offset, long len) throws IOException {
      long blockId = mReadRequestBuilder.getBlockId();
      BufferCachingGrpcDataReader reader;
      try (LockResource ignored = new LockResource(getLock(blockId).writeLock())) {
        reader = BLOCK_READERS.get(blockId);
        if (reader == null) {
          // Even we may only need a portion, create a reader to read the whole block
          ReadRequest cacheRequest = mReadRequestBuilder.setOffset(0).setLength(mBlockSize).build();
          reader = BufferCachingGrpcDataReader.create(mContext, mAddress, cacheRequest);
          BLOCK_READERS.put(blockId, reader);
        }

        reader.ref();
      }
      return new SharedGrpcDataReader(
          mReadRequestBuilder.setOffset(offset).setLength(len).build(), reader);
    }

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

