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.master.block;
  
  import static alluxio.master.block.BlockMasterTestUtils.findWorkerInfo;
  import static alluxio.master.block.BlockMasterTestUtils.verifyBlockNotExisting;
  import static alluxio.master.block.BlockMasterTestUtils.verifyBlockOnWorkers;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertThrows;
  import static org.junit.Assert.assertTrue;
  
  import alluxio.client.block.options.GetWorkerReportOptions;
  import alluxio.clock.ManualClock;
  import alluxio.exception.BlockInfoException;
  import alluxio.grpc.Command;
  import alluxio.grpc.CommandType;
  import alluxio.grpc.RegisterWorkerPOptions;
  import alluxio.grpc.StorageList;
  import alluxio.heartbeat.HeartbeatContext;
  import alluxio.heartbeat.ManuallyScheduleHeartbeat;
  import alluxio.master.AlwaysPrimaryPrimarySelector;
  import alluxio.master.CoreMasterContext;
  import alluxio.master.MasterRegistry;
  import alluxio.master.MasterTestUtils;
  import alluxio.master.journal.noop.NoopJournalSystem;
  import alluxio.master.metrics.MetricsMaster;
  import alluxio.master.metrics.MetricsMasterFactory;
  import alluxio.proto.meta.Block;
  import alluxio.util.ThreadFactoryUtils;
  import alluxio.util.executor.ExecutorServiceFactories;
  import alluxio.wire.BlockInfo;
  import alluxio.wire.BlockLocation;
  import alluxio.wire.WorkerInfo;
  import alluxio.wire.WorkerNetAddress;
  
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.ImmutableMap;
  import org.junit.After;
  import org.junit.Before;
  import org.junit.ClassRule;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.rules.ExpectedException;
  import org.junit.rules.TemporaryFolder;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.Map;
  import java.util.Queue;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ConcurrentLinkedQueue;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  public class ConcurrentBlockMasterTest {
    private static final WorkerNetAddress NET_ADDRESS_1 = new WorkerNetAddress().setHost("localhost")
        .setRpcPort(80).setDataPort(81).setWebPort(82);
    private static final WorkerNetAddress NET_ADDRESS_2 = new WorkerNetAddress().setHost("localhost")
        .setRpcPort(83).setDataPort(84).setWebPort(85);
  
    private static final Map<Block.BlockLocation, List<Long>> NO_BLOCKS_ON_LOCATION =
        ImmutableMap.of();
    private static final Map<String, StorageList> NO_LOST_STORAGE = ImmutableMap.of();
    private static final int CONCURRENT_CLIENT_COUNT = 20;
    private static final long BLOCK1_ID = 1L;
    private static final long BLOCK1_LENGTH = 49L;
    private static final long BLOCK2_ID = 2L;
    private static final long BLOCK2_LENGTH = 59L;
    private static final Map<String, Long> MEM_CAPACITY = ImmutableMap.of("MEM", 100L);
    private static final Map<String, Long> MEM_USAGE_EMPTY = ImmutableMap.of("MEM", 0L);
    private static final Command FREE_BLOCK1_CMD = Command.newBuilder()
        .setCommandType(CommandType.Free).addData(1).build();
    private static final Command EMPTY_CMD = Command.newBuilder()
        .setCommandType(CommandType.Nothing).build();
  
    private SignalBlockMaster mBlockMaster;
    private MasterRegistry mRegistry;
    private ManualClock mClock;
    private ExecutorService mExecutorService;
    private ExecutorService mClientExecutorService;
    private MetricsMaster mMetricsMaster;
    private CoreMasterContext mMasterContext;
  
    /** Rule to create a new temporary folder during each test. */
    @Rule
    public TemporaryFolder mTestFolder = new TemporaryFolder();
  
    /** The exception expected to be thrown. */
    @Rule
    public ExpectedException mThrown = ExpectedException.none();
  
    @ClassRule
    public static ManuallyScheduleHeartbeat sManuallySchedule = new ManuallyScheduleHeartbeat(
            HeartbeatContext.MASTER_LOST_WORKER_DETECTION);
  
    /**
     * Sets up the dependencies before a test runs.
     */
    @Before
    public void before() throws Exception {
      mRegistry = new MasterRegistry();
      mMasterContext = MasterTestUtils.testMasterContext(
          new NoopJournalSystem(), null, new AlwaysPrimaryPrimarySelector()
      );
      mMetricsMaster = new MetricsMasterFactory().create(mRegistry, mMasterContext);
      mClock = new ManualClock();
      mExecutorService =
          Executors.newFixedThreadPool(2, ThreadFactoryUtils.build("TestBlockMaster-%d", true));
      mClientExecutorService = Executors.newFixedThreadPool(2,
          ThreadFactoryUtils.build("TestBlockMaster-%d", true));
      // No one is listening to this latch, so do any preparation for the test
      CountDownLatch voidLatch = new CountDownLatch(1);
      mBlockMaster = new SignalBlockMaster(mMetricsMaster, mMasterContext, mClock,
          ExecutorServiceFactories.constantExecutorServiceFactory(mExecutorService), voidLatch);
      mRegistry.add(BlockMaster.class, mBlockMaster);
      mRegistry.start(true);
    }
  
    /**
     * Stops the master after a test ran.
     */
    @After
    public void after() throws Exception {
      mRegistry.stop();
  
      // When the registry is stopped, the BlockMaster will stop the given ExecutorService
      // We need to manually shutdown this client thread pool
      mClientExecutorService.shutdown();
    }
  
    /**
     * RW contention: Concurrent commit and readers.
     * Signal in commit and the readers inquire the state
     */
    @Test
    public void concurrentCommitWithReaders() throws Exception {
      // Prepare worker
      long worker1 = registerEmptyWorker(NET_ADDRESS_1);
      // Replace the latch used for preparation
      CountDownLatch readerLatch = new CountDownLatch(1);
      mBlockMaster.setLatch(readerLatch);
  
      concurrentWriterWithReaders(
          readerLatch,
          // Writer
          () -> {
            mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
            return null;
          },
          // Reader
          () -> {
            try {
              // If the block is not committed yet, a BlockInfoException will be thrown
              BlockInfo blockInfo = mBlockMaster.getBlockInfo(BLOCK1_ID);
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              WorkerInfo worker = findWorkerInfo(workerInfoList, worker1);
              assertEquals(BLOCK1_LENGTH, worker.getUsedBytes());
  
              BlockLocation blockLocation = new BlockLocation()
                  .setTierAlias("MEM")
                  .setWorkerAddress(NET_ADDRESS_1)
                  .setWorkerId(worker1)
                  .setMediumType("MEM");
              BlockInfo expectedBlockInfo = new BlockInfo()
                  .setBlockId(BLOCK1_ID)
                  .setLength(BLOCK1_LENGTH)
                  .setLocations(ImmutableList.of(blockLocation));
              assertEquals(expectedBlockInfo, blockInfo);
              assertEquals(1, workerInfoList.size());
            } catch (BlockInfoException e) {
              // The reader came in before the writer started the commit
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              assertEquals(1, workerInfoList.size());
              WorkerInfo worker = workerInfoList.get(0);
              // We may just see the result before or after the commit
              // But other values should be illegal
              assertTrue(BLOCK1_LENGTH == worker.getUsedBytes() || 100L == worker.getUsedBytes());
            }
            return null;
          });
    }
  
    /**
     * RW contention: Concurrent remove operation and readers.
     * Readers should read the state either before or after the removal.
     */
    @Test
    public void concurrentRemoveWithReaders() throws Exception {
      for (boolean deleteMetadata : ImmutableList.of(true, false)) {
        // Prepare worker and block 1 on worker
        long worker1 = registerEmptyWorker(NET_ADDRESS_1);
        mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM",
            BLOCK1_ID, BLOCK1_LENGTH);
  
        CountDownLatch readerLatch = new CountDownLatch(1);
        mBlockMaster.setLatch(readerLatch);
  
        concurrentWriterWithReaders(readerLatch,
            // Writer
            () -> {
              List<Long> blocksToRemove = new ArrayList<>();
              blocksToRemove.add(BLOCK1_ID);
              mBlockMaster.removeBlocks(blocksToRemove, deleteMetadata);
              return null;
            },
            // Reader
            () -> {
              try {
                // Even if the block is removed, the worker usage will not be updated
                // until the next worker heartbeat
                List<WorkerInfo> workerInfoList = mBlockMaster.getWorkerReport(
                    GetWorkerReportOptions.defaults());
                assertEquals(1, workerInfoList.size());
                WorkerInfo worker = workerInfoList.get(0);
                assertEquals(BLOCK1_LENGTH, worker.getUsedBytes());
  
                // If the block is removed already, a BlockInfoException will be thrown
                verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, workerInfoList);
              } catch (BlockInfoException e) {
                // If the block has been removed, this exception is expected
                // There is nothing more to test here
              }
              return null;
            });
      }
    }
  
    /**
     *  Concurrently run writer with a bunch of readers.
     *  The readers and the writer use a {@link CountDownLatch} for signal passing.
     *  The writer will release the latch in the middle or the write.
     *  The readers will wait on the latch and start to read on seeing the latch released.
     *
     * @param readerLatch the latch that writer releases and readers wait on
     * @param writer the writer function, nothing will be returned
     * @param reader the reader function, nothing will be returned
     */
    private void concurrentWriterWithReaders(CountDownLatch readerLatch,
                                             Callable writer, Callable reader) throws Exception {
      // This thread count is intentionally larger than the client thread pool
      // In the hope that even if the first batch of clients all read the state before
      // the commit really happens
      // The following batch will capture the state after the commit
      Queue<Throwable> uncaughtThrowables = new ConcurrentLinkedQueue<>();
      CountDownLatch allClientFinished = new CountDownLatch(CONCURRENT_CLIENT_COUNT);
      for (int i = 0; i < CONCURRENT_CLIENT_COUNT; i++) {
        mClientExecutorService.submit(() -> {
          // Wait until the writer enters the critical section and sends a signal
          try {
            readerLatch.await();
          } catch (Throwable t) {
            uncaughtThrowables.add(t);
            // Fail to wait for the signal, just give up
            allClientFinished.countDown();
            return;
          }
          // Trigger the reader
          try {
            reader.call();
          } catch (Throwable t) {
            uncaughtThrowables.add(t);
          } finally {
            allClientFinished.countDown();
          }
        });
      }
  
      // The readers should be waiting for the writer to send the signal
      writer.call();
  
      allClientFinished.await();
      // If any assertion failed, the failed assertion will throw an AssertError
      assertEquals(0, uncaughtThrowables.size());
    }
  
    /**
     * Tests on WW contentions.
     *
     * Write operations are:
     * 1. commit
     * 2. remove
     * 3. workerRegister
     * 4. workerHeartbeat
     *
     * Test W1 race condition with W2 where W1 will send a signal in the middle of run and trigger W2
     * W1 is commit/remove
     * W2 is commit/remove/workerRegister/workerHeartbeat
     *
     * When W1 is operating on block B, if W2 is commit/remove:
     * 1. W2 is on the same block
     * 2. W2 is on a different block
     *
     * When W1 is operating on block B, if W2 is workerRegister/workerHeartbeat,
     * the options are:
     * Opt1: W2 may be from the same worker or a different worker
     * Opt2: W2 may contain the same block or not
     */
    @Test
    public void concurrentCommitWithRegisterNewWorkerSameBlock() throws Exception {
      // Prepare worker
      long worker1 = registerEmptyWorker(NET_ADDRESS_1);
      long worker2 = mBlockMaster.getWorkerId(NET_ADDRESS_2);
      CountDownLatch w1Latch = new CountDownLatch(1);
      mBlockMaster.setLatch(w1Latch);
      concurrentWriterWithWriter(w1Latch,
          // W1
          () -> {
            mBlockMaster.commitBlock(worker1, 49L, "MEM", "MEM",
                BLOCK1_ID, BLOCK1_LENGTH);
            return null;
          },
          // W2
          () -> {
            // W1 will commit the block exclusively before worker 2 registers with the same block
            // So when worker 2 comes in, the block should be committed already
            // So the block on worker 2 should be recognized
            mBlockMaster.workerRegister(worker2, Arrays.asList("MEM"),
                MEM_CAPACITY,
                ImmutableMap.of("MEM", BLOCK1_LENGTH),
                ImmutableMap.of(newBlockLocationOnWorkerMemTier(worker2),
                    ImmutableList.of(BLOCK1_ID)),
                NO_LOST_STORAGE,
                RegisterWorkerPOptions.getDefaultInstance());
            return null;
          },
          // Verifier
          () -> {
            // After registration, verify the worker info
            List<WorkerInfo> workerInfoList = mBlockMaster.getWorkerReport(
                GetWorkerReportOptions.defaults());
            assertEquals(2, workerInfoList.size());
            WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
            assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
            WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
            assertEquals(BLOCK1_LENGTH, worker2Info.getUsedBytes());
  
            verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, workerInfoList);
  
            return null;
          });
    }
  
    @Test
    public void concurrentCommitWithRegisterNewWorkerDifferentBlock() throws Exception {
      // Prepare worker
      long worker1 = registerEmptyWorker(NET_ADDRESS_1);
      long worker2 = mBlockMaster.getWorkerId(NET_ADDRESS_2);
      CountDownLatch w1Latch = new CountDownLatch(1);
      mBlockMaster.setLatch(w1Latch);
  
      concurrentWriterWithWriter(w1Latch,
          // W1
          () -> {
            mBlockMaster.commitBlock(worker1, 49L, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
            return null;
          },
          // W2
          () -> {
            // The new worker contains another block
            // The new block on worker 2 is not recognized and will be ignored by master
            // because the block metadata is not in alluxio
            mBlockMaster.workerRegister(worker2, Arrays.asList("MEM"),
                MEM_CAPACITY,
                ImmutableMap.of("MEM", BLOCK2_LENGTH),
                ImmutableMap.of(newBlockLocationOnWorkerMemTier(worker2),
                    ImmutableList.of(BLOCK2_ID)),
                NO_LOST_STORAGE,
                RegisterWorkerPOptions.getDefaultInstance());
            return null;
          },
          // Verifier
          () -> {
            // After registration, verify the worker info
            List<WorkerInfo> workerInfoList = mBlockMaster
                .getWorkerReport(GetWorkerReportOptions.defaults());
            assertEquals(2, workerInfoList.size());
            WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
            assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
            // Although the new block is not recognized, the worker usage will be taken as-is
            // That new block will be marked orphaned on the worker and will be removed later
            // So that later the worker usage will be rectified
            WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
            assertEquals(BLOCK2_LENGTH, worker2Info.getUsedBytes());
  
            // Verify the block metadata
            verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, Arrays.asList(worker1Info));
            verifyBlockNotExisting(mBlockMaster, BLOCK2_ID);
            return null;
          });
    }
  
    @Test
    public void concurrentCommitWithSameWorkerHeartbeatSameBlock() throws Exception {
      // Prepare worker
      long worker1 = registerEmptyWorker(NET_ADDRESS_1);
      CountDownLatch w1Latch = new CountDownLatch(1);
      mBlockMaster.setLatch(w1Latch);
  
      concurrentWriterWithWriter(w1Latch,
          // W1
          () -> {
            mBlockMaster.commitBlock(worker1, 49L, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
            return null;
          },
          // W2
          () -> {
            // The same block is removed on worker in this heartbeat
            // This should succeed as commit locks the block exclusively and finishes first
            // When the block heartbeat processes the same block, it has been committed
            Command cmd = mBlockMaster.workerHeartbeat(worker1,
                MEM_CAPACITY,
                // 0 used because the block removed on this worker
                MEM_USAGE_EMPTY,
                // list of removed blockIds
                ImmutableList.of(BLOCK1_ID),
                ImmutableMap.of(),
                NO_LOST_STORAGE,
                ImmutableList.of());
  
            // The block has been removed, nothing from command
            assertEquals(EMPTY_CMD, cmd);
  
            return null;
          },
          // Verifier
          () -> {
            // After heartbeat, verify the worker info
            List<WorkerInfo> workerInfoList = mBlockMaster.getWorkerReport(
                GetWorkerReportOptions.defaults());
            assertEquals(1, workerInfoList.size());
            WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
            assertEquals(0L, worker1Info.getUsedBytes());
  
            // The block has no locations now because the last location is removed
            verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, Arrays.asList());
  
            return null;
          });
    }
  
    @Test
    public void concurrentCommitWithSameWorkerHeartbeatDifferentBlock() throws Exception {
      // Prepare worker
      long worker1 = registerEmptyWorker(NET_ADDRESS_1);
      // Block 2 is on worker 1
      mBlockMaster.commitBlock(worker1, BLOCK2_LENGTH, "MEM", "MEM", BLOCK2_ID, BLOCK2_LENGTH);
      CountDownLatch w1Latch = new CountDownLatch(1);
      mBlockMaster.setLatch(w1Latch);
  
      concurrentWriterWithWriter(w1Latch,
          // W1
          () -> {
            mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
            return null;
          },
          // W2
          () -> {
            // A different block is removed on the same worker
            // This should contend on the worker metadata
            Command cmd = mBlockMaster.workerHeartbeat(worker1,
                MEM_CAPACITY,
                // 0 used because the block is already removed
                MEM_USAGE_EMPTY,
                // list of removed blockIds
                ImmutableList.of(BLOCK2_ID),
                ImmutableMap.of(),
                NO_LOST_STORAGE,
                ImmutableList.of());
  
            // The block has been removed, nothing from command
            assertEquals(EMPTY_CMD, cmd);
  
            return null;
          },
          // Verifier
          () -> {
            // After heartbeat, verify the worker info
            List<WorkerInfo> workerInfoList = mBlockMaster.getWorkerReport(
                GetWorkerReportOptions.defaults());
            assertEquals(1, workerInfoList.size());
            WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
            assertEquals(0L, worker1Info.getUsedBytes());
  
            verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, workerInfoList);
            return null;
          });
    }
  
    @Test
    public void concurrentCommitWithDifferentWorkerHeartbeatSameBlock() throws Exception {
      // Prepare worker
      long worker1 = registerEmptyWorker(NET_ADDRESS_1);
      long worker2 = registerEmptyWorker(NET_ADDRESS_2);
      // The block is on worker 2
      mBlockMaster.commitBlock(worker2, BLOCK2_LENGTH, "MEM", "MEM", BLOCK2_ID, BLOCK2_LENGTH);
  
      CountDownLatch w1Latch = new CountDownLatch(1);
      mBlockMaster.setLatch(w1Latch);
  
      concurrentWriterWithWriter(w1Latch,
          // W1
          () -> {
            mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
            return null;
          },
          // W2
          () -> {
            // The same block is removed on another worker
            // This should succeed as commit locks the block exclusively and finishes first
            // When the block heartbeat processes the same block, it has been committed
            Command cmd = mBlockMaster.workerHeartbeat(worker2,
                MEM_CAPACITY,
                // 0 used because the block is already removed
                MEM_USAGE_EMPTY,
                // list of removed blockIds
                ImmutableList.of(BLOCK1_ID),
                ImmutableMap.of(),
                NO_LOST_STORAGE,
                ImmutableList.of());
  
            // The block has been removed, nothing from command
            assertEquals(EMPTY_CMD, cmd);
  
            return null;
          },
          // Verifier
          () -> {
            // After heartbeat, verify the worker info
            List<WorkerInfo> workerInfoList = mBlockMaster
                .getWorkerReport(GetWorkerReportOptions.defaults());
            assertEquals(2, workerInfoList.size());
            WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
            assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
            WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
            assertEquals(0L, worker2Info.getUsedBytes());
  
            // The block has 1 location on worker 1
            verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, Arrays.asList(worker1Info));
  
            return null;
          });
    }
  
    private Block.BlockLocation newBlockLocationOnWorkerMemTier(long workerId) {
      return Block.BlockLocation.newBuilder().setTier("MEM").setMediumType("MEM")
          .setWorkerId(workerId).build();
    }
  
    @Test
    public void concurrentCommitWithDifferentWorkerHeartbeatDifferentBlock() throws Exception {
      // Prepare worker
      long worker1 = registerEmptyWorker(NET_ADDRESS_1);
      long worker2 = registerEmptyWorker(NET_ADDRESS_2);
      // Register with block 2 on both workers
      mBlockMaster.commitBlock(worker1, BLOCK2_LENGTH, "MEM", "MEM", BLOCK2_ID, BLOCK2_LENGTH);
      mBlockMaster.commitBlock(worker2, BLOCK2_LENGTH, "MEM", "MEM", BLOCK2_ID, BLOCK2_LENGTH);
  
      CountDownLatch w1Latch = new CountDownLatch(1);
      mBlockMaster.setLatch(w1Latch);
  
      concurrentWriterWithWriter(w1Latch,
          // W1
          () -> {
            // worker 1 has block 1 and block 2 now
            mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH + BLOCK2_LENGTH,
                "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
            return null;
          },
          // W2
          () -> {
            // A different block is removed on the same worker
            // This should contend on the worker metadata
            Command cmd = mBlockMaster.workerHeartbeat(worker2,
                MEM_CAPACITY,
                // 0 used because the block is already removed
                MEM_USAGE_EMPTY,
                // list of removed blockIds
                ImmutableList.of(BLOCK2_ID),
                ImmutableMap.of(),
                NO_LOST_STORAGE,
                ImmutableList.of());
  
            // The block has been removed, nothing from command
            assertEquals(EMPTY_CMD, cmd);
  
            return null;
          },
          // Verifier
          () -> {
            // After heartbeat, verify the worker info
            List<WorkerInfo> workerInfoList = mBlockMaster
                .getWorkerReport(GetWorkerReportOptions.defaults());
            assertEquals(2, workerInfoList.size());
            WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
            assertEquals(BLOCK1_LENGTH + BLOCK2_LENGTH, worker1Info.getUsedBytes());
            WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
            assertEquals(0L, worker2Info.getUsedBytes());
  
            // Block 1 should exist on master 1
            verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, Arrays.asList(worker1Info));
  
            // Block 2 should exist on master 1
            verifyBlockOnWorkers(mBlockMaster, BLOCK2_ID, BLOCK2_LENGTH, Arrays.asList(worker1Info));
            return null;
          });
    }
  
    @Test
    public void concurrentRemoveWithRegisterNewWorkerSameBlock() throws Exception {
      for (boolean deleteMetadata : ImmutableList.of(true, false)) {
        // Prepare worker
        long worker1 = registerEmptyWorker(NET_ADDRESS_1);
        // Prepare block on the worker
        mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
        CountDownLatch w1Latch = new CountDownLatch(1);
        mBlockMaster.setLatch(w1Latch);
  
        // A new worker as the W2
        long worker2 = mBlockMaster.getWorkerId(NET_ADDRESS_2);
        concurrentWriterWithWriter(w1Latch,
            // W1
            () -> {
              mBlockMaster.removeBlocks(ImmutableList.of(BLOCK1_ID), deleteMetadata);
              return null;
            },
            // W2
            () -> {
              // The new worker contains the block
              // W1 will remove the block exclusively before worker2 registers with the same block
              // So when worker 2 comes in, the block should be removed already
              // So the block on worker 2 should be ignored
              mBlockMaster.workerRegister(worker2, Arrays.asList("MEM"),
                  MEM_CAPACITY,
                  ImmutableMap.of("MEM", BLOCK1_LENGTH),
                  ImmutableMap.of(newBlockLocationOnWorkerMemTier(worker2),
                      ImmutableList.of(BLOCK1_ID)),
                  NO_LOST_STORAGE,
                  RegisterWorkerPOptions.getDefaultInstance());
              return null;
            },
            // Verifier
            () -> {
              // After registration, verify the worker info
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              assertEquals(2, workerInfoList.size());
              WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
              assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
              WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
              assertEquals(BLOCK1_LENGTH, worker2Info.getUsedBytes());
  
              // Verify the block metadata
              if (deleteMetadata) {
                // If the block metadata has been removed, getting that will get an exception
                assertThrows(BlockInfoException.class, () -> {
                  mBlockMaster.getBlockInfo(BLOCK1_ID);
                });
              } else {
                // The master will issue commands to remove blocks on the next heartbeat
                // So now the locations are still there
                verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, workerInfoList);
              }
  
              // Verify the heartbeat from worker will get a command to remove the block
              Command worker1HeartbeatCmd = mBlockMaster.workerHeartbeat(worker1,
                  MEM_CAPACITY,
                  // the block has not yet been removed
                  ImmutableMap.of("MEM", BLOCK1_LENGTH),
                  // an empty list of removed blockIds
                  ImmutableList.of(),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
              assertEquals(FREE_BLOCK1_CMD, worker1HeartbeatCmd);
  
              if (deleteMetadata) {
                // Block on worker 2 will be freed because the block is already removed
                Command worker2HeartbeatCmd = mBlockMaster.workerHeartbeat(worker2,
                    MEM_CAPACITY,
                    // the block has not yet been removed
                    ImmutableMap.of("MEM", BLOCK1_LENGTH),
                    // an empty list of removed blockIds
                    ImmutableList.of(),
                    ImmutableMap.of(),
                    NO_LOST_STORAGE,
                    ImmutableList.of());
                // Block on worker 2 will be freed because the block is already removed
                // Unrecognized blocks will be freed
                assertEquals(FREE_BLOCK1_CMD, worker2HeartbeatCmd);
              } else {
                // Two cases can happen:
                // 1. Worker 2 registers before the free operation checks the block locations
                //    In this case the block on worker 2 will be freed
                // 2. Worker 2 registers after the free operation is complete
                //    In this case the block on worker 2 will not be freed
                Command worker2HeartbeatCmd = mBlockMaster.workerHeartbeat(worker2,
                    MEM_CAPACITY,
                    // the block has not yet been removed
                    ImmutableMap.of("MEM", BLOCK1_LENGTH),
                    // an empty list of removed blockIds
                    ImmutableList.of(),
                    ImmutableMap.of(),
                    NO_LOST_STORAGE,
                    ImmutableList.of());
                assertTrue(worker2HeartbeatCmd.equals(FREE_BLOCK1_CMD)
                    || worker2HeartbeatCmd.equals(EMPTY_CMD));
              }
  
              return null;
            });
      }
    }
  
    @Test
    public void concurrentRemoveWithRegisterNewWorkerDifferentBlock() throws Exception {
      for (boolean deleteMetadata : ImmutableList.of(true, false)) {
        // Prepare worker
        long worker1 = registerEmptyWorker(NET_ADDRESS_1);
        // Prepare block on the worker
        mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
        // Prepare block 2 so it is recognized at worker register
        mBlockMaster.commitBlockInUFS(BLOCK2_ID, BLOCK2_LENGTH);
        CountDownLatch w1Latch = new CountDownLatch(1);
        mBlockMaster.setLatch(w1Latch);
  
        // A new worker as the W2
        long worker2 = mBlockMaster.getWorkerId(NET_ADDRESS_2);
        concurrentWriterWithWriter(w1Latch,
            // W1
            () -> {
              mBlockMaster.removeBlocks(ImmutableList.of(BLOCK1_ID), deleteMetadata);
              return null;
            },
            // W2
            () -> {
              // The new worker contains the block
              // W1 will remove the block exclusively before worker2 registers with the same block
              // So when worker 2 comes in, the block should be removed already
              // So the block on worker 2 should be ignored
              mBlockMaster.workerRegister(worker2, Arrays.asList("MEM"),
                  MEM_CAPACITY,
                  ImmutableMap.of("MEM", BLOCK2_LENGTH),
                  ImmutableMap.of(newBlockLocationOnWorkerMemTier(worker2),
                      ImmutableList.of(BLOCK2_ID)),
                  NO_LOST_STORAGE,
                  RegisterWorkerPOptions.getDefaultInstance());
              return null;
            },
            // Verifier
            () -> {
              // After registration, verify the worker info
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              assertEquals(2, workerInfoList.size());
              // Block 1 has not been removed from the workers yet
              WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
              assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
              WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
              assertEquals(BLOCK2_LENGTH, worker2Info.getUsedBytes());
  
              // Verify the block metadata
              if (deleteMetadata) {
                verifyBlockNotExisting(mBlockMaster, BLOCK1_ID);
              } else {
                // The master will issue commands to remove blocks on the next heartbeat
                // So now the locations are still there
                verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH,
                    ImmutableList.of(worker1Info));
              }
              // Block 2 is unaffected
              verifyBlockOnWorkers(mBlockMaster, BLOCK2_ID, BLOCK2_LENGTH,
                  ImmutableList.of(worker2Info));
  
              // Regardless of whether the metadata is removed, the existing block will be freed
              Command worker1HeartbeatCmd = mBlockMaster.workerHeartbeat(worker1,
                  MEM_CAPACITY,
                  // the block has not yet been removed
                  ImmutableMap.of("MEM", BLOCK1_LENGTH),
                  // an empty list of removed blockIds
                  ImmutableList.of(),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
              assertEquals(FREE_BLOCK1_CMD, worker1HeartbeatCmd);
  
              Command worker2HeartbeatCmd = mBlockMaster.workerHeartbeat(worker2,
                  MEM_CAPACITY,
                  // the block has not yet been removed
                  ImmutableMap.of("MEM", BLOCK1_LENGTH),
                  // an empty list of removed blockIds
                  ImmutableList.of(),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
              // Blocks on worker 2 are unaffected
              assertEquals(EMPTY_CMD, worker2HeartbeatCmd);
              return null;
            });
      }
    }
  
    @Test
    public void concurrentRemoveWithSameWorkerHeartbeatSameBlock() throws Exception {
      for (boolean deleteMetadata : ImmutableList.of(true, false)) {
        // Prepare worker
        long worker1 = registerEmptyWorker(NET_ADDRESS_1);
        // Prepare block in alluxio
        mBlockMaster.commitBlockInUFS(BLOCK1_ID, BLOCK1_LENGTH);
        CountDownLatch w1Latch = new CountDownLatch(1);
        mBlockMaster.setLatch(w1Latch);
  
        concurrentWriterWithWriter(w1Latch,
            // W1
            () -> {
              mBlockMaster.removeBlocks(ImmutableList.of(BLOCK1_ID), deleteMetadata);
              return null;
            },
            // W2
            () -> {
              // The same block is removed on worker in this heartbeat
              // This should succeed as commit locks the block exclusively and finishes first
              // When the block heartbeat processes the same block, it has been committed
              Command cmd = mBlockMaster.workerHeartbeat(worker1,
                  MEM_CAPACITY,
                  // 0 used because the block is already removed
                  MEM_USAGE_EMPTY,
                  // list of removed blockIds
                  ImmutableList.of(BLOCK1_ID),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
  
              // The block has been removed, nothing from command
              assertEquals(EMPTY_CMD, cmd);
  
              return null;
            },
            // Verifier
            () -> {
              // After heartbeat, verify the worker info
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              assertEquals(1, workerInfoList.size());
              WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
              assertEquals(0L, worker1Info.getUsedBytes());
  
              if (deleteMetadata) {
                verifyBlockNotExisting(mBlockMaster, BLOCK1_ID);
              } else {
                // The block has no locations now because the last location is removed
                verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, Arrays.asList());
              }
  
              return null;
            });
      }
    }
  
    private long registerEmptyWorker(WorkerNetAddress address) throws Exception {
      long workerId = mBlockMaster.getWorkerId(address);
      mBlockMaster.workerRegister(workerId, Arrays.asList("MEM"), MEM_CAPACITY,
          MEM_USAGE_EMPTY, NO_BLOCKS_ON_LOCATION, NO_LOST_STORAGE,
          RegisterWorkerPOptions.getDefaultInstance());
      return workerId;
    }
  
    @Test
    public void concurrentRemoveWithSameWorkerHeartbeatDifferentBlock() throws Exception {
      for (boolean deleteMetadata : ImmutableList.of(true)) {
        // Prepare block 1 and 2 on the worker
        long worker1 = registerEmptyWorker(NET_ADDRESS_1);
        mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
        mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH + BLOCK2_LENGTH, "MEM", "MEM",
            BLOCK2_ID, BLOCK2_LENGTH);
        CountDownLatch w1Latch = new CountDownLatch(1);
        mBlockMaster.setLatch(w1Latch);
  
        AtomicBoolean freeCommandSeen = new AtomicBoolean(false);
        concurrentWriterWithWriter(w1Latch,
            // W1
            () -> {
              mBlockMaster.removeBlocks(ImmutableList.of(BLOCK1_ID), deleteMetadata);
              return null;
            },
            // W2
            () -> {
              // A different block is removed on the same worker
              // This should contend on the worker metadata
              Command cmd = mBlockMaster.workerHeartbeat(worker1,
                  MEM_CAPACITY,
                  // Block 2 is removed but 1 is still on the worker
                  ImmutableMap.of("MEM", BLOCK1_LENGTH),
                  // list of removed blockIds
                  ImmutableList.of(BLOCK2_ID),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
  
              // The heartbeat contends on the block lock of block 2, worker usage lock and
              // worker block list lock
              // The remove operation will first remove the block metadata with the block lock,
              // then update the worker to-be-removed list with the block list lock
              // There are two possible outcomes:
              // 1. Remove gets the block list lock first and updates the to-be-removed list
              //    In this case the returned value will be a free command.
              // 2. Worker heartbeat gets the block list lock first before the remove operation
              //    adds to the to-be-removed list. In this case the return command has nothing.
              if (cmd.equals(FREE_BLOCK1_CMD)) {
                freeCommandSeen.set(true);
              } else {
                assertEquals(EMPTY_CMD, cmd);
              }
  
              return null;
            },
            // Verifier
            () -> {
              // After heartbeat, verify the worker info
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              assertEquals(1, workerInfoList.size());
              WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
              assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
  
              if (deleteMetadata) {
                verifyBlockNotExisting(mBlockMaster, BLOCK1_ID);
              } else {
                // All locations of block 1 are freed in metadata
                verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH, workerInfoList);
              }
              verifyBlockOnWorkers(mBlockMaster, BLOCK2_ID, BLOCK2_LENGTH, ImmutableList.of());
  
              // If the 1st heartbeat does not see the free command
              // This heartbeat should definitely see it,
              // because the verifier is run after W1 fully finished
              // and updated the to-be-removed list
              if (!freeCommandSeen.get()) {
                Command cmd = mBlockMaster.workerHeartbeat(worker1,
                    MEM_CAPACITY,
                    // Block 2 is removed but 1 is still on the worker
                    ImmutableMap.of("MEM", BLOCK1_LENGTH),
                    // list of removed blockIds
                    ImmutableList.of(BLOCK2_ID),
                    ImmutableMap.of(),
                    NO_LOST_STORAGE,
                    ImmutableList.of());
                assertEquals(FREE_BLOCK1_CMD, cmd);
              }
              return null;
            });
      }
    }
  
    @Test
    public void concurrentRemoveWithDifferentWorkerHeartbeatSameBlock() throws Exception {
      for (boolean deleteMetadata : ImmutableList.of(true, false)) {
        // Prepare worker
        long worker1 = registerEmptyWorker(NET_ADDRESS_1);
        long worker2 = registerEmptyWorker(NET_ADDRESS_2);
        mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
        mBlockMaster.commitBlock(worker2, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
        CountDownLatch w1Latch = new CountDownLatch(1);
        mBlockMaster.setLatch(w1Latch);
  
        concurrentWriterWithWriter(w1Latch,
            // W1
            () -> {
              mBlockMaster.removeBlocks(ImmutableList.of(BLOCK1_ID), deleteMetadata);
              return null;
            },
            // W2
            () -> {
              // The same block is removed on another worker
              // This should succeed as commit locks the block exclusively and finishes first
              // When the block heartbeat processes the same block, it has been committed
              Command cmd = mBlockMaster.workerHeartbeat(worker2,
                  MEM_CAPACITY,
                  // 0 used because the block is already removed
                  MEM_USAGE_EMPTY,
                  // list of removed blockIds
                  ImmutableList.of(BLOCK1_ID),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
  
              // The block has been removed, nothing from command
              assertEquals(EMPTY_CMD, cmd);
  
              return null;
            },
            // Verifier
            () -> {
              // After heartbeat, verify the worker info
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              assertEquals(2, workerInfoList.size());
              // The block is still on worker 1, will be removed on the next heartbeat
              WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
              assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
              WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
              assertEquals(0L, worker2Info.getUsedBytes());
  
              if (deleteMetadata) {
                verifyBlockNotExisting(mBlockMaster, BLOCK1_ID);
              } else {
                // The location is still on worker 1, until it is removed after the next heartbeat
                verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH,
                    ImmutableList.of(worker1Info));
              }
  
              // On the heartbeat worker 1 block will be removed
              Command cmd = mBlockMaster.workerHeartbeat(worker1,
                  MEM_CAPACITY,
                  // Block 1 is still on worker 1
                  ImmutableMap.of("MEM", BLOCK1_LENGTH),
                  // list of removed blockIds
                  ImmutableList.of(),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
              assertEquals(FREE_BLOCK1_CMD, cmd);
              return null;
            });
      }
    }
  
    @Test
    public void concurrentRemoveWithDifferentWorkerHeartbeatDifferentBlock() throws Exception {
      for (boolean deleteMetadata : ImmutableList.of(true, false)) {
        // Prepare worker
        long worker1 = registerEmptyWorker(NET_ADDRESS_1);
        long worker2 = registerEmptyWorker(NET_ADDRESS_2);
        // Worker 1 has block 1
        mBlockMaster.commitBlock(worker1, BLOCK1_LENGTH, "MEM", "MEM", BLOCK1_ID, BLOCK1_LENGTH);
        // Worker 2 has block 2
        mBlockMaster.commitBlock(worker2, BLOCK2_LENGTH, "MEM", "MEM", BLOCK2_ID, BLOCK2_LENGTH);
  
        CountDownLatch w1Latch = new CountDownLatch(1);
        mBlockMaster.setLatch(w1Latch);
  
        concurrentWriterWithWriter(w1Latch,
            // W1
            () -> {
              mBlockMaster.removeBlocks(ImmutableList.of(BLOCK1_ID), deleteMetadata);
              return null;
            },
            // W2
            () -> {
              // A different block is removed on another worker
              Command cmd = mBlockMaster.workerHeartbeat(worker2,
                  MEM_CAPACITY,
                  // 0 used because the block is already removed
                  MEM_USAGE_EMPTY,
                  // list of removed blockIds
                  ImmutableList.of(BLOCK2_ID),
                  ImmutableMap.of(),
                  NO_LOST_STORAGE,
                  ImmutableList.of());
  
              // Nothing for worker 2 to do because it does not have block 1
              assertEquals(EMPTY_CMD, cmd);
  
              return null;
            },
            // Verifier
            () -> {
              // After heartbeat, verify the worker info
              List<WorkerInfo> workerInfoList = mBlockMaster
                  .getWorkerReport(GetWorkerReportOptions.defaults());
              assertEquals(2, workerInfoList.size());
              WorkerInfo worker1Info = findWorkerInfo(workerInfoList, worker1);
              assertEquals(BLOCK1_LENGTH, worker1Info.getUsedBytes());
              WorkerInfo worker2Info = findWorkerInfo(workerInfoList, worker2);
              assertEquals(0L, worker2Info.getUsedBytes());
  
              if (deleteMetadata) {
                verifyBlockNotExisting(mBlockMaster, BLOCK1_ID);
              } else {
                // Block 1 should still exist on worker 1 until the next heartbeat frees it
                verifyBlockOnWorkers(mBlockMaster, BLOCK1_ID, BLOCK1_LENGTH,
                    Arrays.asList(worker1Info));
              }
  
              // No copies for block 2
              verifyBlockOnWorkers(mBlockMaster, BLOCK2_ID, BLOCK2_LENGTH, Arrays.asList());
              return null;
            });
      }
    }
  
    /**
     * Concurrently runs two writers W1 and W2, and verifies the state after W1 and W2 are complete.
     * W1 runs first, in the middle of execution W1 will release the {@link CountDownLatch}.
     * W2 will be waiting for the latch and then runs the writer operation.
     * The verifier will run last verifying the final state.
     *
     * @param w1Latch W1 will trigger this latch, W2 will wait on this latch
     * @param w1 writer 1
     * @param w2 writer 2
     * @param verifier the verifier of the final state
     */
    private void concurrentWriterWithWriter(
        CountDownLatch w1Latch, Callable w1, Callable w2, Callable verifier) throws Exception {
      // This thread count is intentionally larger than the client thread pool
      // In the hope that even if the first batch of clients all read the state before
      // the commit really happens
      // The following batch will capture the state after the commit
      Queue<Throwable> uncaughtThrowables = new ConcurrentLinkedQueue<>();
      CountDownLatch writerFinished = new CountDownLatch(1);
  
      mClientExecutorService.submit(() -> {
        // Wait until the writer enters the critical section and sends a signal
        try {
          w1Latch.await();
        } catch (Throwable t) {
          uncaughtThrowables.add(t);
          // Fail to wait for the signal, just give up
          writerFinished.countDown();
          return;
        }
        // Trigger the other writer
        try {
          w2.call();
        } catch (Throwable t) {
          uncaughtThrowables.add(t);
        } finally {
          writerFinished.countDown();
        }
      });
  
      // Call W1 in the same thread
      w1.call();
  
      writerFinished.await();
      verifier.call();
      // W2 has finished, verify here
      // If any assertion failed, the failed assertion will throw an AssertError
      assertEquals(0, uncaughtThrowables.size());
    }
  }
  