alluxio-master-2.x/core/client/fs/src/main/java/alluxio/client/block/stream/BlockInStream.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.Seekable;
import alluxio.client.BoundedStream;
import alluxio.client.CanUnbuffer;
import alluxio.client.PositionedReadable;
import alluxio.client.ReadType;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.options.InStreamOptions;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.exception.PreconditionMessage;
import alluxio.exception.status.NotFoundException;
import alluxio.exception.status.OutOfRangeException;
import alluxio.grpc.ReadRequest;
import alluxio.network.protocol.databuffer.DataBuffer;
import alluxio.proto.dataserver.Protocol;
import alluxio.util.LogUtils;
import alluxio.util.io.BufferUtils;
import alluxio.util.network.NettyUtils;
import alluxio.util.network.NetworkAddressUtils;
import alluxio.wire.BlockInfo;
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.io.InputStream;
import java.nio.ByteBuffer;
import java.util.Objects;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * Provides an {@link InputStream} implementation that is based on {@link DataReader}s to
 * stream data chunk by chunk.
 */
@NotThreadSafe
public class BlockInStream extends InputStream implements BoundedStream, Seekable,
    PositionedReadable, CanUnbuffer {
  private static final Logger LOG = LoggerFactory.getLogger(BlockInStream.class);

  /** the source tracking where the block is from. */
  public enum BlockInStreamSource {
    PROCESS_LOCAL, // The block is from a worker in the same process
    NODE_LOCAL, // The block is from a separate worker process on the same node
    REMOTE, // The block is from a remote worker
    UFS // The block is in UFS
  }

  private final WorkerNetAddress mAddress;
  private final BlockInStreamSource mInStreamSource;
  /** The id of the block or UFS file to which this InStream provides access. */
  private final long mId;
  /** The size in bytes of the block. */
  private final long mLength;
  private final byte[] mSingleByte = new byte[1];

  /** Current position of the stream, relative to the start of the block. */
  private long mPos = 0;
  /** The current data chunk. */
  protected DataBuffer mCurrentChunk;

  protected DataReader mDataReader;
  private final DataReader.Factory mDataReaderFactory;

  private boolean mClosed = false;
  private boolean mEOF = false;

  /**
   * Creates a {@link BlockInStream}.
   *
   * One of several read behaviors:
   *
   * 1. Domain socket - if the data source is the local worker and the local worker has a domain
   * socket server
   * 2. Short-Circuit - if the data source is the local worker
   * 3. Local Loopback Read - if the data source is the local worker and short circuit is disabled
   * 4. Read from remote worker - if the data source is a remote worker
   * 5. UFS Read from worker - if the data source is UFS, read from the UFS policy's designated
   * worker (ufs -> local or remote worker -> client)
   *
   * @param context the file system context
   * @param info the block info
   * @param dataSource the Alluxio worker which should read the data
   * @param dataSourceType the source location of the block
   * @param options the InStream options
   * @return the {@link BlockInStream} object
   */
  public static BlockInStream create(FileSystemContext context, BlockInfo info,
      WorkerNetAddress dataSource, BlockInStreamSource dataSourceType, InStreamOptions options)
      throws IOException {
    long blockId = info.getBlockId();
    long blockSize = info.getLength();

    if (dataSourceType == BlockInStreamSource.PROCESS_LOCAL) {
      // Interaction between the current client and the worker it embedded to should
      // go through worker internal communication directly without RPC involves
      LOG.debug("Creating worker process local input stream for block {} @ {}",
          blockId, dataSource);
      return createProcessLocalBlockInStream(context, dataSource, blockId, blockSize, options);
    }

    AlluxioConfiguration alluxioConf = context.getClusterConf();
    boolean shortCircuit = alluxioConf.getBoolean(PropertyKey.USER_SHORT_CIRCUIT_ENABLED);
    boolean shortCircuitPreferred =
        alluxioConf.getBoolean(PropertyKey.USER_SHORT_CIRCUIT_PREFERRED);
    boolean sourceSupportsDomainSocket = NettyUtils.isDomainSocketSupported(dataSource);
    boolean sourceIsLocal = dataSourceType == BlockInStreamSource.NODE_LOCAL;

    // Short circuit is enabled when
    // 1. data source is local node
    // 2. alluxio.user.short.circuit.enabled is true
    // 3. the worker's domain socket is not configured
    //      OR alluxio.user.short.circuit.preferred is true
    if (sourceIsLocal && shortCircuit && (shortCircuitPreferred || !sourceSupportsDomainSocket)) {
      LOG.debug("Creating short circuit input stream for block {} @ {}", blockId, dataSource);
      try {
        return createLocalBlockInStream(context, dataSource, blockId, blockSize, options);
      } catch (NotFoundException e) {
        // Failed to do short circuit read because the block is not available in Alluxio.
        // We will try to read via gRPC. So this exception is ignored.
        LOG.warn("Failed to create short circuit input stream for block {} @ {}. Falling back to "
            + "network transfer", blockId, dataSource);
      }
    }

    // gRPC
    LOG.debug("Creating gRPC input stream for block {} @ {} from client {} reading through {} ("
        + "data locates in the local worker {}, shortCircuitEnabled {}, "
        + "shortCircuitPreferred {}, sourceSupportDomainSocket {})",
        blockId, dataSource, NetworkAddressUtils.getClientHostName(alluxioConf), dataSource,
        sourceIsLocal, shortCircuit, shortCircuitPreferred, sourceSupportsDomainSocket);
    return createGrpcBlockInStream(context, dataSource, dataSourceType, blockId,
        blockSize, options);
  }

  /**
   * Creates a {@link BlockInStream} to read from the worker process-local to this client
   * directly without RPC involves, if the block does not exist in this worker, will read from
   * the UFS storage via this worker.
   *
   * @param context the file system context
   * @param address the network address of the gRPC data server to read from
   * @param blockId the block ID
   * @param length the block length
   * @param options the in stream options
   * @return the {@link BlockInStream} created
   */
  private static BlockInStream createProcessLocalBlockInStream(FileSystemContext context,
      WorkerNetAddress address, long blockId, long length, InStreamOptions options) {
    AlluxioConfiguration conf = context.getClusterConf();
    long chunkSize = conf.getBytes(
        PropertyKey.USER_LOCAL_READER_CHUNK_SIZE_BYTES);
    return new BlockInStream(new BlockWorkerDataReader.Factory(
        context.getProcessLocalWorker().orElseThrow(NullPointerException::new),
        blockId, chunkSize, options),
        address, BlockInStreamSource.PROCESS_LOCAL, blockId, length);
  }

  /**
   * Creates a {@link BlockInStream} to read from a local file.
   *
   * @param context the file system context
   * @param address the network address of the gRPC data server to read from
   * @param blockId the block ID
   * @param length the block length
   * @param options the in stream options
   * @return the {@link BlockInStream} created
   */
  private static BlockInStream createLocalBlockInStream(FileSystemContext context,
      WorkerNetAddress address, long blockId, long length, InStreamOptions options)
      throws IOException {
    AlluxioConfiguration conf = context.getClusterConf();
    long chunkSize = conf.getBytes(
        PropertyKey.USER_LOCAL_READER_CHUNK_SIZE_BYTES);
    return new BlockInStream(
        new LocalFileDataReader.Factory(context, address, blockId, chunkSize, options),
        address, BlockInStreamSource.NODE_LOCAL, blockId, length);
  }

  /**
   * Creates a {@link BlockInStream} to read from a gRPC data server.
   *
   * @param context the file system context
   * @param address the address of the gRPC data server
   * @param blockSource the source location of the block
   * @param blockSize the block size
   * @param blockId the block id
   * @return the {@link BlockInStream} created
   */
  private static BlockInStream createGrpcBlockInStream(FileSystemContext context,
      WorkerNetAddress address, BlockInStreamSource blockSource,
      long blockId, long blockSize, InStreamOptions options) {
    AlluxioConfiguration conf = context.getClusterConf();
    long chunkSize = conf.getBytes(
        PropertyKey.USER_STREAMING_READER_CHUNK_SIZE_BYTES);
    // Construct the partial read request
    ReadRequest.Builder builder = ReadRequest.newBuilder()
        .setBlockId(blockId)
        .setPromote(ReadType.fromProto(options.getOptions().getReadType()).isPromote())
        .setOpenUfsBlockOptions(options.getOpenUfsBlockOptions(blockId)) // Add UFS fallback options
        .setPositionShort(options.getPositionShort())
        .setChunkSize(chunkSize);
    DataReader.Factory factory;
    if (context.getClusterConf().getBoolean(PropertyKey.FUSE_SHARED_CACHING_READER_ENABLED)
        && blockSize > chunkSize * 4) {
      // Heuristic to resolve issues/12146, guarded by alluxio.fuse.shared.caching.reader.enabled
      // GrpcDataReader instances are shared across FileInStreams to mitigate seek cost
      factory = new SharedGrpcDataReader.Factory(context, address, builder, blockSize);
    } else {
      factory = new GrpcDataReader.Factory(context, address, builder);
    }
    return new BlockInStream(factory, address, blockSource, blockId, blockSize);
  }

  /**
   * Creates a {@link BlockInStream} to read from a specific remote server. Should only be used
   * in cases where the data source and method of reading is known, i.e. worker - worker
   * communication.
   *
   * @param context the file system context
   * @param blockId the block id
   * @param address the address of the gRPC data server
   * @param blockSource the source location of the block
   * @param blockSize the size of the block
   * @param ufsOptions the ufs read options
   * @return the {@link BlockInStream} created
   */
  public static BlockInStream createRemoteBlockInStream(FileSystemContext context, long blockId,
      WorkerNetAddress address, BlockInStreamSource blockSource, long blockSize,
      Protocol.OpenUfsBlockOptions ufsOptions) {
    AlluxioConfiguration conf = context.getClusterConf();
    long chunkSize = conf.getBytes(
        PropertyKey.USER_STREAMING_READER_CHUNK_SIZE_BYTES);
    ReadRequest readRequest = ReadRequest.newBuilder().setBlockId(blockId)
        .setOpenUfsBlockOptions(ufsOptions).setChunkSize(chunkSize).buildPartial();
    DataReader.Factory factory = new GrpcDataReader.Factory(context, address,
        readRequest.toBuilder());
    return new BlockInStream(factory, address, blockSource, blockId, blockSize);
  }

  /**
   * Creates an instance of {@link BlockInStream}.
   *
   * @param dataReaderFactory the data reader factory
   * @param address the address of the gRPC data server
   * @param blockSource the source location of the block
   * @param id the ID (either block ID or UFS file ID)
   * @param length the length
   */
  @VisibleForTesting
  protected BlockInStream(DataReader.Factory dataReaderFactory,
      WorkerNetAddress address, BlockInStreamSource blockSource, long id, long length) {
    mDataReaderFactory = dataReaderFactory;
    mAddress = address;
    mInStreamSource = blockSource;
    mId = id;
    mLength = length;
  }

  @Override
  public long getPos() {
    return mPos;
  }

  @Override
  public int read() throws IOException {
    int bytesRead = read(mSingleByte);
    if (bytesRead == -1) {
      return -1;
    }
    Preconditions.checkState(bytesRead == 1);
    return BufferUtils.byteToInt(mSingleByte[0]);
  }

  @Override
  public int read(byte[] b) throws IOException {
    return read(b, 0, b.length);
  }

  @Override
  public int read(byte[] b, int off, int len) throws IOException {
    Objects.requireNonNull(b, "Read buffer cannot be null");
    return read(ByteBuffer.wrap(b), off, len);
  }

  /**
   * Reads up to len bytes of data from the input stream into the byte buffer.
   *
   * @param byteBuffer the buffer into which the data is read
   * @param off the start offset in the buffer at which the data is written
   * @param len the maximum number of bytes to read
   * @return the total number of bytes read into the buffer, or -1 if there is no more data because
   *         the end of the stream has been reached
   */
  public int read(ByteBuffer byteBuffer, int off, int len) throws IOException {
    Preconditions.checkArgument(off >= 0 && len >= 0 && len + off <= byteBuffer.capacity(),
        PreconditionMessage.ERR_BUFFER_STATE.toString(), byteBuffer.capacity(), off, len);
    checkIfClosed();
    if (len == 0) {

      return 0;
    }
    if (mPos == mLength) {
      return -1;
    }
    readChunk();
    if (mCurrentChunk == null) {
      mEOF = true;
    }
    if (mEOF) {
      closeDataReader();
      if (mPos < mLength) {
        throw new OutOfRangeException(String.format("Block %s is expected to be %s bytes, "
            + "but only %s bytes are available in the UFS. "
            + "Please retry the read and on the next access, "
            + "Alluxio will sync with the UFS and fetch the updated file content.",
            mId, mLength, mPos));
      }
      return -1;
    }
    int toRead = Math.min(len, mCurrentChunk.readableBytes());
    byteBuffer.position(off).limit(off + toRead);
    mCurrentChunk.readBytes(byteBuffer);
    mPos += toRead;
    if (mPos == mLength) {
      // a performance improvement introduced by https://github.com/Alluxio/alluxio/issues/14020
      closeDataReader();
    }
    return toRead;
  }

  @Override
  public int positionedRead(long pos, byte[] b, int off, int len) throws IOException {
    if (len == 0) {
      return 0;
    }
    if (pos < 0 || pos >= mLength) {
      return -1;
    }

    int lenCopy = len;
    try (DataReader reader = mDataReaderFactory.create(pos, len)) {
      // We try to read len bytes instead of returning after reading one chunk because
      // it is not free to create/close a DataReader.
      while (len > 0) {
        DataBuffer dataBuffer = null;
        try {
          dataBuffer = reader.readChunk();
          if (dataBuffer == null) {
            break;
          }
          Preconditions.checkState(dataBuffer.readableBytes() <= len);
          int toRead = dataBuffer.readableBytes();
          dataBuffer.readBytes(b, off, toRead);
          len -= toRead;
          off += toRead;
        } finally {
          if (dataBuffer != null) {
            dataBuffer.release();
          }
        }
      }
    }
    if (lenCopy == len) {
      return -1;
    }
    return lenCopy - len;
  }

  @Override
  public long remaining() {
    return mEOF ? 0 : mLength - mPos;
  }

  @Override
  public void seek(long pos) throws IOException {
    checkIfClosed();
    Preconditions.checkArgument(pos >= 0, PreconditionMessage.ERR_SEEK_NEGATIVE.toString(), pos);
    Preconditions.checkArgument(pos <= mLength,
        "Seek position past the end of the read region (block or file).", mId);
    if (pos == mPos) {
      return;
    }
    // When alluxio.fuse.shared.caching.reader.enabled is on (to resolve issues/12146),
    // use the heuristic to improve seek performance with fewer data reader close.
    if (mDataReader instanceof SharedGrpcDataReader) {
      seekForSharedGrpcDataReader(pos);
      return;
    }
    if (pos < mPos) {
      mEOF = false;
    }
    closeDataReader();
    mPos = pos;
  }

  private void seekForSharedGrpcDataReader(long pos) throws IOException {
    if (pos < mPos) {
      mEOF = false;
      // because the reader is shared, let's not close it but simply seek
      ((SharedGrpcDataReader) mDataReader).seek(pos);
      if (mCurrentChunk != null) {
        mCurrentChunk.release();
        mCurrentChunk = null;
      }
    } else {
      // TODO(lu) combine the original seek logic and the following general improvements
      // that are helpful in both fuse and non-fuse scenarios
      // Try to read data already received but haven't processed
      long curPos = mPos;
      while (mCurrentChunk != null && curPos < pos) {
        long nextPos = curPos + mCurrentChunk.readableBytes();
        if (nextPos <= pos) {
          curPos = nextPos;
          mCurrentChunk.release();
          mCurrentChunk = mDataReader.readChunk();
        } else {
          // TODO(chaowang) introduce seek in DataBuffer
          int toRead = (int) (pos - curPos);
          final byte[] b = new byte[toRead];
          mCurrentChunk.readBytes(b, 0, toRead);
          curPos = pos;
        }
      }

      if (curPos < pos) {
        // Not enough data in queue, close the data reader
        closeDataReader();
      }
    }
    mPos = pos;
  }

  @Override
  public long skip(long n) throws IOException {
    checkIfClosed();
    if (n <= 0) {
      return 0;
    }

    long toSkip = Math.min(remaining(), n);
    mPos += toSkip;

    closeDataReader();
    return toSkip;
  }

  @Override
  public void close() throws IOException {
    if (mClosed) {
      return;
    }
    try {
      closeDataReader();
    } finally {
      mDataReaderFactory.close();
    }
    mClosed = true;
  }

  /**
   * @return the underlying data reader factory
   */
  @VisibleForTesting
  public DataReader.Factory getDataReaderFactory() {
    return mDataReaderFactory;
  }

  /**
   * Reads a new chunk from the channel if all of the current chunk is read.
   */
  private void readChunk() throws IOException {
    if (mDataReader == null) {
      mDataReader = mDataReaderFactory.create(mPos, mLength - mPos);
    }

    if (mCurrentChunk != null && mCurrentChunk.readableBytes() == 0) {
      mCurrentChunk.release();
      mCurrentChunk = null;
    }
    if (mCurrentChunk == null) {
      mCurrentChunk = mDataReader.readChunk();
    }
  }

  /**
   * Close the current data reader.
   */
  private void closeDataReader() throws IOException {
    if (mCurrentChunk != null) {
      mCurrentChunk.release();
      mCurrentChunk = null;
    }
    if (mDataReader != null) {
      mDataReader.close();
    }
    mDataReader = null;
  }

  @Override
  public void unbuffer() {
    try {
      closeDataReader();
    } catch (IOException e) {
      LogUtils.warnWithException(LOG, "failed to unbuffer the block stream", e);
    }
  }

  /**
   * Convenience method to ensure the stream is not closed.
   */
  private void checkIfClosed() {
    Preconditions.checkState(!mClosed, "Cannot do operations on a closed BlockInStream");
  }

  /**
   * @return the address of the data server
   */
  public WorkerNetAddress getAddress() {
    return mAddress;
  }

  /**
   * @return the source of the block location
   */
  public BlockInStreamSource getSource() {
    return mInStreamSource;
  }

  /**
   * @return the block ID
   */
  public long getId() {
    return mId;
  }
}
