src/main/java/org/apache/aurora/scheduler/updater/JobDiff.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.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)));
  }
}