src/main/java/org/apache/aurora/scheduler/updater/UpdateFactory.java - aurora - Git at Google

 /**
  * 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;
     }
   }
 }
	/**
	* 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;
	}
	}
	}