alluxio-master-2.x/core/server/worker/src/test/java/alluxio/worker/block/management/tier/BlockTransferPartitionerTest.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.worker.block.management.tier;

import alluxio.worker.block.BlockStoreLocation;
import alluxio.worker.block.evictor.BlockTransferInfo;
import alluxio.worker.block.management.BlockTransferPartitioner;

import org.junit.Assert;
import org.junit.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

public class BlockTransferPartitionerTest {
  private BlockTransferPartitioner mPartitioner = new BlockTransferPartitioner();

  @Test
  public void testEmptyList() throws Exception {
    validatePartitions(mPartitioner.partitionTransfers(Collections.emptyList(), 1), 1, 0);
    validatePartitions(mPartitioner.partitionTransfers(Collections.emptyList(), 2), 1, 0);
  }

  @Test
  public void testPartitioning() throws Exception {
    // Test partitioning for when there are 2 distinct locations with "2 transfers each".
    for (List<BlockTransferInfo> transfers : generateTransferLists(4, new int[] {2, 2})) {
      validatePartitions(mPartitioner.partitionTransfers(transfers, 1), 1, 4);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 2), 2, 2, 2);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 3), 2, 2, 2);
    }

    // Test partitioning for when there are 2 distinct locations with "3 and 1" transfers.
    for (List<BlockTransferInfo> transfers : generateTransferLists(4, new int[] {3, 1})) {
      validatePartitions(mPartitioner.partitionTransfers(transfers, 1), 1, 4);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 2), 2, 3, 1);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 3), 2, 3, 1);
    }

    // Test partitioning for when there are 3 distinct locations with "1, 1 and 2" transfers.
    for (List<BlockTransferInfo> transfers : generateTransferLists(4, new int[] {1, 1, 2})) {
      validatePartitions(mPartitioner.partitionTransfers(transfers, 1), 1, 4);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 2), 2, 3, 1);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 3), 3, 1, 1, 2);
    }

    // Test partitioning for when there are 4 distinct locations with "a transfer for each".
    for (List<BlockTransferInfo> transfers : generateTransferLists(4, new int[] {1, 1, 1, 1})) {
      validatePartitions(mPartitioner.partitionTransfers(transfers, 1), 1, 4);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 2), 2, 2, 2);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 3), 3, 2, 1, 1);
      validatePartitions(mPartitioner.partitionTransfers(transfers, 4), 4, 1, 1, 1, 1);
    }
  }

  /**
   * Generates transfer-lists to validate for when SRC/DST or both are
   * allocated(fully identified) locations.
   */
  private List<List<BlockTransferInfo>> generateTransferLists(int count, int[] locDistribution) {
    List<List<BlockTransferInfo>> transfersToValidate = new LinkedList<>();
    transfersToValidate.add(generateTransfers(count, true, false, locDistribution, new int[] {4}));
    transfersToValidate.add(generateTransfers(count, false, true, new int[] {4}, locDistribution));
    transfersToValidate.add(generateTransfers(count, true, true, locDistribution, locDistribution));
    return transfersToValidate;
  }

  /**
   * Generates transfers with controls to specify if src/dst is allocated.
   * It also acceptslocation-group parameter to specify how many transfers will
   * each unique location contain.
   */
  List<BlockTransferInfo> generateTransfers(int count, boolean srcAllocated, boolean dstAllocated,
      int[] srcLocGroups, int[] dstLocGroups) {
    Assert.assertEquals(count, Arrays.stream(srcLocGroups).sum());
    Assert.assertEquals(count, Arrays.stream(dstLocGroups).sum());

    List<BlockTransferInfo> transferInfos = new ArrayList<>(count);
    int[] srcLocGroupsCopy = Arrays.copyOf(srcLocGroups, srcLocGroups.length);
    int[] dstLocGroupsCopy = Arrays.copyOf(dstLocGroups, dstLocGroups.length);
    int srcTierOrdinal = 0;
    int dstTierOrdinal = 0;
    int srcDirIndex = 0;
    int dstDirIndex = 0;
    int currLocationGroupSrc = 0;
    int currLocationGroupDst = 0;
    for (int blockId = 0; blockId < count; blockId++) {
      // Generate the src location.
      BlockStoreLocation srcLoc;
      if (!srcAllocated) {
        srcLoc = BlockStoreLocation.anyDirInTier(Integer.toString(srcTierOrdinal));
      } else {
        srcLoc = new BlockStoreLocation(Integer.toString(srcTierOrdinal), srcDirIndex);
      }

      // Generate the dst location.
      BlockStoreLocation dstLoc;
      if (!dstAllocated) {
        dstLoc = BlockStoreLocation.anyDirInTier(Integer.toString(dstTierOrdinal));
      } else {
        dstLoc = new BlockStoreLocation(Integer.toString(dstTierOrdinal), dstDirIndex);
      }

      // Create the transfer info.
      transferInfos.add(BlockTransferInfo.createMove(srcLoc, blockId, dstLoc));

      // Advance counters for the next location based on requested location groups.
      if (--srcLocGroupsCopy[currLocationGroupSrc] == 0) {
        currLocationGroupSrc++;
        srcTierOrdinal++;
        srcDirIndex++;
      }
      if (--dstLocGroupsCopy[currLocationGroupDst] == 0) {
        currLocationGroupDst++;
        dstTierOrdinal++;
        dstDirIndex++;
      }
    }

    return transferInfos;
  }

  /**
   * Validates given generated partitions.
   *
   * @param transferPartitions partitions list
   * @param expectedPartitionCount number of partitions
   * @param expectedPartitionSizes sizes for each partition
   */
  private void validatePartitions(List<List<BlockTransferInfo>> transferPartitions,
      int expectedPartitionCount, int... expectedPartitionSizes) {
    Assert.assertEquals(expectedPartitionCount, transferPartitions.size());

    for (int i = 0; i < transferPartitions.size(); i++) {
      Assert.assertEquals(expectedPartitionSizes[i], transferPartitions.get(i).size());
    }
  }
}
