/*
 * 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.metastore.heap;

import alluxio.collections.TwoKeyConcurrentMap;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.master.metastore.BlockMetaStore;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.proto.meta.Block.BlockLocation;
import alluxio.proto.meta.Block.BlockMeta;
import alluxio.resource.CloseableIterator;
import alluxio.util.ObjectSizeCalculator;

import com.google.common.collect.ImmutableSet;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import javax.annotation.concurrent.ThreadSafe;

/**
 * This class requires external synchronization for operations on the same block id. Operations on
 * different block ids can be performed concurrently.
 */
@ThreadSafe
public class HeapBlockMetaStore implements BlockMetaStore {
  // Map from block id to block metadata.
  public final Map<Long, BlockMeta> mBlocks = new ConcurrentHashMap<>();
  // Map from block id to block locations.
  public final TwoKeyConcurrentMap<Long, Long, BlockLocation, Map<Long, BlockLocation>>
      mBlockLocations = new TwoKeyConcurrentMap<>(() -> new HashMap<>(4));

  /**
   * constructor a HeapBlockStore.
   *
   */
  public HeapBlockMetaStore() {
    super();
    if (Configuration.getBoolean(PropertyKey.MASTER_METRICS_HEAP_ENABLED)) {
      MetricsSystem.registerCachedGaugeIfAbsent(MetricKey.MASTER_BLOCK_HEAP_SIZE.getName(),
          () -> ObjectSizeCalculator.getObjectSize(mBlocks,
              ImmutableSet.of(Long.class, BlockMeta.class)));
    }
  }

  @Override
  public Optional<BlockMeta> getBlock(long id) {
    return Optional.ofNullable(mBlocks.get(id));
  }

  @Override
  public void putBlock(long id, BlockMeta meta) {
    mBlocks.put(id, meta);
  }

  @Override
  public void removeBlock(long id) {
    mBlocks.remove(id);
  }

  @Override
  public CloseableIterator<Block> getCloseableIterator() {
    Iterator<Entry<Long, BlockMeta>> it = mBlocks.entrySet().iterator();
    return CloseableIterator.noopCloseable(new Iterator<Block>() {
      @Override
      public boolean hasNext() {
        return it.hasNext();
      }

      @Override
      public Block next() {
        Entry<Long, BlockMeta> entry = it.next();
        return new Block(entry.getKey(), entry.getValue());
      }
    });
  }

  @Override
  public void clear() {
    mBlocks.clear();
  }

  @Override
  public void close() {
    // Nothing to close for HEAP store.
  }

  @Override
  public long size() {
    return mBlocks.size();
  }

  @Override
  public List<BlockLocation> getLocations(long blockid) {
    if (!mBlockLocations.containsKey(blockid)) {
      return Collections.emptyList();
    }
    return new ArrayList<>(mBlockLocations.get(blockid).values());
  }

  @Override
  public void addLocation(long blockId, BlockLocation location) {
    mBlockLocations.addInnerValue(blockId, location.getWorkerId(), location);
  }

  @Override
  public void removeLocation(long blockId, long workerId) {
    mBlockLocations.removeInnerValue(blockId, workerId);
  }
}
