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

import static alluxio.util.CommonUtils.waitFor;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyMap;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;

import alluxio.AlluxioURI;
import alluxio.ConfigurationRule;
import alluxio.Constants;
import alluxio.Sessions;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.CacheRequest;
import alluxio.grpc.Command;
import alluxio.grpc.CommandType;
import alluxio.master.NoopUfsManager;
import alluxio.proto.dataserver.Protocol;
import alluxio.underfs.UfsManager;
import alluxio.underfs.UnderFileSystemConfiguration;
import alluxio.util.WaitForOptions;
import alluxio.util.io.BufferUtils;
import alluxio.util.network.NetworkAddressUtils;
import alluxio.wire.WorkerNetAddress;
import alluxio.worker.block.io.BlockReader;
import alluxio.worker.file.FileSystemMasterClient;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import org.junit.Before;
import org.junit.Rule;
import org.junit.rules.TemporaryFolder;

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

/**
 * Test base for block worker related tests.
 */
public class DefaultBlockWorkerTestBase {
  protected static final int BLOCK_SIZE = 128;

  TieredBlockStore mTieredBlockStore;
  // worker configurations
  protected static final long WORKER_ID = 30L;
  // ufs for fallback read
  protected static final long UFS_MOUNT_ID = 1L;
  // ufs for batch load
  protected static final long UFS_LOAD_MOUNT_ID = 2L;
  protected static final WorkerNetAddress WORKER_ADDRESS =
      new WorkerNetAddress().setHost("localhost").setRpcPort(20001);

  // invalid initial worker id
  protected static final long INVALID_WORKER_ID = -1L;

  // test subject
  protected DefaultBlockWorker mBlockWorker;

  // mocked dependencies of DefaultBlockWorker
  protected BlockMasterClient mBlockMasterClient;
  protected FileSystemMasterClient mFileSystemMasterClient;

  protected final Random mRandom = new Random();

  @Rule
  public TemporaryFolder mTestFolder = new TemporaryFolder();
  // worker's local storage directories
  protected String mMemDir;
  protected String mHddDir;
  // ufs file for fallback read
  protected File mTestUfsFile;

  // ufs root path for batch load
  protected String mRootUfs;
  // ufs file for batch load
  protected String mTestLoadFilePath;
  protected BlockMasterClientPool mBlockMasterClientPool;

  @Rule
  public ConfigurationRule mConfigurationRule =
      new ConfigurationRule(new ImmutableMap.Builder<PropertyKey, Object>()
          .put(PropertyKey.WORKER_TIERED_STORE_LEVELS, 2)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_ALIAS, Constants.MEDIUM_MEM)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_MEDIUMTYPE, Constants.MEDIUM_MEM)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_QUOTA, "1GB")
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL1_ALIAS, Constants.MEDIUM_HDD)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL1_DIRS_MEDIUMTYPE, Constants.MEDIUM_HDD)
          .put(PropertyKey.WORKER_TIERED_STORE_LEVEL1_DIRS_QUOTA, "2GB")
          .put(PropertyKey.WORKER_RPC_PORT, 0)
          .put(PropertyKey.WORKER_MANAGEMENT_TIER_ALIGN_RESERVED_BYTES, "0")
          .put(PropertyKey.WORKER_BLOCK_HEARTBEAT_INTERVAL_MS, "10ms")
          .build(), Configuration.modifiableGlobal());
  protected BlockStore mBlockStore;

  /**
   * Sets up all dependencies before a test runs.
   */
  @Before
  public void before() throws Exception {
    // set up storage directories
    mMemDir = mTestFolder.newFolder().getAbsolutePath();
    mHddDir = mTestFolder.newFolder().getAbsolutePath();
    mConfigurationRule.set(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_PATH, mMemDir);
    mConfigurationRule.set(PropertyKey.WORKER_TIERED_STORE_LEVEL1_DIRS_PATH, mHddDir);

    // set up BlockMasterClient
    mBlockMasterClient = createMockBlockMasterClient();
    mBlockMasterClientPool = spy(new BlockMasterClientPool());
    doReturn(mBlockMasterClient).when(mBlockMasterClientPool).createNewResource();

    mTieredBlockStore = spy(new TieredBlockStore());
    UfsManager ufsManager = new NoopUfsManager();
    AtomicReference<Long> workerId = new AtomicReference<>(INVALID_WORKER_ID);
    mBlockStore =
        spy(new MonoBlockStore(mTieredBlockStore, mBlockMasterClientPool, ufsManager, workerId));

    mFileSystemMasterClient = createMockFileSystemMasterClient();

    Sessions sessions = mock(Sessions.class);

    // set up a ufs directory for batch load jobs
    mRootUfs = mTestFolder.newFolder("DefaultBlockWorkerTest").getAbsolutePath();
    mConfigurationRule.set(PropertyKey.MASTER_MOUNT_TABLE_ROOT_UFS, mRootUfs);
    ufsManager.addMount(UFS_LOAD_MOUNT_ID,
        new AlluxioURI(mRootUfs),
        UnderFileSystemConfiguration.defaults(Configuration.global()));
    // Write an actual file to UFS
    mTestLoadFilePath = mTestFolder.newFile("temp").getAbsolutePath();
    byte[] buffer = BufferUtils.getIncreasingByteArray((int) (BLOCK_SIZE * 1.5));
    BufferUtils.writeBufferToFile(mTestLoadFilePath, buffer);

    // set up ufs directory for fallback reading
    mTestUfsFile = mTestFolder.newFile();
    // mount test file to UFS_MOUNT_ID
    ufsManager.addMount(
        UFS_MOUNT_ID,
        new AlluxioURI(mTestUfsFile.getAbsolutePath()),
        UnderFileSystemConfiguration.defaults(Configuration.global())
    );

    mBlockWorker = new DefaultBlockWorker(mBlockMasterClientPool, mFileSystemMasterClient,
        sessions, mBlockStore, workerId);
  }

  protected void cacheBlock(boolean async) throws Exception {
    // flush 1MB random data to ufs so that caching will take a while
    long ufsBlockSize = 1024 * 1024;
    byte[] data = new byte[(int) ufsBlockSize];
    mRandom.nextBytes(data);

    try (FileOutputStream fileOut = new FileOutputStream(mTestUfsFile);
         BufferedOutputStream bufOut = new BufferedOutputStream(fileOut)) {
      bufOut.write(data);
      bufOut.flush();
    }

    // ufs options: delegate to the ufs mounted at UFS_MOUNT_ID
    // with path to our test file
    long blockId = mRandom.nextLong();
    Protocol.OpenUfsBlockOptions options = Protocol.OpenUfsBlockOptions
        .newBuilder()
        .setBlockSize(ufsBlockSize)
        .setUfsPath(mTestUfsFile.getAbsolutePath())
        .setMountId(UFS_MOUNT_ID)
        .setNoCache(false)
        .setOffsetInFile(0)
        .build();

    // cache request:
    // delegate to local ufs client rather than remote worker
    CacheRequest request = CacheRequest
        .newBuilder()
        .setSourceHost(NetworkAddressUtils.getLocalHostName(500))
        .setBlockId(blockId)
        .setLength(ufsBlockSize)
        .setAsync(async)
        .setOpenUfsBlockOptions(options)
        .build();

    mBlockWorker.cache(request);

    // check that the block metadata is present
    if (async) {
      assertFalse(mBlockWorker.getBlockStore().hasBlockMeta(blockId));
      waitFor(
          "Wait for async cache",
          () -> mBlockWorker.getBlockStore().hasBlockMeta(blockId),
          WaitForOptions.defaults().setInterval(10).setTimeoutMs(2000));
    } else {
      assertTrue(mBlockWorker.getBlockStore().hasBlockMeta(blockId));
    }

    long sessionId = mRandom.nextLong();
    // check that we can read the block locally
    // note: this time we use an OpenUfsOption without ufsPath and blockInUfsTier so
    // that the worker can't fall back to ufs read.
    Protocol.OpenUfsBlockOptions noFallbackOptions = Protocol.OpenUfsBlockOptions.newBuilder()
        .setBlockInUfsTier(false).build();
    try (BlockReader reader = mBlockWorker.createBlockReader(
        sessionId, blockId, 0, false, noFallbackOptions)) {
      ByteBuffer buf = reader.read(0, ufsBlockSize);
      // alert: LocalFileBlockReader uses a MappedByteBuffer, which does not
      // support the array operation. So we need to compare ByteBuffer manually
      assertEquals(0, buf.compareTo(ByteBuffer.wrap(data)));
    }
  }

  // create a BlockMasterClient that simulates reasonable default
  // interactions with the block master
  protected BlockMasterClient createMockBlockMasterClient() throws Exception {
    BlockMasterClient client = mock(BlockMasterClient.class);

    // return designated worker id
    doReturn(WORKER_ID)
        .when(client)
        .getId(any(WorkerNetAddress.class));

    // return Command.Nothing for heartbeat
    doReturn(Command.newBuilder().setCommandType(CommandType.Nothing).build())
        .when(client)
        .heartbeat(
            anyLong(),
            anyMap(),
            anyMap(),
            anyList(),
            anyMap(),
            anyMap(),
            anyList()
        );
    return client;
  }

  // create a mocked FileSystemMasterClient that simulates reasonable default
  // interactions with file system master
  protected FileSystemMasterClient createMockFileSystemMasterClient() throws Exception {
    FileSystemMasterClient client = mock(FileSystemMasterClient.class);
    doReturn(ImmutableSet.of())
        .when(client)
        .getPinList();
    return client;
  }
}
