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 com.google.common.collect.ImmutableList.toImmutableList;
  
  import alluxio.Constants;
  import alluxio.conf.Configuration;
  import alluxio.conf.PropertyKey;
  import alluxio.util.io.BufferUtils;
  import alluxio.util.io.FileUtils;
  import alluxio.util.io.PathUtils;
  import alluxio.worker.block.annotator.BlockIterator;
  import alluxio.worker.block.evictor.Evictor;
  import alluxio.worker.block.io.BlockWriter;
  import alluxio.worker.block.io.LocalFileBlockWriter;
  import alluxio.worker.block.meta.DefaultTempBlockMeta;
  import alluxio.worker.block.meta.StorageDir;
  import alluxio.worker.block.meta.TempBlockMeta;
  
  import com.google.common.base.Preconditions;
  import com.google.common.primitives.Ints;
  
  import java.util.Arrays;
  import java.util.Collections;
  
  /**
   * Utility methods for setting and testing {@link TieredBlockStore}.
   */
  public final class TieredBlockStoreTestUtils {
    /**
     * Default configurations of a TieredBlockStore for use in {@link #defaultMetadataManager}. They
     * represent a block store with a MEM tier and a SSD tier, there are two directories with capacity
     * 2000 bytes and 3000 bytes separately in the MEM tier and three directories with capacity 10000,
     * 20000, 30000 bytes separately in the SSD tier.
     */
    public static final int[] TIER_ORDINAL = {0, 1};
    public static final String[] TIER_ALIAS = {Constants.MEDIUM_MEM, Constants.MEDIUM_SSD};
    public static final String[][] TIER_PATH = {{"/mem/0", "/mem/1"}, {"/ssd/0", "/ssd/1", "/ssd/2"}};
    public static final long[][] TIER_CAPACITY_BYTES = {{2000, 3000}, {10000, 20000, 30000}};
    public static final String[][] TIER_MEDIA_TYPES =
        {{Constants.MEDIUM_MEM, Constants.MEDIUM_MEM},
            {Constants.MEDIUM_SSD, Constants.MEDIUM_SSD, Constants.MEDIUM_SSD}};
    public static final String WORKER_DATA_FOLDER = "/alluxioworker/";
  
    /**
     * Sets up a {@link Configuration} for a {@link TieredBlockStore} with several tiers
     * configured by the parameters. For simplicity, you can use {@link #setupDefaultConf(String)}
     * which calls this method with default values.
     *
     * @param baseDir the directory path as prefix for all the paths of directories in the tiered
     *        storage; when specified, the directory needs to exist before calling this method
     * @param tierOrdinal like {@link #TIER_ORDINAL}, length must be &gt; 0
     * @param tierAlias like {@link #TIER_ALIAS}, each corresponds to an element in tierLevel
     * @param tierPath like {@link #TIER_PATH}, each list represents directories of the tier with the
     *        same list index in tierAlias
     * @param tierCapacity like {@link #TIER_CAPACITY_BYTES}, should be in the same dimension with
     *        tierPath, each element is the capacity of the corresponding dir in tierPath
     * @param tierMediumType should be in the same dimension with tierPath, each element is the
     *                       medium type of the corresponding dir in tierPath
     * @param workerDataFolder when specified it sets up the alluxio.worker.data.folder property
     */
    public static void setupConfWithMultiTier(String baseDir, int[] tierOrdinal, String[] tierAlias,
        String[][] tierPath, long[][] tierCapacity, String[][] tierMediumType,
        String workerDataFolder) throws Exception {
      // make sure dimensions are legal
      Preconditions.checkNotNull(tierOrdinal, "tierOrdinal");
      Preconditions.checkNotNull(tierAlias, "tierAlias");
      Preconditions.checkNotNull(tierPath, "tierPath");
      Preconditions.checkNotNull(tierCapacity, "tierCapacity");
      Preconditions.checkNotNull(tierMediumType, "tierMediumType");
  
      Preconditions.checkArgument(tierOrdinal.length > 0, "length of tierLevel should be > 0");
      Preconditions.checkArgument(tierOrdinal.length == tierAlias.length,
          "tierAlias and tierLevel should have the same length");
      Preconditions.checkArgument(tierOrdinal.length == tierPath.length,
          "tierPath and tierLevel should have the same length");
      Preconditions.checkArgument(tierOrdinal.length == tierCapacity.length,
          "tierCapacity and tierLevel should have the same length");
      Preconditions.checkArgument(tierOrdinal.length == tierMediumType.length,
          "tierMediumType and tierLevel should have the same length");
      int nTier = tierOrdinal.length;
  
      tierPath = createDirHierarchy(baseDir, tierPath);
      if (workerDataFolder != null) {
        Configuration.set(PropertyKey.WORKER_DATA_FOLDER, workerDataFolder);
      }
      Configuration.set(PropertyKey.WORKER_TIERED_STORE_LEVELS, nTier);
  
      // sets up each tier in turn
      for (int i = 0; i < nTier; i++) {
        setupConfTier(tierOrdinal[i], tierAlias[i], tierPath[i], tierCapacity[i], tierMediumType[i]);
      }
    }
  
    /**
     * Sets up a {@link Configuration} for a {@link TieredBlockStore} with only *one tier*
     * configured by the parameters. For simplicity, you can use {@link #setupDefaultConf(String)}
     * which sets up the tierBlockStore with default values.
     *
     * This method modifies the configuration, so be sure to reset it when done.
     *
     * @param baseDir the directory path as prefix for all the paths of directories in the tiered
     *        storage; when specified, the directory needs to exist before calling this method
     * @param tierOrdinal ordinal of this tier
     * @param tierAlias alias of this tier
     * @param tierPath path of this tier; when `baseDir` is specified, the actual test tierPath turns
     *        into `baseDir/tierPath`
     * @param tierCapacity capacity of this tier
     * @param workerDataFolder when specified it sets up the alluxio.worker.data.folder property
     */
    public static void setupConfWithSingleTier(String baseDir, int tierOrdinal, String tierAlias,
        String[] tierPath, long[] tierCapacity, String[] tierMedia, String workerDataFolder)
        throws Exception {
      if (baseDir != null) {
        tierPath = createDirHierarchy(baseDir, tierPath);
      }
      if (workerDataFolder != null) {
        Configuration.set(PropertyKey.WORKER_DATA_FOLDER, workerDataFolder);
      }
      Configuration.set(PropertyKey.WORKER_TIERED_STORE_LEVELS, 1);
      setupConfTier(tierOrdinal, tierAlias, tierPath, tierCapacity, tierMedia);
    }
  
    /**
     * Sets up a specific tier's {@link Configuration} for a {@link TieredBlockStore}.
     *
     * @param tierAlias alias of the tier
     * @param tierPath absolute path of the tier
     * @param tierCapacity capacity of the tier
     */
    private static void setupConfTier(int ordinal, String tierAlias, String[] tierPath,
        long[] tierCapacity, String[] tierMediumType) {
      Preconditions.checkNotNull(tierPath, "tierPath");
      Preconditions.checkNotNull(tierCapacity, "tierCapacity");
      Preconditions.checkNotNull(tierMediumType, "tierMediumType");
      Preconditions.checkArgument(tierPath.length == tierCapacity.length,
          "tierPath and tierCapacity should have the same length");
  
      Configuration
          .set(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_ALIAS.format(ordinal),
              tierAlias);
  
      Configuration
          .set(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_PATH.format(ordinal),
              Arrays.stream(tierPath).collect(toImmutableList()));
  
      Configuration
          .set(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_QUOTA.format(ordinal),
              Arrays.stream(tierCapacity).boxed().map(String::valueOf).collect(toImmutableList()));
  
      Configuration
          .set(PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_MEDIUMTYPE.format(ordinal),
              Arrays.stream(tierMediumType).collect(toImmutableList()));
    }
  
    /**
     * Joins baseDir with all the paths listed in the array and then create the new generated path.
     *
     * @param baseDir the directory path as prefix for all the paths in the array 'dirs'
     * @param dirs 2-D array of directory paths
     * @return new joined and created paths array
     */
    private static String[][] createDirHierarchy(String baseDir, final String[][] dirs)
        throws Exception {
      if (baseDir == null) {
        return dirs;
      }
      String[][] newDirs = new String[dirs.length][];
      for (int i = 0; i < dirs.length; i++) {
        newDirs[i] = createDirHierarchy(baseDir, dirs[i]);
      }
      return newDirs;
    }
  
    /**
     * Joins baseDir with all the paths listed in the array and then create the new generated path.
     *
     * @param baseDir the directory path as prefix for all the paths in the array 'dirs'
     * @param dirs 1-D array of directory paths
     * @return new joined and created paths array
     */
    private static String[] createDirHierarchy(String baseDir, final String[] dirs) throws Exception {
      if (baseDir == null) {
        return dirs;
      }
      String[] newDirs = new String[dirs.length];
      for (int i = 0; i < dirs.length; i++) {
        newDirs[i] = PathUtils.concatPath(baseDir, dirs[i]);
        FileUtils.createDir(newDirs[i]);
      }
      return newDirs;
    }
  
    /**
     * Creates a BlockMetadataManager with {@link #setupDefaultConf(String)}.
     *
     * @param baseDir the directory path as prefix for paths of directories in the tiered storage; the
     *        directory needs to exist before calling this method
     * @return the created metadata manager
     */
    public static BlockMetadataManager defaultMetadataManager(String baseDir) throws Exception {
      setupDefaultConf(baseDir);
      return BlockMetadataManager.createBlockMetadataManager();
    }
  
    /**
     * Creates a {@link BlockMetadataEvictorView} with {@link #setupDefaultConf(String)}.
     *
     * @param baseDir the directory path as prefix for paths of directories in the tiered storage; the
     *        directory needs to exist before calling this method
     * @return the created metadata evictor view
     */
    public static BlockMetadataEvictorView defaultMetadataManagerView(String baseDir)
        throws Exception {
      BlockMetadataManager metaManager = TieredBlockStoreTestUtils.defaultMetadataManager(baseDir);
      return new BlockMetadataEvictorView(metaManager, Collections.<Long>emptySet(),
          Collections.<Long>emptySet());
    }
  
    /**
     * Sets up a {@link Configuration} with default values of {@link #TIER_ORDINAL},
     * {@link #TIER_ALIAS}, {@link #TIER_PATH} with the baseDir as path prefix,
     * {@link #TIER_CAPACITY_BYTES}.
     *
     * @param baseDir the directory path as prefix for paths of directories in the tiered storage; the
     *        directory needs to exist before calling this method
     */
    public static void setupDefaultConf(String baseDir) throws Exception {
      setupConfWithMultiTier(baseDir, TIER_ORDINAL, TIER_ALIAS, TIER_PATH, TIER_CAPACITY_BYTES,
          TIER_MEDIA_TYPES, WORKER_DATA_FOLDER);
    }
  
    /**
     * Caches bytes into {@link StorageDir}.
     *
     * @param sessionId session who caches the data
     * @param blockId id of the cached block
     * @param bytes size of the block in bytes
     * @param dir the {@link StorageDir} the block resides in
     * @param meta the metadata manager to update meta of the block
     * @param evictor the evictor to be informed of the new block
     */
    public static void cache(long sessionId, long blockId, long bytes, StorageDir dir,
        BlockMetadataManager meta, Evictor evictor) throws Exception {
      BlockStoreEventListener listener = null;
      if (evictor instanceof BlockStoreEventListener) {
        listener = (BlockStoreEventListener) evictor;
      }
      cache2(sessionId, blockId, bytes, dir, meta, listener);
    }
  
    /**
     * Caches bytes into {@link LocalBlockStore} at specific location.
     *
     * @param sessionId session who caches the data
     * @param blockId id of the cached block
     * @param bytes size of the block in bytes
     * @param blockStore block store that the block is written into
     * @param location the location where the block resides
     * @param pinOnCreate whether to pin block on create
     */
    public static void cache(long sessionId, long blockId, long bytes, LocalBlockStore blockStore,
        BlockStoreLocation location, boolean pinOnCreate) throws Exception {
      TempBlockMeta tempBlockMeta = blockStore.createBlock(sessionId, blockId,
          AllocateOptions.forCreate(bytes, location));
      // write data
      BlockWriter writer = new LocalFileBlockWriter(tempBlockMeta.getPath());
      writer.append(BufferUtils.getIncreasingByteBuffer(Ints.checkedCast(bytes)));
      writer.close();
  
      // commit block
      blockStore.commitBlock(sessionId, blockId, pinOnCreate);
    }
  
    /**
     * Caches bytes into {@link LocalBlockStore} at specific location.
     *
     * @param sessionId session who caches the data
     * @param blockId id of the cached block
     * @param options allocation options
     * @param blockStore block store that the block is written into
     * @param pinOnCreate whether to pin block on create
     */
    public static void cache(long sessionId, long blockId, AllocateOptions options,
        LocalBlockStore blockStore, boolean pinOnCreate) throws Exception {
      TempBlockMeta tempBlockMeta = blockStore.createBlock(sessionId, blockId, options);
      // write data
      BlockWriter writer = new LocalFileBlockWriter(tempBlockMeta.getPath());
      writer.append(BufferUtils.getIncreasingByteBuffer(Ints.checkedCast(options.getSize())));
      writer.close();
  
      // commit block
      blockStore.commitBlock(sessionId, blockId, pinOnCreate);
    }
  
    /**
     * Caches bytes into {@link StorageDir}.
     *
     * @param sessionId session who caches the data
     * @param blockId id of the cached block
     * @param bytes size of the block in bytes
     * @param tierLevel tier level of the {@link StorageDir} the block resides in
     * @param dirIndex index of directory in the tierLevel the block resides in
     * @param meta the metadata manager to update meta of the block
     * @param evictor the evictor to be informed of the new block
     */
    public static void cache(long sessionId, long blockId, long bytes, int tierLevel, int dirIndex,
        BlockMetadataManager meta, Evictor evictor) throws Exception {
      StorageDir dir = meta.getTiers().get(tierLevel).getDir(dirIndex);
      cache(sessionId, blockId, bytes, dir, meta, evictor);
    }
  
    /**
     * Caches bytes into {@link StorageDir}.
     *
     * @param sessionId session who caches the data
     * @param blockId id of the cached block
     * @param bytes size of the block in bytes
     * @param dir the {@link StorageDir} the block resides in
     * @param meta the metadata manager to update meta of the block
     * @param iterator the iterator to be informed of the new block
     */
    public static void cache2(long sessionId, long blockId, long bytes, StorageDir dir,
        BlockMetadataManager meta, BlockIterator iterator) throws Exception {
      cache2(sessionId, blockId, bytes, dir, meta, (BlockStoreEventListener) null);
      if (iterator != null) {
        for (BlockStoreEventListener listener : iterator.getListeners()) {
          listener.onCommitBlockToLocal(blockId, dir.toBlockStoreLocation());
        }
      }
    }
  
    /**
     * Caches bytes into {@link StorageDir}.
     *
     * @param sessionId session who caches the data
     * @param blockId id of the cached block
     * @param bytes size of the block in bytes
     * @param dir the {@link StorageDir} the block resides in
     * @param meta the metadata manager to update meta of the block
     * @param listener the listener to be informed of the new block
     */
    public static void cache2(long sessionId, long blockId, long bytes, StorageDir dir,
        BlockMetadataManager meta, BlockStoreEventListener listener) throws Exception {
      TempBlockMeta tempBlockMeta = createTempBlock(sessionId, blockId, bytes, dir);
  
      // commit block.
      FileUtils.move(tempBlockMeta.getPath(), tempBlockMeta.getCommitPath());
      meta.commitTempBlockMeta(tempBlockMeta);
  
      // update iterator if a listener.
      if (listener != null) {
        listener.onCommitBlockToLocal(blockId, dir.toBlockStoreLocation());
      }
    }
  
    /**
     * Makes a temp block of a given size in {@link StorageDir}.
     *
     * @param sessionId session who caches the data
     * @param blockId id of the cached block
     * @param bytes size of the block in bytes
     * @param dir the {@link StorageDir} the block resides in
     * @return the temp block meta
     */
    public static TempBlockMeta createTempBlock(long sessionId, long blockId, long bytes,
        StorageDir dir) throws Exception {
      // prepare temp block
      TempBlockMeta tempBlockMeta = new DefaultTempBlockMeta(sessionId, blockId, bytes, dir);
      dir.addTempBlockMeta(tempBlockMeta);
      // write data
      FileUtils.createFile(tempBlockMeta.getPath());
      BlockWriter writer = new LocalFileBlockWriter(tempBlockMeta.getPath());
      writer.append(BufferUtils.getIncreasingByteBuffer(Ints.checkedCast(bytes)));
      writer.close();
      return tempBlockMeta;
    }
  
    /**
     * Gets the total capacity of all tiers in bytes.
     *
     * @return total capacity of all tiers in bytes
     */
    public static long getDefaultTotalCapacityBytes() {
      long totalCapacity = 0;
      for (long[] tierCapacityBytes : TIER_CAPACITY_BYTES) {
        for (long tierCapacityByte : tierCapacityBytes) {
          totalCapacity += tierCapacityByte;
        }
      }
      return totalCapacity;
    }
  
    /**
     * Gets the number of testing directories of all tiers.
     *
     * @return number of testing directories of all tiers
     */
    public static long getDefaultDirNum() {
      int dirNum = 0;
      for (String[] tierPath : TIER_PATH) {
        dirNum += tierPath.length;
      }
      return dirNum;
    }
  }
  