src/main/java/org/apache/aurora/scheduler/quota/QuotaManager.java

/**
 * Licensed 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.aurora.scheduler.quota;

import java.util.EnumSet;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicate;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Multimap;
import com.google.common.collect.RangeSet;

import org.apache.aurora.gen.JobUpdateQuery;
import org.apache.aurora.scheduler.base.JobKeys;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.configuration.ConfigurationManager;
import org.apache.aurora.scheduler.resources.ResourceBag;
import org.apache.aurora.scheduler.resources.ResourceManager;
import org.apache.aurora.scheduler.resources.ResourceType;
import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
import org.apache.aurora.scheduler.storage.entities.IInstanceTaskConfig;
import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateInstructions;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateQuery;
import org.apache.aurora.scheduler.storage.entities.IRange;
import org.apache.aurora.scheduler.storage.entities.IResource;
import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.apache.aurora.scheduler.updater.Updates;

import static java.util.Objects.requireNonNull;

import static com.google.common.base.Predicates.and;
import static com.google.common.base.Predicates.compose;
import static com.google.common.base.Predicates.equalTo;
import static com.google.common.base.Predicates.in;
import static com.google.common.base.Predicates.not;
import static com.google.common.base.Predicates.or;

import static org.apache.aurora.scheduler.quota.QuotaCheckResult.Result.SUFFICIENT_QUOTA;
import static org.apache.aurora.scheduler.resources.ResourceBag.EMPTY;
import static org.apache.aurora.scheduler.resources.ResourceBag.IS_NEGATIVE;
import static org.apache.aurora.scheduler.resources.ResourceManager.bagFromAggregate;
import static org.apache.aurora.scheduler.resources.ResourceManager.bagFromResources;
import static org.apache.aurora.scheduler.resources.ResourceManager.getTaskResources;
import static org.apache.aurora.scheduler.resources.ResourceType.CPUS;
import static org.apache.aurora.scheduler.resources.ResourceType.DISK_MB;
import static org.apache.aurora.scheduler.resources.ResourceType.RAM_MB;
import static org.apache.aurora.scheduler.updater.Updates.getInstanceIds;

/**
 * Allows access to resource quotas, and tracks quota consumption.
 */
public interface QuotaManager {
  Predicate<ITaskConfig> PROD = ITaskConfig::isProduction;
  Predicate<ITaskConfig> DEDICATED =
      e -> ConfigurationManager.isDedicated(e.getConstraints());
  Predicate<ITaskConfig> PROD_SHARED = and(PROD, not(DEDICATED));
  Predicate<ITaskConfig> PROD_DEDICATED = and(PROD, DEDICATED);
  Predicate<ITaskConfig> NON_PROD_SHARED = and(not(PROD), not(DEDICATED));
  Predicate<ITaskConfig> NON_PROD_DEDICATED = and(not(PROD), DEDICATED);

  EnumSet<ResourceType> QUOTA_RESOURCE_TYPES = EnumSet.of(CPUS, RAM_MB, DISK_MB);

  Function<ITaskConfig, ResourceBag> QUOTA_RESOURCES =
      config -> bagFromResources(getTaskResources(config, QUOTA_RESOURCE_TYPES));

  /**
   * Saves a new quota for the provided role or overrides the existing one.
   *
   * @param role Quota owner.
   * @param quota Quota to save.
   * @param storeProvider A store provider to access quota and other data.
   * @throws QuotaException If provided quota specification is invalid.
   */
  void saveQuota(
      String role,
      IResourceAggregate quota,
      MutableStoreProvider storeProvider) throws QuotaException;

  /**
   * Gets {@code QuotaInfo} for the specified role.
   *
   * @param role Quota owner.
   * @param storeProvider A store provider to access quota data.
   * @return quota usage information for the given role.
   */
  QuotaInfo getQuotaInfo(String role, StoreProvider storeProvider);

  /**
   * Checks if there is enough resource quota available for adding {@code instances} of
   * {@code template} tasks provided resources consumed by {@code releasedTemplates} tasks
   * are released. The quota is defined at the task owner (role) level.
   *
   * @param template Task resource requirement.
   * @param instances Number of additional instances requested.
   * @param storeProvider A store provider to access quota data.
   * @return quota check result details.
   */
  QuotaCheckResult checkInstanceAddition(
      ITaskConfig template,
      int instances,
      StoreProvider storeProvider);

  /**
   * Checks if there is enough resource quota available for performing a job update represented
   * by the {@code jobUpdate}. The quota is defined at the task owner (role) level.
   *
   * @param jobUpdate Job update to check quota for.
   * @param storeProvider A store provider to access quota data.
   * @return quota check result details.
   */
  QuotaCheckResult checkJobUpdate(IJobUpdate jobUpdate, StoreProvider storeProvider);

  /**
   * Check if there is enough resource quota available for creating or updating a cron job
   * represented by the {@code cronConfig}. The quota is defined at the task owner (role) level.
   *
   * @param cronConfig Cron job configuration.
   * @param storeProvider A store provider to access quota data.
   * @return quota check result details.
   */
  QuotaCheckResult checkCronUpdate(IJobConfiguration cronConfig, StoreProvider storeProvider);

  /**
   * Thrown when quota related operation failed.
   */
  class QuotaException extends Exception {
    public QuotaException(String msg) {
      super(msg);
    }
  }

  /**
   * Quota provider that stores quotas in the canonical store.
   */
  class QuotaManagerImpl implements QuotaManager {
    private static final Predicate<ITaskConfig> NO_QUOTA_CHECK = or(PROD_DEDICATED, not(PROD));

    @Override
    public void saveQuota(
        final String ownerRole,
        final IResourceAggregate quota,
        MutableStoreProvider storeProvider) throws QuotaException {

      for (IResource resource : quota.getResources()) {
        if (resource.isSetNumCpus() && resource.getNumCpus() < 0.0
            || resource.isSetRamMb() && resource.getRamMb() < 0
            || resource.isSetDiskMb() && resource.getDiskMb() < 0) {
          throw new QuotaException("Negative values in: " + quota.toString());
        }
      }

      QuotaInfo info = getQuotaInfo(ownerRole, Optional.empty(), storeProvider);
      ResourceBag prodConsumption = info.getProdSharedConsumption();
      ResourceBag overage = bagFromAggregate(quota).subtract(prodConsumption);
      if (!overage.filter(IS_NEGATIVE).getResourceVectors().isEmpty()) {
        throw new QuotaException(String.format(
            "Quota: %s is less then current prod reservation: %s",
            quota.toString(),
            prodConsumption.toString()));
      }

      storeProvider.getQuotaStore().saveQuota(ownerRole, quota);
    }

    @Override
    public QuotaInfo getQuotaInfo(String role, StoreProvider storeProvider) {
      return getQuotaInfo(role, Optional.empty(), storeProvider);
    }

    @Override
    public QuotaCheckResult checkInstanceAddition(
        ITaskConfig template,
        int instances,
        StoreProvider storeProvider) {

      Preconditions.checkArgument(instances >= 0);
      if (NO_QUOTA_CHECK.apply(template)) {
        return new QuotaCheckResult(SUFFICIENT_QUOTA);
      }

      QuotaInfo quotaInfo = getQuotaInfo(template.getJob().getRole(), storeProvider);
      ResourceBag requestedTotal =
          quotaInfo.getProdSharedConsumption().add(scale(template, instances));

      return QuotaCheckResult.greaterOrEqual(quotaInfo.getQuota(), requestedTotal);
    }

    @Override
    public QuotaCheckResult checkJobUpdate(IJobUpdate jobUpdate, StoreProvider storeProvider) {
      requireNonNull(jobUpdate);
      if (!jobUpdate.getInstructions().isSetDesiredState()
          || NO_QUOTA_CHECK.apply(jobUpdate.getInstructions().getDesiredState().getTask())) {

        return new QuotaCheckResult(SUFFICIENT_QUOTA);
      }

      QuotaInfo quotaInfo = getQuotaInfo(
          jobUpdate.getSummary().getKey().getJob().getRole(),
          Optional.of(jobUpdate),
          storeProvider);

      return QuotaCheckResult.greaterOrEqual(
          quotaInfo.getQuota(),
          quotaInfo.getProdSharedConsumption());
    }

    @Override
    public QuotaCheckResult checkCronUpdate(
        IJobConfiguration cronConfig,
        StoreProvider storeProvider) {

      if (!cronConfig.getTaskConfig().isProduction()) {
        return new QuotaCheckResult(SUFFICIENT_QUOTA);
      }

      QuotaInfo quotaInfo =
          getQuotaInfo(cronConfig.getKey().getRole(), Optional.empty(), storeProvider);

      Optional<IJobConfiguration> oldCron =
          storeProvider.getCronJobStore().fetchJob(cronConfig.getKey());

      ResourceBag oldResource = oldCron.isPresent() ? scale(oldCron.get()) : EMPTY;

      // Calculate requested total as a sum of current prod consumption and a delta between
      // new and old cron templates.
      ResourceBag requestedTotal =
          quotaInfo.getProdSharedConsumption().add(scale(cronConfig).subtract(oldResource));

      return QuotaCheckResult.greaterOrEqual(quotaInfo.getQuota(), requestedTotal);
    }

    /**
     * Gets QuotaInfo with currently allocated quota and actual consumption data.
     * <p>
     * In case an optional {@code requestedUpdate} is specified, the consumption returned also
     * includes an estimated resources share of that update as if it was already in progress.
     *
     * @param role Role to get quota info for.
     * @param requestedUpdate An optional {@code IJobUpdate} to forecast the consumption.
     * @param storeProvider A store provider to access quota data.
     * @return {@code QuotaInfo} with quota and consumption details.
     */
    private QuotaInfo getQuotaInfo(
        String role,
        Optional<IJobUpdate> requestedUpdate,
        StoreProvider storeProvider) {

      FluentIterable<IAssignedTask> tasks = FluentIterable
          .from(storeProvider.getTaskStore().fetchTasks(Query.roleScoped(role).active()))
          .transform(IScheduledTask::getAssignedTask);

      // Relies on the invariant of at-most-one active update per job.
      Map<IJobKey, IJobUpdateInstructions> updates = storeProvider.getJobUpdateStore()
          .fetchJobUpdates(updateQuery(role))
          .stream()
          .collect(Collectors.toMap(
              u -> u.getUpdate().getSummary().getKey().getJob(),
              u -> u.getUpdate().getInstructions()));

      // Mix in a requested job update (if present) to correctly calculate consumption.
      // This would be an update that is not saved in the store yet (i.e. the one quota is
      // checked for).
      if (requestedUpdate.isPresent()) {
        updates.put(
            requestedUpdate.get().getSummary().getKey().getJob(),
            requestedUpdate.get().getInstructions());
      }

      Map<IJobKey, IJobConfiguration> cronTemplates =
          FluentIterable.from(storeProvider.getCronJobStore().fetchJobs())
              .filter(compose(equalTo(role), JobKeys::getRole))
              .uniqueIndex(IJobConfiguration::getKey);

      return new QuotaInfo(
          storeProvider.getQuotaStore().fetchQuota(role)
              .map(ResourceManager::bagFromAggregate)
              .orElse(EMPTY),
          getConsumption(tasks, updates, cronTemplates, PROD_SHARED),
          getConsumption(tasks, updates, cronTemplates, PROD_DEDICATED),
          getConsumption(tasks, updates, cronTemplates, NON_PROD_SHARED),
          getConsumption(tasks, updates, cronTemplates, NON_PROD_DEDICATED));
    }

    private ResourceBag getConsumption(
        FluentIterable<IAssignedTask> tasks,
        Map<IJobKey, IJobUpdateInstructions> updatesByKey,
        Map<IJobKey, IJobConfiguration> cronTemplatesByKey,
        Predicate<ITaskConfig> filter) {

      FluentIterable<IAssignedTask> filteredTasks =
          tasks.filter(compose(filter, IAssignedTask::getTask));

      Predicate<IAssignedTask> excludeCron = compose(
          not(in(cronTemplatesByKey.keySet())),
          Tasks::getJob);

      ResourceBag nonCronConsumption = getNonCronConsumption(
          updatesByKey,
          filteredTasks.filter(excludeCron),
          filter);

      ResourceBag cronConsumption = getCronConsumption(
          Iterables.filter(
              cronTemplatesByKey.values(),
              compose(filter, IJobConfiguration::getTaskConfig)),
          filteredTasks.transform(IAssignedTask::getTask));

      return nonCronConsumption.add(cronConsumption);
    }

    private static ResourceBag getNonCronConsumption(
        Map<IJobKey, IJobUpdateInstructions> updatesByKey,
        FluentIterable<IAssignedTask> tasks,
        final Predicate<ITaskConfig> configFilter) {

      // 1. Get all active tasks that belong to jobs without active updates OR unaffected by an
      //    active update working set. An example of the latter would be instances not updated by
      //    the update due to being already in desired state or outside of update range (e.g.
      //    not in JobUpdateInstructions.updateOnlyTheseInstances). Calculate consumed resources
      //    as "nonUpdateConsumption".
      //
      // 2. Calculate consumed resources from instances affected by the active job updates as
      //    "updateConsumption".
      //
      // 3. Add up the two to yield total consumption.

      ResourceBag nonUpdateConsumption = fromTasks(tasks
          .filter(buildNonUpdatingTasksFilter(updatesByKey))
          .transform(IAssignedTask::getTask));

      final Predicate<IInstanceTaskConfig> instanceFilter =
          compose(configFilter, IInstanceTaskConfig::getTask);

      ResourceBag updateConsumption =
          addAll(Iterables.transform(updatesByKey.values(), updateResources(instanceFilter)));

      return nonUpdateConsumption.add(updateConsumption);
    }

    private static ResourceBag getCronConsumption(
        Iterable<IJobConfiguration> cronTemplates,
        FluentIterable<ITaskConfig> tasks) {

      // Calculate the overall cron consumption as MAX between cron template resources and active
      // cron tasks. This is required to account for a case when a running cron task has higher
      // resource requirements than its updated template.
      //
      // While this is the "worst case" calculation that does not account for a possible "staggered"
      // cron scheduling, it's the simplest approach possible given the system constraints (e.g.:
      // lack of enforcement on a cron job run duration).

      final Multimap<IJobKey, ITaskConfig> taskConfigsByKey = tasks.index(ITaskConfig::getJob);
      return addAll(Iterables.transform(
          cronTemplates,
          config ->
              scale(config.getTaskConfig(), config.getInstanceCount())
                  .max(fromTasks(taskConfigsByKey.get(config.getKey())))));
    }

    private static Predicate<IAssignedTask> buildNonUpdatingTasksFilter(
        final Map<IJobKey, IJobUpdateInstructions> roleJobUpdates) {

      return task -> {
        Optional<IJobUpdateInstructions> update = Optional.ofNullable(
            roleJobUpdates.get(task.getTask().getJob()));

        if (update.isPresent()) {
          IJobUpdateInstructions instructions = update.get();
          RangeSet<Integer> initialInstances = getInstanceIds(instructions.getInitialState());
          RangeSet<Integer> desiredInstances = getInstanceIds(instructions.isSetDesiredState()
              ? ImmutableSet.of(instructions.getDesiredState())
              : ImmutableSet.of());

          int instanceId = task.getInstanceId();
          return !initialInstances.contains(instanceId) && !desiredInstances.contains(instanceId);
        }
        return true;
      };
    }

    @VisibleForTesting
    static IJobUpdateQuery updateQuery(String role) {
      return IJobUpdateQuery.build(new JobUpdateQuery()
          .setRole(role)
          .setUpdateStatuses(Updates.ACTIVE_JOB_UPDATE_STATES));
    }

    private static final Function<IInstanceTaskConfig, ResourceBag> INSTANCE_RESOURCES =
        config -> scale(config.getTask(), getUpdateInstanceCount(config.getInstances()));

    private static ResourceBag instructionsToResources(
        Iterable<IInstanceTaskConfig> instructions) {

      return addAll(FluentIterable.from(instructions).transform(INSTANCE_RESOURCES));
    }

    /**
     * Calculates max aggregate resources consumed by the job update
     * {@code instructions}. The max is calculated between existing and desired task configs on per
     * resource basis. This means max CPU, RAM and DISK values are computed individually and may
     * come from different task configurations. While it may not be the most accurate
     * representation of job update resources during the update, it does guarantee none of the
     * individual resource values is exceeded during the forward/back roll.
     * <p/>
     * NOTE: In case of a job update converting the job production bit (i.e. prod -> non-prod or
     *       non-prod -> prod), only the matching state is counted towards consumption. For example,
     *       prod -> non-prod AND {@code prodSharedConsumption=True}: only the initial state
     *       is accounted.
     */
    private static Function<IJobUpdateInstructions, ResourceBag> updateResources(
        final Predicate<IInstanceTaskConfig> instanceFilter) {

      return instructions -> {
        Iterable<IInstanceTaskConfig> initialState =
            Iterables.filter(instructions.getInitialState(), instanceFilter);
        Iterable<IInstanceTaskConfig> desiredState = Iterables.filter(
            Optional.ofNullable(instructions.getDesiredState())
                .map(ImmutableSet::of)
                .orElse(ImmutableSet.of()),
            instanceFilter);

        // Calculate result as max(existing, desired) per resource type.
        return instructionsToResources(initialState).max(instructionsToResources(desiredState));
      };
    }

    private static ResourceBag addAll(Iterable<ResourceBag> aggregates) {
      return StreamSupport.stream(aggregates.spliterator(), false)
          .reduce((l, r) -> l.add(r))
          .orElse(EMPTY);
    }

    private static ResourceBag scale(ITaskConfig taskConfig, int instanceCount) {
      return QUOTA_RESOURCES.apply(taskConfig).scale(instanceCount);
    }

    private static ResourceBag scale(IJobConfiguration jobConfiguration) {
      return scale(jobConfiguration.getTaskConfig(), jobConfiguration.getInstanceCount());
    }

    private static ResourceBag fromTasks(Iterable<ITaskConfig> tasks) {
      return addAll(Iterables.transform(tasks, QUOTA_RESOURCES));
    }

    private static int getUpdateInstanceCount(Set<IRange> ranges) {
      int instanceCount = 0;
      for (IRange range : ranges) {
        instanceCount += range.getLast() - range.getFirst() + 1;
      }

      return instanceCount;
    }
  }
}