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

import static org.junit.Assert.assertFalse;

import alluxio.grpc.LocationBlockIdListEntry;
import alluxio.proto.meta.Block;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.junit.Test;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class BlockMapIteratorTest {
  @Test
  public void convertStream() {
    List<Integer> memList = ImmutableList.of(1000, 1000);
    List<Integer> ssdList = ImmutableList.of(10000, 10000);
    List<Integer> hddList = ImmutableList.of(20000, 20000);
    Map<String, List<Integer>> blockConfig =
        ImmutableMap.of("MEM", memList, "SSD", ssdList, "HDD", hddList);

    Map<BlockStoreLocation, List<Long>> blockMap = generateBlockIdOnTiers(blockConfig);

    BlockMapIterator iter = new BlockMapIterator(blockMap);

    while (iter.hasNext()) {
      List<LocationBlockIdListEntry> entries = iter.next();

      // Each batch should not contain duplicate locations
      Set<Block.BlockLocation> locations = new HashSet<>();
      for (LocationBlockIdListEntry e : entries) {
        // No duplicate
        Block.BlockLocation loc = Block.BlockLocation.newBuilder()
            .setTier(e.getKey().getTierAlias())
            .setMediumType(e.getKey().getMediumType())
            .build();
        assertFalse(locations.contains(loc));
        locations.add(loc);
      }
    }
  }

  @Test
  public void convertStreamMergedTiers() {
    List<Integer> memList = ImmutableList.of(100, 100);
    List<Integer> ssdList = ImmutableList.of(101, 22);
    List<Integer> hddList = ImmutableList.of(10, 20);
    Map<String, List<Integer>> blockConfig =
        ImmutableMap.of("MEM", memList, "SSD", ssdList, "HDD", hddList);

    Map<BlockStoreLocation, List<Long>> blockMap = generateBlockIdOnTiers(blockConfig);

    BlockMapIterator iter = new BlockMapIterator(blockMap);

    while (iter.hasNext()) {
      List<LocationBlockIdListEntry> entries = iter.next();

      // Each batch should not contain duplicate locations
      Set<Block.BlockLocation> locations = new HashSet<>();
      for (LocationBlockIdListEntry e : entries) {
        // No duplicate
        Block.BlockLocation loc = Block.BlockLocation.newBuilder()
                .setTier(e.getKey().getTierAlias())
                .setMediumType(e.getKey().getMediumType())
                .build();
        assertFalse(locations.contains(loc));
        locations.add(loc);
      }
    }
  }

  private static Map<BlockStoreLocation, List<Long>> generateBlockIdOnTiers(
          Map<String, List<Integer>> tiersConfig) {
    Map<BlockStoreLocation, List<Long>> blockMap = new HashMap<>();

    long blockIdStart = Long.MAX_VALUE;
    for (Map.Entry<String, List<Integer>> tierConfig : tiersConfig.entrySet()) {
      List<Integer> dirConfigs = tierConfig.getValue();
      for (int i = 0; i < dirConfigs.size(); i++) {
        int dirNumBlocks = dirConfigs.get(i);
        BlockStoreLocation loc = new BlockStoreLocation(tierConfig.getKey(), i);
        List<Long> blockIds = generateDecreasingNumbers(blockIdStart, dirNumBlocks);
        blockMap.put(loc, blockIds);
        blockIdStart -= dirNumBlocks;
      }
    }
    return blockMap;
  }

  private static List<Long> generateDecreasingNumbers(long start, int count) {
    List<Long> list = new ArrayList<>(count);
    for (long i = 0; i < count; i++) {
      list.add(start - i);
    }
    return list;
  }
}
