alluxio-master-2.x/microbench/src/main/java/alluxio/worker/BlockStoreBase.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;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import alluxio.AlluxioTestDirectory;
import alluxio.AlluxioURI;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.master.NoopUfsManager;
import alluxio.underfs.UfsManager;
import alluxio.underfs.UnderFileSystemConfiguration;
import alluxio.worker.block.BlockMasterClient;
import alluxio.worker.block.BlockMasterClientPool;
import alluxio.worker.block.BlockStore;
import alluxio.worker.block.CreateBlockOptions;
import alluxio.worker.block.MonoBlockStore;
import alluxio.worker.block.TieredBlockStore;
import alluxio.worker.block.io.BlockWriter;
import alluxio.worker.page.PagedBlockStore;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.nio.ByteBuffer;
import java.util.concurrent.atomic.AtomicReference;

/**
 * The class contains necessary objects to perform benchmarks on {@link BlockStore}.
 * It is intended for use case of comparing performance of different BlockStore implementations,
 * where only one of the stores will be actually used. For this reason, UFS is shared among all
 * the objects. Users should be careful about this for different kinds of use cases.
 */
public class BlockStoreBase implements AutoCloseable {
  /**
   * {@link MonoBlockStore} backed by {@link TieredBlockStore} as local storage.
   */
  public MonoBlockStore mMonoBlockStore;

  /**
   * {@link PagedBlockStore} that uses page cache as local storage.
   */
  public PagedBlockStore mPagedBlockStore;

  /**
   * The shared {@link UfsManager} used for block stores.
   */
  public UfsManager mUfsManager;

  /**
   * Create a new instance with initialized block store objects.
   *
   * @return the {@link BlockStoreBase}
   */
  public static BlockStoreBase create() {
    UfsManager ufs = new NoopUfsManager();
    MonoBlockStore monoBlockStore = createMonoBlockStore(ufs);
    PagedBlockStore pagedBlockStore = createPagedBlockStore(ufs);
    return new BlockStoreBase(monoBlockStore, pagedBlockStore, ufs);
  }

  private static MonoBlockStore createMonoBlockStore(UfsManager ufsManager) {
    // set up configurations
    Configuration.set(PropertyKey.WORKER_TIERED_STORE_LEVELS, 1);
    String tieredStoreDir = AlluxioTestDirectory
        .createTemporaryDirectory("worker_local_store").getAbsolutePath();
    Configuration.set(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_PATH, tieredStoreDir);

    // mock a BlockMasterClientPool as we won't use it
    BlockMasterClient mockClient = mock(BlockMasterClient.class);
    BlockMasterClientPool mockPool = mock(BlockMasterClientPool.class);
    when(mockPool.acquire()).thenReturn(mockClient);

    TieredBlockStore tieredBlockStore = new TieredBlockStore();

    return new MonoBlockStore(tieredBlockStore, mockPool, ufsManager, new AtomicReference<>(1L));
  }

  private static PagedBlockStore createPagedBlockStore(UfsManager ufsManager) {
    // set up configurations
    String cacheDir = AlluxioTestDirectory
        .createTemporaryDirectory("worker_cache").getAbsolutePath();
    Configuration.set(PropertyKey.USER_CLIENT_CACHE_DIRS, cacheDir);
    BlockMasterClientPool blockMasterClientPool = new BlockMasterClientPool();
    AtomicReference<Long> workerId = new AtomicReference<>(-1L);

    return PagedBlockStore.create(ufsManager, blockMasterClientPool, workerId);
  }

  public BlockStoreBase(MonoBlockStore monoBlockStore,
                        PagedBlockStore pagedBlockStore, UfsManager ufs) {
    mMonoBlockStore = monoBlockStore;
    mPagedBlockStore = pagedBlockStore;
    mUfsManager = ufs;
  }

  /**
   * Create a local block to both Mono and Paged block store.
   * Flush some data to the block and commit it for later use.
   *
   * @param blockId block id
   * @param blockSize size of block
   * @param data read-only data to flush to block
   * @throws Exception if I/O error happens
   */
  public void prepareLocalBlock(long blockId, long blockSize, byte[] data)
      throws Exception {
    mMonoBlockStore.createBlock(1, blockId, 0,
                    new CreateBlockOptions(null, null, blockSize));
    try (BlockWriter writer = mMonoBlockStore
            .createBlockWriter(1, blockId)) {
      writer.append(ByteBuffer.wrap(data));
    }
    mMonoBlockStore.commitBlock(1, blockId, false);

    mPagedBlockStore.createBlock(1, blockId, 0,
            new CreateBlockOptions(null, null, blockSize));
    try (BlockWriter writer = mPagedBlockStore.createBlockWriter(1, blockId)) {
      writer.append(ByteBuffer.wrap(data));
    }
    mPagedBlockStore.commitBlock(1, blockId, false);
  }

  /**
   * Mount an ufs to {@link UfsManager}, which is shared for both stores.
   *
   * @param mountId mount id
   * @param rootDir directory path of ufs
   */
  public void mountUfs(long mountId, String rootDir) {
    mUfsManager.addMount(mountId, new AlluxioURI(rootDir),
        UnderFileSystemConfiguration.defaults(Configuration.global()));
  }

  /**
   * Flush some data to the local file located at ufsFilePath.
   *
   * @param ufsFilePath file path
   * @param data ufs file data
   * @throws Exception if error occurs
   */
  public void prepareUfsFile(String ufsFilePath, byte[] data) throws Exception {
    File ufsFile = new File(ufsFilePath);
    if (!ufsFile.createNewFile()) {
      throw new IllegalStateException(String.format("UFS file %s already exists", ufsFilePath));
    }

    try (FileOutputStream out = new FileOutputStream(ufsFile);
         BufferedOutputStream bout = new BufferedOutputStream(out)) {
      bout.write(data);
      bout.flush();
    }
  }

  @Override
  public void close() throws Exception {
    mMonoBlockStore.close();
    mPagedBlockStore.close();
    mUfsManager.close();
    Configuration.reloadProperties();
  }
}
