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.io.Serializable;
  import java.util.Optional;
  import java.util.Set;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.DiscreteDomain;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.Ordering;
  import com.google.common.collect.Sets;
  import com.google.inject.Inject;
  
  import org.apache.aurora.common.quantity.Amount;
  import org.apache.aurora.common.quantity.Time;
  import org.apache.aurora.common.util.Clock;
  import org.apache.aurora.gen.JobUpdateStatus;
  import org.apache.aurora.scheduler.storage.entities.IInstanceTaskConfig;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateInstructions;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateSettings;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
  import org.apache.aurora.scheduler.updater.strategy.BatchStrategy;
  import org.apache.aurora.scheduler.updater.strategy.QueueStrategy;
  import org.apache.aurora.scheduler.updater.strategy.UpdateStrategy;
  import org.apache.aurora.scheduler.updater.strategy.VariableBatchStrategy;
  
  import static java.util.Objects.requireNonNull;
  
  import static com.google.common.base.Preconditions.checkArgument;
  
  import static org.apache.aurora.scheduler.updater.Updates.getConfig;
  
  /**
   * A factory that produces job updaters based on a job update configuration.
   * <p>
   * TODO(wfarner): Use AssistedInject to inject this (github.com/google/guice/wiki/AssistedInject).
   */
  interface UpdateFactory {
  
    /**
     * Creates a one-way job updater that will execute the job update configuration in the direction
     * specified by {@code rollingForward}.
     *
     * @param configuration Configuration to act on.
     * @param rollingForward {@code true} if this is a job update, {@code false} if it is a rollback.
     * @return An updater that will execute the job update as specified in the
     *         {@code configuration}.
     */
    Update newUpdate(
        IJobUpdateInstructions configuration,
        boolean rollingForward);
  
    class UpdateFactoryImpl implements UpdateFactory {
      private final Clock clock;
  
      @Inject
      UpdateFactoryImpl(Clock clock) {
        this.clock = requireNonNull(clock);
      }
  
      @Override
      public Update newUpdate(IJobUpdateInstructions instructions, boolean rollingForward) {
        requireNonNull(instructions);
        IJobUpdateSettings settings = instructions.getSettings();
  
        checkArgument(
            settings.getMinWaitInInstanceRunningMs() >= 0,
            "Min wait in running must be non-negative.");
  
        if (settings.getUpdateStrategy().isSetBatchStrategy()) {
          checkArgument(
              settings.getUpdateStrategy().getBatchStrategy().getGroupSize() > 0,
              GROUP_SIZES_INVALID);
        } else if (settings.getUpdateStrategy().isSetVarBatchStrategy()) {
          checkArgument(
              settings.getUpdateStrategy().
                  getVarBatchStrategy().
                  getGroupSizes().
                  stream().
                  reduce(0, Integer::sum) > 0,
              GROUP_SIZES_INVALID);
  
        } else {
          checkArgument(
              settings.getUpdateStrategy().getQueueStrategy().getGroupSize() > 0,
              GROUP_SIZES_INVALID);
        }
  
        Set<Integer> currentInstances = expandInstanceIds(instructions.getInitialState());
        Set<Integer> desiredInstances = instructions.isSetDesiredState()
            ? expandInstanceIds(ImmutableSet.of(instructions.getDesiredState()))
            : ImmutableSet.of();
  
        Set<Integer> instances = ImmutableSet.copyOf(Sets.union(currentInstances, desiredInstances));
  
        ImmutableMap.Builder<Integer, StateEvaluator<Optional<IScheduledTask>>> evaluators =
            ImmutableMap.builder();
        for (int instanceId : instances) {
          Optional<ITaskConfig> desiredStateConfig;
          if (rollingForward) {
            desiredStateConfig = desiredInstances.contains(instanceId)
                ? Optional.of(instructions.getDesiredState().getTask())
                : Optional.empty();
          } else {
            desiredStateConfig = getConfig(instanceId, instructions.getInitialState());
          }
  
          evaluators.put(
              instanceId,
              new InstanceUpdater(
                  desiredStateConfig,
                  settings.getMaxPerInstanceFailures(),
                  Amount.of((long) settings.getMinWaitInInstanceRunningMs(), Time.MILLISECONDS),
                  clock));
        }
  
        Ordering<Integer> updateOrdering = new UpdateOrdering(currentInstances, desiredInstances);
        Ordering<Integer> updateOrder = rollingForward
            ? updateOrdering
            : updateOrdering.reverse();
  
        UpdateStrategy<Integer> strategy;
  
        // Note: Verification that the update strategy exists and is valid has already taken
        // place when the scheduler receives the thrift call.
        // TODO(rdelvalle): Consider combining Batch Update and Variable Batch update strategies.
        if (settings.getUpdateStrategy().isSetBatchStrategy()) {
          strategy = new BatchStrategy<>(
              updateOrder,
              settings.getUpdateStrategy().getBatchStrategy().getGroupSize());
        } else if (settings.getUpdateStrategy().isSetVarBatchStrategy()) {
          strategy = new VariableBatchStrategy<>(
              updateOrder,
              settings.getUpdateStrategy().getVarBatchStrategy().getGroupSizes(),
              rollingForward);
        } else {
          strategy = new QueueStrategy<>(
              updateOrder,
              settings.getUpdateStrategy().getQueueStrategy().getGroupSize());
        }
  
        JobUpdateStatus successStatus =
            rollingForward ? JobUpdateStatus.ROLLED_FORWARD : JobUpdateStatus.ROLLED_BACK;
        JobUpdateStatus failureStatus = rollingForward && settings.isRollbackOnFailure()
            ? JobUpdateStatus.ROLLING_BACK
            : JobUpdateStatus.FAILED;
  
        return new Update(
            new OneWayJobUpdater<>(
                strategy,
                settings.getMaxFailedInstances(),
                evaluators.build()),
            successStatus,
            failureStatus);
      }
  
      @VisibleForTesting
      static Set<Integer> expandInstanceIds(Set<IInstanceTaskConfig> instanceGroups) {
        return Updates.getInstanceIds(instanceGroups).asSet(DiscreteDomain.integers());
      }
  
      static final String GROUP_SIZES_INVALID = "Update group size(s) must be positive.";
    }
  
    /**
     * An instance ID ordering that prefers to create new instances first, then update existing
     * instances, and finally kill instances.
     */
    @VisibleForTesting
    class UpdateOrdering extends Ordering<Integer> implements Serializable {
      /**
       * Associates an instance ID to an action (create, update, or kill) priority.
       */
      private final ImmutableMap<Integer, Integer> instanceToActionPriority;
  
      /**
       * Creates an {@link UpdateOrdering}. Determines the action of the instance (create, update, or
       * kill) by comparing the current instance IDs against the desired instance IDs after the
       * update.
       *
       * @param currentInstances The current instance IDs.
       * @param desiredInstances The desired instance IDs after the update.
       */
      UpdateOrdering(Set<Integer> currentInstances, Set<Integer> desiredInstances) {
        requireNonNull(desiredInstances);
        requireNonNull(currentInstances);
  
        Set<Integer> toCreate = Sets.difference(desiredInstances, currentInstances);
        Set<Integer> toUpdate = Sets.intersection(desiredInstances, currentInstances);
        Set<Integer> toKill = Sets.difference(currentInstances, desiredInstances);
  
        // Build a mapping of ordering priority (lower is more important) to the instance action
        // group. Then, we invert it for easy lookup of instance ID to priority.
        ImmutableMap.Builder<Integer, Integer> builder = new ImmutableMap.Builder<>();
        ImmutableMap.of(
            1, toCreate,
            2, toUpdate,
            3, toKill
        ).forEach((priority, instances) -> instances.forEach(id -> builder.put(id, priority)));
        this.instanceToActionPriority = builder.build();
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public int compare(Integer a, Integer b) {
        Integer aActionPriority = instanceToActionPriority.get(a);
        Integer bActionPriority = instanceToActionPriority.get(b);
  
        // Try to order by the instance's action.
        if (!aActionPriority.equals(bActionPriority)) {
          return Integer.compare(aActionPriority, bActionPriority);
        }
  
        // If it is the same action, order the IDs numerically.
        return Integer.compare(a, b);
      }
    }
  
    class Update {
      private final OneWayJobUpdater<Integer, Optional<IScheduledTask>> updater;
      private final JobUpdateStatus successStatus;
      private final JobUpdateStatus failureStatus;
  
      Update(
          OneWayJobUpdater<Integer, Optional<IScheduledTask>> updater,
          JobUpdateStatus successStatus,
          JobUpdateStatus failureStatus) {
  
        this.updater = requireNonNull(updater);
        this.successStatus = requireNonNull(successStatus);
        this.failureStatus = requireNonNull(failureStatus);
      }
  
      OneWayJobUpdater<Integer, Optional<IScheduledTask>> getUpdater() {
        return updater;
      }
  
      JobUpdateStatus getSuccessStatus() {
        return successStatus;
      }
  
      JobUpdateStatus getFailureStatus() {
        return failureStatus;
      }
    }
  }
  