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

import static alluxio.grpc.BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceStub;
import static alluxio.server.block.RegisterStreamTestUtils.BATCH_SIZE;
import static alluxio.server.block.RegisterStreamTestUtils.CAPACITY_MAP;
import static alluxio.server.block.RegisterStreamTestUtils.EMPTY_CONFIG;
import static alluxio.server.block.RegisterStreamTestUtils.LOST_STORAGE;
import static alluxio.server.block.RegisterStreamTestUtils.MEM_CAPACITY;
import static alluxio.server.block.RegisterStreamTestUtils.MEM_USAGE_EMPTY;
import static alluxio.server.block.RegisterStreamTestUtils.NET_ADDRESS_1;
import static alluxio.server.block.RegisterStreamTestUtils.TIER_BLOCK_TOTAL;
import static alluxio.server.block.RegisterStreamTestUtils.TIER_CONFIG;
import static alluxio.server.block.RegisterStreamTestUtils.USAGE_MAP;
import static alluxio.server.block.RegisterStreamTestUtils.findFirstBlock;
import static alluxio.server.block.RegisterStreamTestUtils.getTierAliases;
import static alluxio.server.block.RegisterStreamTestUtils.parseTierConfig;
import static alluxio.stress.cli.RpcBenchPreparationUtils.CAPACITY;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import alluxio.AlluxioTestDirectory;
import alluxio.ClientContext;
import alluxio.ConfigurationRule;
import alluxio.Constants;
import alluxio.Sessions;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.status.DeadlineExceededException;
import alluxio.exception.status.InternalException;
import alluxio.exception.status.NotFoundException;
import alluxio.grpc.BlockMasterWorkerServiceGrpc;
import alluxio.grpc.Command;
import alluxio.grpc.ConfigProperty;
import alluxio.grpc.GrpcExceptionUtils;
import alluxio.grpc.LocationBlockIdListEntry;
import alluxio.grpc.Metric;
import alluxio.grpc.RegisterWorkerPRequest;
import alluxio.grpc.RegisterWorkerPResponse;
import alluxio.grpc.StorageList;
import alluxio.master.MasterClientContext;
import alluxio.stress.cli.RpcBenchPreparationUtils;
import alluxio.stress.rpc.TierAlias;
import alluxio.underfs.UfsManager;
import alluxio.util.ThreadFactoryUtils;
import alluxio.worker.block.BlockMasterClient;
import alluxio.worker.block.BlockMasterClientPool;
import alluxio.worker.block.BlockMasterSync;
import alluxio.worker.block.BlockStoreLocation;
import alluxio.worker.block.CreateBlockOptions;
import alluxio.worker.block.DefaultBlockWorker;
import alluxio.worker.block.MonoBlockStore;
import alluxio.worker.block.RegisterStreamer;
import alluxio.worker.block.TieredBlockStore;
import alluxio.worker.file.FileSystemMasterClient;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import io.grpc.StatusException;
import io.grpc.stub.StreamObserver;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
import org.powermock.api.mockito.PowerMockito;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;

import java.io.IOException;
import java.lang.reflect.Field;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Integration tests for the client-side logic for the register stream.
 */
@RunWith(PowerMockRunner.class)
@PrepareForTest({BlockMasterWorkerServiceGrpc.BlockMasterWorkerServiceStub.class})
public class BlockWorkerRegisterStreamIntegrationTest {
  private static final long WORKER_ID = 1L;

  // Some fields for made-up data
  private List<String> mTierAliases;
  private Map<String, Long> mCapacityMap;
  private Map<String, Long> mUsedMap;
  private Map<BlockStoreLocation, List<Long>> mBlockMap;
  private Map<String, Integer> mTierToIndex =
      ImmutableMap.of("MEM", 0, "SSD", 1, "HDD", 2);

  // Used to mock the block worker
  private BlockMasterClient mBlockMasterClient;
  private BlockMasterClientPool mBlockMasterClientPool;
  private DefaultBlockWorker mBlockWorker;
  private String mMemDir =
      AlluxioTestDirectory.createTemporaryDirectory(Constants.MEDIUM_MEM).getAbsolutePath();
  private String mHddDir =
      AlluxioTestDirectory.createTemporaryDirectory(Constants.MEDIUM_HDD).getAbsolutePath();

  @Rule
  public TemporaryFolder mTestFolder = new TemporaryFolder();
  @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_LEVEL0_DIRS_PATH, mMemDir)
          .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_TIERED_STORE_LEVEL1_DIRS_PATH, mHddDir)
          .put(PropertyKey.WORKER_RPC_PORT, 0)
          .put(PropertyKey.WORKER_MANAGEMENT_TIER_ALIGN_RESERVED_BYTES, "0")
          .put(PropertyKey.MASTER_WORKER_REGISTER_LEASE_ENABLED, false)
          .build(), Configuration.modifiableGlobal());

  private ExecutorService mExecutorService;

  /**
   * Sets up the dependencies before a test runs.
   */
  @Before
  public void before() throws Exception {
    // Set the config properties
    Configuration.set(PropertyKey.WORKER_REGISTER_STREAM_ENABLED, true);
    Configuration.set(PropertyKey.WORKER_REGISTER_STREAM_BATCH_SIZE, BATCH_SIZE);
    Configuration.set(PropertyKey.WORKER_REGISTER_STREAM_RESPONSE_TIMEOUT, "1s");
    Configuration.set(PropertyKey.WORKER_REGISTER_STREAM_COMPLETE_TIMEOUT, "3s");

    // Prepare test data
    mTierAliases = getTierAliases(parseTierConfig(TIER_CONFIG));
    mCapacityMap = Maps.toMap(mTierAliases, (tier) -> CAPACITY);
    mUsedMap = Maps.toMap(mTierAliases, (tier) -> 0L);
    mBlockMap = RpcBenchPreparationUtils.generateBlockIdOnTiers(parseTierConfig(TIER_CONFIG));

    mExecutorService =
        Executors.newFixedThreadPool(2, ThreadFactoryUtils.build("TestBlockWorker-%d", true));
  }

  @After
  public void after() throws Exception {
    mExecutorService.shutdown();
  }

  public void initBlockWorker() throws Exception {
    // Prepare a block worker
    mBlockMasterClientPool = Mockito.spy(new BlockMasterClientPool());
    when(mBlockMasterClientPool.createNewResource()).thenReturn(mBlockMasterClient);
    when(mBlockMasterClientPool.acquire()).thenReturn(mBlockMasterClient);
    TieredBlockStore tieredBlockStore = new TieredBlockStore();
    UfsManager ufsManager = mock(UfsManager.class);
    AtomicReference<Long> workerId = new AtomicReference<>(-1L);
    MonoBlockStore blockStore =
        new MonoBlockStore(tieredBlockStore, mBlockMasterClientPool, ufsManager, workerId);
    FileSystemMasterClient fileSystemMasterClient = mock(FileSystemMasterClient.class);
    Sessions sessions = mock(Sessions.class);

    mBlockWorker = new DefaultBlockWorker(mBlockMasterClientPool, fileSystemMasterClient,
        sessions, blockStore, workerId);
  }

  /**
   * Tests below cover the most normal cases.
   */
  @Test
  public void requestsForEmptyWorker() throws Exception {
    List<RegisterWorkerPRequest> requestChunks =
        RegisterStreamTestUtils.generateRegisterStreamForEmptyWorker(WORKER_ID);

    // Verify the size and content of the requests
    assertEquals(1, requestChunks.size());
    RegisterWorkerPRequest request = requestChunks.get(0);
    assertEquals(WORKER_ID, request.getWorkerId());
    assertEquals(MEM_USAGE_EMPTY, request.getUsedBytesOnTiersMap());
    assertEquals(MEM_CAPACITY, request.getTotalBytesOnTiersMap());
    Map<String, StorageList> lostMap = request.getLostStorageMap();
    assertEquals(1, lostMap.size());
    assertEquals(StorageList.newBuilder().build(), lostMap.get("MEM"));
    assertEquals(ImmutableList.of("MEM"), request.getStorageTiersList());

    List<LocationBlockIdListEntry> entries = request.getCurrentBlocksList();
    assertEquals(0, entries.size());
  }

  @Test
  public void requestsForWorker() throws Exception {
    List<RegisterWorkerPRequest> requestChunks =
        RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(WORKER_ID);

    // Verify the size and content of the requests
    int expectedBatchCount = (int) Math.ceil((TIER_BLOCK_TOTAL) / (double) BATCH_SIZE);
    Set<Long> containedBlockIds = new HashSet<>();
    assertEquals(expectedBatchCount, requestChunks.size());
    for (int i = 0; i < expectedBatchCount; i++) {
      RegisterWorkerPRequest request = requestChunks.get(i);
      assertEquals(WORKER_ID, request.getWorkerId());
      List<LocationBlockIdListEntry> entries = request.getCurrentBlocksList();

      int totalSize = 0;
      for (LocationBlockIdListEntry entry : entries) {
        totalSize += entry.getValue().getBlockIdCount();
        containedBlockIds.addAll(entry.getValue().getBlockIdList());
      }
      if (i != expectedBatchCount - 1) {
        assertEquals(BATCH_SIZE, totalSize);
      }

      // The 1st request contains metadata but the following do not
      if (i == 0) {
        assertEquals(USAGE_MAP, request.getUsedBytesOnTiersMap());
        assertEquals(CAPACITY_MAP, request.getTotalBytesOnTiersMap());
        Map<String, StorageList> lostMap = request.getLostStorageMap();
        assertEquals(1, lostMap.size());
        assertEquals(StorageList.newBuilder().build(), lostMap.get("MEM"));
        assertEquals(ImmutableList.of("MEM", "SSD", "HDD"), request.getStorageTiersList());
      } else {
        assertEquals(0, request.getStorageTiersCount());
        assertEquals(ImmutableMap.of(), request.getUsedBytesOnTiersMap());
        assertEquals(ImmutableMap.of(), request.getTotalBytesOnTiersMap());
        Map<String, StorageList> lostMap = request.getLostStorageMap();
        assertEquals(0, lostMap.size());
      }
    }
    assertEquals(containedBlockIds.size(), TIER_BLOCK_TOTAL);
  }

  @Test
  public void normalFlow() throws Exception {
    // Create the stream and control the request/response observers
    BlockMasterWorkerServiceStub asyncClient =
        PowerMockito.mock(BlockMasterWorkerServiceStub.class);
    when(asyncClient.withDeadlineAfter(anyLong(), any())).thenReturn(asyncClient);
    AtomicInteger receivedCount = new AtomicInteger();
    AtomicBoolean completed = new AtomicBoolean();
    ManualRegisterStreamObserver requestObserver =
        new ManualRegisterStreamObserver(
            MasterMode.COMPLETE_AND_REPORT, receivedCount, completed);
    when(asyncClient.registerWorkerStream(any())).thenReturn(requestObserver);
    RegisterStreamer registerStreamer = new RegisterStreamer(asyncClient,
        WORKER_ID, mTierAliases, mCapacityMap, mUsedMap, mBlockMap, LOST_STORAGE, EMPTY_CONFIG);
    StreamObserver<RegisterWorkerPResponse> responseObserver =
        getResponseObserver(registerStreamer);
    requestObserver.setResponseObserver(responseObserver);

    registerStreamer.registerWithMaster();

    int expectedBatchCount = (int) Math.ceil((TIER_BLOCK_TOTAL) / (double) BATCH_SIZE);
    assertEquals(expectedBatchCount, receivedCount.get());
    assertEquals(true, completed.get());
  }

  // The field is private so we have to use reflection to bypass the permission control
  private StreamObserver<RegisterWorkerPResponse> getResponseObserver(
      RegisterStreamer stream) throws Exception {
    Field privateField
        = RegisterStreamer.class.getDeclaredField("mMasterResponseObserver");
    privateField.setAccessible(true);
    return (StreamObserver<RegisterWorkerPResponse>) privateField.get(stream);
  }

  /**
   * Tests below cover various failure cases.
   */
  @Test
  public void registerWorkerErrorAtStreamStart() throws Exception {
    // Create the stream and control the request/response observers
    BlockMasterWorkerServiceStub asyncClient =
        PowerMockito.mock(BlockMasterWorkerServiceStub.class);
    when(asyncClient.withDeadlineAfter(anyLong(), any())).thenReturn(asyncClient);
    ManualRegisterStreamObserver requestObserver =
        new ManualRegisterStreamObserver(MasterMode.FIRST_REQUEST);
    when(asyncClient.registerWorkerStream(any())).thenReturn(requestObserver);
    RegisterStreamer registerStreamer = new RegisterStreamer(asyncClient,
        WORKER_ID, mTierAliases, mCapacityMap, mUsedMap, mBlockMap, LOST_STORAGE, EMPTY_CONFIG);
    StreamObserver<RegisterWorkerPResponse> responseObserver =
        getResponseObserver(registerStreamer);
    requestObserver.setResponseObserver(responseObserver);

    assertThrows(InternalException.class, () -> {
      // An error received from the master side will be InternalException
      registerStreamer.registerWithMaster();
    });
  }

  @Test
  public void registerWorkerErrorDuringStream() throws Exception {
    // Create the stream and control the request/response observers
    BlockMasterWorkerServiceStub asyncClient =
        PowerMockito.mock(BlockMasterWorkerServiceStub.class);
    when(asyncClient.withDeadlineAfter(anyLong(), any())).thenReturn(asyncClient);
    ManualRegisterStreamObserver requestObserver =
        new ManualRegisterStreamObserver(MasterMode.SECOND_REQUEST);
    when(asyncClient.registerWorkerStream(any())).thenReturn(requestObserver);
    RegisterStreamer registerStreamer = new RegisterStreamer(asyncClient,
        WORKER_ID, mTierAliases, mCapacityMap, mUsedMap, mBlockMap, LOST_STORAGE, EMPTY_CONFIG);
    StreamObserver<RegisterWorkerPResponse> responseObserver =
        getResponseObserver(registerStreamer);
    requestObserver.setResponseObserver(responseObserver);

    assertThrows(InternalException.class, () -> {
      // An error received from the master side will be InternalException
      registerStreamer.registerWithMaster();
    });
  }

  @Test
  public void registerWorkerErrorAtCompletion() throws Exception {
    // Create the stream and control the request/response observers
    BlockMasterWorkerServiceStub asyncClient =
        PowerMockito.mock(BlockMasterWorkerServiceStub.class);
    when(asyncClient.withDeadlineAfter(anyLong(), any())).thenReturn(asyncClient);
    ManualRegisterStreamObserver requestObserver =
        new ManualRegisterStreamObserver(MasterMode.ON_COMPLETED);
    when(asyncClient.registerWorkerStream(any())).thenReturn(requestObserver);
    RegisterStreamer registerStreamer = new RegisterStreamer(asyncClient,
        WORKER_ID, mTierAliases, mCapacityMap, mUsedMap, mBlockMap, LOST_STORAGE, EMPTY_CONFIG);
    StreamObserver<RegisterWorkerPResponse> responseObserver =
        getResponseObserver(registerStreamer);
    requestObserver.setResponseObserver(responseObserver);

    assertThrows(InternalException.class, () -> {
      // An error received from the master side will be InternalException
      registerStreamer.registerWithMaster();
    });
  }

  @Test
  public void masterHangsInStream() throws Exception {
    // Create the stream and control the request/response observers
    BlockMasterWorkerServiceStub asyncClient =
        PowerMockito.mock(BlockMasterWorkerServiceStub.class);
    when(asyncClient.withDeadlineAfter(anyLong(), any())).thenReturn(asyncClient);
    ManualRegisterStreamObserver requestObserver =
        new ManualRegisterStreamObserver(MasterMode.HANG_IN_STREAM);
    when(asyncClient.registerWorkerStream(any())).thenReturn(requestObserver);
    RegisterStreamer registerStreamer = new RegisterStreamer(asyncClient,
        WORKER_ID, mTierAliases, mCapacityMap, mUsedMap, mBlockMap, LOST_STORAGE, EMPTY_CONFIG);
    StreamObserver<RegisterWorkerPResponse> responseObserver =
        getResponseObserver(registerStreamer);
    requestObserver.setResponseObserver(responseObserver);

    // The master will hang during the stream which exceeds the deadline between messages
    assertThrows(DeadlineExceededException.class, () -> {
      registerStreamer.registerWithMaster();
    });
  }

  @Test
  public void masterHangsOnCompleted() throws Exception {
    // Create the stream and control the request/response observers
    BlockMasterWorkerServiceStub asyncClient =
        PowerMockito.mock(BlockMasterWorkerServiceStub.class);
    when(asyncClient.withDeadlineAfter(anyLong(), any())).thenReturn(asyncClient);
    ManualRegisterStreamObserver requestObserver =
        new ManualRegisterStreamObserver(MasterMode.HANG_ON_COMPLETED);
    when(asyncClient.registerWorkerStream(any())).thenReturn(requestObserver);
    RegisterStreamer registerStreamer = new RegisterStreamer(asyncClient,
        WORKER_ID, mTierAliases, mCapacityMap, mUsedMap, mBlockMap, LOST_STORAGE, EMPTY_CONFIG);
    StreamObserver<RegisterWorkerPResponse> responseObserver =
        getResponseObserver(registerStreamer);
    requestObserver.setResponseObserver(responseObserver);

    // The master will hang at the completion step which exceeds the deadline between messages
    assertThrows(DeadlineExceededException.class, () -> {
      registerStreamer.registerWithMaster();
    });
  }

  /**
   * Tests below cover the race conditions during concurrent executions.
   *
   * If a worker is new to the cluster, no clients should know this worker.
   * Therefore there is no concurrent client-incurred write operations on this worker.
   * The races happen typically when the worker re-registers with the master,
   * where some clients already know this worker and can direct invoke writes on the worker.
   *
   * Tests here verify the integrity of the worker-side metadata.
   * In other words, even a delete/move happens on the worker during the register stream,
   * the change should be successful and the update should be recorded correctly.
   * The update should later be reported to the master.
   */
  @Test
  public void deleteDuringRegisterStream() throws Exception {
    // Generate a request stream of blocks
    List<RegisterWorkerPRequest> requestChunks =
        RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(WORKER_ID);
    // Select a block to remove concurrent with the stream
    long blockToRemove = findFirstBlock(requestChunks);

    // Manually create a stream and control the request/response observers
    BlockMasterWorkerServiceStub asyncClient =
        PowerMockito.mock(BlockMasterWorkerServiceStub.class);
    when(asyncClient.withDeadlineAfter(anyLong(), any())).thenReturn(asyncClient);
    CountDownLatch signalLatch = new CountDownLatch(1);
    ManualRegisterStreamObserver requestObserver =
        new ManualRegisterStreamObserver(MasterMode.SIGNAL_IN_STREAM, signalLatch);
    when(asyncClient.registerWorkerStream(any())).thenReturn(requestObserver);
    RegisterStreamer registerStreamer = new RegisterStreamer(asyncClient,
        WORKER_ID, mTierAliases, mCapacityMap, mUsedMap, mBlockMap, LOST_STORAGE, EMPTY_CONFIG);
    StreamObserver<RegisterWorkerPResponse> responseObserver =
        getResponseObserver(registerStreamer);
    requestObserver.setResponseObserver(responseObserver);

    // Prepare the block worker to use the overriden stream
    MasterClientContext context = MasterClientContext
        .newBuilder(ClientContext.create(Configuration.global())).build();
    // On heartbeat, the expected values will be checked against
    List<Long> expectedLostBlocks = ImmutableList.of(blockToRemove);
    Map<BlockStoreLocation, List<Long>> expectedAddedBlocks = ImmutableMap.of();
    mBlockMasterClient = new TestBlockMasterClient(context, registerStreamer,
        expectedLostBlocks, expectedAddedBlocks);
    initBlockWorker();
    // Prepare the blocks in the BlockWorker storage that match the register stream
    prepareBlocksOnWorker(TIER_CONFIG);

    // Let a delete job wait on the signal
    AtomicReference<Throwable> error = new AtomicReference<>();
    Future f = mExecutorService.submit(() -> {
      try {
        signalLatch.await();
        mBlockWorker.removeBlock(1L, blockToRemove);
      } catch (Exception e) {
        error.set(e);
      }
    });

    // Kick off the register stream
    AtomicReference<Long> workerId = new AtomicReference<>(WORKER_ID);
    BlockMasterSync sync = new BlockMasterSync(
        mBlockWorker, workerId, NET_ADDRESS_1, mBlockMasterClientPool);

    // Check the next heartbeat to be sent to the master
    f.get();
    assertNull(error.get());
    // Validation will happen on the heartbeat
    sync.heartbeat(Long.MAX_VALUE);
  }

  // TODO(jiacheng): an internal block movement happens during register stream
  //  verify the next heartbeat

  // An overriden BlockMasterClient
  class TestBlockMasterClient extends BlockMasterClient {
    RegisterStreamer mStreamer;
    List<Long> mRemovedBlocks;
    Map<BlockStoreLocation, List<Long>> mAddedBlocks;

    public TestBlockMasterClient(MasterClientContext conf, RegisterStreamer streamer,
         final List<Long> removedBlocks, final Map<BlockStoreLocation, List<Long>> addedBlocks) {
      super(conf);
      mStreamer = streamer;
      mRemovedBlocks = removedBlocks;
      mAddedBlocks = addedBlocks;
    }

    // Just use the prepared stream
    @Override
    public void registerWithStream(
        final long workerId, final List<String> storageTierAliases,
        final Map<String, Long> totalBytesOnTiers, final Map<String, Long> usedBytesOnTiers,
        final Map<BlockStoreLocation, List<Long>> currentBlocksOnLocation,
        final Map<String, List<String>> lostStorage,
        final List<ConfigProperty> configList) throws IOException {
      AtomicReference<IOException> ioe = new AtomicReference<>();
      try {
        mStreamer.registerWithMaster();
      } catch (IOException e) {
        ioe.set(e);
      } catch (InterruptedException e) {
        ioe.set(new IOException(e));
      }
      if (ioe.get() != null) {
        throw ioe.get();
      }
      return;
    }

    @Override
    public synchronized Command heartbeat(
        final long workerId, final Map<String, Long> capacityBytesOnTiers,
        final Map<String, Long> usedBytesOnTiers,
        final List<Long> removedBlocks, final Map<BlockStoreLocation, List<Long>> addedBlocks,
        final Map<String, List<String>> lostStorage, final List<Metric> metrics) {
      assertEquals(mRemovedBlocks, removedBlocks);
      assertEquals(mAddedBlocks, addedBlocks);
      return Command.getDefaultInstance();
    }

    @Override
    public void commitBlock(final long workerId, final long usedBytesOnTier,
        final String tierAlias, final String mediumType,
        final long blockId, final long length) {
      // Noop because there is no master
    }
  }

  private void prepareBlocksOnWorker(String tierConfig) {
    List<String> tierAliases = getTierAliases(parseTierConfig(tierConfig));
    // Generate block IDs heuristically
    Map<TierAlias, List<Integer>> tierConfigMap = parseTierConfig(tierConfig);
    Map<BlockStoreLocation, List<Long>> blockMap =
        RpcBenchPreparationUtils.generateBlockIdOnTiers(tierConfigMap);

    for (Map.Entry<BlockStoreLocation, List<Long>> entry : blockMap.entrySet()) {
      BlockStoreLocation loc = entry.getKey();
      int tierIndex = mTierToIndex.get(loc.tierAlias());
      for (long blockId : entry.getValue()) {
        mBlockWorker.createBlock(1L, blockId, tierIndex,
            new CreateBlockOptions(null, loc.tierAlias(), 1));
        mBlockWorker.commitBlock(1L, blockId, false);
      }
    }
  }

  // Controls the behavior of the manually overriden master-side register stream observer
  enum MasterMode {
    FIRST_REQUEST,
    SECOND_REQUEST,
    ON_COMPLETED,
    HANG_IN_STREAM,
    HANG_ON_COMPLETED,
    COMPLETE_AND_REPORT,
    SIGNAL_IN_STREAM
  }

  // A master side register stream handler that may return error on certain calls
  class ManualRegisterStreamObserver implements StreamObserver<RegisterWorkerPRequest> {
    private int mBatch = 0;
    MasterMode mMasterMode;
    StreamObserver<RegisterWorkerPResponse> mResponseObserver;
    // Used to pass a signal back to the tester
    AtomicInteger mReceivedCount = null;
    AtomicBoolean mCompleted = null;
    CountDownLatch mSignalLatch = null;

    ManualRegisterStreamObserver(MasterMode masterMode) {
      mMasterMode = masterMode;
    }

    ManualRegisterStreamObserver(MasterMode masterMode,
        AtomicInteger receivedCount, AtomicBoolean completed) {
      this(masterMode);
      mReceivedCount = receivedCount;
      mCompleted = completed;
    }

    ManualRegisterStreamObserver(MasterMode masterMode, CountDownLatch signalLatch) {
      this(masterMode);
      mSignalLatch = signalLatch;
    }

    void setResponseObserver(StreamObserver<RegisterWorkerPResponse> responseObserver) {
      mResponseObserver = responseObserver;
    }

    @Override
    public void onNext(alluxio.grpc.RegisterWorkerPRequest chunk) {
      if (mMasterMode == MasterMode.HANG_IN_STREAM) {
        return;
      }
      if (mBatch == 0 && mMasterMode == MasterMode.FIRST_REQUEST) {
        // Throw a checked exception that is the most likely at this stage
        mResponseObserver.onError(GrpcExceptionUtils.fromThrowable(
            new NotFoundException("Simulate worker is not found")));
        mBatch++;
        return;
      }
      if (mBatch == 1 && mMasterMode == MasterMode.SECOND_REQUEST) {
        // There is no checked exception after the 1st request
        // It is probably because something is wrong on the master side
        StatusException x = new InternalException(
            new RuntimeException("Error on the server side")).toGrpcStatusException();
        mResponseObserver.onError(GrpcExceptionUtils.fromThrowable(x));
        mBatch++;
        return;
      }
      mResponseObserver.onNext(RegisterWorkerPResponse.getDefaultInstance());
      if (mBatch == 0 && mMasterMode == MasterMode.SIGNAL_IN_STREAM) {
        // Send a signal and continue
        mSignalLatch.countDown();
      }
      mBatch++;
    }

    @Override
    public void onError(Throwable t) {
      mResponseObserver.onError(t);
    }

    @Override
    public void onCompleted() {
      if (mMasterMode == MasterMode.HANG_ON_COMPLETED) {
        // No response is equal to infinite waiting
        return;
      }
      if (mMasterMode == MasterMode.ON_COMPLETED) {
        // There is no checked exception after the 1st request
        // It is probably because something is wrong on the master side
        StatusException x = new InternalException(
            new RuntimeException("Error on completing the server side")).toGrpcStatusException();
        mResponseObserver.onError(GrpcExceptionUtils.fromThrowable(x));
        return;
      }
      if (mMasterMode == MasterMode.COMPLETE_AND_REPORT) {
        mReceivedCount.set(mBatch);
        mCompleted.set(true);
      }
      mResponseObserver.onCompleted();
    }
  }
}
