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

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.anyLong;
import static org.mockito.Mockito.anyString;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import alluxio.Constants;
import alluxio.client.file.cache.CacheManager;
import alluxio.client.file.cache.CacheManagerOptions;
import alluxio.client.file.cache.evictor.CacheEvictorOptions;
import alluxio.client.file.cache.evictor.FIFOCacheEvictor;
import alluxio.client.file.cache.store.PageStoreDir;
import alluxio.client.file.cache.store.PageStoreOptions;
import alluxio.client.file.cache.store.PageStoreType;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.Configuration;
import alluxio.conf.InstancedConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.exception.status.AlluxioStatusException;
import alluxio.master.NoopUfsManager;
import alluxio.underfs.UfsManager;
import alluxio.util.CommonUtils;
import alluxio.worker.block.AbstractBlockStoreEventListener;
import alluxio.worker.block.BlockMasterClient;
import alluxio.worker.block.BlockMasterClientPool;
import alluxio.worker.block.BlockStoreEventListener;
import alluxio.worker.block.BlockStoreLocation;
import alluxio.worker.block.CreateBlockOptions;
import alluxio.worker.block.io.BlockWriter;

import com.google.common.collect.ImmutableList;
import io.grpc.Status;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;

public class PagedBlockStoreCommitBlockTest {
  BlockStoreEventListener mListener;
  UfsManager mUfs;
  AlluxioConfiguration mConf;
  CacheManagerOptions mCacheManagerOptions;
  PagedBlockMetaStore mPageMetaStore;
  List<PagedBlockStoreDir> mDirs;
  PagedBlockStore mPagedBlockStore;
  BlockMasterClientPool mBlockMasterClientPool;
  BlockMasterClient mMockedBlockMasterClient;
  AtomicReference<Long> mWorkerId;

  CacheManager mCacheManager;

  private static final int DIR_INDEX = 0;

  private static final Long SESSION_ID = 1L;
  private static final long BLOCK_ID = 2L;
  final int mBlockSize = 64;

  public int mPageSize = 2;

  private static final int OFFSET = 0;

  @Rule
  public TemporaryFolder mTempFolder = new TemporaryFolder();

  @Before
  public void setup() throws Exception {
    List<PageStoreDir> pageStoreDirs;
    InstancedConfiguration cacheManagerConf = Configuration.copyGlobal();

    Path dirPath = mTempFolder.newFolder().toPath();
    InstancedConfiguration dirConf = Configuration.modifiableGlobal();
    dirConf.set(PropertyKey.WORKER_PAGE_STORE_DIRS, ImmutableList.of(dirPath));
    dirConf.set(PropertyKey.WORKER_PAGE_STORE_SIZES, ImmutableList.of(Constants.MB));
    dirConf.set(PropertyKey.WORKER_PAGE_STORE_TYPE, PageStoreType.LOCAL);
    PageStoreDir pageStoreDir =
            PageStoreDir.createPageStoreDir(
                    new CacheEvictorOptions().setEvictorClass(FIFOCacheEvictor.class),
                    PageStoreOptions.createForWorkerPageStore(dirConf).get(DIR_INDEX));

    mUfs = new NoopUfsManager();
    mConf = Configuration.global();
    cacheManagerConf.set(PropertyKey.WORKER_PAGE_STORE_PAGE_SIZE, mPageSize);
    cacheManagerConf.set(PropertyKey.WORKER_PAGE_STORE_DIRS, ImmutableList.of(dirPath));

    // Here mock BlockMasterClientPool and BlockMasterClient since I have no idea
    // about how to override them.
    // mockedPool will return a mocked BlockMasterClient when require() is called,
    // and do nothing when releasing, maybe add some action later on.
    mBlockMasterClientPool = mock(BlockMasterClientPool.class);
    mMockedBlockMasterClient = mock(BlockMasterClient.class);
    when(mBlockMasterClientPool.acquire()).thenReturn(mMockedBlockMasterClient);
    doNothing().when(mBlockMasterClientPool).release(any());
    mWorkerId = new AtomicReference<>(-1L);
    mCacheManagerOptions = CacheManagerOptions.createForWorker(cacheManagerConf);
    pageStoreDirs = new ArrayList<PageStoreDir>();
    pageStoreDirs.add(pageStoreDir);
    mDirs = PagedBlockStoreDir.fromPageStoreDirs(pageStoreDirs);

    mListener = spy(new AbstractBlockStoreEventListener() {
      @Override
      public void onCommitBlockToLocal(long blockId, BlockStoreLocation location) {
        assertEquals(BLOCK_ID, blockId);
      }

      @Override
      public void onCommitBlockToMaster(long blockId, BlockStoreLocation location) {
        assertEquals(BLOCK_ID, blockId);
      }
    });
  }

  @After
  public void tearDown() throws IOException {
    mPagedBlockStore.close();
  }

  // This Test case success both to commit, no Exception should be thrown,
  // and both onCommit method should be called
  @Test
  public void localCommitAndMasterCommitBothSuccess()
         throws IOException, InterruptedException, TimeoutException {
    mPageMetaStore = new PagedBlockMetaStore(mDirs);
    mCacheManager = CacheManager.Factory.create(mConf, mCacheManagerOptions, mPageMetaStore);

    mPagedBlockStore = new PagedBlockStore(mCacheManager, mUfs, mBlockMasterClientPool, mWorkerId,
            mPageMetaStore, mCacheManagerOptions.getPageSize());

    prepareBlockStore();

    mPagedBlockStore.commitBlock(SESSION_ID, BLOCK_ID, false);
    verify(mListener).onCommitBlockToLocal(anyLong(), any(BlockStoreLocation.class));
    verify(mListener).onCommitBlockToMaster(anyLong(), any(BlockStoreLocation.class));
  }

  // This Test case success commitToMaster, expecting one exception,

  @Test
  public void localCommitFailAndMasterCommitSuccess()
          throws IOException, InterruptedException, TimeoutException {
    mPageMetaStore = new PagedBlockMetaStore(mDirs) {
        // here commit always throw Exception
        @Override
        public PagedBlockMeta commit(long BLOCK_ID) {
            throw new RuntimeException();
        }
    };
    mCacheManager = CacheManager.Factory.create(mConf, mCacheManagerOptions, mPageMetaStore);

    mPagedBlockStore = new PagedBlockStore(mCacheManager, mUfs, mBlockMasterClientPool,
            mWorkerId, mPageMetaStore, mCacheManagerOptions.getPageSize());

    prepareBlockStore();

    assertThrows(RuntimeException.class, () -> {
      mPagedBlockStore.commitBlock(SESSION_ID, BLOCK_ID, false);
    });

    verify(mListener, never()).onCommitBlockToLocal(anyLong(), any(BlockStoreLocation.class));
    verify(mListener, never()).onCommitBlockToMaster(anyLong(), any(BlockStoreLocation.class));
  }

  // This Test case success commitToLocal, expecting one exception,
  // and only one onCommit method should be called.
  @Test
  public void localCommitSuccessAndMasterCommitFail()
          throws IOException, InterruptedException, TimeoutException {
    doAnswer((i) -> {
      throw new AlluxioStatusException(Status.UNAVAILABLE);
    }).when(mMockedBlockMasterClient).commitBlock(anyLong(), anyLong(), anyString(),
             anyString(), anyLong(), anyLong());
    mPageMetaStore = new PagedBlockMetaStore(mDirs);
    mCacheManager = CacheManager.Factory.create(mConf, mCacheManagerOptions, mPageMetaStore);

    mPagedBlockStore = new PagedBlockStore(mCacheManager, mUfs, mBlockMasterClientPool, mWorkerId,
             mPageMetaStore, mCacheManagerOptions.getPageSize());

    prepareBlockStore();

    assertThrows(RuntimeException.class, () -> {
      mPagedBlockStore.commitBlock(SESSION_ID, BLOCK_ID, false);
    });
    verify(mListener).onCommitBlockToLocal(anyLong(), any(BlockStoreLocation.class));
    verify(mListener, never()).onCommitBlockToMaster(anyLong(), any(BlockStoreLocation.class));
  }

  // Prepare PageBlockStore and creat a temp block for following test
  public void prepareBlockStore() throws IOException, InterruptedException, TimeoutException {
    PagedBlockStoreDir dir =
            (PagedBlockStoreDir) mPageMetaStore.allocate(BlockPageId.tempFileIdOf(BLOCK_ID), 1);

    dir.putTempFile(BlockPageId.tempFileIdOf(BLOCK_ID));
    PagedTempBlockMeta blockMeta = new PagedTempBlockMeta(BLOCK_ID, dir);
    mPagedBlockStore.createBlock(SESSION_ID, BLOCK_ID, OFFSET,
             new CreateBlockOptions(null, null, mBlockSize));
    byte[] data = new byte[mBlockSize];
    Arrays.fill(data, (byte) 1);
    ByteBuffer buf = ByteBuffer.wrap(data);
    BlockWriter writer = mPagedBlockStore.createBlockWriter(SESSION_ID, BLOCK_ID);
    CommonUtils.waitFor("writer initiation complete",
            () -> mPagedBlockStore.getCacheManagerState() == CacheManager.State.READ_WRITE);
    writer.append(buf);

    mPagedBlockStore.registerBlockStoreEventListener(mListener);
  }
}
