content: /*
   * 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.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_BYTES;
  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.getErrorCapturingResponseObserver;
  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.assertThrows;
  
  import alluxio.clock.ManualClock;
  import alluxio.conf.Configuration;
  import alluxio.conf.PropertyKey;
  import alluxio.exception.BlockInfoException;
  import alluxio.exception.runtime.FailedPreconditionRuntimeException;
  import alluxio.exception.status.UnavailableException;
  import alluxio.grpc.Command;
  import alluxio.grpc.CommandType;
  import alluxio.grpc.LocationBlockIdListEntry;
  import alluxio.grpc.RegisterWorkerPRequest;
  import alluxio.grpc.RegisterWorkerPResponse;
  import alluxio.grpc.StorageList;
  import alluxio.master.AlwaysPrimaryPrimarySelector;
  import alluxio.master.CoreMasterContext;
  import alluxio.master.MasterRegistry;
  import alluxio.master.MasterTestUtils;
  import alluxio.master.block.BlockMaster;
  import alluxio.master.block.BlockMasterTestUtils;
  import alluxio.master.block.BlockMasterWorkerServiceHandler;
  import alluxio.master.block.DefaultBlockMaster;
  import alluxio.master.block.RegisterStreamObserver;
  import alluxio.master.block.WorkerRegisterContext;
  import alluxio.master.block.meta.MasterWorkerInfo;
  import alluxio.master.journal.noop.NoopJournalSystem;
  import alluxio.master.metrics.MetricsMaster;
  import alluxio.master.metrics.MetricsMasterFactory;
  import alluxio.stress.cli.RpcBenchPreparationUtils;
  import alluxio.stress.rpc.TierAlias;
  import alluxio.util.CommonUtils;
  import alluxio.util.ThreadFactoryUtils;
  import alluxio.util.executor.ExecutorServiceFactories;
  import alluxio.util.executor.ExecutorServiceFactory;
  import alluxio.wire.BlockInfo;
  import alluxio.wire.WorkerInfo;
  import alluxio.worker.block.BlockStoreLocation;
  import alluxio.worker.block.RegisterStreamer;
  
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.Maps;
  import io.grpc.stub.StreamObserver;
  import org.junit.After;
  import org.junit.Before;
  import org.junit.Test;
  
  import java.io.IOException;
  import java.time.Clock;
  import java.time.Duration;
  import java.time.temporal.ChronoUnit;
  import java.util.Collection;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Queue;
  import java.util.Set;
  import java.util.concurrent.ConcurrentLinkedQueue;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.stream.Collectors;
  
  /**
   * Integration tests for the server-side logic for the register stream.
   */
  public class BlockMasterRegisterStreamIntegrationTest {
    private BlockMaster mBlockMaster;
    private CoreMasterContext mMasterContext;
    private MasterRegistry mRegistry;
    private ManualClock mClock;
    private ExecutorService mExecutorService;
    private MetricsMaster mMetricsMaster;
    private BlockMasterWorkerServiceHandler mHandler;
  
    private static final int MASTER_WORKER_TIMEOUT = 1_000_000;
    private static final long BLOCK_SIZE =
        Configuration.getBytes(PropertyKey.USER_BLOCK_SIZE_BYTES_DEFAULT);
    private static final Map<String, StorageList> NO_LOST_STORAGE = ImmutableMap.of();
    private static final Command EMPTY_CMD =
        Command.newBuilder().setCommandType(CommandType.Nothing).build();
  
    /**
     * 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.MASTER_WORKER_TIMEOUT_MS, MASTER_WORKER_TIMEOUT);
      Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_STREAM_RESPONSE_TIMEOUT, "1s");
      Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_ENABLED, false);
  
      mRegistry = new MasterRegistry();
      mMasterContext = MasterTestUtils.testMasterContext(
          new NoopJournalSystem(), null, new AlwaysPrimaryPrimarySelector()
      );
      mMetricsMaster = new MetricsMasterFactory().create(mRegistry, mMasterContext);
      mRegistry.add(MetricsMaster.class, mMetricsMaster);
      mClock = new ManualClock();
  
      mExecutorService =
          Executors.newFixedThreadPool(10, ThreadFactoryUtils.build("TestBlockMaster-%d", true));
      mBlockMaster = new DefaultBlockMaster(mMetricsMaster, mMasterContext, mClock,
          ExecutorServiceFactories.constantExecutorServiceFactory(mExecutorService));
      mRegistry.add(BlockMaster.class, mBlockMaster);
      mRegistry.start(true);
      mHandler = new BlockMasterWorkerServiceHandler(mBlockMaster);
    }
  
    @After
    public void after() throws Exception {
      mRegistry.stop();
    }
  
    /**
     * Tests below cover the most normal cases.
     */
    @Test
    public void registerEmptyWorkerStream() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForEmptyWorker(workerId);
  
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
  
      // Verify the worker is registered
      assertEquals(0, errorQueue.size());
      assertEquals(1, mBlockMaster.getWorkerCount());
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      assertEquals(0, worker.getBlockCount());
      assertEquals(0, worker.getToRemoveBlockCount());
      assertEquals(MEM_CAPACITY_BYTES, worker.getCapacityBytes());
      assertEquals(MEM_CAPACITY_BYTES, worker.getAvailableBytes());
      assertEquals(0L, worker.getUsedBytes());
  
      // Verify the worker is readable and writable
      verifyWorkerWritable(workerId);
    }
  
    @Test
    public void registerWorkerStream() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
  
      // Verify the worker is registered
      assertEquals(0, errorQueue.size());
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      assertEquals(TIER_BLOCK_TOTAL, worker.getBlockCount());
      assertEquals(0, worker.getToRemoveBlockCount());
      assertEquals(1, mBlockMaster.getWorkerCount());
  
      // Verify the worker is readable and writable
      verifyWorkerWritable(workerId);
    }
  
    @Test
    // The worker once registered with the master and has lost heartbeat.
    // The master has marked the worker as lost.
    public void registerLostWorker() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      // The worker registers to the master
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
      // Verify the worker has been registered
      assertEquals(0, errorQueue.size());
      assertEquals(1, mBlockMaster.getWorkerCount());
  
      // The worker has lost heartbeat and been forgotten
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      long newTimeMs = worker.getLastUpdatedTimeMs() + MASTER_WORKER_TIMEOUT + 1;
      mClock.setTimeMs(newTimeMs);
      DefaultBlockMaster.LostWorkerDetectionHeartbeatExecutor lostWorkerDetector =
          ((DefaultBlockMaster) mBlockMaster).new LostWorkerDetectionHeartbeatExecutor();
      lostWorkerDetector.heartbeat(Long.MAX_VALUE);
  
      // Verify the worker has been forgotten
      assertEquals(0, mBlockMaster.getWorkerCount());
  
      // Register again
      Queue<Throwable> newErrorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(newErrorQueue));
  
      // Verify the worker is registered again
      assertEquals(0, errorQueue.size());
      MasterWorkerInfo updatedWorker = mBlockMaster.getWorker(workerId);
      assertEquals(TIER_BLOCK_TOTAL, updatedWorker.getBlockCount());
      assertEquals(0, updatedWorker.getToRemoveBlockCount());
      assertEquals(1, mBlockMaster.getWorkerCount());
  
      // Verify the worker is readable and writable
      verifyWorkerWritable(workerId);
    }
  
    @Test
    // The worker has not yet been forgotten by the master but attempts to register again.
    // This can happen when a worker process is restarted.
    public void registerExistingWorker() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
      assertEquals(0, errorQueue.size());
      // Verify the worker has registered
      assertEquals(1, mBlockMaster.getWorkerCount());
  
      // Register again
      Queue<Throwable> newErrorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(newErrorQueue));
      assertEquals(0, newErrorQueue.size());
  
      // Verify the worker is registered
      MasterWorkerInfo updatedWorker = mBlockMaster.getWorker(workerId);
      assertEquals(TIER_BLOCK_TOTAL, updatedWorker.getBlockCount());
      assertEquals(0, updatedWorker.getToRemoveBlockCount());
      assertEquals(1, mBlockMaster.getWorkerCount());
  
      // Verify the worker is readable and writable
      verifyWorkerWritable(workerId);
    }
  
    @Test
    public void registerExistingWorkerBlocksLost() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      // Register the worker for the 1st time
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
      assertEquals(0, errorQueue.size());
      // Verify the worker has registered
      assertEquals(1, mBlockMaster.getWorkerCount());
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      assertEquals(TIER_BLOCK_TOTAL, worker.getBlockCount());
      assertEquals(0, worker.getToRemoveBlockCount());
  
      // Manually generate the blocks again and remove some
      List<String> tierAliases = getTierAliases(parseTierConfig(TIER_CONFIG));
      Map<BlockStoreLocation, List<Long>> blockMap =
          RpcBenchPreparationUtils.generateBlockIdOnTiers(parseTierConfig(TIER_CONFIG));
      Set<Long> lostBlocks = removeSomeBlocks(blockMap);
      // Regenerate the requests
      RegisterStreamer newRegisterStreamer = new RegisterStreamer(null,
          workerId, tierAliases, CAPACITY_MAP, USAGE_MAP, blockMap, LOST_STORAGE, EMPTY_CONFIG);
      List<RegisterWorkerPRequest> newRequestChunks = ImmutableList.copyOf(newRegisterStreamer);
      int newExpectedBatchCount =
          (int) Math.ceil((TIER_BLOCK_TOTAL - lostBlocks.size()) / (double) BATCH_SIZE);
      assertEquals(newExpectedBatchCount, newRequestChunks.size());
  
      // Register again with the updated stream
      Queue<Throwable> newErrorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, newRequestChunks,
          getErrorCapturingResponseObserver(newErrorQueue));
      assertEquals(0, newErrorQueue.size());
  
      // Verify the worker is registered
      assertEquals(1, mBlockMaster.getWorkerCount());
      MasterWorkerInfo updatedWorker = mBlockMaster.getWorker(workerId);
      assertEquals(TIER_BLOCK_TOTAL - lostBlocks.size(), updatedWorker.getBlockCount());
      // The master will mark the lost blocks as to be removed
      // This is to ensure the unrecognized blocks do no live on the worker anymore
      assertEquals(lostBlocks.size(), updatedWorker.getToRemoveBlockCount());
  
      // The update is received during the registration so no command to send to the worker
      Command command = sendHeartbeatToMaster(workerId);
      assertEquals(CommandType.Free, command.getCommandType());
      assertEquals(lostBlocks, new HashSet<>(command.getDataList()));
  
      // Verify the worker is readable and writable
      verifyWorkerWritable(workerId);
    }
  
    @Test
    public void registerExistingWorkerBlocksAdded() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      // Register the worker for the 1st time
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
      assertEquals(0, errorQueue.size());
      // Verify the worker has registered
      assertEquals(1, mBlockMaster.getWorkerCount());
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      assertEquals(TIER_BLOCK_TOTAL, worker.getBlockCount());
      assertEquals(0, worker.getToRemoveBlockCount());
  
      // Generate block IDs in the same way but add some more
      Map<BlockStoreLocation, List<Long>> blockMap =
          RpcBenchPreparationUtils.generateBlockIdOnTiers(parseTierConfig(TIER_CONFIG));
      Set<Long> addedBlocks = addSomeBlocks(blockMap);
      // Make the master accept these blocks
      prepareBlocksOnMaster(addedBlocks);
  
      // Re-generate the request
      List<String> tierAliases = getTierAliases(parseTierConfig(TIER_CONFIG));
      Map<String, Long> capacityMap = Maps.toMap(tierAliases, (tier) -> CAPACITY);
      Map<String, Long> usedMap = Maps.toMap(tierAliases, (tier) -> 0L);
      RegisterStreamer newRegisterStreamer = new RegisterStreamer(null,
          workerId, tierAliases, capacityMap, usedMap, blockMap, LOST_STORAGE, EMPTY_CONFIG);
      List<RegisterWorkerPRequest> newRequestChunks = ImmutableList.copyOf(newRegisterStreamer);
      int newExpectedBatchCount =
          (int) Math.ceil((TIER_BLOCK_TOTAL + addedBlocks.size()) / (double) BATCH_SIZE);
      assertEquals(newExpectedBatchCount, newRequestChunks.size());
  
      // Register again with the new request stream
      Queue<Throwable> newErrorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, newRequestChunks,
          getErrorCapturingResponseObserver(newErrorQueue));
      assertEquals(0, newErrorQueue.size());
  
      // Verify the worker is registered
      assertEquals(1, mBlockMaster.getWorkerCount());
      MasterWorkerInfo updatedWorker = mBlockMaster.getWorker(workerId);
      assertEquals(TIER_BLOCK_TOTAL + addedBlocks.size(), updatedWorker.getBlockCount());
      assertEquals(0, updatedWorker.getToRemoveBlockCount());
  
      // No command from the master because the update is received during registration
      assertEquals(EMPTY_CMD, sendHeartbeatToMaster(workerId));
  
      // Verify the worker is readable and writable
      verifyWorkerWritable(workerId);
    }
  
    /**
     * Tests below cover various failure cases.
     */
    @Test
    public void hangingWorkerSessionRecycled() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
  
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      StreamObserver<RegisterWorkerPRequest> requestObserver =
          mHandler.registerWorkerStream(getErrorCapturingResponseObserver(errorQueue));
  
      // Feed the chunks into the requestObserver
      for (int i = 0; i < requestChunks.size(); i++) {
        RegisterWorkerPRequest chunk = requestChunks.get(i);
        // From the 2nd request on, the request will be rejected
        requestObserver.onNext(chunk);
        if (i == 0) {
          // Progress the clock so the worker stream expired
          mClock.addTime(Duration.of(100_000_000, ChronoUnit.MILLIS));
          // Sleep and wait for the stream recycler thread heartbeat
          CommonUtils.sleepMs(3000);
        }
      }
      // This will be rejected too
      requestObserver.onCompleted();
      // The 5 requests after expiry will be rejected
      // And the complete message will be rejected too
      // -1 because the 1st request was accepted
      // +1 because master sends the TimeoutException to worker on timeout
      // +count because all following requests are rejected
      // +1 because the onComplete() is also rejected
      assertEquals(requestChunks.size() - 1 + 1 + 1, errorQueue.size());
  
      // Verify the session is recycled
      assertEquals(0, mBlockMaster.getWorkerCount());
  
      // Verify the worker can re-register and be operated, so the locks are managed correctly
      verifyWorkerCanReregister(workerId, requestChunks, TIER_BLOCK_TOTAL);
      verifyWorkerWritable(workerId);
    }
  
    @Test
    public void workerSendsErrorOnStart() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
  
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      StreamObserver<RegisterWorkerPRequest> requestObserver =
          mHandler.registerWorkerStream(
              getErrorCapturingResponseObserver(errorQueue));
  
      // Instead of sending requests to the master, the worker is interrupted
      // around the beginning of the stream. The error propagated to the master.
      Exception e = new InterruptedException("Worker is interrupted");
      assertThrows(IllegalStateException.class, () -> {
        requestObserver.onError(e);
      });
  
      // After the requestObserver.onError(), no requests
      // shall be accepted by the master.
      for (RegisterWorkerPRequest chunk : requestChunks) {
        requestObserver.onNext(chunk);
      }
      requestObserver.onCompleted();
      assertEquals(requestChunks.size() + 1, errorQueue.size());
  
      // verify the worker is not registered
      assertEquals(0, mBlockMaster.getWorkerCount());
  
      verifyWorkerCanReregister(workerId, requestChunks, TIER_BLOCK_TOTAL);
    }
  
    @Test
    public void workerSendsErrorInStream() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
  
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      StreamObserver<RegisterWorkerPRequest> requestObserver =
          mHandler.registerWorkerStream(
              getErrorCapturingResponseObserver(errorQueue));
  
      // An error took place in the worker during the stream
      RegisterWorkerPRequest first = requestChunks.get(0);
      requestObserver.onNext(first);
      Exception e = new IOException("An exception occurred during the stream");
      requestObserver.onError(e);
  
      // The following requests will be rejected as the stream has been closed
      // on the master side.
      RegisterWorkerPRequest second = requestChunks.get(1);
      requestObserver.onNext(second);
      requestObserver.onCompleted();
      assertEquals(2, errorQueue.size());
  
      // verify the worker is not registered
      assertEquals(0, mBlockMaster.getWorkerCount());
  
      verifyWorkerCanReregister(workerId, requestChunks, TIER_BLOCK_TOTAL);
    }
  
    @Test
    public void workerSendsErrorBeforeCompleting() throws Exception {
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
  
      // Send the requests to the master
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      StreamObserver<RegisterWorkerPRequest> requestObserver =
          mHandler.registerWorkerStream(
              getErrorCapturingResponseObserver(errorQueue));
      for (RegisterWorkerPRequest chunk : requestChunks) {
        requestObserver.onNext(chunk);
      }
  
      // An error took place on the worker side after all the requests have been sent
      Exception e = new IOException("Exception on the worker side");
      requestObserver.onError(e);
      // The complete should be rejected
      requestObserver.onCompleted();
      assertEquals(1, errorQueue.size());
      // verify the worker is not registered
      assertEquals(0, mBlockMaster.getWorkerCount());
  
      verifyWorkerCanReregister(workerId, requestChunks, TIER_BLOCK_TOTAL);
    }
  
    @Test
    public void workerRegisterStartThrowsError() throws Exception {
      // Hijack the block master so it throws errors
      BrokenBlockMaster brokenBlockMaster = new BrokenBlockMaster(
          mMetricsMaster, mMasterContext, mClock,
          ExecutorServiceFactories.constantExecutorServiceFactory(mExecutorService),
          WorkerRegisterMode.ERROR_START);
      mBlockMaster = brokenBlockMaster;
      mHandler = new BlockMasterWorkerServiceHandler(brokenBlockMaster);
  
      long workerId = brokenBlockMaster.getWorkerId(NET_ADDRESS_1);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamObserver streamOb =
          new RegisterStreamObserver(brokenBlockMaster,
              getErrorCapturingResponseObserver(errorQueue));
  
      // Generate requests
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      for (RegisterWorkerPRequest chunk : requestChunks) {
        streamOb.onNext(chunk);
      }
      streamOb.onCompleted();
      // All requests + onCompleted should receive an error
      assertEquals(requestChunks.size() + 1, errorQueue.size());
  
      // The BlockMaster is not throwing error this time
      brokenBlockMaster.setCorrect();
      verifyWorkerCanReregister(workerId, requestChunks, TIER_BLOCK_TOTAL);
    }
  
    @Test
    public void workerRegisterBatchThrowsError() throws Exception {
      // Hijack the block master so it throws errors
      BrokenBlockMaster brokenBlockMaster = new BrokenBlockMaster(
          mMetricsMaster, mMasterContext, mClock,
          ExecutorServiceFactories.constantExecutorServiceFactory(mExecutorService),
          WorkerRegisterMode.ERROR_STREAM);
      mBlockMaster = brokenBlockMaster;
      mHandler = new BlockMasterWorkerServiceHandler(brokenBlockMaster);
  
      long workerId = brokenBlockMaster.getWorkerId(NET_ADDRESS_1);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamObserver streamOb =
          new RegisterStreamObserver(brokenBlockMaster,
              getErrorCapturingResponseObserver(errorQueue));
  
      // Generate requests
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      for (RegisterWorkerPRequest chunk : requestChunks) {
        streamOb.onNext(chunk);
      }
      streamOb.onCompleted();
      // All requests except the 1st one + onCompleted should receive an error
      assertEquals(requestChunks.size() - 1 + 1, errorQueue.size());
  
      // The BlockMaster is not throwing error this time
      brokenBlockMaster.setCorrect();
      verifyWorkerCanReregister(workerId, requestChunks, TIER_BLOCK_TOTAL);
    }
  
    @Test
    public void workerRegisterCompleteThrowsError() throws Exception {
      // Hijack the block master so it throws errors
      BrokenBlockMaster brokenBlockMaster = new BrokenBlockMaster(
          mMetricsMaster, mMasterContext, mClock,
          ExecutorServiceFactories.constantExecutorServiceFactory(mExecutorService),
          WorkerRegisterMode.ERROR_COMPLETE);
      mBlockMaster = brokenBlockMaster;
      mHandler = new BlockMasterWorkerServiceHandler(brokenBlockMaster);
  
      long workerId = brokenBlockMaster.getWorkerId(NET_ADDRESS_1);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamObserver streamOb =
          new RegisterStreamObserver(brokenBlockMaster,
              getErrorCapturingResponseObserver(errorQueue));
  
      // Generate requests
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      for (RegisterWorkerPRequest chunk : requestChunks) {
        streamOb.onNext(chunk);
      }
      streamOb.onCompleted();
      // Only onCompleted should receive an error
      assertEquals(1, errorQueue.size());
  
      // The BlockMaster is not throwing error this time
      brokenBlockMaster.setCorrect();
      verifyWorkerCanReregister(workerId, requestChunks, TIER_BLOCK_TOTAL);
    }
  
    // Different modes that defines how the manually overriden block master work
    enum WorkerRegisterMode {
      ERROR_START,
      ERROR_STREAM,
      ERROR_COMPLETE,
      OK
    }
  
    // A BlockMaster that can throws error on certain calls
    class BrokenBlockMaster extends DefaultBlockMaster {
      private WorkerRegisterMode mMode;
  
      public BrokenBlockMaster(MetricsMaster metricsMaster, CoreMasterContext masterContext,
           Clock clock, ExecutorServiceFactory executorServiceFactory, WorkerRegisterMode mode) {
        super(metricsMaster, masterContext, clock, executorServiceFactory);
        mMode = mode;
      }
  
      public void setCorrect() {
        mMode = WorkerRegisterMode.OK;
      }
  
      @Override
      public void workerRegisterStart(WorkerRegisterContext context, RegisterWorkerPRequest chunk) {
        if (mMode == WorkerRegisterMode.ERROR_START) {
          throw new FailedPreconditionRuntimeException("Error from workerRegisterStart");
        }
        super.workerRegisterStart(context, chunk);
      }
  
      @Override
      public void workerRegisterBatch(WorkerRegisterContext context, RegisterWorkerPRequest chunk) {
        if (mMode == WorkerRegisterMode.ERROR_STREAM) {
          throw new FailedPreconditionRuntimeException("Error from workerRegisterBatch");
        }
        super.workerRegisterBatch(context, chunk);
      }
  
      @Override
      public void workerRegisterFinish(WorkerRegisterContext context) {
        if (mMode == WorkerRegisterMode.ERROR_COMPLETE) {
          throw new FailedPreconditionRuntimeException("Error from workerRegisterBatch");
        }
        super.workerRegisterFinish(context);
      }
    }
  
    /**
     * Tests below cover the race conditions during concurrent executions.
     *
     * When the worker registers for the 1st time, 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 master-side metadata.
     * In other words, we assume those writers succeed on the worker, and the subsequent
     * update on the master-side metadata should also succeed and be correct.
     */
    @Test
    public void reregisterWithDelete() throws Exception {
      // Register the worker so the worker is marked active in master
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
      assertEquals(0, errorQueue.size());
      assertEquals(1, mBlockMaster.getWorkerCount());
  
      // Find a block to remove
      long blockToRemove = RegisterStreamTestUtils.findFirstBlock(requestChunks);
  
      // Register again
      CountDownLatch latch = new CountDownLatch(1);
      Queue<Throwable> newErrorQueue = new ConcurrentLinkedQueue<>();
      Future f = mExecutorService.submit(() -> {
        sendStreamToMasterAndSignal(requestChunks,
            getErrorCapturingResponseObserver(newErrorQueue), latch);
      });
  
      // During the register stream, trigger a delete on worker
      latch.await();
      mBlockMaster.removeBlocks(ImmutableList.of(blockToRemove), true);
  
      // Wait for the register to finish
      f.get();
  
      assertThrows(BlockInfoException.class, () -> {
        mBlockMaster.getBlockInfo(blockToRemove);
      });
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      assertEquals(1, mBlockMaster.getWorkerCount());
      assertEquals(TIER_BLOCK_TOTAL - 1, worker.getBlockCount());
  
      // BlockMaster.removeBlocks() will first remove the block from master metadata
      // (with block lock) then update the block locations (with worker lock).
      // The worker lock is being held by the registering worker, but the 1st part
      // will likely succeed.
      // So during registration when checking on the block, the block is not recognized
      // any more and will remain in MasterWorkerInfo.mToRemoveBlocks.
      // In the next heartbeat the master will issue a command to remove the block
      // from the worker.
      // Even if the block is already removed on the worker it is fine,
      // because deletion of a not-found block is a noop.
      Command command = sendHeartbeatToMaster(workerId);
      assertEquals(Command.newBuilder()
          .addData(blockToRemove).setCommandType(CommandType.Free).build(), command);
    }
  
    @Test
    public void reregisterWithFree() throws Exception {
      // Register the worker so the worker is marked active in master
      long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
      List<RegisterWorkerPRequest> requestChunks =
          RegisterStreamTestUtils.generateRegisterStreamForWorkerWithTiers(workerId);
      prepareBlocksOnMaster(requestChunks);
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
      assertEquals(0, errorQueue.size());
      assertEquals(1, mBlockMaster.getWorkerCount());
  
      // Find a block to free
      long blockToRemove = RegisterStreamTestUtils.findFirstBlock(requestChunks);
  
      // Register again
      CountDownLatch latch = new CountDownLatch(1);
      Queue<Throwable> newErrorQueue = new ConcurrentLinkedQueue<>();
      mExecutorService.submit(() -> {
        sendStreamToMasterAndSignal(requestChunks,
            getErrorCapturingResponseObserver(newErrorQueue), latch);
      });
  
      // During the register stream, trigger a delete on worker
      latch.await();
      mBlockMaster.removeBlocks(ImmutableList.of(blockToRemove), false);
  
      BlockInfo info = mBlockMaster.getBlockInfo(blockToRemove);
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      assertEquals(0, newErrorQueue.size());
      assertEquals(1, mBlockMaster.getWorkerCount());
      // The block still exists on the worker but a command will be issued to remove it
      assertEquals(TIER_BLOCK_TOTAL, worker.getBlockCount());
  
      Command command = sendHeartbeatToMaster(workerId);
      assertEquals(Command.newBuilder()
          .setCommandType(CommandType.Free).addData(blockToRemove).build(), command);
    }
  
    private void prepareBlocksOnMaster(Collection<Long> blockIds) throws UnavailableException {
      for (long id : blockIds) {
        mBlockMaster.commitBlockInUFS(id, BLOCK_SIZE);
      }
    }
  
    private void prepareBlocksOnMaster(List<RegisterWorkerPRequest> requestChunks) throws Exception {
      for (RegisterWorkerPRequest chunk : requestChunks) {
        List<LocationBlockIdListEntry> entries = chunk.getCurrentBlocksList();
        for (LocationBlockIdListEntry entry : entries) {
          for (long blockId : entry.getValue().getBlockIdList()) {
            mBlockMaster.commitBlockInUFS(blockId, BLOCK_SIZE);
          }
        }
      }
    }
  
    private void sendStreamToMasterAndSignal(
        List<RegisterWorkerPRequest> requestChunks,
        StreamObserver<RegisterWorkerPResponse> responseObserver,
        CountDownLatch latch) {
      StreamObserver<RegisterWorkerPRequest> requestObserver =
              mHandler.registerWorkerStream(responseObserver);
      for (int i = 0; i < requestChunks.size(); i++) {
        RegisterWorkerPRequest chunk = requestChunks.get(i);
        requestObserver.onNext(chunk);
        // Signal after the 1st request has been sent
        if (i == 0) {
          latch.countDown();
        }
      }
      requestObserver.onCompleted();
    }
  
    // Removes blocks from the current worker blocks.
    // Returns the set of blocks that are removed.
    private Set<Long> removeSomeBlocks(Map<BlockStoreLocation, List<Long>> blockMap) {
      Set<Long> toRemove = new HashSet<>();
      // 1 block goes missing from each location
      for (Map.Entry<BlockStoreLocation, List<Long>> entry : blockMap.entrySet()) {
        List<Long> blocks = entry.getValue();
        if (blocks.isEmpty()) {
          continue;
        }
        toRemove.add(blocks.get(0));
        blocks.remove(0);
      }
      return toRemove;
    }
  
    // Add some blocks to the current worker blocks.
    // Returns the set of blocks that are added.
    private Set<Long> addSomeBlocks(Map<BlockStoreLocation, List<Long>> blockMap) {
      long blockId = 100L;
      Set<Long> added = new HashSet<>();
      // 1 block is added into each location
      for (Map.Entry<BlockStoreLocation, List<Long>> entry : blockMap.entrySet()) {
        List<Long> blocks = entry.getValue();
        added.add(blockId);
        blocks.add(blockId);
        blockId++;
      }
      return added;
    }
  
    // Verify the worker is writable by trying to commit a block in it
    private void verifyWorkerWritable(long workerId) throws Exception {
      // First see the worker usage
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      long used = worker.getUsedBytes();
      // Commit a new block to the worker
      long newBlockId = 999L;
      long newBlockSize = 1024L;
  
      mBlockMaster.commitBlock(workerId, used + newBlockSize,
          "MEM", "MEM", newBlockId, newBlockSize);
  
      List<WorkerInfo> workers = mBlockMaster.getWorkerInfoList();
      WorkerInfo workerInfo = BlockMasterTestUtils.findWorkerInfo(workers, workerId);
      BlockMasterTestUtils.verifyBlockOnWorkers(
          mBlockMaster, newBlockId, newBlockSize, ImmutableList.of(workerInfo));
    }
  
    private static List<String> getTierAliases(Map<TierAlias, List<Integer>> tierConfig) {
      return tierConfig.keySet().stream().map(TierAlias::toString).collect(Collectors.toList());
    }
  
    // Verify a worker can reregister and have the correct final blocks
    private void verifyWorkerCanReregister(long workerId,
        List<RegisterWorkerPRequest> requestChunks, int expectedBlockCount) throws Exception {
      Queue<Throwable> errorQueue = new ConcurrentLinkedQueue<>();
      RegisterStreamTestUtils.sendStreamToMaster(mHandler, requestChunks,
          getErrorCapturingResponseObserver(errorQueue));
  
      assertEquals(errorQueue.toString(), 0, errorQueue.size());
      MasterWorkerInfo worker = mBlockMaster.getWorker(workerId);
      assertEquals(expectedBlockCount, worker.getBlockCount());
      assertEquals(0, worker.getToRemoveBlockCount());
      assertEquals(1, mBlockMaster.getWorkerCount());
    }
  
    private Command sendHeartbeatToMaster(long workerId) {
      return mBlockMaster.workerHeartbeat(
          workerId,
          CAPACITY_MAP,
          USAGE_MAP,
          // list of removed blockIds
          ImmutableList.of(),
          // list of added blockIds
          ImmutableMap.of(),
          NO_LOST_STORAGE,
          ImmutableList.of());
    }
  }
  