content: /**
   * 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.thrift;
  
  import java.util.Collection;
  import java.util.Date;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Optional;
  import java.util.Set;
  import java.util.stream.Collectors;
  
  import javax.annotation.Nullable;
  import javax.inject.Inject;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.base.Function;
  import com.google.common.base.Functions;
  import com.google.common.base.Joiner;
  import com.google.common.base.Predicate;
  import com.google.common.base.Predicates;
  import com.google.common.collect.FluentIterable;
  import com.google.common.collect.HashMultimap;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.Iterables;
  import com.google.common.collect.Lists;
  import com.google.common.collect.Maps;
  import com.google.common.collect.Multimap;
  import com.google.common.collect.Multimaps;
  import com.google.common.collect.Sets;
  
  import org.apache.aurora.GuavaUtils;
  import org.apache.aurora.common.base.MorePreconditions;
  import org.apache.aurora.gen.ConfigGroup;
  import org.apache.aurora.gen.ConfigSummary;
  import org.apache.aurora.gen.ConfigSummaryResult;
  import org.apache.aurora.gen.GetJobUpdateDetailsResult;
  import org.apache.aurora.gen.GetJobUpdateDiffResult;
  import org.apache.aurora.gen.GetJobUpdateSummariesResult;
  import org.apache.aurora.gen.GetJobsResult;
  import org.apache.aurora.gen.GetPendingReasonResult;
  import org.apache.aurora.gen.GetQuotaResult;
  import org.apache.aurora.gen.GetTierConfigResult;
  import org.apache.aurora.gen.JobConfiguration;
  import org.apache.aurora.gen.JobKey;
  import org.apache.aurora.gen.JobSummary;
  import org.apache.aurora.gen.JobSummaryResult;
  import org.apache.aurora.gen.JobUpdateKey;
  import org.apache.aurora.gen.JobUpdateQuery;
  import org.apache.aurora.gen.JobUpdateRequest;
  import org.apache.aurora.gen.PendingReason;
  import org.apache.aurora.gen.PopulateJobResult;
  import org.apache.aurora.gen.ReadOnlyScheduler;
  import org.apache.aurora.gen.Response;
  import org.apache.aurora.gen.Result;
  import org.apache.aurora.gen.RoleSummary;
  import org.apache.aurora.gen.RoleSummaryResult;
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.gen.ScheduleStatusResult;
  import org.apache.aurora.gen.ScheduledTask;
  import org.apache.aurora.gen.TaskConfig;
  import org.apache.aurora.gen.TaskQuery;
  import org.apache.aurora.gen.TierConfig;
  import org.apache.aurora.scheduler.TierManager;
  import org.apache.aurora.scheduler.base.JobKeys;
  import org.apache.aurora.scheduler.base.Jobs;
  import org.apache.aurora.scheduler.base.Query;
  import org.apache.aurora.scheduler.base.TaskGroupKey;
  import org.apache.aurora.scheduler.base.Tasks;
  import org.apache.aurora.scheduler.configuration.ConfigurationManager;
  import org.apache.aurora.scheduler.configuration.ConfigurationManager.TaskDescriptionException;
  import org.apache.aurora.scheduler.configuration.SanitizedConfiguration;
  import org.apache.aurora.scheduler.cron.CronPredictor;
  import org.apache.aurora.scheduler.cron.CrontabEntry;
  import org.apache.aurora.scheduler.filter.SchedulingFilter.Veto;
  import org.apache.aurora.scheduler.metadata.NearestFit;
  import org.apache.aurora.scheduler.quota.QuotaInfo;
  import org.apache.aurora.scheduler.quota.QuotaManager;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
  import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateDetails;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateQuery;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateRequest;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateSummary;
  import org.apache.aurora.scheduler.storage.entities.IRange;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
  import org.apache.aurora.scheduler.updater.JobDiff;
  import org.apache.thrift.TException;
  
  import static java.util.Objects.requireNonNull;
  
  import static org.apache.aurora.gen.ResponseCode.INVALID_REQUEST;
  import static org.apache.aurora.scheduler.base.Numbers.convertRanges;
  import static org.apache.aurora.scheduler.base.Numbers.toRanges;
  import static org.apache.aurora.scheduler.resources.ResourceManager.aggregateFromBag;
  import static org.apache.aurora.scheduler.thrift.Responses.addMessage;
  import static org.apache.aurora.scheduler.thrift.Responses.error;
  import static org.apache.aurora.scheduler.thrift.Responses.invalidRequest;
  import static org.apache.aurora.scheduler.thrift.Responses.ok;
  
  class ReadOnlySchedulerImpl implements ReadOnlyScheduler.Iface {
    private static final Function<Entry<ITaskConfig, Collection<Integer>>, ConfigGroup> TO_GROUP =
        input -> new ConfigGroup()
            .setConfig(input.getKey().newBuilder())
            .setInstances(IRange.toBuildersSet(convertRanges(toRanges(input.getValue()))));
  
    private final ConfigurationManager configurationManager;
    private final Storage storage;
    private final NearestFit nearestFit;
    private final CronPredictor cronPredictor;
    private final QuotaManager quotaManager;
    private final TierManager tierManager;
  
    @Inject
    ReadOnlySchedulerImpl(
        ConfigurationManager configurationManager,
        Storage storage,
        NearestFit nearestFit,
        CronPredictor cronPredictor,
        QuotaManager quotaManager,
        TierManager tierManager) {
  
      this.configurationManager = requireNonNull(configurationManager);
      this.storage = requireNonNull(storage);
      this.nearestFit = requireNonNull(nearestFit);
      this.cronPredictor = requireNonNull(cronPredictor);
      this.quotaManager = requireNonNull(quotaManager);
      this.tierManager = requireNonNull(tierManager);
    }
  
    @Override
    public Response populateJobConfig(JobConfiguration description) {
      requireNonNull(description);
  
      try {
        ITaskConfig populatedTaskConfig = SanitizedConfiguration.fromUnsanitized(
            configurationManager,
            IJobConfiguration.build(description)).getJobConfig().getTaskConfig();
        return ok(Result.populateJobResult(
            new PopulateJobResult().setTaskConfig(populatedTaskConfig.newBuilder())));
      } catch (TaskDescriptionException e) {
        return invalidRequest("Invalid configuration: " + e.getMessage());
      }
    }
  
    // TODO(William Farner): Provide status information about cron jobs here.
    @Override
    public Response getTasksStatus(TaskQuery query) {
      return ok(Result.scheduleStatusResult(
          new ScheduleStatusResult().setTasks(getTasks(query))));
    }
  
    @Override
    public Response getTasksWithoutConfigs(TaskQuery query) {
      List<ScheduledTask> tasks = Lists.transform(
          getTasks(query),
          task -> {
            task.getAssignedTask().getTask().getExecutorConfig().unsetData();
            return task;
          });
  
      return ok(Result.scheduleStatusResult(new ScheduleStatusResult().setTasks(tasks)));
    }
  
    @Override
    public Response getPendingReason(TaskQuery query) throws TException {
      requireNonNull(query);
  
      if (query.isSetSlaveHosts() && !query.getSlaveHosts().isEmpty()) {
        return invalidRequest(
            "SlaveHosts are not supported in " + query.toString());
      }
      if (query.isSetStatuses() && !query.getStatuses().isEmpty()) {
        return invalidRequest(
            "Statuses is not supported in " + query.toString());
      }
  
      // Only PENDING tasks should be considered.
      query.setStatuses(ImmutableSet.of(ScheduleStatus.PENDING));
  
      Set<PendingReason> reasons = FluentIterable.from(getTasks(query))
          .transform(scheduledTask -> {
            TaskGroupKey groupKey = TaskGroupKey.from(
                ITaskConfig.build(scheduledTask.getAssignedTask().getTask()));
  
            String reason = Joiner.on(',').join(Iterables.transform(
                nearestFit.getNearestFit(groupKey),
                Veto::getReason));
  
            return new PendingReason()
                .setTaskId(scheduledTask.getAssignedTask().getTaskId())
                .setReason(reason);
          }).toSet();
  
      return ok(Result.getPendingReasonResult(new GetPendingReasonResult(reasons)));
    }
  
    @Override
    public Response getConfigSummary(JobKey job) throws TException {
      IJobKey jobKey = JobKeys.assertValid(IJobKey.build(job));
  
      Iterable<IAssignedTask> assignedTasks = Iterables.transform(
          Storage.Util.fetchTasks(storage, Query.jobScoped(jobKey).active()),
          IScheduledTask::getAssignedTask);
      Map<Integer, ITaskConfig> tasksByInstance = Maps.transformValues(
          Maps.uniqueIndex(assignedTasks, IAssignedTask::getInstanceId),
          IAssignedTask::getTask);
      Set<ConfigGroup> groups = instancesToConfigGroups(tasksByInstance);
  
      return ok(Result.configSummaryResult(
          new ConfigSummaryResult().setSummary(new ConfigSummary(job, groups))));
    }
  
    @Override
    public Response getRoleSummary() {
      Multimap<String, IJobKey> jobsByRole = storage.read(
          storeProvider ->
              Multimaps.index(storeProvider.getTaskStore().getJobKeys(), IJobKey::getRole));
  
      Multimap<String, IJobKey> cronJobsByRole = Multimaps.index(
          Iterables.transform(Storage.Util.fetchCronJobs(storage), IJobConfiguration::getKey),
          IJobKey::getRole);
  
      Set<RoleSummary> summaries = FluentIterable.from(
          Sets.union(jobsByRole.keySet(), cronJobsByRole.keySet()))
          .transform(role -> new RoleSummary(
              role,
              jobsByRole.get(role).size(),
              cronJobsByRole.get(role).size()))
          .toSet();
  
      return ok(Result.roleSummaryResult(new RoleSummaryResult(summaries)));
    }
  
    @Override
    public Response getJobSummary(@Nullable String maybeNullRole) {
      Optional<String> ownerRole = Optional.ofNullable(maybeNullRole);
  
      Multimap<IJobKey, IScheduledTask> tasks = getTasks(maybeRoleScoped(ownerRole));
      Map<IJobKey, IJobConfiguration> jobs = getJobs(ownerRole, tasks);
  
      Function<IJobKey, JobSummary> makeJobSummary = jobKey -> {
        IJobConfiguration job = jobs.get(jobKey);
        JobSummary summary = new JobSummary()
            .setJob(job.newBuilder())
            .setStats(Jobs.getJobStats(tasks.get(jobKey)).newBuilder());
  
        if (job.isSetCronSchedule()) {
          CrontabEntry crontabEntry = CrontabEntry.parse(job.getCronSchedule());
          Optional<Date> nextRun = cronPredictor.predictNextRun(crontabEntry);
          return nextRun.map(date -> summary.setNextCronRunMs(date.getTime())).orElse(summary);
        } else {
          return summary;
        }
      };
  
      ImmutableSet<JobSummary> jobSummaries =
          FluentIterable.from(jobs.keySet()).transform(makeJobSummary).toSet();
  
      return ok(Result.jobSummaryResult(new JobSummaryResult().setSummaries(jobSummaries)));
    }
  
    @Override
    public Response getJobs(@Nullable String maybeNullRole) {
      Optional<String> ownerRole = Optional.ofNullable(maybeNullRole);
  
      return ok(Result.getJobsResult(
          new GetJobsResult()
              .setConfigs(IJobConfiguration.toBuildersSet(
                  getJobs(ownerRole, getTasks(maybeRoleScoped(ownerRole).active())).values()))));
    }
  
    @Override
    public Response getQuota(String ownerRole) {
      MorePreconditions.checkNotBlank(ownerRole);
      return storage.read(storeProvider -> {
        QuotaInfo quotaInfo = quotaManager.getQuotaInfo(ownerRole, storeProvider);
        GetQuotaResult result = new GetQuotaResult()
            .setQuota(aggregateFromBag(quotaInfo.getQuota()).newBuilder())
            .setProdSharedConsumption(aggregateFromBag(
                quotaInfo.getProdSharedConsumption()).newBuilder())
            .setProdDedicatedConsumption(aggregateFromBag(
                quotaInfo.getProdDedicatedConsumption()).newBuilder())
            .setNonProdSharedConsumption(aggregateFromBag(
                quotaInfo.getNonProdSharedConsumption()).newBuilder())
            .setNonProdDedicatedConsumption(aggregateFromBag(
                quotaInfo.getNonProdDedicatedConsumption()).newBuilder());
  
        return ok(Result.getQuotaResult(result));
      });
    }
  
    @Override
    public Response getJobUpdateSummaries(JobUpdateQuery mutableQuery) {
      IJobUpdateQuery query = IJobUpdateQuery.build(requireNonNull(mutableQuery));
  
      List<IJobUpdateSummary> summaries = storage.read(
          storeProvider -> storeProvider.getJobUpdateStore()
              .fetchJobUpdates(query)
              .stream()
              .map(u -> u.getUpdate().getSummary()).collect(Collectors.toList()));
  
      return ok(Result.getJobUpdateSummariesResult(new GetJobUpdateSummariesResult()
          .setUpdateSummaries(IJobUpdateSummary.toBuildersList(summaries))));
    }
  
    @Override
    public Response getJobUpdateDetails(JobUpdateKey mutableKey, JobUpdateQuery mutableQuery) {
      if (mutableKey == null && mutableQuery == null)  {
        return error("Either key or query must be set.");
      }
  
      if (mutableQuery != null) {
        IJobUpdateQuery query = IJobUpdateQuery.build(mutableQuery);
  
        List<IJobUpdateDetails> details =
            storage.read(storeProvider -> storeProvider.getJobUpdateStore().fetchJobUpdates(query));
  
        return ok(Result.getJobUpdateDetailsResult(new GetJobUpdateDetailsResult()
            .setDetailsList(IJobUpdateDetails.toBuildersList(details))));
      }
  
      // TODO(zmanji): Remove this code once `mutableKey` is removed in AURORA-1765
      IJobUpdateKey key = IJobUpdateKey.build(mutableKey);
      Optional<IJobUpdateDetails> details = storage.read(storeProvider ->
          storeProvider.getJobUpdateStore().fetchJobUpdate(key));
  
      if (details.isPresent()) {
        return addMessage(ok(Result.getJobUpdateDetailsResult(
            new GetJobUpdateDetailsResult().setDetails(details.get().newBuilder()))),
            "The key argument is deprecated, use the query argument instead");
      } else {
        return invalidRequest("Invalid update: " + key);
      }
    }
  
    @Override
    public Response getJobUpdateDiff(JobUpdateRequest mutableRequest) {
      IJobUpdateRequest request;
      try {
        request = IJobUpdateRequest.build(
            new JobUpdateRequest(mutableRequest)
                .setTaskConfig(configurationManager
                    .validateAndPopulate(
                        ITaskConfig.build(mutableRequest.getTaskConfig()),
                        mutableRequest.getInstanceCount())
                    .newBuilder()));
      } catch (TaskDescriptionException e) {
        return error(INVALID_REQUEST, e);
      }
  
      IJobKey job = request.getTaskConfig().getJob();
  
      return storage.read(storeProvider -> {
        if (storeProvider.getCronJobStore().fetchJob(job).isPresent()) {
          return invalidRequest(NO_CRON);
        }
  
        JobDiff diff = JobDiff.compute(
            storeProvider.getTaskStore(),
            job,
            JobDiff.asMap(request.getTaskConfig(), request.getInstanceCount()),
            request.getSettings().getUpdateOnlyTheseInstances());
  
        Map<Integer, ITaskConfig> replaced = diff.getReplacedInstances();
        Map<Integer, ITaskConfig> replacements = Maps.asMap(
            diff.getReplacementInstances(),
            Functions.constant(request.getTaskConfig()));
  
        Map<Integer, ITaskConfig> add = Maps.filterKeys(
            replacements,
            Predicates.in(Sets.difference(replacements.keySet(), replaced.keySet())));
        Map<Integer, ITaskConfig> remove = Maps.filterKeys(
            replaced,
            Predicates.in(Sets.difference(replaced.keySet(), replacements.keySet())));
        Map<Integer, ITaskConfig> update = Maps.filterKeys(
            replaced,
            Predicates.in(Sets.intersection(replaced.keySet(), replacements.keySet())));
  
        return ok(Result.getJobUpdateDiffResult(new GetJobUpdateDiffResult()
            .setAdd(instancesToConfigGroups(add))
            .setRemove(instancesToConfigGroups(remove))
            .setUpdate(instancesToConfigGroups(update))
            .setUnchanged(instancesToConfigGroups(diff.getUnchangedInstances()))));
      });
    }
  
    @Override
    public Response getTierConfigs() throws TException {
      return ok(Result.getTierConfigResult(
          new GetTierConfigResult(
              tierManager.getDefaultTierName(),
              tierManager.getTiers()
                  .entrySet()
                  .stream()
                  .map(entry -> new TierConfig(entry.getKey(), entry.getValue().toMap()))
                  .collect(GuavaUtils.toImmutableSet()))));
    }
  
    private static Set<ConfigGroup> instancesToConfigGroups(Map<Integer, ITaskConfig> tasks) {
      Multimap<ITaskConfig, Integer> instancesByDetails = Multimaps.invertFrom(
          Multimaps.forMap(tasks),
          HashMultimap.create());
      return ImmutableSet.copyOf(
          Iterables.transform(instancesByDetails.asMap().entrySet(), TO_GROUP));
    }
  
    private List<ScheduledTask> getTasks(TaskQuery query) {
      requireNonNull(query);
  
      Iterable<IScheduledTask> tasks = Storage.Util.fetchTasks(storage, Query.arbitrary(query));
      if (query.getOffset() > 0) {
        tasks = Iterables.skip(tasks, query.getOffset());
      }
      if (query.getLimit() > 0) {
        tasks = Iterables.limit(tasks, query.getLimit());
      }
  
      return IScheduledTask.toBuildersList(tasks);
    }
  
    private Query.Builder maybeRoleScoped(Optional<String> ownerRole) {
      return ownerRole.isPresent()
          ? Query.roleScoped(ownerRole.get())
          : Query.unscoped();
    }
  
    private Map<IJobKey, IJobConfiguration> getJobs(
        Optional<String> ownerRole,
        Multimap<IJobKey, IScheduledTask> tasks) {
  
      // We need to synthesize the JobConfiguration from the the current tasks because the
      // ImmediateJobManager doesn't store jobs directly and ImmediateJobManager#getJobs always
      // returns an empty Collection.
      Map<IJobKey, IJobConfiguration> jobs = Maps.newHashMap();
  
      jobs.putAll(Maps.transformEntries(tasks.asMap(),
          (jobKey, tasks1) -> {
  
            // Pick the latest transitioned task for each immediate job since the job can be in the
            // middle of an update or some shards have been selectively created.
            TaskConfig mostRecentTaskConfig =
                Tasks.getLatestActiveTask(tasks1).getAssignedTask().getTask().newBuilder();
  
            return IJobConfiguration.build(new JobConfiguration()
                .setKey(jobKey.newBuilder())
                .setOwner(mostRecentTaskConfig.getOwner())
                .setTaskConfig(mostRecentTaskConfig)
                .setInstanceCount(tasks1.size()));
          }));
  
      // Get cron jobs directly from the manager. Do this after querying the task store so the real
      // template JobConfiguration for a cron job will overwrite the synthesized one that could have
      // been created above.
      Predicate<IJobConfiguration> configFilter = ownerRole.isPresent()
          ? Predicates.compose(Predicates.equalTo(ownerRole.get()), JobKeys::getRole)
          : Predicates.alwaysTrue();
      jobs.putAll(Maps.uniqueIndex(
          FluentIterable.from(Storage.Util.fetchCronJobs(storage)).filter(configFilter),
          IJobConfiguration::getKey));
  
      return jobs;
    }
  
    private Multimap<IJobKey, IScheduledTask> getTasks(Query.Builder query) {
      return Tasks.byJobKey(Storage.Util.fetchTasks(storage, query));
    }
  
    @VisibleForTesting
    static final String NO_CRON = "Cron jobs are not supported.";
  }
  