OneWayJobUpdater.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.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicates;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;

import org.apache.aurora.common.util.StateMachine;
import org.apache.aurora.scheduler.updater.strategy.UpdateStrategy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

import static org.apache.aurora.scheduler.updater.SideEffect.InstanceUpdateStatus.FAILED;
import static org.apache.aurora.scheduler.updater.SideEffect.InstanceUpdateStatus.IDLE;
import static org.apache.aurora.scheduler.updater.SideEffect.InstanceUpdateStatus.SUCCEEDED;
import static org.apache.aurora.scheduler.updater.SideEffect.InstanceUpdateStatus.WORKING;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result;

/**
 * Controller for a one-way job update (i.e. no rollbacks).  The controller will coordinate updates
 * of all instances within the job, and roll up the results of the individual updates into the
 * result of the job update.
 *
 * @param <K> Type used to uniquely identify instances.
 * @param <T> Instance data type.
 */
class OneWayJobUpdater<K, T> {
  private static final Logger LOG = LoggerFactory.getLogger(OneWayJobUpdater.class);

  private final UpdateStrategy<K> strategy;
  private final int maxFailedInstances;
  private final Map<K, InstanceUpdate<T>> instances;
  private final StateMachine<OneWayStatus> stateMachine =
      StateMachine.<OneWayStatus>builder("job_update")
          .initialState(OneWayStatus.IDLE)
          .addState(OneWayStatus.IDLE, OneWayStatus.WORKING)
          .addState(OneWayStatus.WORKING, OneWayStatus.SUCCEEDED, OneWayStatus.FAILED)
          .addState(OneWayStatus.SUCCEEDED, OneWayStatus.SUCCEEDED)
          .addState(OneWayStatus.FAILED, OneWayStatus.FAILED)
          .throwOnBadTransition(true)
          .build();

  /**
   * Creates a new one-way updater.
   *
   * @param strategy The strategy to decide which instances to update after a state change.
   * @param maxFailedInstances Maximum tolerated failures before the update is considered failed.
   * @param instanceEvaluators Evaluate the state of individual instances, and decide what actions
   *                           must be taken to update them.
   */
  OneWayJobUpdater(
      UpdateStrategy<K> strategy,
      int maxFailedInstances,
      Map<K, StateEvaluator<T>> instanceEvaluators) {

    this.strategy = requireNonNull(strategy);
    this.maxFailedInstances = maxFailedInstances;
    requireNonNull(instanceEvaluators);

    this.instances = ImmutableMap.copyOf(Maps.transformEntries(
        instanceEvaluators,
        (key, value) -> new InstanceUpdate<>("Instance " + key, value)));
  }

  private static final Function<InstanceUpdate<?>, SideEffect.InstanceUpdateStatus> GET_STATE =
      InstanceUpdate::getState;

  private static <K, T> Set<K> filterByStatus(
      Map<K, InstanceUpdate<T>> instances,
      SideEffect.InstanceUpdateStatus status) {

    return ImmutableSet.copyOf(
        Maps.filterValues(instances, Predicates.compose(Predicates.equalTo(status), GET_STATE))
            .keySet());
  }

  @VisibleForTesting
  Set<K> getInstances() {
    return ImmutableSet.copyOf(instances.keySet());
  }

  /**
   * Checks whether an instance is in scope for this update.
   *
   * @param instanceId Instance id to check.
   * @return {@code true} if the instance is part of the update, as originally specified by
   *         {@code instanceEvaluators}.
   */
  boolean containsInstance(K instanceId) {
    return instances.containsKey(instanceId);
  }

  /**
   * Performs an evaluation of the job.  An evaluation would normally be triggered to initiate the
   * update, as a result of a state change relevant to the update, or due to a
   * {@link InstanceAction#WATCH_TASK requested} instance re-evaluation.
   *
   * @param instancesNeedingUpdate Instances triggering the event, if any.
   * @param stateProvider Provider to fetch state of instances, and pass to
   *                      {@link StateEvaluator#evaluate(Object)}.
   * @return The outcome of the evaluation, including the state of the job update and actions the
   *         caller should perform on individual instances.
   * @throws IllegalStateException if the job updater is not currently
   *         {@link OneWayStatus#WORKING working} state, as indicated by a previous evaluation.
   */
  synchronized EvaluationResult<K> evaluate(
      Map<K, T> instancesNeedingUpdate,
      InstanceStateProvider<K, T> stateProvider) {

    if (stateMachine.getState() == OneWayStatus.IDLE) {
      stateMachine.transition(OneWayStatus.WORKING);
    }
    Preconditions.checkState(
        stateMachine.getState() == OneWayStatus.WORKING,
        "Attempted to evaluate an inactive job updater.");

    // Call order is important here: update on-demand instances, evaluate new instances, compute
    // job update state.
    ImmutableMap.Builder<K, SideEffect> actions = ImmutableMap.<K, SideEffect>builder()
        // Re-evaluate instances that are in need of update.
        .putAll(evaluateInstances(instancesNeedingUpdate));

    if (computeJobUpdateStatus() == OneWayStatus.WORKING) {
      // If ready to begin updating more instances, evaluate those as well.
      actions.putAll(startNextInstanceGroup(stateProvider));
    }

    return new EvaluationResult<>(computeJobUpdateStatus(), actions.build());
  }

  private Map<K, SideEffect> evaluateInstances(Map<K, T> updatedInstances) {
    ImmutableMap.Builder<K, SideEffect> sideEffects = ImmutableMap.builder();
    for (Map.Entry<K, T> entry : updatedInstances.entrySet()) {
      K instanceId = entry.getKey();
      InstanceUpdate<T> update = instances.get(instanceId);
      // Suppress state changes for updates that are not in-progress.
      if (update.getState() == WORKING) {
        sideEffects.put(instanceId, update.evaluate(entry.getValue()));
      } else {
        LOG.info("Ignoring state change for instance outside working set: " + instanceId);
      }
    }

    return sideEffects.build();
  }

  private Map<K, SideEffect> startNextInstanceGroup(InstanceStateProvider<K, T> stateProvider) {
    Set<K> idle = filterByStatus(instances, IDLE);
    if (idle.isEmpty()) {
      return ImmutableMap.of();
    } else {
      ImmutableMap.Builder<K, SideEffect> builder = ImmutableMap.builder();
      Set<K> working = filterByStatus(instances, WORKING);
      Set<K> nextGroup = strategy.getNextGroup(idle, working);
      if (!nextGroup.isEmpty()) {
        for (K instance : nextGroup) {
          builder.put(instance, instances.get(instance).evaluate(stateProvider.getState(instance)));
        }
        LOG.info("Changed working set for update to "
            + filterByStatus(instances, WORKING));
      }

      Map<K, SideEffect> sideEffects = builder.build();
      if (!idle.isEmpty() && working.isEmpty() && !SideEffect.hasActions(sideEffects.values())) {
        // There's no in-flight instances, and no actions - so there's nothing left to initiate more
        // work on this job. Try to find more work, or converge.
        return builder.putAll(startNextInstanceGroup(stateProvider)).build();
      } else {
        return sideEffects;
      }
    }
  }

  private OneWayStatus computeJobUpdateStatus() {
    Set<K> idle = filterByStatus(instances, IDLE);
    Set<K> working = filterByStatus(instances, WORKING);
    Set<K> failed = filterByStatus(instances, FAILED);
    if (failed.size() > maxFailedInstances) {
      stateMachine.transition(OneWayStatus.FAILED);
    } else if (working.isEmpty() && idle.isEmpty()) {
      stateMachine.transition(OneWayStatus.SUCCEEDED);
    }

    return stateMachine.getState();
  }

  /**
   * Container and state for the update of an individual instance.
   */
  private static class InstanceUpdate<T> {
    private final StateEvaluator<T> evaluator;
    private final StateMachine<SideEffect.InstanceUpdateStatus> stateMachine;

    InstanceUpdate(String name, StateEvaluator<T> evaluator) {
      this.evaluator = requireNonNull(evaluator);
      stateMachine = StateMachine.<SideEffect.InstanceUpdateStatus>builder(name)
          .initialState(IDLE)
          .addState(IDLE, WORKING)
          .addState(WORKING, SUCCEEDED, FAILED)
          .addState(SUCCEEDED)
          .addState(FAILED)
          .throwOnBadTransition(true)
          .logTransitions()
          .build();
    }

    SideEffect.InstanceUpdateStatus getState() {
      return stateMachine.getState();
    }

    SideEffect evaluate(T actualState) {
      ImmutableSet.Builder<SideEffect.InstanceUpdateStatus> statusChanges = ImmutableSet.builder();

      if (stateMachine.getState() == IDLE) {
        stateMachine.transition(WORKING);
        statusChanges.add(WORKING);
      }

      Result result = evaluator.evaluate(actualState);
      if (TERMINAL_RESULT_TO_STATUS.containsKey(result)) {
        SideEffect.InstanceUpdateStatus status = TERMINAL_RESULT_TO_STATUS.get(result);
        stateMachine.transition(status);
        statusChanges.add(status);
      }

      return new SideEffect(result.getAction(), statusChanges.build(), result.getFailure());
    }
  }

  private static final Map<Result, SideEffect.InstanceUpdateStatus> TERMINAL_RESULT_TO_STATUS =
      ImmutableMap.of(
          Result.SUCCEEDED, SUCCEEDED,
          Result.FAILED_TERMINATED, FAILED
  );

  /**
   * Status of the job update.
   */
  enum OneWayStatus {
    IDLE,
    WORKING,
    SUCCEEDED,
    FAILED
  }

  /**
   * Result of an evaluation round.
   */
  static class EvaluationResult<K> {
    private final OneWayStatus status;
    private final Map<K, SideEffect> sideEffects;

    EvaluationResult(OneWayStatus jobStatus, Map<K, SideEffect> sideEffects) {
      this.status = requireNonNull(jobStatus);
      this.sideEffects = requireNonNull(sideEffects);
    }

    public OneWayStatus getStatus() {
      return status;
    }

    public Map<K, SideEffect> getSideEffects() {
      return sideEffects;
    }

    @Override
    public boolean equals(Object obj) {
      if (!(obj instanceof EvaluationResult)) {
        return false;
      }
      @SuppressWarnings("unchecked")
      EvaluationResult<K> other = (EvaluationResult<K>) obj;
      return other.getStatus().equals(this.getStatus())
          && other.getSideEffects().equals(this.getSideEffects());
    }

    @Override
    public int hashCode() {
      return Objects.hash(getStatus(), getSideEffects());
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(this)
          .add("status", getStatus())
          .add("sideEffects", getSideEffects())
          .toString();
    }
  }
}