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.updater;
  
  import java.util.Map;
  import java.util.Objects;
  import java.util.Set;
  
  import com.google.common.base.Function;
  import com.google.common.base.Functions;
  import com.google.common.base.MoreObjects;
  import com.google.common.base.Predicates;
  import com.google.common.collect.ContiguousSet;
  import com.google.common.collect.DiscreteDomain;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.MapDifference;
  import com.google.common.collect.MapDifference.ValueDifference;
  import com.google.common.collect.Maps;
  import com.google.common.collect.Range;
  import com.google.common.collect.Sets;
  
  import org.apache.aurora.scheduler.base.Numbers;
  import org.apache.aurora.scheduler.base.Query;
  import org.apache.aurora.scheduler.base.Tasks;
  import org.apache.aurora.scheduler.storage.TaskStore;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IRange;
  import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * The difference between two states of a job.
   */
  public final class JobDiff {
    private final Map<Integer, ITaskConfig> replacedInstances;
    private final Set<Integer> replacementInstances;
    private final Map<Integer, ITaskConfig> unchangedInstances;
  
    /**
     * Creates a job diff containing the instances to be replaced (original state), and instances
     * replacing them (target state). An instance may be absent in either side, which means that
     * it was absent in that state.
     *
     * @param replacedInstances Instances being replaced.
     * @param replacementInstances Instances to replace {@code replacedInstances}.
     * @param unchangedInstances Instances unchanged by the update.
     */
    public JobDiff(
        Map<Integer, ITaskConfig> replacedInstances,
        Set<Integer> replacementInstances,
        Map<Integer, ITaskConfig> unchangedInstances) {
  
      this.replacedInstances = requireNonNull(replacedInstances);
      this.replacementInstances = requireNonNull(replacementInstances);
      this.unchangedInstances = requireNonNull(unchangedInstances);
    }
  
    public Map<Integer, ITaskConfig> getReplacedInstances() {
      return replacedInstances;
    }
  
    public Set<Integer> getReplacementInstances() {
      return replacementInstances;
    }
  
    public Map<Integer, ITaskConfig> getUnchangedInstances() {
      return unchangedInstances;
    }
  
    /**
     * Gets instances contained in a {@code scope} that are not a part of the job diff.
     *
     * @param scope Scope to search within.
     * @return Instance IDs in {@code scope} that are not in this job diff.
     */
    public Set<Integer> getOutOfScopeInstances(Set<Integer> scope) {
      Set<Integer> allValidInstances = ImmutableSet.<Integer>builder()
          .addAll(getReplacedInstances().keySet())
          .addAll(getReplacementInstances())
          .addAll(getUnchangedInstances().keySet())
          .build();
      return ImmutableSet.copyOf(Sets.difference(scope, allValidInstances));
    }
  
    /**
     * Checks whether this diff contains no work to be done.
     *
     * @return {@code true} if this diff is a no-op, otherwise {@code false}.
     */
    public boolean isNoop() {
      return replacedInstances.isEmpty() && replacementInstances.isEmpty();
    }
  
    @Override
    public int hashCode() {
      return Objects.hash(replacedInstances, replacementInstances, unchangedInstances);
    }
  
    @Override
    public boolean equals(Object o) {
      if (!(o instanceof JobDiff)) {
        return false;
      }
  
      JobDiff other = (JobDiff) o;
      return Objects.equals(getReplacedInstances(), other.getReplacedInstances())
          && Objects.equals(getReplacementInstances(), other.getReplacementInstances())
          && Objects.equals(getUnchangedInstances(), other.getUnchangedInstances());
    }
  
    @Override
    public String toString() {
      return MoreObjects.toStringHelper(this)
          .add("replacedInstances", getReplacedInstances())
          .add("replacementInstances", getReplacementInstances())
          .add("unchangedInstances", getUnchangedInstances())
          .toString();
    }
  
    private static <V> Function<MapDifference.ValueDifference<V>, V> leftValue() {
      return ValueDifference::leftValue;
    }
  
    private static JobDiff computeUnscoped(
        Map<Integer, ITaskConfig> currentState,
        IJobKey job,
        Map<Integer, ITaskConfig> proposedState) {
  
      requireNonNull(job);
      requireNonNull(proposedState);
  
      MapDifference<Integer, ITaskConfig> diff = Maps.difference(currentState, proposedState);
  
      Map<Integer, ITaskConfig> removedInstances = ImmutableMap.<Integer, ITaskConfig>builder()
          .putAll(diff.entriesOnlyOnLeft())
          .putAll(Maps.transformValues(diff.entriesDiffering(), JobDiff.leftValue()))
          .build();
  
      Set<Integer> addedInstances = ImmutableSet.<Integer>builder()
          .addAll(diff.entriesOnlyOnRight().keySet())
          .addAll(diff.entriesDiffering().keySet())
          .build();
  
      return new JobDiff(
          removedInstances,
          addedInstances,
          ImmutableMap.copyOf(diff.entriesInCommon()));
    }
  
    /**
     * Calculates the diff necessary to change the current state of a job to the proposed state.
     *
     * @param taskStore Store to fetch the job's current state from.
     * @param job Job being diffed.
     * @param proposedState Proposed state to move the job towards.
     * @param scope Instances to limit the diff to.
     * @return A diff of the current state compared with {@code proposedState}, within {@code scope}.
     */
    public static JobDiff compute(
        TaskStore taskStore,
        IJobKey job,
        Map<Integer, ITaskConfig> proposedState,
        Set<IRange> scope) {
  
      Map<Integer, ITaskConfig> currentState = ImmutableMap.copyOf(
          Maps.transformValues(
              Maps.uniqueIndex(
                  taskStore.fetchTasks(Query.jobScoped(job).active()),
                  Tasks::getInstanceId),
              Tasks::getConfig));
  
      JobDiff diff = computeUnscoped(currentState, job, proposedState);
      if (scope.isEmpty()) {
        return diff;
      } else {
        Set<Integer> limit = Numbers.rangesToInstanceIds(scope);
        Map<Integer, ITaskConfig> replaced =
            ImmutableMap.copyOf(Maps.filterKeys(diff.getReplacedInstances(), Predicates.in(limit)));
        Set<Integer> replacements =
            ImmutableSet.copyOf(Sets.intersection(diff.getReplacementInstances(), limit));
  
        Set<Integer> unchangedIds = ImmutableSet.copyOf(
            Sets.difference(
                ImmutableSet.copyOf(Sets.difference(currentState.keySet(), replaced.keySet())),
                replacements));
        Map<Integer, ITaskConfig> unchanged =
            ImmutableMap.copyOf(Maps.filterKeys(currentState, Predicates.in(unchangedIds)));
  
        return new JobDiff(replaced, replacements, unchanged);
      }
    }
  
    /**
     * Creates a map of {@code instanceCount} copies of {@code config}.
     *
     * @param config Configuration to generate an instance mapping for.
     * @param instanceCount Number of instances to represent.
     * @return A map of instance IDs (from 0 to {@code instanceCount - 1}) to {@code config}.
     */
    public static Map<Integer, ITaskConfig> asMap(ITaskConfig config, int instanceCount) {
      requireNonNull(config);
  
      Set<Integer> desiredInstances = ContiguousSet.create(
          Range.closedOpen(0, instanceCount),
          DiscreteDomain.integers());
      return ImmutableMap.copyOf(Maps.asMap(desiredInstances, Functions.constant(config)));
    }
  }
  