alluxio-master-2.x/core/server/master/src/test/java/alluxio/master/file/replication/ReplicationCheckerTest.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.master.file.replication;

import static org.mockito.Mockito.mock;

import alluxio.AlluxioURI;
import alluxio.Constants;
import alluxio.collections.Pair;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.CreateFilePOptions;
import alluxio.grpc.RegisterWorkerPOptions;
import alluxio.grpc.StorageList;
import alluxio.job.plan.replicate.ReplicationHandler;
import alluxio.job.wire.Status;
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.BlockMasterFactory;
import alluxio.master.file.RpcContext;
import alluxio.master.file.contexts.CreateFileContext;
import alluxio.master.file.contexts.CreatePathContext;
import alluxio.master.file.meta.Inode;
import alluxio.master.file.meta.InodeDirectoryIdGenerator;
import alluxio.master.file.meta.InodeLockManager;
import alluxio.master.file.meta.InodeTree;
import alluxio.master.file.meta.InodeTree.LockPattern;
import alluxio.master.file.meta.LockedInodePath;
import alluxio.master.file.meta.MountTable;
import alluxio.master.file.meta.MutableInodeFile;
import alluxio.master.file.meta.options.MountInfo;
import alluxio.master.journal.JournalSystem;
import alluxio.master.journal.JournalTestUtils;
import alluxio.master.journal.JournalType;
import alluxio.master.journal.NoopJournalContext;
import alluxio.master.metastore.InodeStore;
import alluxio.master.metrics.MetricsMasterFactory;
import alluxio.metrics.Metric;
import alluxio.proto.meta.Block;
import alluxio.security.authorization.Mode;
import alluxio.underfs.UfsManager;
import alluxio.wire.WorkerNetAddress;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.IOException;
import java.time.Clock;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.concurrent.ThreadSafe;

/**
 * Unit tests for {@link ReplicationChecker}.
 */
public final class ReplicationCheckerTest {
  private static final String TEST_OWNER = "user1";
  private static final String TEST_GROUP = "";
  private static final Mode TEST_MODE = new Mode((short) 0755);
  private static final AlluxioURI TEST_FILE_1 = new AlluxioURI("/test1");
  private static final AlluxioURI TEST_FILE_2 = new AlluxioURI("/test2");
  private static final AlluxioURI TEST_FILE_3 = new AlluxioURI("/test3");
  private static final List<Long> NO_BLOCKS = ImmutableList.of();
  private static final List<Metric> NO_METRICS = ImmutableList.of();
  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 Map EMPTY = ImmutableMap.of();

  /**
   * A mock class of AdjustReplicationHandler, used to test the output of ReplicationChecker.
   */
  @ThreadSafe
  private static class MockHandler implements ReplicationHandler {
    private final Map<Long, Integer> mSetReplicaRequests = Maps.newHashMap();
    private final Map<Long, Pair<String, String>>
        mMigrateRequests = Maps.newHashMap();
    private final List<Long> mJobStatusRequests = Lists.newArrayList();
    private final Map<Long, Status> mJobStatus = Maps.newHashMap();

    private long mNextJobId = 0;

    public void setJobStatus(long jobId, Status status) {
      mJobStatus.put(jobId, status);
    }

    @Override
    public Status getJobStatus(long jobId) throws IOException {
      mJobStatusRequests.add(jobId);
      if (mJobStatus.containsKey(jobId)) {
        return mJobStatus.get(jobId);
      }
      return Status.RUNNING;
    }

    @Override
    public List<Long> findJobs(String jobName, Set<Status> statusList) throws IOException {
      return mJobStatus.entrySet().stream()
          .filter(x -> statusList.isEmpty() || statusList.contains(x.getValue()))
          .map(x -> x.getKey()).collect(Collectors.toList());
    }

    @Override
    public long setReplica(AlluxioURI uri, long blockId, int numReplicas) {
      mSetReplicaRequests.put(blockId, numReplicas);
      return ++mNextJobId;
    }

    @Override
    public long migrate(AlluxioURI uri, long blockId, String workerHost, String mediumType) {
      mMigrateRequests.put(blockId, new Pair<>(workerHost, mediumType));
      return ++mNextJobId;
    }

    public Map<Long, Integer> getSetReplicaRequests() {
      return mSetReplicaRequests;
    }

    public Map<Long, Pair<String, String>> getMigrateRequests() {
      return mMigrateRequests;
    }
  }

  private CoreMasterContext mContext;
  private InodeStore mInodeStore;
  private InodeTree mInodeTree;
  private BlockMaster mBlockMaster;
  private ReplicationChecker mReplicationChecker;
  private MockHandler mMockReplicationHandler;
  private CreateFileContext mFileContext =
      CreateFileContext.mergeFrom(CreateFilePOptions.newBuilder().setBlockSizeBytes(Constants.KB)
          .setMode(TEST_MODE.toProto())).setOwner(TEST_OWNER).setGroup(TEST_GROUP);
  private Set<Long> mKnownWorkers = Sets.newHashSet();

  /** Rule to create a new temporary folder during each test. */
  @Rule
  public TemporaryFolder mTestFolder = new TemporaryFolder();

  @Before
  public void before() throws Exception {
    Configuration.set(PropertyKey.MASTER_JOURNAL_TYPE, JournalType.UFS);
    MasterRegistry registry = new MasterRegistry();
    JournalSystem journalSystem = JournalTestUtils.createJournalSystem(mTestFolder);
    mContext = MasterTestUtils.testMasterContext(journalSystem,
        null, new AlwaysPrimaryPrimarySelector());
    new MetricsMasterFactory().create(registry, mContext);
    mBlockMaster = new BlockMasterFactory().create(registry, mContext);
    InodeDirectoryIdGenerator directoryIdGenerator = new InodeDirectoryIdGenerator(mBlockMaster);
    UfsManager manager = mock(UfsManager.class);
    MountTable mountTable = new MountTable(manager, mock(MountInfo.class), Clock.systemUTC());
    InodeLockManager lockManager = new InodeLockManager();
    mInodeStore = mContext.getInodeStoreFactory().apply(lockManager);
    mInodeTree =
        new InodeTree(mInodeStore, mBlockMaster, directoryIdGenerator, mountTable, lockManager);

    journalSystem.start();
    journalSystem.gainPrimacy();
    mBlockMaster.start(true);

    Configuration.set(PropertyKey.TEST_MODE, true);
    Configuration.set(PropertyKey.SECURITY_AUTHORIZATION_PERMISSION_ENABLED, true);
    Configuration
        .set(PropertyKey.SECURITY_AUTHORIZATION_PERMISSION_SUPERGROUP, "test-supergroup");
    mInodeTree.initializeRoot(TEST_OWNER, TEST_GROUP, TEST_MODE, NoopJournalContext.INSTANCE);

    mMockReplicationHandler = new MockHandler();
    mReplicationChecker = new ReplicationChecker(mInodeTree, mBlockMaster,
        mContext.getSafeModeManager(), mMockReplicationHandler);
  }

  @After
  public void after() {
    Configuration.reloadProperties();
  }

  /**
   * Helper to create a file with a single block.
   *
   * @param path Alluxio path of the file
   * @param context context to create the file
   * @param pinLocation
   * @return the block ID
   */
  private long createBlockHelper(AlluxioURI path, CreatePathContext<?, ?> context,
      String pinLocation) throws Exception {
    try (LockedInodePath inodePath = mInodeTree.lockInodePath(
        path, LockPattern.WRITE_EDGE, NoopJournalContext.INSTANCE)
    ) {
      List<Inode> created = mInodeTree.createPath(RpcContext.NOOP, inodePath, context);
      if (!pinLocation.equals("")) {
        mInodeTree.setPinned(RpcContext.NOOP, inodePath, true, ImmutableList.of(pinLocation), 0);
      }
      MutableInodeFile inodeFile = mInodeStore.getMutable(created.get(0).getId()).get().asFile();
      inodeFile.setBlockSizeBytes(1);
      inodeFile.setBlockIds(Arrays.asList(inodeFile.getNewBlockId()));
      inodeFile.setCompleted(true);
      mInodeStore.writeInode(inodeFile);
      return inodeFile.getBlockIdByIndex(0);
    }
  }

  /**
   * Helper to create and add a given number of locations for block ID.
   *
   * @param blockId ID of the block to add location
   * @param numLocations number of locations to add
   */
  private void addBlockLocationHelper(long blockId, int numLocations) throws Exception {
    // Commit blockId to the first worker.
    mBlockMaster.commitBlock(createWorkerHelper(0), 50L,
        Constants.MEDIUM_MEM, Constants.MEDIUM_MEM, blockId, 20L);

    // Send a heartbeat from other workers saying that it's added blockId.
    for (int i = 1; i < numLocations; i++) {
      heartbeatToAddLocationHelper(blockId, createWorkerHelper(i));
    }
  }

  /**
   * Helper to register a new worker.
   *
   * @param workerIndex the index of the worker in all workers
   * @return the created worker ID
   */
  private long createWorkerHelper(int workerIndex) throws Exception {
    WorkerNetAddress address = new WorkerNetAddress().setHost("host" + workerIndex).setRpcPort(1000)
        .setDataPort(2000).setWebPort(3000);
    long workerId = mBlockMaster.getWorkerId(address);
    if (!mKnownWorkers.contains(workerId)) {
      // Do not re-register works, otherwise added block will be removed
      mBlockMaster.workerRegister(workerId, ImmutableList.of(Constants.MEDIUM_MEM),
          ImmutableMap.of(Constants.MEDIUM_MEM, 100L),
          ImmutableMap.of(Constants.MEDIUM_MEM, 0L), NO_BLOCKS_ON_LOCATION, NO_LOST_STORAGE,
          RegisterWorkerPOptions.getDefaultInstance());
      mKnownWorkers.add(workerId);
    }
    return workerId;
  }

  /**
   * Helper to heartbeat to a worker and report a newly added block.
   *
   * @param blockId ID of the block to add location
   * @param workerId ID of the worker to heartbeat
   */
  private void heartbeatToAddLocationHelper(long blockId, long workerId) throws Exception {
    List<Long> addedBlocks = ImmutableList.of(blockId);
    Block.BlockLocation blockLocation =
        Block.BlockLocation.newBuilder().setWorkerId(workerId)
            .setTier(Constants.MEDIUM_MEM).setMediumType(Constants.MEDIUM_MEM).build();

    mBlockMaster.workerHeartbeat(workerId, null,
        ImmutableMap.of(Constants.MEDIUM_MEM, 0L), NO_BLOCKS,
        ImmutableMap.of(blockLocation, addedBlocks), NO_LOST_STORAGE, NO_METRICS);
  }

  @Test
  public void heartbeatWhenTreeIsEmpty() throws Exception {
    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatFileWithinRange() throws Exception {
    mFileContext.getOptions().setReplicationMin(1).setReplicationMax(3);
    long blockId =
        createBlockHelper(TEST_FILE_1, mFileContext, "");
    // One replica, meeting replication min
    addBlockLocationHelper(blockId, 1);
    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getSetReplicaRequests());

    // Two replicas, good
    heartbeatToAddLocationHelper(blockId, createWorkerHelper(1));
    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getSetReplicaRequests());

    // Three replicas, meeting replication max, still good
    heartbeatToAddLocationHelper(blockId, createWorkerHelper(2));
    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatFileUnderReplicatedBy1() throws Exception {
    mFileContext.getOptions().setReplicationMin(1);
    long blockId = createBlockHelper(TEST_FILE_1, mFileContext, "");

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Integer> expected = ImmutableMap.of(blockId, 1);
    Assert.assertEquals(expected, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatFileNeedsMove() throws Exception {
    mFileContext.getOptions().setReplicationMin(1);
    long blockId = createBlockHelper(TEST_FILE_1, mFileContext, Constants.MEDIUM_SSD);
    addBlockLocationHelper(blockId, 1);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Pair<String, String>> expected =
        ImmutableMap.of(blockId, new Pair<>("host0", Constants.MEDIUM_SSD));
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getSetReplicaRequests());
    Assert.assertEquals(expected, mMockReplicationHandler.getMigrateRequests());
  }

  @Test
  public void heartbeatFileDoesnotNeedMove() throws Exception {
    mFileContext.getOptions().setReplicationMin(1);
    long blockId = createBlockHelper(TEST_FILE_1, mFileContext, Constants.MEDIUM_MEM);
    addBlockLocationHelper(blockId, 1);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getSetReplicaRequests());
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getMigrateRequests());
  }

  @Test
  public void heartbeatFileUnderReplicatedBy10() throws Exception {
    mFileContext.getOptions().setReplicationMin(10);
    long blockId = createBlockHelper(TEST_FILE_1, mFileContext, "");

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Integer> expected = ImmutableMap.of(blockId, 10);
    Assert.assertEquals(expected, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatMultipleFilesUnderReplicated() throws Exception {
    mFileContext.getOptions().setReplicationMin(1);
    long blockId1 = createBlockHelper(TEST_FILE_1, mFileContext, "");
    mFileContext.getOptions().setReplicationMin(2);
    long blockId2 = createBlockHelper(TEST_FILE_2, mFileContext, "");

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Integer> expected = ImmutableMap.of(blockId1, 1, blockId2, 2);
    Assert.assertEquals(expected, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatFileUnderReplicatedAndLost() throws Exception {
    mFileContext.getOptions().setReplicationMin(2);
    long blockId = createBlockHelper(TEST_FILE_1, mFileContext, "");

    // Create a worker.
    long workerId = mBlockMaster.getWorkerId(new WorkerNetAddress().setHost("localhost")
        .setRpcPort(80).setDataPort(81).setWebPort(82));
    mBlockMaster.workerRegister(workerId, Collections.singletonList(Constants.MEDIUM_MEM),
        ImmutableMap.of(Constants.MEDIUM_MEM, 100L),
        ImmutableMap.of(Constants.MEDIUM_MEM, 0L), NO_BLOCKS_ON_LOCATION,
        NO_LOST_STORAGE, RegisterWorkerPOptions.getDefaultInstance());
    mBlockMaster.commitBlock(workerId, 50L,
        Constants.MEDIUM_MEM, Constants.MEDIUM_MEM, blockId, 20L);

    // Indicate that blockId is removed on the worker.
    mBlockMaster.workerHeartbeat(workerId, null,
        ImmutableMap.of(Constants.MEDIUM_MEM, 0L),
        ImmutableList.of(blockId), NO_BLOCKS_ON_LOCATION, NO_LOST_STORAGE, NO_METRICS);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(EMPTY, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatFileOverReplicatedBy1() throws Exception {
    mFileContext.getOptions().setReplicationMax(1);
    long blockId = createBlockHelper(TEST_FILE_1, mFileContext, "");
    addBlockLocationHelper(blockId, 2);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Integer> expected = ImmutableMap.of(blockId, 1);
    Assert.assertEquals(expected, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatFileOverReplicatedBy10() throws Exception {
    mFileContext.getOptions().setReplicationMax(1);
    long blockId = createBlockHelper(TEST_FILE_1, mFileContext, "");
    addBlockLocationHelper(blockId, 11);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Integer> expected = ImmutableMap.of(blockId, 1);
    Assert.assertEquals(expected, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatMultipleFilesOverReplicated() throws Exception {
    mFileContext.getOptions().setReplicationMax(1);
    long blockId1 = createBlockHelper(TEST_FILE_1, mFileContext, "");
    mFileContext.getOptions().setReplicationMax(2);
    long blockId2 = createBlockHelper(TEST_FILE_2, mFileContext, "");
    addBlockLocationHelper(blockId1, 2);
    addBlockLocationHelper(blockId2, 4);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Integer> expected = ImmutableMap.of(blockId1, 1, blockId2, 2);
    Assert.assertEquals(expected, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatFilesUnderAndOverReplicated() throws Exception {
    mFileContext.getOptions().setReplicationMin(2).setReplicationMax(-1);
    long blockId1 = createBlockHelper(TEST_FILE_1, mFileContext, "");
    mFileContext.getOptions().setReplicationMin(0).setReplicationMax(3);
    long blockId2 = createBlockHelper(TEST_FILE_2, mFileContext, "");
    addBlockLocationHelper(blockId1, 1);
    addBlockLocationHelper(blockId2, 5);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Map<Long, Integer> expected1 = ImmutableMap.of(blockId1, 2, blockId2, 3);
    Assert.assertEquals(expected1, mMockReplicationHandler.getSetReplicaRequests());
  }

  @Test
  public void heartbeatPartial() throws Exception {
    Configuration.set(PropertyKey.JOB_MASTER_JOB_CAPACITY, 20);
    mReplicationChecker = new ReplicationChecker(mInodeTree, mBlockMaster,
        mContext.getSafeModeManager(), mMockReplicationHandler);
    mFileContext.getOptions().setReplicationMin(3).setReplicationMax(-1);
    long blockId1 = createBlockHelper(TEST_FILE_1, mFileContext, "");
    long blockId2 = createBlockHelper(TEST_FILE_2, mFileContext, "");
    long blockId3 = createBlockHelper(TEST_FILE_3, mFileContext, "");
    addBlockLocationHelper(blockId1, 1);
    addBlockLocationHelper(blockId2, 1);
    addBlockLocationHelper(blockId3, 1);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    final Map<Long, Integer> replicateRequests = mMockReplicationHandler.getSetReplicaRequests();
    System.out.println(replicateRequests);
    Assert.assertEquals(2, replicateRequests.size());
    Assert.assertEquals(3, replicateRequests.values().toArray()[0]);
    Assert.assertEquals(3, replicateRequests.values().toArray()[1]);
    replicateRequests.clear();

    mMockReplicationHandler.setJobStatus(1, Status.RUNNING);
    mMockReplicationHandler.setJobStatus(2, Status.RUNNING);
    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(0, replicateRequests.size());

    mMockReplicationHandler.setJobStatus(1, Status.FAILED);
    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(1, replicateRequests.size());
    Assert.assertEquals(3, replicateRequests.values().toArray()[0]);

    replicateRequests.clear();
    addBlockLocationHelper(blockId1, 3);
    addBlockLocationHelper(blockId2, 3);
    mMockReplicationHandler.setJobStatus(2, Status.COMPLETED);
    mMockReplicationHandler.setJobStatus(3, Status.COMPLETED);

    mReplicationChecker.heartbeat(Long.MAX_VALUE);
    Assert.assertEquals(1, replicateRequests.size());
    Assert.assertTrue(replicateRequests.containsKey(blockId3));
    Assert.assertEquals(3, replicateRequests.values().toArray()[0]);
  }
}
