alluxio-master-2.x/core/server/worker/src/main/java/alluxio/worker/block/BlockMapIterator.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 alluxio.conf.AlluxioConfiguration;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.BlockIdList;
import alluxio.grpc.BlockStoreLocationProto;
import alluxio.grpc.LocationBlockIdListEntry;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * Each iteration returns a list of {@link LocationBlockIdListEntry} which consists of
 * blocks for one {@link alluxio.grpc.RegisterWorkerPRequest}.
 * The number of blocks included in one iteration is specified by
 * {@code PropertyKey.WORKER_REGISTER_STREAM_BATCH_SIZE}.
 */
public class BlockMapIterator implements Iterator<List<LocationBlockIdListEntry>> {
  private static final Logger LOG = LoggerFactory.getLogger(BlockMapIterator.class);

  private final int mBatchSize;
  private final int mBlockCount;
  // Keeps the order of iteration
  private final List<BlockStoreLocationProto> mBlockStoreLocationProtoList;
  private final Map<BlockStoreLocationProto, Iterator<Long>> mBlockLocationIteratorMap;
  // Iteration states
  private int mCurrentBlockLocationIndex;
  private final Iterator<Long> mCurrentIterator;
  // A global counter of how many blocks have been traversed
  private int mCounter = 0;

  /**
   * Constructor.
   *
   * @param blockLocationMap the block lists for each location
   */
  public BlockMapIterator(Map<BlockStoreLocation, List<Long>> blockLocationMap) {
    this(blockLocationMap, Configuration.global());
  }

  /**
   * Constructor.
   *
   * @param blockLocationMap the block lists for each location
   * @param conf configuration properties
   */
  public BlockMapIterator(
      Map<BlockStoreLocation, List<Long>> blockLocationMap, AlluxioConfiguration conf) {
    mBatchSize = conf.getInt(PropertyKey.WORKER_REGISTER_STREAM_BATCH_SIZE);
    LOG.info("Worker register stream batchSize={}", mBatchSize);

    // The worker will merge the block lists from dirs on the same tier
    // because the master only wants one list for each tier.
    // TODO(jiacheng): try to avoid this merge copy
    Map<BlockStoreLocationProto, List<Long>> tierToBlocks = new HashMap<>();
    for (Map.Entry<BlockStoreLocation, List<Long>> entry : blockLocationMap.entrySet()) {
      BlockStoreLocation loc = entry.getKey();
      BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder()
          .setTierAlias(loc.tierAlias()).setMediumType(loc.mediumType()).build();
      if (tierToBlocks.containsKey(locationProto)) {
        tierToBlocks.get(locationProto).addAll(entry.getValue());
      } else {
        List<Long> blockList = new ArrayList<>(entry.getValue());
        tierToBlocks.put(locationProto, blockList);
      }
    }
    LOG.debug("Found blocks for tiers: {}", tierToBlocks.keySet());

    // Initialize the iteration statuses
    mBlockStoreLocationProtoList = new ArrayList<>(tierToBlocks.keySet());
    mCurrentBlockLocationIndex = 0; // Corresponding to index in mBlockLocationKeyList

    // Find the iterator of each tier and calculate total block count
    mBlockLocationIteratorMap = new HashMap<>();
    int totalCount = 0;
    for (Map.Entry<BlockStoreLocationProto, List<Long>> entry : tierToBlocks.entrySet()) {
      totalCount += entry.getValue().size();
      mBlockLocationIteratorMap.put(entry.getKey(), entry.getValue().iterator());
    }
    mBlockCount = totalCount;

    // Initialize the iteration status
    if (mBlockStoreLocationProtoList.size() == 0) {
      mCurrentIterator = Collections.emptyIterator();
    } else {
      mCurrentIterator = mBlockLocationIteratorMap.get(mBlockStoreLocationProtoList.get(0));
    }
  }

  @Override
  public boolean hasNext() {
    return mCurrentIterator.hasNext()
        || mCurrentBlockLocationIndex < mBlockStoreLocationProtoList.size();
  }

  private LocationBlockIdListEntry nextBatchFromTier(
      BlockStoreLocationProto currentLoc, Iterator<Long> currentIterator, int spaceLeft) {
    // Generate the next batch
    List<Long> blockIdBatch = new ArrayList<>(spaceLeft);
    while (blockIdBatch.size() < spaceLeft && currentIterator.hasNext()) {
      blockIdBatch.add(currentIterator.next());
      mCounter++;
    }
    BlockIdList blockIdList = BlockIdList.newBuilder().addAllBlockId(blockIdBatch).build();
    return LocationBlockIdListEntry.newBuilder()
        .setKey(currentLoc).setValue(blockIdList).build();
  }

  @Override
  public List<LocationBlockIdListEntry> next() {
    List<LocationBlockIdListEntry> result = new ArrayList<>();
    int currentCounter = mCounter;
    int targetCounter = currentCounter + mBatchSize;

    while (mCounter < targetCounter) {
      // Find the BlockStoreLocation
      BlockStoreLocationProto currentLoc =
          mBlockStoreLocationProtoList.get(mCurrentBlockLocationIndex);
      Iterator<Long> currentIterator = mBlockLocationIteratorMap.get(currentLoc);

      // Generate the next batch
      // This method does NOT progress the pointers
      int spaceLeft = targetCounter - mCounter;
      LocationBlockIdListEntry batchFromThisTier =
          nextBatchFromTier(currentLoc, currentIterator, spaceLeft);
      result.add(batchFromThisTier);

      // Progress the iterator based on the break condition
      if (!currentIterator.hasNext()) {
        // We keep filling in using the next tier
        mCurrentBlockLocationIndex++;
        if (mCurrentBlockLocationIndex >= mBlockStoreLocationProtoList.size()) {
          // We break out of the loop when the current iterator is exhausted
          // and all iterators have been exhausted
          return result;
        }
        continue;
      } else {
        return result;
      }
    }

    // Should never reach here
    return result;
  }

  /**
   * Gets the number of batches.
   *
   * @return the count
   */
  public int getBatchCount() {
    return (int) Math.ceil(mBlockCount / (double) mBatchSize);
  }
}
