alluxio-master-2.x/core/server/worker/src/test/java/alluxio/worker/block/management/tier/BaseTierManagementTaskTest.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.worker.block.management.tier;

import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.worker.block.AllocateOptions;
import alluxio.worker.block.BlockMetadataManager;
import alluxio.worker.block.BlockStoreLocation;
import alluxio.worker.block.TieredBlockStore;
import alluxio.worker.block.TieredBlockStoreTestUtils;
import alluxio.worker.block.annotator.BlockIterator;
import alluxio.worker.block.annotator.LRUAnnotator;
import alluxio.worker.block.io.BlockWriter;
import alluxio.worker.block.meta.StorageDir;

import org.junit.Rule;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.lang.reflect.Field;

/**
 * Provides base functionality for tier-management tests.
 */
public abstract class BaseTierManagementTaskTest {
  protected static final String FIRST_TIER_ALIAS = TieredBlockStoreTestUtils.TIER_ALIAS[0];
  protected static final String SECOND_TIER_ALIAS = TieredBlockStoreTestUtils.TIER_ALIAS[1];
  protected static final long SIMULATE_LOAD_SESSION_ID = 1;
  protected static final long SIMULATE_LOAD_BLOCK_ID = 1;
  protected static final long SMALL_BLOCK_SIZE = 10;
  protected static final long BLOCK_SIZE = 100;

  @Rule
  public TemporaryFolder mTestFolder = new TemporaryFolder();

  protected TieredBlockStore mBlockStore;
  protected BlockMetadataManager mMetaManager;
  protected BlockIterator mBlockIterator;

  protected StorageDir mTestDir1;
  protected StorageDir mTestDir2;
  protected StorageDir mTestDir3;
  protected StorageDir mTestDir4;

  protected BlockWriter mSimulateWriter;

  /**
   * Sets up all dependencies before a test runs.
   */
  protected void init() throws Exception {
    // Disable reviewer to make sure the allocator behavior stays deterministic
    Configuration.set(PropertyKey.WORKER_REVIEWER_CLASS,
            "alluxio.worker.block.reviewer.AcceptingReviewer");
    // Use LRU for stronger overlap guarantee.
    Configuration.set(PropertyKey.WORKER_BLOCK_ANNOTATOR_CLASS, LRUAnnotator.class.getName());
    Configuration.set(PropertyKey.USER_BLOCK_SIZE_BYTES_DEFAULT, BLOCK_SIZE);
    // Set timeout for faster task execution.
    Configuration.set(PropertyKey.WORKER_MANAGEMENT_LOAD_DETECTION_COOL_DOWN_TIME, "100ms");

    File tempFolder = mTestFolder.newFolder();
    TieredBlockStoreTestUtils.setupDefaultConf(tempFolder.getAbsolutePath());
    mBlockStore = new TieredBlockStore();
    Field field = mBlockStore.getClass().getDeclaredField("mMetaManager");
    field.setAccessible(true);
    mMetaManager = (BlockMetadataManager) field.get(mBlockStore);
    mBlockIterator = mMetaManager.getBlockIterator();

    mTestDir1 = mMetaManager.getTier(FIRST_TIER_ALIAS).getDir(0);
    mTestDir2 = mMetaManager.getTier(FIRST_TIER_ALIAS).getDir(1);
    mTestDir3 = mMetaManager.getTier(SECOND_TIER_ALIAS).getDir(1);
    mTestDir4 = mMetaManager.getTier(SECOND_TIER_ALIAS).getDir(2);
  }

  /**
   * Stars simulating load on the worker.
   */
  protected void startSimulateLoad() throws Exception {
    mBlockStore.createBlock(SIMULATE_LOAD_SESSION_ID, SIMULATE_LOAD_BLOCK_ID,
        AllocateOptions.forCreate(0, BlockStoreLocation.anyTier()));
    mSimulateWriter =
        mBlockStore.createBlockWriter(SIMULATE_LOAD_SESSION_ID, SIMULATE_LOAD_BLOCK_ID);
  }

  /**
   * Stops simulating load on the worker.
   */
  protected void stopSimulateLoad() throws Exception {
    mBlockStore.abortBlock(SIMULATE_LOAD_SESSION_ID, SIMULATE_LOAD_BLOCK_ID);
    mSimulateWriter.close();
  }
}
