/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.gobblin.yarn;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.DataOutputBuffer;
import org.apache.hadoop.security.Credentials;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.yarn.api.ApplicationConstants;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse;
import org.apache.hadoop.yarn.api.records.ApplicationAccessType;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerExitStatus;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.ContainerLaunchContext;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
import org.apache.hadoop.yarn.api.records.LocalResource;
import org.apache.hadoop.yarn.api.records.LocalResourceType;
import org.apache.hadoop.yarn.api.records.NodeReport;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.api.records.ResourceRequest;
import org.apache.hadoop.yarn.client.api.AMRMClient;
import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync;
import org.apache.hadoop.yarn.client.api.async.NMClientAsync;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
import org.apache.hadoop.yarn.util.Records;
import org.apache.helix.HelixAdmin;
import org.apache.helix.HelixManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.base.Throwables;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.google.common.eventbus.EventBus;
import com.google.common.eventbus.Subscribe;
import com.google.common.io.Closer;
import com.google.common.util.concurrent.AbstractIdleService;
import com.typesafe.config.Config;

import lombok.AccessLevel;
import lombok.Getter;

import org.apache.gobblin.cluster.GobblinClusterConfigurationKeys;
import org.apache.gobblin.cluster.GobblinClusterMetricTagNames;
import org.apache.gobblin.cluster.GobblinClusterUtils;
import org.apache.gobblin.cluster.HelixUtils;
import org.apache.gobblin.cluster.event.ClusterManagerShutdownRequest;
import org.apache.gobblin.configuration.ConfigurationKeys;
import org.apache.gobblin.metrics.GobblinMetrics;
import org.apache.gobblin.metrics.MetricReporterException;
import org.apache.gobblin.metrics.MultiReporterException;
import org.apache.gobblin.metrics.Tag;
import org.apache.gobblin.metrics.event.EventSubmitter;
import org.apache.gobblin.util.ConfigUtils;
import org.apache.gobblin.util.ExecutorsUtils;
import org.apache.gobblin.util.JvmUtils;
import org.apache.gobblin.util.executors.ScalingThreadPoolExecutor;
import org.apache.gobblin.yarn.event.ContainerReleaseRequest;
import org.apache.gobblin.yarn.event.ContainerShutdownRequest;
import org.apache.gobblin.yarn.event.NewContainerRequest;
import org.apache.gobblin.yarn.helix.HelixInstancePurgerWithMetrics;

import static org.apache.gobblin.yarn.GobblinYarnTaskRunner.HELIX_YARN_INSTANCE_NAME_PREFIX;

/**
 * This class is responsible for all Yarn-related stuffs including ApplicationMaster registration,
 * ApplicationMaster un-registration, Yarn container management, etc.
 *
 * @author Yinan Li
 */
public class YarnService extends AbstractIdleService {

  private static final Logger LOGGER = LoggerFactory.getLogger(YarnService.class);

  private static final String UNKNOWN_HELIX_INSTANCE = "UNKNOWN";

  private final String applicationName;
  private final String applicationId;
  private final String appViewAcl;
  //Default helix instance tag derived from cluster level config
  private final String helixInstanceTags;

  private final Config config;

  @Getter
  private final EventBus eventBus;

  private final Configuration yarnConfiguration;
  private final FileSystem fs;

  private final Optional<GobblinMetrics> gobblinMetrics;
  @Getter
  private final Optional<EventSubmitter> eventSubmitter;

  @VisibleForTesting
  @Getter(AccessLevel.PROTECTED)
  private final AMRMClientAsync<AMRMClient.ContainerRequest> amrmClientAsync;
  private final NMClientAsync nmClientAsync;
  private final ExecutorService containerLaunchExecutor;

  private final int initialContainers;
  private final int requestedContainerMemoryMbs;
  private final int requestedContainerCores;
  private final int jvmMemoryOverheadMbs;
  private final double jvmMemoryXmxRatio;
  private final boolean containerHostAffinityEnabled;

  private final int helixInstanceMaxRetries;

  private final Optional<String> containerJvmArgs;
  private final String containerTimezone;
  private final HelixManager helixManager;
  private final HelixAdmin helixAdmin;

  @Getter(AccessLevel.PROTECTED)
  private volatile Optional<Resource> maxResourceCapacity = Optional.absent();

  // Security tokens for accessing HDFS
  private ByteBuffer tokens;

  private final Closer closer = Closer.create();

  private final Object allContainersStopped = new Object();

  // A map from container IDs to Container instances, Helix participant IDs of the containers and Helix Tag
  @VisibleForTesting
  @Getter(AccessLevel.PROTECTED)
  private final ConcurrentMap<ContainerId, ContainerInfo> containerMap = Maps.newConcurrentMap();

  // A cache of the containers with an outstanding container release request.
  // This is a cache instead of a set to get the automatic cleanup in case a container completes before the requested
  // release.
  @VisibleForTesting
  @Getter(AccessLevel.PROTECTED)
  private final Cache<ContainerId, String> releasedContainerCache;

  // A generator for an integer ID of a Helix instance (participant)
  private final AtomicInteger helixInstanceIdGenerator = new AtomicInteger(0);

  // A map from Helix instance names to the number times the instances are retried to be started
  private final ConcurrentMap<String, AtomicInteger> helixInstanceRetryCount = Maps.newConcurrentMap();

  // A concurrent HashSet of unused Helix instance names. An unused Helix instance name gets put
  // into the set if the container running the instance completes. Unused Helix
  // instance names get picked up when replacement containers get allocated.
  private final Set<String> unusedHelixInstanceNames = ConcurrentHashMap.newKeySet();

  // The map from helix tag to allocated container count
  private final ConcurrentMap<String, AtomicInteger> allocatedContainerCountMap = Maps.newConcurrentMap();

  private final boolean isPurgingOfflineHelixInstancesEnabled;
  private final long helixPurgeLaggingThresholdMs;
  private final long helixPurgeStatusPollingRateMs;
  private final ConcurrentMap<ContainerId, Long> containerIdleSince = Maps.newConcurrentMap();
  private final ConcurrentMap<ContainerId, String> removedContainerID = Maps.newConcurrentMap();

  private volatile YarnContainerRequestBundle yarnContainerRequest;
  private final AtomicInteger priorityNumGenerator = new AtomicInteger(0);
  private final Map<String, Integer> resourcePriorityMap = new HashMap<>();

  private volatile boolean shutdownInProgress = false;
  private volatile boolean startupInProgress = true;

  public YarnService(Config config, String applicationName, String applicationId, YarnConfiguration yarnConfiguration,
      FileSystem fs, EventBus eventBus, HelixManager helixManager, HelixAdmin helixAdmin) throws Exception {
    this.applicationName = applicationName;
    this.applicationId = applicationId;

    this.config = config;

    this.eventBus = eventBus;

    this.helixManager = helixManager;
    this.helixAdmin = helixAdmin;

    this.gobblinMetrics = config.getBoolean(ConfigurationKeys.METRICS_ENABLED_KEY) ?
        Optional.of(buildGobblinMetrics()) : Optional.<GobblinMetrics>absent();

    this.eventSubmitter = config.getBoolean(ConfigurationKeys.METRICS_ENABLED_KEY) ?
        Optional.of(buildEventSubmitter()) : Optional.<EventSubmitter>absent();

    this.yarnConfiguration = yarnConfiguration;
    this.fs = fs;

    int amRmHeartbeatIntervalMillis = Long.valueOf(TimeUnit.SECONDS.toMillis(
        ConfigUtils.getInt(config, GobblinYarnConfigurationKeys.AMRM_HEARTBEAT_INTERVAL_SECS,
            GobblinYarnConfigurationKeys.DEFAULT_AMRM_HEARTBEAT_INTERVAL_SECS))).intValue();
    this.amrmClientAsync = closer.register(
        AMRMClientAsync.createAMRMClientAsync(amRmHeartbeatIntervalMillis, new AMRMClientCallbackHandler()));
    this.amrmClientAsync.init(this.yarnConfiguration);
    this.nmClientAsync = closer.register(NMClientAsync.createNMClientAsync(getNMClientCallbackHandler()));
    this.nmClientAsync.init(this.yarnConfiguration);

    this.initialContainers = config.getInt(GobblinYarnConfigurationKeys.INITIAL_CONTAINERS_KEY);
    this.requestedContainerMemoryMbs = config.getInt(GobblinYarnConfigurationKeys.CONTAINER_MEMORY_MBS_KEY);
    this.requestedContainerCores = config.getInt(GobblinYarnConfigurationKeys.CONTAINER_CORES_KEY);
    this.containerHostAffinityEnabled = config.getBoolean(GobblinYarnConfigurationKeys.CONTAINER_HOST_AFFINITY_ENABLED);

    this.helixInstanceMaxRetries = config.getInt(GobblinYarnConfigurationKeys.HELIX_INSTANCE_MAX_RETRIES);
    this.helixInstanceTags = ConfigUtils.getString(config,
        GobblinClusterConfigurationKeys.HELIX_INSTANCE_TAGS_KEY, GobblinClusterConfigurationKeys.HELIX_DEFAULT_TAG);
    this.isPurgingOfflineHelixInstancesEnabled = ConfigUtils.getBoolean(config,
        GobblinYarnConfigurationKeys.HELIX_PURGE_OFFLINE_INSTANCES_ENABLED,
        GobblinYarnConfigurationKeys.DEFAULT_HELIX_PURGE_OFFLINE_INSTANCES_ENABLED);
    this.helixPurgeLaggingThresholdMs = ConfigUtils.getLong(config,
        GobblinYarnConfigurationKeys.HELIX_PURGE_LAGGING_THRESHOLD_MILLIS,
        GobblinYarnConfigurationKeys.DEFAULT_HELIX_PURGE_LAGGING_THRESHOLD_MILLIS);
    this.helixPurgeStatusPollingRateMs = ConfigUtils.getLong(config,
        GobblinYarnConfigurationKeys.HELIX_PURGE_POLLING_RATE_MILLIS,
        GobblinYarnConfigurationKeys.DEFAULT_HELIX_PURGE_POLLING_RATE_MILLIS);

    this.containerJvmArgs = config.hasPath(GobblinYarnConfigurationKeys.CONTAINER_JVM_ARGS_KEY) ?
        Optional.of(config.getString(GobblinYarnConfigurationKeys.CONTAINER_JVM_ARGS_KEY)) :
        Optional.<String>absent();

    int numContainerLaunchThreads =
        ConfigUtils.getInt(config, GobblinYarnConfigurationKeys.MAX_CONTAINER_LAUNCH_THREADS_KEY,
            GobblinYarnConfigurationKeys.DEFAULT_MAX_CONTAINER_LAUNCH_THREADS);
    this.containerLaunchExecutor = ScalingThreadPoolExecutor.newScalingThreadPool(5, numContainerLaunchThreads, 0L,
        ExecutorsUtils.newThreadFactory(Optional.of(LOGGER), Optional.of("ContainerLaunchExecutor")));

    this.tokens = getSecurityTokens();

    this.releasedContainerCache = CacheBuilder.newBuilder().expireAfterAccess(ConfigUtils.getInt(config,
        GobblinYarnConfigurationKeys.RELEASED_CONTAINERS_CACHE_EXPIRY_SECS,
        GobblinYarnConfigurationKeys.DEFAULT_RELEASED_CONTAINERS_CACHE_EXPIRY_SECS), TimeUnit.SECONDS).build();

    this.jvmMemoryXmxRatio = ConfigUtils.getDouble(this.config,
        GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_XMX_RATIO_KEY,
        GobblinYarnConfigurationKeys.DEFAULT_CONTAINER_JVM_MEMORY_XMX_RATIO);

    Preconditions.checkArgument(this.jvmMemoryXmxRatio >= 0 && this.jvmMemoryXmxRatio <= 1,
        GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_XMX_RATIO_KEY + " must be between 0 and 1 inclusive");

    this.jvmMemoryOverheadMbs = ConfigUtils.getInt(this.config,
        GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_OVERHEAD_MBS_KEY,
        GobblinYarnConfigurationKeys.DEFAULT_CONTAINER_JVM_MEMORY_OVERHEAD_MBS);

    Preconditions.checkArgument(this.jvmMemoryOverheadMbs < this.requestedContainerMemoryMbs * this.jvmMemoryXmxRatio,
        GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_OVERHEAD_MBS_KEY + " cannot be more than "
            + GobblinYarnConfigurationKeys.CONTAINER_MEMORY_MBS_KEY + " * "
            + GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_XMX_RATIO_KEY);

    this.appViewAcl = ConfigUtils.getString(this.config, GobblinYarnConfigurationKeys.APP_VIEW_ACL,
        GobblinYarnConfigurationKeys.DEFAULT_APP_VIEW_ACL);
    this.containerTimezone = ConfigUtils.getString(this.config, GobblinYarnConfigurationKeys.GOBBLIN_YARN_CONTAINER_TIMEZONE,
        GobblinYarnConfigurationKeys.DEFAULT_GOBBLIN_YARN_CONTAINER_TIMEZONE);
  }

  @SuppressWarnings("unused")
  @Subscribe
  public void handleNewContainerRequest(NewContainerRequest newContainerRequest) {
    if (!this.maxResourceCapacity.isPresent()) {
      LOGGER.error(String.format(
          "Unable to handle new container request as maximum resource capacity is not available: "
              + "[memory (MBs) requested = %d, vcores requested = %d]", this.requestedContainerMemoryMbs,
          this.requestedContainerCores));
      return;
    }
    requestContainer(newContainerRequest.getReplacedContainer().transform(container -> container.getNodeId().getHost()),
        newContainerRequest.getResource());
  }

  /**
   *  getContainerInfoGivenHelixParticipant returns the container of the given helixParticipant if it exists else
   *  return Optional<Container>
   * @param helixParticipant
   * @return Container
   */
  public Optional<Container> getContainerInfoGivenHelixParticipant(final String helixParticipant) {
    for (ContainerInfo containerInfo : this.containerMap.values()) {
      if (containerInfo.getHelixParticipantId().equals(helixParticipant)) {
        return Optional.fromNullable(containerInfo.getContainer());
      }
    }
    return Optional.absent();
  }

  protected NMClientCallbackHandler getNMClientCallbackHandler() {
    return new NMClientCallbackHandler();
  }

  @SuppressWarnings("unused")
  @Subscribe
  public void handleContainerShutdownRequest(ContainerShutdownRequest containerShutdownRequest) {
    for (Container container : containerShutdownRequest.getContainers()) {
      LOGGER.info(String.format("Stopping container %s running on %s", container.getId(), container.getNodeId()));
      this.nmClientAsync.stopContainerAsync(container.getId(), container.getNodeId());
    }
  }

  /**
   * Request the Resource Manager to release the container
   * @param containerReleaseRequest containers to release
   */
  @Subscribe
  public void handleContainerReleaseRequest(ContainerReleaseRequest containerReleaseRequest) {
    for (Container container : containerReleaseRequest.getContainers()) {
      LOGGER.info(String.format("Releasing container %s running on %s", container.getId(), container.getNodeId()));

      if(!containerReleaseRequest.isShouldSpinUpReplacementContainers()) {
        // Record that this container was explicitly released so that a new one is not spawned to replace it
        // Put the container id in the releasedContainerCache before releasing it so that handleContainerCompletion()
        // can check for the container id and skip spawning a replacement container.
        // Note that this is the best effort since these are asynchronous operations and a container may abort concurrently
        // with the release call. So in some cases a replacement container may have already been spawned before
        // the container is put into the black list.
        this.releasedContainerCache.put(container.getId(), "");
      }
      this.amrmClientAsync.releaseAssignedContainer(container.getId());
    }
  }

  @Override
  protected synchronized void startUp() throws Exception {
    LOGGER.info("Starting the YarnService");

    // Register itself with the EventBus for container-related requests
    this.eventBus.register(this);

    this.amrmClientAsync.start();
    this.nmClientAsync.start();

    // The ApplicationMaster registration response is used to determine the maximum resource capacity of the cluster
    RegisterApplicationMasterResponse response = this.amrmClientAsync.registerApplicationMaster(
        GobblinClusterUtils.getHostname(), -1, "");
    LOGGER.info("ApplicationMaster registration response: " + response);
    this.maxResourceCapacity = Optional.of(response.getMaximumResourceCapability());

    if (this.isPurgingOfflineHelixInstancesEnabled) {
      purgeHelixOfflineInstances(this.helixPurgeLaggingThresholdMs);
    }

    LOGGER.info("Requesting initial containers");
    requestInitialContainers(this.initialContainers);
    startupInProgress = false;
  }

  private void purgeHelixOfflineInstances(long laggingThresholdMs) {
    LOGGER.info("Purging offline helix instances before allocating containers for helixClusterName={}, connectionString={}, helixPurgeStatusPollingRateMs={}",
        helixManager.getClusterName(), helixManager.getMetadataStoreConnectionString(), this.helixPurgeStatusPollingRateMs);
    HelixInstancePurgerWithMetrics purger = new HelixInstancePurgerWithMetrics(this.eventSubmitter.orNull(),
        this.helixPurgeStatusPollingRateMs);
    Map<String, String> gteMetadata = ImmutableMap.of(
        "connectionString", this.helixManager.getMetadataStoreConnectionString(),
        "clusterName", this.helixManager.getClusterName()
    );
    purger.purgeAllOfflineInstances(this.helixAdmin, this.helixManager.getClusterName(), laggingThresholdMs, gteMetadata);
  }

  @Override
  protected void shutDown() throws IOException {
    LOGGER.info("Stopping the YarnService");

    this.shutdownInProgress = true;

    try {
      ExecutorsUtils.shutdownExecutorService(this.containerLaunchExecutor, Optional.of(LOGGER));

      // Stop the running containers
      for (ContainerInfo containerInfo : this.containerMap.values()) {
        LOGGER.info("Stopping container {} running participant {}", containerInfo.getContainer().getId(),
            containerInfo.getHelixParticipantId());
        this.nmClientAsync.stopContainerAsync(containerInfo.getContainer().getId(), containerInfo.getContainer().getNodeId());
      }

      if (!this.containerMap.isEmpty()) {
        synchronized (this.allContainersStopped) {
          try {
            // Wait 5 minutes for the containers to stop
            Duration waitTimeout = Duration.ofMinutes(5);
            this.allContainersStopped.wait(waitTimeout.toMillis());
            LOGGER.info("All of the containers have been stopped");
          } catch (InterruptedException ie) {
            Thread.currentThread().interrupt();
          }
        }
      }

      this.amrmClientAsync.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, null, null);
    } catch (IOException | YarnException e) {
      LOGGER.error("Failed to unregister the ApplicationMaster", e);
    } finally {
      try {
        this.closer.close();
      } finally {
        if (this.gobblinMetrics.isPresent()) {
          this.gobblinMetrics.get().stopMetricsReporting();
        }
      }
    }
  }

  public void updateToken() throws IOException{
    this.tokens = getSecurityTokens();
  }

  private GobblinMetrics buildGobblinMetrics() {
    // Create tags list
    ImmutableList.Builder<Tag<?>> tags = new ImmutableList.Builder<>();
    tags.add(new Tag<>(GobblinClusterMetricTagNames.APPLICATION_ID, this.applicationId));
    tags.add(new Tag<>(GobblinClusterMetricTagNames.APPLICATION_NAME, this.applicationName));

    // Intialize Gobblin metrics and start reporters
    GobblinMetrics gobblinMetrics = GobblinMetrics.get(this.applicationId, null, tags.build());
    try {
      gobblinMetrics.startMetricReporting(ConfigUtils.configToProperties(config));
    } catch (MultiReporterException ex) {
      for (MetricReporterException e: ex.getExceptions()) {
        LOGGER.error("Failed to start {} {} reporter.", e.getSinkType().name(), e.getReporterType().name(), e);
      }
    }

    return gobblinMetrics;
  }

  private EventSubmitter buildEventSubmitter() {
    return new EventSubmitter.Builder(this.gobblinMetrics.get().getMetricContext(),
        GobblinYarnEventConstants.EVENT_NAMESPACE)
        .build();
  }

  /**
   * Request an allocation of containers. If numTargetContainers is larger than the max of current and expected number
   * of containers then additional containers are requested.
   * <p>
   * If numTargetContainers is less than the current number of allocated containers then release free containers.
   * Shrinking is relative to the number of currently allocated containers since it takes time for containers
   * to be allocated and assigned work and we want to avoid releasing a container prematurely before it is assigned
   * work. This means that a container may not be released even though numTargetContainers is less than the requested
   * number of containers. The intended usage is for the caller of this method to make periodic calls to attempt to
   * adjust the cluster towards the desired number of containers.
   *
   * @param yarnContainerRequestBundle the desired containers information, including numbers, resource and helix tag
   * @param inUseInstances  a set of in use instances
   * @return whether successfully requested the target number of containers
   */
  public synchronized boolean requestTargetNumberOfContainers(YarnContainerRequestBundle yarnContainerRequestBundle, Set<String> inUseInstances) {
    LOGGER.info("Trying to set numTargetContainers={}, in-use helix instances count is {}, container map size is {}",
        yarnContainerRequestBundle.getTotalContainers(), inUseInstances.size(), this.containerMap.size());
    if (startupInProgress) {
      LOGGER.warn("YarnService is still starting up. Unable to request containers from yarn until YarnService is finished starting up.");
      return false;
    }

    //Correct the containerMap first as there is cases that handleContainerCompletion() is called before onContainersAllocated()
    for (ContainerId removedId : this.removedContainerID.keySet()) {
      ContainerInfo containerInfo = this.containerMap.remove(removedId);
      if (containerInfo != null) {
        String helixTag = containerInfo.getHelixTag();
        allocatedContainerCountMap.putIfAbsent(helixTag, new AtomicInteger(0));
        this.allocatedContainerCountMap.get(helixTag).decrementAndGet();
        this.removedContainerID.remove(removedId);
      }
    }

    int numTargetContainers = yarnContainerRequestBundle.getTotalContainers();
    // YARN can allocate more than the requested number of containers, compute additional allocations and deallocations
    // based on the max of the requested and actual allocated counts
    // Represents the number of containers allocated for across all helix tags
    int totalAllocatedContainers = this.containerMap.size();
    int totalContainersInContainerCountMap = 0;
    for (AtomicInteger count: allocatedContainerCountMap.values()) {
      totalContainersInContainerCountMap += count.get();
    }
    if (totalContainersInContainerCountMap != totalAllocatedContainers) {
      LOGGER.warn(String.format("Container number mismatch in containerMap and allocatedContainerCountMap, "
          + "we have %s containers in containerMap while %s in allocatedContainerCountMap", totalAllocatedContainers, totalContainersInContainerCountMap));
    }

    // Request additional containers if the desired count is higher than the max of the current allocation or previously
    // requested amount. Note that there may be in-flight or additional allocations after numContainers has been computed
    // so overshooting can occur, but periodic calls to this method will make adjustments towards the target.
    for (Map.Entry<String, Integer> entry : yarnContainerRequestBundle.getHelixTagContainerCountMap().entrySet()) {
      String currentHelixTag = entry.getKey();
      int desiredContainerCount = entry.getValue();
      Resource resourceForHelixTag = yarnContainerRequestBundle.getHelixTagResourceMap().get(currentHelixTag);

      // Calculate requested container count based on adding allocated count and outstanding ContainerRequests in Yarn
      allocatedContainerCountMap.putIfAbsent(currentHelixTag, new AtomicInteger(0));
      int allocatedContainersForHelixTag = allocatedContainerCountMap.get(currentHelixTag).get();
      int outstandingContainerRequests = getMatchingRequestsCount(resourceForHelixTag);
      int requestedContainerCount = allocatedContainersForHelixTag + outstandingContainerRequests;
      int numContainersNeeded = desiredContainerCount - requestedContainerCount;
      LOGGER.info("Container counts for helixTag={} (allocatedContainers={}, outstandingContainerRequests={}, desiredContainerCount={}, numContainersNeeded={})",
          currentHelixTag, allocatedContainersForHelixTag, outstandingContainerRequests, desiredContainerCount, numContainersNeeded);

      if (numContainersNeeded > 0) {
        requestContainers(numContainersNeeded, resourceForHelixTag);
      }
    }

    //Iterate through all containers allocated and check whether the corresponding helix instance is still LIVE within the helix cluster.
    // A container that has a bad connection to zookeeper will be dropped from the Helix cluster if the disconnection is greater than the specified timeout.
    // In these cases, we want to release the container to get a new container because these containers won't be assigned tasks by Helix

    List<Container> containersToRelease = new ArrayList<>();
    HashSet<ContainerId> idleContainerIdsToRelease = new HashSet<>();
    for (Map.Entry<ContainerId, ContainerInfo> entry : this.containerMap.entrySet()) {
      ContainerInfo containerInfo = entry.getValue();
      if (!HelixUtils.isInstanceLive(helixManager, containerInfo.getHelixParticipantId())) {
        containerIdleSince.putIfAbsent(entry.getKey(), System.currentTimeMillis());
        if (System.currentTimeMillis() - containerIdleSince.get(entry.getKey())
            >= TimeUnit.MINUTES.toMillis(YarnAutoScalingManager.DEFAULT_MAX_CONTAINER_IDLE_TIME_BEFORE_SCALING_DOWN_MINUTES)) {
          LOGGER.info("Releasing Container {} because the assigned participant {} has been in-active for more than {} minutes",
              entry.getKey(), containerInfo.getHelixParticipantId(), YarnAutoScalingManager.DEFAULT_MAX_CONTAINER_IDLE_TIME_BEFORE_SCALING_DOWN_MINUTES);
          containersToRelease.add(containerInfo.getContainer());
          idleContainerIdsToRelease.add(entry.getKey());
        }
      } else {
        containerIdleSince.remove(entry.getKey());
      }
    }

    // If the total desired is lower than the currently allocated amount then release free containers.
    // This is based on the currently allocated amount since containers may still be in the process of being allocated
    // and assigned work. Resizing based on numRequestedContainers at this point may release a container right before
    // or soon after it is assigned work.
    if (numTargetContainers < totalAllocatedContainers - idleContainerIdsToRelease.size()) {
      int numToShutdown = totalAllocatedContainers - numTargetContainers;

      LOGGER.info("Shrinking number of containers by {} because numTargetContainers < totalAllocatedContainers - idleContainersToRelease ({} < {} - {})",
          totalAllocatedContainers - numTargetContainers - idleContainerIdsToRelease.size(), numTargetContainers, totalAllocatedContainers, idleContainerIdsToRelease.size());

      // Look for eligible containers to release. If a container is in use then it is not released.
      for (Map.Entry<ContainerId, ContainerInfo> entry : this.containerMap.entrySet()) {
        ContainerInfo containerInfo = entry.getValue();
        if (!inUseInstances.contains(containerInfo.getHelixParticipantId()) && !idleContainerIdsToRelease.contains(entry.getKey())) {
          containersToRelease.add(containerInfo.getContainer());
        }

        if (containersToRelease.size() >= numToShutdown) {
          break;
        }
      }

      LOGGER.info("Shutting down {} containers. containersToRelease={}", containersToRelease.size(), containersToRelease);
    }

    if (!containersToRelease.isEmpty()) {
      this.eventBus.post(new ContainerReleaseRequest(containersToRelease));
    }
    this.yarnContainerRequest = yarnContainerRequestBundle;
    LOGGER.info("Current tag-container desired count:{}, tag-container allocated: {}",
        yarnContainerRequestBundle.getHelixTagContainerCountMap(), this.allocatedContainerCountMap);
    return true;
  }

  // Request initial containers with default resource and helix tag
  private void requestInitialContainers(int containersRequested) {
    YarnContainerRequestBundle initialYarnContainerRequest = new YarnContainerRequestBundle();
    Resource capability = Resource.newInstance(this.requestedContainerMemoryMbs, this.requestedContainerCores);
    initialYarnContainerRequest.add(this.helixInstanceTags, containersRequested, capability);
    requestTargetNumberOfContainers(initialYarnContainerRequest, Collections.EMPTY_SET);
  }

  private void requestContainer(Optional<String> preferredNode, Optional<Resource> resourceOptional) {
    Resource desiredResource = resourceOptional.or(Resource.newInstance(
        this.requestedContainerMemoryMbs, this.requestedContainerCores));
    requestContainer(preferredNode, desiredResource);
  }

  /**
   * Request {@param numContainers} from yarn with the specified resource. Resources will be allocated without a preferred
   * node
   * @param numContainers
   * @param resource
   */
  private void requestContainers(int numContainers, Resource resource) {
    LOGGER.info("Requesting {} containers with resource={}", numContainers, resource);
    IntStream.range(0, numContainers)
        .forEach(i -> requestContainer(Optional.absent(), resource));
  }

  // Request containers with specific resource requirement
  private void requestContainer(Optional<String> preferredNode, Resource resource) {
    // Fail if Yarn cannot meet container resource requirements
    Preconditions.checkArgument(resource.getMemory() <= this.maxResourceCapacity.get().getMemory() &&
            resource.getVirtualCores() <= this.maxResourceCapacity.get().getVirtualCores(),
        "Resource requirement must less than the max resource capacity. Requested resource" + resource.toString()
            + " exceed the max resource limit " + this.maxResourceCapacity.get().toString());

    // Due to YARN-314, different resource capacity needs different priority, otherwise Yarn will not allocate container
    Priority priority = Records.newRecord(Priority.class);
    if(!resourcePriorityMap.containsKey(resource.toString())) {
      resourcePriorityMap.put(resource.toString(), priorityNumGenerator.getAndIncrement());
    }
    int priorityNum = resourcePriorityMap.get(resource.toString());
