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

import static alluxio.AlluxioMockUtil.setStaticInternalState;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import alluxio.ConfigurationRule;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.OpenLocalBlockRequest;
import alluxio.grpc.OpenLocalBlockResponse;
import alluxio.underfs.UfsManager;
import alluxio.util.io.PathUtils;
import alluxio.worker.block.BlockMasterClient;
import alluxio.worker.block.BlockMasterClientPool;
import alluxio.worker.block.BlockStore;
import alluxio.worker.block.BlockStoreType;
import alluxio.worker.block.CreateBlockOptions;
import alluxio.worker.block.MonoBlockStore;
import alluxio.worker.block.TieredBlockStore;
import alluxio.worker.block.io.BlockWriter;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import io.grpc.Status;
import io.grpc.StatusException;
import io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

public class ShortCircuitBlockReadHandlerTest {
  // templates used in current configuration
  private static final PropertyKey.Template ALIAS_TEMPLATE =
      PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_ALIAS;
  private static final PropertyKey.Template PATH_TEMPLATE =
      PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_PATH;
  private static final PropertyKey.Template QUOTA_TEMPLATE =
      PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_QUOTA;
  private static final PropertyKey.Template MEDIUM_TEMPLATE =
      PropertyKey.Template.WORKER_TIERED_STORE_LEVEL_DIRS_MEDIUMTYPE;

  // test configuration values
  private static final String TIER_0_ALIAS = "ALIAS_0";
  private static final String TIER_1_ALIAS = "ALIAS_1";
  private static final String WORKER_DATA_FOLDER = "TestWorker";
  private static final String WORKER_TMP_FOLDER = "TestTmpWorker";

  // default session id
  private static final long SESSION_ID = 1L;

  // default block id
  private static final long BLOCK_ID = 2L;

  // first tier id
  private static final int FIRST_TIER = 0;

  // second tier id
  private static final int SECOND_TIER = 1;

  @Rule
  public ConfigurationRule mConfiguration = new ConfigurationRule(
      new ImmutableMap.Builder<PropertyKey, Object>()
          // disable paging for this test
          .put(PropertyKey.WORKER_BLOCK_STORE_TYPE, BlockStoreType.FILE)
          // use 2 tiers to test possible moves of blocks
          .put(PropertyKey.WORKER_TIERED_STORE_LEVELS, 2)
          // disable alignment so that no space needs to be reserved
          // on directories
          .put(PropertyKey.WORKER_MANAGEMENT_TIER_ALIGN_ENABLED, false)
          .put(PropertyKey.WORKER_DATA_FOLDER, WORKER_DATA_FOLDER)
          .put(PropertyKey.WORKER_DATA_TMP_FOLDER, WORKER_TMP_FOLDER)
          // tier configurations
          .put(ALIAS_TEMPLATE.format(0), TIER_0_ALIAS)
          .put(ALIAS_TEMPLATE.format(1), TIER_1_ALIAS)
          .put(QUOTA_TEMPLATE.format(0), "1024")
          .put(QUOTA_TEMPLATE.format(1), "1024")
          .put(MEDIUM_TEMPLATE.format(0), "MEM")
          .put(MEDIUM_TEMPLATE.format(1), "SSD")
          // the default reviewer is ProbabilisticBufferReviewer,
          // it stops accepting new blocks into it when free space is under a threshold
          .put(PropertyKey.WORKER_REVIEWER_CLASS,
              "alluxio.worker.block.reviewer.AcceptingReviewer")
          .build(),
      Configuration.modifiableGlobal()
  );

  @Rule
  public TemporaryFolder mFolder = new TemporaryFolder();

  // each tier's storage directory
  private File mTier0Dir;
  private File mTier1Dir;

  // mocked response observer to interact with the handler
  private final TestResponseObserver mResponseObserver = new TestResponseObserver();

  // local store that backs the short circuit read
  private BlockStore mBlockStore;
  private ShortCircuitBlockReadHandler mTestHandler;

  @Before
  public void before() throws Exception {
    // set up storage tier directories
    mTier0Dir = mFolder.newFolder();
    mTier1Dir = mFolder.newFolder();
    mConfiguration.set(PATH_TEMPLATE.format(0), mTier0Dir.getAbsolutePath());
    mConfiguration.set(PATH_TEMPLATE.format(1), mTier1Dir.getAbsolutePath());

    // set up local storage
    // todo(yangchen): BlockStore should not worry about the logic of reporting to master,
    // could leave that work to BlockWorker so that we can get rid of these irrelevant mocks
    BlockMasterClientPool pool = mock(BlockMasterClientPool.class);
    BlockMasterClient cli = mock(BlockMasterClient.class);
    when(pool.createNewResource()).thenReturn(cli);
    when(pool.acquire()).thenReturn(cli);

    setStaticInternalState(TieredBlockStore.class, "REMOVE_BLOCK_TIMEOUT_MS", 1000L);
    TieredBlockStore tieredBlockStore = new TieredBlockStore();

    mBlockStore = new MonoBlockStore(tieredBlockStore, pool,
        mock(UfsManager.class), new AtomicReference<>(1L));

    mTestHandler = new ShortCircuitBlockReadHandler(mBlockStore, mResponseObserver);
  }

  @Test
  public void accessNonExistentBlock() {
    // we started with empty directories without any block
    OpenLocalBlockRequest request = createRequest(BLOCK_ID, false);
    mTestHandler.onNext(request);

    // check that we get a proper error and no response
    assertTrue(mResponseObserver.getResponses().isEmpty());
    assertFalse(mResponseObserver.isCompleted());
    Throwable t = mResponseObserver.getError();
    assertTrue(t instanceof StatusException
            && ((StatusException) t).getStatus().getCode() == Status.Code.NOT_FOUND);
  }

  @Test
  public void accessBlockNoPromote() throws Exception {
    accessBlock(false);
  }

  @Test
  public void accessBlockPromote() throws Exception {
    accessBlock(true);
  }

  @Test
  public void accessBlockPinned() throws Exception {
    createLocalBlock(SESSION_ID, BLOCK_ID, FIRST_TIER);

    // accessing a block through short circuit will pin the block
    OpenLocalBlockRequest request = createRequest(BLOCK_ID, false);
    mTestHandler.onNext(request);

    // block 2 should be pinned now
    // and an attempt to remove the block from local store should fail
    assertThrows(Exception.class, () -> mBlockStore.removeBlock(3L, BLOCK_ID));
  }

  @Test(timeout = 2000)
  public void unpinBlockOnError() throws Exception {
    createLocalBlock(SESSION_ID, BLOCK_ID, FIRST_TIER);

    OpenLocalBlockRequest request = createRequest(BLOCK_ID, false);
    mTestHandler.onNext(request);
    mTestHandler.onError(new RuntimeException());

    // block should be unpinned and another session can remove it
    long anotherSessionId = 4L;
    mBlockStore.removeBlock(anotherSessionId, BLOCK_ID);
  }

  @Test(timeout = 2000)
  public void errorReAccessingBlock() throws Exception {
    createLocalBlock(SESSION_ID, BLOCK_ID, FIRST_TIER);

    OpenLocalBlockRequest request = createRequest(BLOCK_ID, false);

    // first call should be ok
    mTestHandler.onNext(request);
    assertNull(mResponseObserver.getError());

    // second call should cause an error
    mTestHandler.onNext(request);

    assertFalse(mResponseObserver.isCompleted());
    assertNotNull(mResponseObserver.getError());

    // the block should be unpinned
    long anotherSessionId = 3L;
    mBlockStore.removeBlock(anotherSessionId, BLOCK_ID);
  }

  private static OpenLocalBlockRequest createRequest(long blockId, boolean promote) {
    return OpenLocalBlockRequest.newBuilder().setBlockId(blockId).setPromote(promote).build();
  }

  private void accessBlock(boolean promote) throws Exception {
    // create a block in second tier(tier1) so that promotion will move block
    createLocalBlock(SESSION_ID, BLOCK_ID, SECOND_TIER);

    // access the new block via short circuit read
    OpenLocalBlockRequest request = createRequest(BLOCK_ID, promote);
    mTestHandler.onNext(request);
    mTestHandler.onCompleted();

    // if promote, the block should be moved to tier 1
    // real block path is <tier_dir_path>/<WORKER_DATA_FOLDER>/<blockId>
    String expectedRootDirPath = (promote ? mTier0Dir : mTier1Dir).getAbsolutePath();
    String expectedBlockPath = PathUtils.concatPath(
        expectedRootDirPath, WORKER_DATA_FOLDER, BLOCK_ID);

    // check that the block file is present on disk
    assertTrue(Files.exists(Paths.get(expectedBlockPath)));

    // check the response is correct
    assertNull(mResponseObserver.getError());
    assertTrue(mResponseObserver.isCompleted());

    List<OpenLocalBlockResponse> responses = mResponseObserver.getResponses();
    assertEquals(1, responses.size());
    OpenLocalBlockResponse response = responses.get(0);

    // check that we get the expected path where the block lies
    assertEquals(expectedBlockPath, response.getPath());
  }

  // create and commit a block in LocalStorage in the specified location
  // so that further request can read the block via short circuit
  private void createLocalBlock(
      long sessionId, long blockId, int tier) throws Exception {
    mBlockStore.createBlock(
        sessionId,
        blockId,
        tier,
        new CreateBlockOptions(null, null, 64));

    byte[] data = new byte[64];
    Arrays.fill(data, (byte) 1);
    ByteBuffer buf = ByteBuffer.wrap(data);
    try (BlockWriter writer = mBlockStore.createBlockWriter(sessionId, blockId)) {
      writer.append(buf);
    }

    mBlockStore.commitBlock(sessionId, blockId, false);
  }

  // a testing response observer that keeps track of the interactions with
  // the handler
  private static class TestResponseObserver implements StreamObserver<OpenLocalBlockResponse> {
    // keeps track of received response
    private final List<OpenLocalBlockResponse> mResponses = new ArrayList<>();
    // keeps track of received error
    private Throwable mError = null;
    // keeps track of completion
    private boolean mCompleted = false;

    @Override
    public void onNext(OpenLocalBlockResponse value) {
      mResponses.add(value);
    }

    @Override
    public void onError(Throwable t) {
      mError = t;
    }

    @Override
    public void onCompleted() {
      mCompleted = true;
    }

    public List<OpenLocalBlockResponse> getResponses() {
      return ImmutableList.copyOf(mResponses);
    }

    public Throwable getError() {
      return mError;
    }

    public boolean isCompleted() {
      return mCompleted;
    }
  }
}
