content: /*
   * 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.block;
  
  import static alluxio.worker.block.BlockMetadataManager.WORKER_STORAGE_TIER_ASSOC;
  
  import alluxio.exception.runtime.AlluxioRuntimeException;
  import alluxio.exception.status.AlluxioStatusException;
  import alluxio.metrics.MetricKey;
  import alluxio.metrics.MetricsSystem;
  import alluxio.resource.CloseableResource;
  import alluxio.underfs.UfsManager;
  import alluxio.underfs.UnderFileSystem;
  import alluxio.underfs.options.OpenOptions;
  import alluxio.util.IdUtils;
  import alluxio.worker.block.io.BlockReader;
  import alluxio.worker.block.io.BlockWriter;
  import alluxio.worker.block.meta.UnderFileSystemBlockMeta;
  
  import com.codahale.metrics.Counter;
  import com.codahale.metrics.Meter;
  import com.google.common.base.Preconditions;
  import io.netty.buffer.ByteBuf;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.nio.ByteBuffer;
  import java.nio.channels.ReadableByteChannel;
  import javax.annotation.concurrent.NotThreadSafe;
  
  /**
   * This class implements a {@link BlockReader} to read a block directly from UFS, and
   * optionally cache the block to the Alluxio worker if the whole block it is read.
   */
  @NotThreadSafe
  public final class UnderFileSystemBlockReader extends BlockReader {
    private static final Logger LOG = LoggerFactory.getLogger(UnderFileSystemBlockReader.class);
  
    /** Metrics. */
    private static final Counter BLOCKS_READ_UFS =
        MetricsSystem.counter(MetricKey.WORKER_BLOCKS_READ_UFS.getName());
  
    private final Counter mUfsBytesRead;
    private final Meter mUfsBytesReadThroughput;
  
    /** The initial size of the block allocated in Alluxio storage when the block is cached. */
    private final long mInitialBlockSize;
    /** The block metadata for the UFS block. */
    private final UnderFileSystemBlockMeta mBlockMeta;
    /** The Local block store. It is used to interact with Alluxio. */
    private final LocalBlockStore mLocalBlockStore;
    /** The cache for all ufs instream. */
    private final UfsInputStreamCache mUfsInstreamCache;
    /** The ufs client resource. */
    private final CloseableResource<UnderFileSystem> mUfsResource;
    private final boolean mIsPositionShort;
  
    /** The input stream to read from UFS. */
    private InputStream mUnderFileSystemInputStream;
    /** The block writer to write the block to Alluxio. */
    private BlockWriter mBlockWriter;
    /** If set, the reader is closed and should not be used afterwards. */
    private boolean mClosed;
  
    /**
     * The position of mUnderFileSystemInputStream (if not null) is blockStart + mInStreamPos.
     * When mUnderFileSystemInputStream is not set, this is set to -1 (an invalid state) when
     * mUnderFileSystemInputStream is null. Check mUnderFileSystemInputStream directly to see whether
     * that is valid instead of relying on this invalid state of the position to be safe.
     */
    private long mInStreamPos;
  
    /**
     * Creates an instance of {@link UnderFileSystemBlockReader} and initializes it with a reading
     * offset.
     *
     * @param blockMeta the block meta
     * @param offset the position within the block to start the read
     * @param localBlockStore the Local block store
     * @param ufsClient the manager of ufs
     * @param positionShort whether the client op is a positioned read to a small buffer
     * @param ufsInStreamCache the UFS in stream cache
     * @param ufsBytesRead counter metric to track ufs bytes read
     * @param ufsBytesReadThroughput meter metric to track bytes read throughput
     * @return the block reader
     */
    public static UnderFileSystemBlockReader create(UnderFileSystemBlockMeta blockMeta, long offset,
        boolean positionShort, LocalBlockStore localBlockStore, UfsManager.UfsClient ufsClient,
        UfsInputStreamCache ufsInStreamCache, Counter ufsBytesRead, Meter ufsBytesReadThroughput)
        throws IOException {
      UnderFileSystemBlockReader ufsBlockReader =
          new UnderFileSystemBlockReader(blockMeta, positionShort, localBlockStore, ufsClient,
              ufsInStreamCache, ufsBytesRead, ufsBytesReadThroughput);
      ufsBlockReader.init(offset);
      return ufsBlockReader;
    }
  
    /**
     * Creates an instance of {@link UnderFileSystemBlockReader}.
     *
     * @param blockMeta the block meta
     * @param localBlockStore the Local block store
     * @param ufsClient the ufs client
     * @param positionShort whether the client op is a positioned read to a small buffer
     * @param ufsInStreamCache the UFS in stream cache
     * @param ufsBytesRead counter metric to track ufs bytes read
     * @param ufsBytesReadThroughput meter metric to track bytes read throughput
     */
    private UnderFileSystemBlockReader(UnderFileSystemBlockMeta blockMeta, boolean positionShort,
        LocalBlockStore localBlockStore, UfsManager.UfsClient ufsClient,
        UfsInputStreamCache ufsInStreamCache, Counter ufsBytesRead, Meter ufsBytesReadThroughput) {
      mInitialBlockSize = blockMeta.getBlockSize();
      mBlockMeta = blockMeta;
      mLocalBlockStore = localBlockStore;
      mInStreamPos = -1;
      mUfsInstreamCache = ufsInStreamCache;
      mUfsResource = ufsClient.acquireUfsResource();
      mIsPositionShort = positionShort;
      mUfsBytesRead = ufsBytesRead;
      mUfsBytesReadThroughput = ufsBytesReadThroughput;
    }
  
    /**
     * Initializes the reader. This is only called in the factory method.
     *
     * @param offset the position within the block to start the read
     */
    private void init(long offset) throws IOException {
      updateUnderFileSystemInputStream(offset);
      updateBlockWriter(offset);
    }
  
    @Override
    public ReadableByteChannel getChannel() {
      throw new UnsupportedOperationException("UFSFileBlockReader#getChannel is not supported");
    }
  
    @Override
    public long getLength() {
      return mBlockMeta.getBlockSize();
    }
  
    @Override
    public ByteBuffer read(long offset, long length) throws IOException {
      Preconditions.checkState(!mClosed);
      updateUnderFileSystemInputStream(offset);
      updateBlockWriter(offset);
  
      long bytesToRead = Math.min(length, mBlockMeta.getBlockSize() - offset);
      if (bytesToRead <= 0) {
        return ByteBuffer.allocate(0);
      }
      byte[] data = new byte[(int) bytesToRead];
      int bytesRead = 0;
      Preconditions.checkNotNull(mUnderFileSystemInputStream, "mUnderFileSystemInputStream");
      while (bytesRead < bytesToRead) {
        int read;
        try {
          read = mUnderFileSystemInputStream.read(data, bytesRead, (int) (bytesToRead - bytesRead));
        } catch (IOException e) {
          throw AlluxioStatusException.fromIOException(e);
        }
        if (read == -1) {
          break;
        }
        bytesRead += read;
      }
      mInStreamPos += bytesRead;
  
      // We should always read the number of bytes as expected since the UFS file length (hence block
      // size) should be always accurate.
      Preconditions.checkState(bytesRead == bytesToRead,
          "Not enough bytes have been read [bytesRead: %s, bytesToRead: %s] from the UFS file: %s.",
          bytesRead, bytesToRead, mBlockMeta.getUnderFileSystemPath());
      if (mBlockWriter != null && mBlockWriter.getPosition() < mInStreamPos) {
        try {
          Preconditions.checkState(mBlockWriter.getPosition() >= offset);
          mLocalBlockStore.requestSpace(mBlockMeta.getSessionId(), mBlockMeta.getBlockId(),
              mInStreamPos - mBlockWriter.getPosition());
          ByteBuffer buffer = ByteBuffer.wrap(data, (int) (mBlockWriter.getPosition() - offset),
              (int) (mInStreamPos - mBlockWriter.getPosition()));
          mBlockWriter.append(buffer.duplicate());
        } catch (Exception e) {
          LOG.warn("Failed to cache data read from UFS (on read()): {}", e.toString());
          try {
            cancelBlockWriter();
          } catch (IOException ee) {
            LOG.error("Failed to cancel block writer:", ee);
          }
        }
      }
      mUfsBytesRead.inc(bytesRead);
      mUfsBytesReadThroughput.mark(bytesRead);
      return ByteBuffer.wrap(data, 0, bytesRead);
    }
  
    /**
     * This interface is supposed to be used for sequence block reads.
     *
     * @param buf the byte buffer
     * @return the number of bytes read, -1 if it reaches EOF and none was read
     */
    @Override
    public int transferTo(ByteBuf buf) throws IOException {
      Preconditions.checkState(!mClosed);
      if (mUnderFileSystemInputStream == null) {
        return -1;
      }
      if (mBlockMeta.getBlockSize() <= mInStreamPos) {
        return -1;
      }
     // Make a copy of the state to keep track of what we have read in this transferTo call.
      ByteBuf bufCopy = null;
      if (mBlockWriter != null) {
        bufCopy = buf.duplicate();
        bufCopy.readerIndex(bufCopy.writerIndex());
      }
      int bytesToRead =
          (int) Math.min(buf.writableBytes(), mBlockMeta.getBlockSize() - mInStreamPos);
      int bytesRead = buf.writeBytes(mUnderFileSystemInputStream, bytesToRead);
      if (bytesRead <= 0) {
        return bytesRead;
      }
  
      mInStreamPos += bytesRead;
  
      if (mBlockWriter != null && bufCopy != null) {
        try {
          bufCopy.writerIndex(buf.writerIndex());
          while (bufCopy.readableBytes() > 0) {
            mLocalBlockStore.requestSpace(mBlockMeta.getSessionId(), mBlockMeta.getBlockId(),
                mInStreamPos - mBlockWriter.getPosition());
            mBlockWriter.append(bufCopy);
          }
        } catch (Exception e) {
          LOG.warn("Failed to cache data read from UFS (on transferTo()): {}", e.toString());
          cancelBlockWriter();
        }
      }
      mUfsBytesRead.inc(bytesRead);
      mUfsBytesReadThroughput.mark(bytesRead);
      return bytesRead;
    }
  
    /**
     * Closes the block reader. After this, this block reader should not be used anymore.
     * This is recommended to be called after the client finishes reading the block. It is usually
     * triggered when the client unlocks the block.
     */
    @Override
    public void close() throws IOException {
      if (mClosed) {
        return;
      }
  
      super.close();
      try {
        // This aborts the block if the block is not fully read.
        updateBlockWriter(mBlockMeta.getBlockSize());
  
        if (mUnderFileSystemInputStream != null) {
          mUfsInstreamCache.release(mUnderFileSystemInputStream);
          mUnderFileSystemInputStream = null;
        }
  
        if (mBlockWriter != null) {
          mBlockWriter.close();
        }
  
        mUfsResource.close();
      } finally {
        mClosed = true;
        BLOCKS_READ_UFS.inc();
      }
    }
  
    @Override
    public boolean isClosed() {
      return mClosed;
    }
  
    @Override
    public String getLocation() {
      return mBlockMeta.getUnderFileSystemPath();
    }
  
    /**
     * Updates the UFS input stream given an offset to read.
     *
     * @param offset the read offset within the block
     */
    private void updateUnderFileSystemInputStream(long offset) throws IOException {
      if ((mUnderFileSystemInputStream != null) && offset != mInStreamPos) {
        mUfsInstreamCache.release(mUnderFileSystemInputStream);
        mUnderFileSystemInputStream = null;
        mInStreamPos = -1;
      }
  
      if (mUnderFileSystemInputStream == null && offset < mBlockMeta.getBlockSize()) {
        UnderFileSystem ufs = mUfsResource.get();
        mUnderFileSystemInputStream = mUfsInstreamCache
            .acquire(ufs, mBlockMeta.getUnderFileSystemPath(),
                IdUtils.fileIdFromBlockId(mBlockMeta.getBlockId()),
                OpenOptions.defaults().setOffset(mBlockMeta.getOffset() + offset)
                    .setPositionShort(mIsPositionShort));
        mInStreamPos = offset;
      }
    }
  
    /**
     * Closes the current block writer, cleans up its temp block and sets it to null.
     */
    private void cancelBlockWriter() throws IOException {
      if (mBlockWriter == null) {
        return;
      }
      try {
        mBlockWriter.close();
        mBlockWriter = null;
        mLocalBlockStore.abortBlock(mBlockMeta.getSessionId(), mBlockMeta.getBlockId());
      } catch (Exception e) {
        // We cannot skip the exception here because we need to make sure that the user of this
        // reader does not commit the block if it fails to abort the block.
        throw AlluxioStatusException.fromThrowable(e);
      }
    }
  
    /**
     * Updates the block writer given an offset to read. If the offset is beyond the current
     * position of the block writer, the block writer will be aborted.
     *
     * @param offset the read offset
     */
    private void updateBlockWriter(long offset) throws IOException {
      if (mBlockWriter != null && offset > mBlockWriter.getPosition()) {
        cancelBlockWriter();
      }
      try {
        if (mBlockWriter == null && offset == 0 && !mBlockMeta.isNoCache()) {
          BlockStoreLocation loc = BlockStoreLocation.anyDirInTier(
              WORKER_STORAGE_TIER_ASSOC.getAlias(0));
          mLocalBlockStore.createBlock(mBlockMeta.getSessionId(), mBlockMeta.getBlockId(),
              AllocateOptions.forCreate(mInitialBlockSize, loc));
          mBlockWriter = mLocalBlockStore.createBlockWriter(
              mBlockMeta.getSessionId(), mBlockMeta.getBlockId());
        }
      } catch (AlluxioRuntimeException e) {
        LOG.warn(
            "Failed to update block writer for UFS block [blockId: {}, ufsPath: {}, offset: {}]: {}",
            mBlockMeta.getBlockId(), mBlockMeta.getUnderFileSystemPath(), offset, e.toString());
        mBlockWriter = null;
      } catch (IllegalStateException e) {
        // This can happen when there are concurrent UFS readers who are all trying to cache to block.
        LOG.debug(
            "Failed to update block writer for UFS block [blockId: {}, ufsPath: {}, offset: {}]."
                + "Concurrent UFS readers may be caching the same block.",
            mBlockMeta.getBlockId(), mBlockMeta.getUnderFileSystemPath(), offset, e);
        mBlockWriter = null;
      }
    }
  }
  