Merge pull request Alluxio#4546 from jsimsa/tier

[SMALLFIX] Generalizing the write tier index API.

core/client/src/main/java/alluxio/client/block/BlockWorkerClient.java

@@ -100,10 +100,12 @@ public interface BlockWorkerClient extends Closeable {
    *
    * @param blockId the ID of the block
    * @param initialBytes the initial size bytes allocated for the block
+   * @param tier the target tier
    * @return the temporary path of the block
    * @throws IOException if a non-Alluxio exception occurs
    */
-  String requestBlockLocation(final long blockId, final long initialBytes) throws IOException;
+  String requestBlockLocation(final long blockId, final long initialBytes, final int tier)
+      throws IOException;
 
   /**
    * Requests space for some block from worker.

core/client/src/main/java/alluxio/client/block/LocalBlockOutStream.java

@@ -70,7 +70,8 @@ public LocalBlockOutStream(long blockId,
     try {
       mBlockWorkerClient = mCloser.register(context.createBlockWorkerClient(workerNetAddress));
       long initialSize = Configuration.getBytes(PropertyKey.USER_FILE_BUFFER_BYTES);
-      String blockPath = mBlockWorkerClient.requestBlockLocation(mBlockId, initialSize);
+      String blockPath =
+          mBlockWorkerClient.requestBlockLocation(mBlockId, initialSize, options.getWriteTier());
       mReservedBytes += initialSize;
       mWriter = new LocalFileBlockWriter(blockPath);
       mCloser.register(mWriter);

core/client/src/main/java/alluxio/client/block/RetryHandlingBlockWorkerClient.java

@@ -16,15 +16,13 @@
 import alluxio.Constants;
 import alluxio.PropertyKey;
 import alluxio.RuntimeConstants;
-import alluxio.client.WriteTier;
 import alluxio.exception.AlluxioException;
 import alluxio.exception.ExceptionMessage;
 import alluxio.exception.WorkerOutOfSpaceException;
 import alluxio.metrics.MetricsSystem;
 import alluxio.thrift.AlluxioTException;
 import alluxio.thrift.BlockWorkerClientService;
 import alluxio.thrift.ThriftIOException;
-import alluxio.thrift.TWriteTier;
 import alluxio.util.ThreadFactoryUtils;
 import alluxio.util.network.NetworkAddressUtils;
 import alluxio.wire.LockBlockResult;
@@ -67,8 +65,6 @@ public final class RetryHandlingBlockWorkerClient
       ThreadFactoryUtils.build("block-worker-heartbeat-cancel-%d", true));
   private final BlockWorkerThriftClientPool mClientPool;
   private final BlockWorkerThriftClientPool mClientHeartbeatPool;
-  private static final TWriteTier WRITE_TIER = WriteTier.toThrift(
-      alluxio.Configuration.getEnum(PropertyKey.USER_FILE_WRITE_TIER_DEFAULT, WriteTier.class));
   // Tracks the number of active heartbeat close requests.
   private static final AtomicInteger NUM_ACTIVE_SESSIONS = new AtomicInteger(0);
 
@@ -249,15 +245,15 @@ public Void call(BlockWorkerClientService.Client client)
   }
 
   @Override
-  public String requestBlockLocation(final long blockId, final long initialBytes)
-      throws IOException {
+  public String requestBlockLocation(final long blockId, final long initialBytes,
+      final int writeTier) throws IOException {
     try {
       return retryRPC(
           new RpcCallableThrowsAlluxioTException<String, BlockWorkerClientService.Client>() {
             @Override
             public String call(BlockWorkerClientService.Client client)
                 throws AlluxioTException, TException {
-              return client.requestBlockLocation(getSessionId(), blockId, initialBytes, WRITE_TIER);
+              return client.requestBlockLocation(getSessionId(), blockId, initialBytes, writeTier);
             }
           });
     } catch (WorkerOutOfSpaceException e) {

core/client/src/main/java/alluxio/client/block/stream/BlockOutStream.java

@@ -63,8 +63,8 @@ public static BlockOutStream createLocalBlockOutStream(long blockId, long blockS
     Closer closer = Closer.create();
     try {
       BlockWorkerClient client = closer.register(context.createBlockWorkerClient(workerNetAddress));
-      PacketOutStream outStream =
-          PacketOutStream.createLocalPacketOutStream(client, blockId, blockSize);
+      PacketOutStream outStream = PacketOutStream
+          .createLocalPacketOutStream(client, blockId, blockSize, options.getWriteTier());
       closer.register(outStream);
       return new BlockOutStream(outStream, blockId, blockSize, client, options);
     } catch (IOException e) {

core/client/src/main/java/alluxio/client/block/stream/LocalFilePacketWriter.java

@@ -48,12 +48,13 @@ public final class LocalFilePacketWriter implements PacketWriter {
    *
    * @param blockWorkerClient the block worker client, not owned by this class
    * @param blockId the block ID
+   * @param tier the target tier
    * @throws IOException if it fails to create the packet writer
    * @return the {@link LocalFilePacketWriter} created
    */
   public static LocalFilePacketWriter create(BlockWorkerClient blockWorkerClient,
-      long blockId) throws IOException {
-    return new LocalFilePacketWriter(blockWorkerClient, blockId);
+      long blockId, int tier) throws IOException {
+    return new LocalFilePacketWriter(blockWorkerClient, blockId, tier);
   }
 
   @Override
@@ -104,11 +105,13 @@ public void close() throws IOException {
    *
    * @param blockWorkerClient the block worker client, not owned by this class
    * @param blockId the block ID
+   * @param tier the target tier
    * @throws IOException if it fails to create the packet writer
    */
-  private LocalFilePacketWriter(BlockWorkerClient blockWorkerClient, long blockId)
+  private LocalFilePacketWriter(BlockWorkerClient blockWorkerClient, long blockId, int tier)
       throws IOException {
-    String blockPath = blockWorkerClient.requestBlockLocation(blockId, FILE_BUFFER_BYTES);
+    String blockPath =
+        blockWorkerClient.requestBlockLocation(blockId, FILE_BUFFER_BYTES, tier);
     mWriter = new LocalFileBlockWriter(blockPath);
     mPosReserved += FILE_BUFFER_BYTES;
     mBlockId = blockId;

core/client/src/main/java/alluxio/client/block/stream/PacketOutStream.java

@@ -52,12 +52,13 @@ public final class PacketOutStream extends OutputStream implements BoundedStream
    * @param client the block worker client
    * @param id the ID
    * @param length the block or file length
+   * @param tier the target tier
    * @return the {@link PacketOutStream} created
    * @throws IOException if it fails to create the object
    */
   public static PacketOutStream createLocalPacketOutStream(BlockWorkerClient client,
-      long id, long length) throws IOException {
-    PacketWriter packetWriter = LocalFilePacketWriter.create(client, id);
+      long id, long length, int tier) throws IOException {
+    PacketWriter packetWriter = LocalFilePacketWriter.create(client, id, tier);
     return new PacketOutStream(packetWriter, length);
   }
 
@@ -88,19 +89,20 @@ public static PacketOutStream createNettyPacketOutStream(FileSystemContext conte
    * @param clients a list of block worker clients
    * @param id the ID (block ID or UFS file ID)
    * @param length the block or file length
+   * @param tier the target tier
    * @param type the request type (either block write or UFS file write)
    * @return the {@link PacketOutStream} created
    * @throws IOException if it fails to create the object
    */
   public static PacketOutStream createReplicatedPacketOutStream(FileSystemContext context,
-      List<BlockWorkerClient> clients, long id, long length,
+      List<BlockWorkerClient> clients, long id, long length, int tier,
       Protocol.RequestType type) throws IOException {
     String localHost = NetworkAddressUtils.getLocalHostName();
 
     List<PacketWriter> packetWriters = new ArrayList<>();
     for (BlockWorkerClient client : clients) {
       if (client.getWorkerNetAddress().getHost().equals(localHost)) {
-        packetWriters.add(LocalFilePacketWriter.create(client, id));
+        packetWriters.add(LocalFilePacketWriter.create(client, id, tier));
       } else {
         packetWriters.add(new NettyPacketWriter(context, client.getDataServerAddress(), id, length,
             client.getSessionId(), type));

core/client/src/main/java/alluxio/client/file/options/CreateFileOptions.java

@@ -43,6 +43,7 @@ public final class CreateFileOptions {
   private long mTtl;
   private TtlAction mTtlAction;
   private Mode mMode; // null if creating the file using system default mode
+  private int mWriteTier;
   private WriteType mWriteType;
 
   /**
@@ -62,6 +63,7 @@ private CreateFileOptions() {
     } catch (Exception e) {
       throw Throwables.propagate(e);
     }
+    mWriteTier = Configuration.getInt(PropertyKey.USER_FILE_WRITE_TIER_DEFAULT);
     mWriteType = Configuration.getEnum(PropertyKey.USER_FILE_WRITE_TYPE_DEFAULT, WriteType.class);
     mTtl = Constants.NO_TTL;
     mTtlAction = TtlAction.DELETE;
@@ -112,6 +114,13 @@ public Mode getMode() {
     return mMode;
   }
 
+  /**
+   * @return the write tier
+   */
+  public int getWriteTier() {
+    return mWriteTier;
+  }
+
   /**
    * @return the write type
    */
@@ -199,6 +208,15 @@ public CreateFileOptions setTtlAction(TtlAction ttlAction) {
     return this;
   }
 
+  /**
+   * @param writeTier the write tier to use for this operation
+   * @return the updated options object
+   */
+  public CreateFileOptions setWriteTier(int writeTier) {
+    mWriteTier = writeTier;
+    return this;
+  }
+
   /**
    * @param writeType the {@link WriteType} to use for this operation. This will override both the
    *        {@link AlluxioStorageType} and {@link UnderStorageType}.
@@ -219,6 +237,7 @@ public OutStreamOptions toOutStreamOptions() {
         .setMode(mMode)
         .setTtl(mTtl)
         .setTtlAction(mTtlAction)
+        .setWriteTier(mWriteTier)
         .setWriteType(mWriteType);
   }
 
@@ -237,13 +256,14 @@ public boolean equals(Object o) {
         && Objects.equal(mMode, that.mMode)
         && Objects.equal(mTtl, that.mTtl)
         && Objects.equal(mTtlAction, that.mTtlAction)
+        && mWriteTier == that.mWriteTier
         && Objects.equal(mWriteType, that.mWriteType);
   }
 
   @Override
   public int hashCode() {
     return Objects.hashCode(mRecursive, mBlockSizeBytes, mLocationPolicy, mMode, mTtl,
-        mTtlAction, mWriteType);
+        mTtlAction, mWriteTier, mWriteType);
   }
 
   @Override
@@ -255,6 +275,7 @@ public String toString() {
         .add("mode", mMode)
         .add("ttl", mTtl)
         .add("ttlAction", mTtlAction)
+        .add("writeTier", mWriteTier)
         .add("writeType", mWriteType)
         .toString();
   }

core/client/src/main/java/alluxio/client/file/options/OutStreamOptions.java

@@ -41,6 +41,7 @@ public final class OutStreamOptions {
   private long mTtl;
   private TtlAction mTtlAction;
   private FileWriteLocationPolicy mLocationPolicy;
+  private int mWriteTier;
   private WriteType mWriteType;
   private Permission mPermission;
   private String mUfsPath;
@@ -64,6 +65,7 @@ private OutStreamOptions() {
     } catch (Exception e) {
       throw Throwables.propagate(e);
     }
+    mWriteTier = Configuration.getInt(PropertyKey.USER_FILE_WRITE_TIER_DEFAULT);
     mWriteType = Configuration.getEnum(PropertyKey.USER_FILE_WRITE_TYPE_DEFAULT, WriteType.class);
     mPermission = Permission.defaults();
     try {
@@ -131,6 +133,13 @@ public String getUfsPath() {
     return mUfsPath;
   }
 
+  /**
+   * @return the write tier
+   */
+  public int getWriteTier() {
+    return mWriteTier;
+  }
+
   /**
    * @return the write type
    */
@@ -180,6 +189,17 @@ public OutStreamOptions setLocationPolicy(FileWriteLocationPolicy locationPolicy
     return this;
   }
 
+  /**
+   * Sets the write tier.
+   *
+   * @param writeTier the write tier to use for this operation
+   * @return the updated options object
+   */
+  public OutStreamOptions setWriteTier(int writeTier) {
+    mWriteTier = writeTier;
+    return this;
+  }
+
   /**
    * Sets the {@link WriteType}.
    *
@@ -239,6 +259,7 @@ public boolean equals(Object o) {
         && Objects.equal(mTtl, that.mTtl)
         && Objects.equal(mTtlAction, that.mTtlAction)
         && Objects.equal(mLocationPolicy, that.mLocationPolicy)
+        && mWriteTier == that.mWriteTier
         && Objects.equal(mWriteType, that.mWriteType)
         && Objects.equal(mUfsPath, that.mUfsPath)
         && Objects.equal(mPermission, that.mPermission);
@@ -250,6 +271,7 @@ public int hashCode() {
         mTtl,
         mTtlAction,
         mLocationPolicy,
+        mWriteTier,
         mWriteType,
         mUfsPath,
         mPermission);
@@ -260,8 +282,9 @@ public String toString() {
     return Objects.toStringHelper(this)
         .add("blockSizeBytes", mBlockSizeBytes)
         .add("ttl", mTtl)
-        .add("mTtlAction", mTtlAction)
+        .add("ttlAction", mTtlAction)
         .add("locationPolicy", mLocationPolicy)
+        .add("writeTier", mWriteTier)
         .add("writeType", mWriteType)
         .add("permission", mPermission)
         .add("ufsPath", mUfsPath)

core/client/src/test/java/alluxio/client/block/AlluxioBlockStoreTest.java

@@ -218,7 +218,8 @@ public void getOutStreamMissingLocationPolicy() throws IOException {
   @Test
   public void getOutStreamLocal() throws Exception {
     File tmp = mTestFolder.newFile();
-    Mockito.when(mBlockWorkerClient.requestBlockLocation(Matchers.eq(BLOCK_ID), Matchers.anyLong()))
+    Mockito.when(mBlockWorkerClient
+        .requestBlockLocation(Matchers.eq(BLOCK_ID), Matchers.anyLong(), Matchers.anyInt()))
         .thenReturn(tmp.getAbsolutePath());
 
     OutStreamOptions options = OutStreamOptions.defaults().setBlockSizeBytes(BLOCK_LENGTH)