add the block source to decide if triggering the caching

core/client/fs/src/main/java/alluxio/client/block/AlluxioBlockStore.java

@@ -14,6 +14,7 @@
 import alluxio.client.block.policy.BlockLocationPolicy;
 import alluxio.client.block.policy.options.GetWorkerOptions;
 import alluxio.client.block.stream.BlockInStream;
+import alluxio.client.block.stream.BlockInStream.BlockInStreamSource;
 import alluxio.client.block.stream.BlockOutStream;
 import alluxio.client.file.FileSystemContext;
 import alluxio.client.file.options.InStreamOptions;
@@ -133,6 +134,7 @@ public BlockInStream getInStream(long blockId, Protocol.OpenUfsBlockOptions open
       blockInfo = masterClientResource.get().getBlockInfo(blockId);
     }
 
+    BlockInStreamSource source = BlockInStreamSource.UFS;
     if (blockInfo.getLocations().isEmpty() && openUfsBlockOptions == null) {
       throw new NotFoundException("Block " + blockId + " is unavailable in both Alluxio and UFS.");
     }
@@ -157,6 +159,7 @@ public BlockInStream getInStream(long blockId, Protocol.OpenUfsBlockOptions open
       WorkerNetAddress workerNetAddress = location.getWorkerAddress();
       if (workerNetAddress.getHost().equals(mLocalHostName)) {
         address = workerNetAddress;
+        source = BlockInStreamSource.LOCAL;
         break;
       }
     }
@@ -165,9 +168,10 @@ public BlockInStream getInStream(long blockId, Protocol.OpenUfsBlockOptions open
       // only randomize among locations in the highest tier, or have the master randomize the order.
       List<BlockLocation> locations = blockInfo.getLocations();
       address = locations.get(mRandom.nextInt(locations.size())).getWorkerAddress();
+      source = BlockInStreamSource.REMOTE;
     }
-    return BlockInStream
-        .create(mContext, blockId, blockInfo.getLength(), address, openUfsBlockOptions, options);
+    return BlockInStream.create(mContext, blockId, blockInfo.getLength(), address, source,
+        openUfsBlockOptions, options);
   }
 
   /**

core/client/fs/src/main/java/alluxio/client/block/stream/BlockInStream.java

@@ -42,13 +42,18 @@
 @NotThreadSafe
 public class BlockInStream extends InputStream implements BoundedStream, Seekable,
     PositionedReadable, Locatable {
+  /** the source tracking where the block is from. */
+  public enum BlockInStreamSource {
+    LOCAL, REMOTE, UFS
+  }
+
   /** 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];
-  private final boolean mLocal;
+  private final BlockInStreamSource mInStreamSource;
   private final WorkerNetAddress mAddress;
 
   /** Current position of the stream, relative to the start of the block. */
@@ -69,17 +74,19 @@ public class BlockInStream extends InputStream implements BoundedStream, Seekabl
    * @param blockId the block ID
    * @param blockSize the block size in bytes
    * @param address the Alluxio worker address
+   * @param blockSource the source location of the block
    * @param openUfsBlockOptions the options to open a UFS block, set to null if this is block is
    *        not persisted in UFS
    * @param options the in stream options
    * @return the {@link InputStream} object
    */
   public static BlockInStream create(FileSystemContext context, long blockId, long blockSize,
-      WorkerNetAddress address, Protocol.OpenUfsBlockOptions openUfsBlockOptions,
-      InStreamOptions options) throws IOException {
-    if (CommonUtils.isLocalHost(address) && Configuration
-        .getBoolean(PropertyKey.USER_SHORT_CIRCUIT_ENABLED) && !NettyUtils
-        .isDomainSocketSupported(address)) {
+      WorkerNetAddress address, BlockInStreamSource blockSource,
+      Protocol.OpenUfsBlockOptions openUfsBlockOptions, InStreamOptions options)
+          throws IOException {
+    if (Configuration.getBoolean(PropertyKey.USER_SHORT_CIRCUIT_ENABLED)
+        && !NettyUtils.isDomainSocketSupported(address)
+        && blockSource == BlockInStreamSource.LOCAL) {
       try {
         return createLocalBlockInStream(context, address, blockId, blockSize, options);
       } catch (NotFoundException e) {
@@ -93,7 +100,8 @@ public static BlockInStream create(FileSystemContext context, long blockId, long
       builder.setOpenUfsBlockOptions(openUfsBlockOptions);
     }
 
-    return createNettyBlockInStream(context, address, builder.buildPartial(), blockSize, options);
+    return createNettyBlockInStream(context, address, blockSource, builder.buildPartial(),
+        blockSize, options);
   }
 
   /**
@@ -112,44 +120,47 @@ private static BlockInStream createLocalBlockInStream(FileSystemContext context,
     long packetSize = Configuration.getBytes(PropertyKey.USER_LOCAL_READER_PACKET_SIZE_BYTES);
     return new BlockInStream(
         new LocalFilePacketReader.Factory(context, address, blockId, packetSize, options), address,
-        blockId, length);
+        BlockInStreamSource.LOCAL, blockId, length);
   }
 
   /**
    * Creates a {@link BlockInStream} to read from a netty data server.
    *
    * @param context the file system context
    * @param address the address of the netty data server
+   * @param blockSource the source location of the block
    * @param blockSize the block size
    * @param readRequestPartial the partial read request
    * @param options the in stream options
    * @return the {@link BlockInStream} created
    */
   private static BlockInStream createNettyBlockInStream(FileSystemContext context,
-      WorkerNetAddress address, Protocol.ReadRequest readRequestPartial, long blockSize,
-      InStreamOptions options) {
+      WorkerNetAddress address, BlockInStreamSource blockSource,
+      Protocol.ReadRequest readRequestPartial, long blockSize, InStreamOptions options) {
     long packetSize =
         Configuration.getBytes(PropertyKey.USER_NETWORK_NETTY_READER_PACKET_SIZE_BYTES);
     PacketReader.Factory factory = new NettyPacketReader.Factory(context, address,
         readRequestPartial.toBuilder().setPacketSize(packetSize).buildPartial(), options);
-    return new BlockInStream(factory, address, readRequestPartial.getBlockId(), blockSize);
+    return new BlockInStream(factory, address, blockSource, readRequestPartial.getBlockId(),
+        blockSize);
   }
 
   /**
    * Creates an instance of {@link BlockInStream}.
    *
    * @param packetReaderFactory the packet reader factory
    * @param address the worker network address
+   * @param blockSource the source location of the block
    * @param id the ID (either block ID or UFS file ID)
    * @param length the length
    */
   protected BlockInStream(PacketReader.Factory packetReaderFactory, WorkerNetAddress address,
-      long id, long length) {
+      BlockInStreamSource blockSource, long id, long length) {
     mPacketReaderFactory = packetReaderFactory;
     mId = id;
     mLength = length;
     mAddress = address;
-    mLocal = CommonUtils.isLocalHost(mAddress);
+    mInStreamSource = blockSource;
   }
 
   @Override
@@ -328,6 +339,13 @@ public WorkerNetAddress location() {
 
   @Override
   public boolean isLocal() {
-    return mLocal;
+    return CommonUtils.isLocalHost(mAddress);
+  }
+
+  /**
+   * @return the source of the block location
+   */
+  public BlockInStreamSource Source() {
+    return mInStreamSource;
   }
 }

core/client/fs/src/main/java/alluxio/client/file/FileInStream.java

@@ -21,6 +21,7 @@
 import alluxio.client.PositionedReadable;
 import alluxio.client.block.AlluxioBlockStore;
 import alluxio.client.block.stream.BlockInStream;
+import alluxio.client.block.stream.BlockInStream.BlockInStreamSource;
 import alluxio.client.block.stream.BlockOutStream;
 import alluxio.client.file.options.InStreamOptions;
 import alluxio.client.file.options.OutStreamOptions;
@@ -85,7 +86,7 @@ public class FileInStream extends InputStream implements BoundedStream, Seekable
    * Caches the entire block even if only a portion of the block is read. Only valid when
    * mShouldCache is true.
    */
-  private final boolean mCachePartiallyReadBlockEnabled;
+  private final boolean mCachePartiallyReadBlock;
   /** Whether to cache blocks in this file into Alluxio. */
   private final boolean mShouldCache;
 
@@ -131,7 +132,7 @@ protected FileInStream(URIStatus status, InStreamOptions options, FileSystemCont
     mContext = context;
     mAlluxioStorageType = options.getAlluxioStorageType();
     mShouldCache = mAlluxioStorageType.isStore();
-    mCachePartiallyReadBlockEnabled = options.isCachePartiallyReadBlock();
+    mCachePartiallyReadBlock = options.isCachePartiallyReadBlock();
     mClosed = false;
     if (mShouldCache) {
       Preconditions.checkNotNull(options.getCacheLocationPolicy(),
@@ -312,7 +313,7 @@ public void seek(long pos) throws IOException {
    * @return if the partially-read block should be cached to the local worker
    */
   private boolean shouldCachePartiallyReadBlock() {
-    return mShouldCache && mCachePartiallyReadBlockEnabled && mCurrentCacheStream != null;
+    return mShouldCache && mCachePartiallyReadBlock;
   }
 
   @Override
@@ -476,20 +477,17 @@ private void updateStreams() throws IOException {
 
   /**
    * Updates {@link #mCurrentCacheStream}. When {@code mShouldCache} is true, {@code FileInStream}
-   * will create an {@code BlockOutStream} to cache the data read only if
-   * <ol>
-   *   <li>the file is read from under storage, or</li>
-   *   <li>the file is read from a remote worker and we have an available local worker.</li>
-   * </ol>
+   * will create an {@code BlockOutStream} to cache the data read only if the file is read from a
+   * remote worker or UFS and we have an available local worker.
+   *
    * The following preconditions are checked inside:
    * <ol>
-   *   <li>{@link #mCurrentCacheStream} is either done or null.</li>
-   *   <li>EOF is reached or {@link #mCurrentBlockInStream} must be valid.</li>
+   * <li>{@link #mCurrentCacheStream} is either done or null.</li>
+   * <li>EOF is reached or {@link #mCurrentBlockInStream} must be valid.</li>
    * </ol>
    * After this call, {@link #mCurrentCacheStream} is either null or freshly created.
-   * {@link #mCurrentCacheStream} is created only if the block is not cached in a chosen machine
-   * and mPos is at the beginning of a block.
-   * This function is only called by {@link #updateStreams()}.
+   * {@link #mCurrentCacheStream} is created only if the block is not cached in a chosen machine and
+   * mPos is at the beginning of a block. This function is only called by {@link #updateStreams()}.
    *
    * @param blockId cached result of {@link #getCurrentBlockId()}
    */
@@ -504,7 +502,9 @@ private void updateCacheStream(long blockId) {
       return;
     }
     Preconditions.checkNotNull(mCurrentBlockInStream, "mCurrentBlockInStream");
-    if (!mShouldCache || mCurrentBlockInStream.isLocal()) {
+
+    // do not create cache stream when the block is in local worker
+    if (!mShouldCache || mCurrentBlockInStream.Source() == BlockInStreamSource.LOCAL) {
       return;
     }
 
@@ -595,22 +595,24 @@ private void seekInternalWithCachingPartiallyReadBlock(long pos) throws IOExcept
     // Precompute this because mPos will be updated several times in this function.
     final boolean isInCurrentBlock = pos / mBlockSize == mPos / mBlockSize;
 
-    // Make sure that mCurrentBlockInStream and mCurrentCacheStream is updated.
-    // mPos is not updated here.
-    updateStreams();
+    if (!(mPos == 0 && mCurrentBlockInStream == null && !isInCurrentBlock)) {
+      // Make sure that mCurrentBlockInStream and mCurrentCacheStream is updated.
+      // mPos is not updated here.
+      updateStreams();
 
-    // Cache till pos if seeking forward within the current block. Otherwise cache the whole
-    // block.
-    cacheCurrentBlockToPos(pos > mPos ? pos : Long.MAX_VALUE);
+      // Cache till pos if seeking forward within the current block. Otherwise cache the whole
+      // block.
+      cacheCurrentBlockToPos(pos > mPos ? pos : Long.MAX_VALUE);
 
-    // Early return if we are at pos already. This happens if we seek forward with caching
-    // enabled for this block.
-    if (mPos == pos) {
-      return;
+      // Early return if we are at pos already. This happens if we seek forward with caching
+      // enabled for this block.
+      if (mPos == pos) {
+        return;
+      }
+      // The early return above guarantees that we won't close an incomplete cache stream.
+      Preconditions.checkState(mCurrentCacheStream == null || mCurrentCacheStream.remaining() == 0);
+      closeOrCancelCacheStream();
     }
-    // The early return above guarantees that we won't close an incomplete cache stream.
-    Preconditions.checkState(mCurrentCacheStream == null || mCurrentCacheStream.remaining() == 0);
-    closeOrCancelCacheStream();
 
     // If seeks within the current block, directly seeks to pos if we are not yet there.
     // If seeks outside the current block, seek to the beginning of that block first, then
@@ -627,7 +629,13 @@ private void seekInternalWithCachingPartiallyReadBlock(long pos) throws IOExcept
     } else {
       mPos = pos / mBlockSize * mBlockSize;
       updateStreams();
-      cacheCurrentBlockToPos(pos);
+      if (mCurrentCacheStream != null) {
+        cacheCurrentBlockToPos(pos);
+      } else if (mCurrentBlockInStream != null) {
+        seekInternal(pos);
+      } else {
+        Preconditions.checkState(remaining() == 0);
+      }
     }
   }
 
@@ -645,10 +653,6 @@ private void cacheCurrentBlockToPos(long pos) throws IOException {
     if (len <= 0) {
       return;
     }
-    if (mPos % mBlockSize == 0 && pos - mPos >= mBlockSize) {
-      closeOrCancelCacheStream();
-      return;
-    }
 
     do {
       // Account for the last read which might be less than mSeekBufferSizeBytes bytes.

core/client/fs/src/test/java/alluxio/client/block/stream/TestBlockInStream.java

@@ -17,8 +17,10 @@
  * A {@link BlockInStream} which reads from the given byte array.
  */
 public class TestBlockInStream extends BlockInStream {
-  public TestBlockInStream(byte[] mData, long id, long length, boolean shortCircuit) {
-    super(new Factory(mData, shortCircuit), new WorkerNetAddress().setHost("local"), id, length);
+  public TestBlockInStream(byte[] mData, long id, long length, boolean shortCircuit,
+      BlockInStreamSource source) {
+    super(new Factory(mData, shortCircuit), new WorkerNetAddress().setHost("local"), source, id,
+        length);
   }
 
   /**