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.Arrays;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.stream.Collectors;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.base.Function;
  import com.google.common.base.Functions;
  import com.google.common.cache.CacheBuilder;
  import com.google.common.cache.CacheLoader;
  import com.google.common.cache.LoadingCache;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.Iterables;
  import com.google.common.collect.Maps;
  import com.google.common.collect.Ordering;
  import com.google.inject.Inject;
  
  import org.apache.aurora.common.application.Lifecycle;
  import org.apache.aurora.common.collections.Pair;
  import org.apache.aurora.common.quantity.Amount;
  import org.apache.aurora.common.quantity.Time;
  import org.apache.aurora.common.stats.StatsProvider;
  import org.apache.aurora.common.util.Clock;
  import org.apache.aurora.gen.JobInstanceUpdateEvent;
  import org.apache.aurora.gen.JobUpdateAction;
  import org.apache.aurora.gen.JobUpdateEvent;
  import org.apache.aurora.gen.JobUpdatePulseStatus;
  import org.apache.aurora.gen.JobUpdateQuery;
  import org.apache.aurora.gen.JobUpdateStatus;
  import org.apache.aurora.scheduler.BatchWorker;
  import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
  import org.apache.aurora.scheduler.base.InstanceKeys;
  import org.apache.aurora.scheduler.base.JobKeys;
  import org.apache.aurora.scheduler.base.Query;
  import org.apache.aurora.scheduler.state.StateManager;
  import org.apache.aurora.scheduler.storage.JobUpdateStore;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
  import org.apache.aurora.scheduler.storage.TaskStore;
  import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
  import org.apache.aurora.scheduler.storage.entities.IJobInstanceUpdateEvent;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateDetails;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateEvent;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateInstructions;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateQuery;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateStrategy;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateSummary;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.aurora.scheduler.updater.StateEvaluator.Failure;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  import static com.google.common.base.Preconditions.checkState;
  
  import static org.apache.aurora.gen.JobUpdateStatus.ABORTED;
  import static org.apache.aurora.gen.JobUpdateStatus.ERROR;
  import static org.apache.aurora.gen.JobUpdateStatus.ROLLING_BACK;
  import static org.apache.aurora.gen.JobUpdateStatus.ROLLING_FORWARD;
  import static org.apache.aurora.gen.JobUpdateStatus.ROLL_FORWARD_AWAITING_PULSE;
  import static org.apache.aurora.scheduler.base.AsyncUtil.shutdownOnError;
  import static org.apache.aurora.scheduler.base.Jobs.AWAITING_PULSE_STATES;
  import static org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
  import static org.apache.aurora.scheduler.storage.Util.jobUpdateActionStatName;
  import static org.apache.aurora.scheduler.storage.Util.jobUpdateStatusStatName;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.ACTIVE_QUERY;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.AUTO_RESUME_STATES;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.GET_ACTIVE_RESUME_STATE;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.GET_BLOCKED_RESUME_STATE;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.GET_PAUSE_STATE;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.GET_UNBLOCKED_STATE;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.MonitorAction;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.MonitorAction.ROLL_BACK;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.MonitorAction.ROLL_FORWARD;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.MonitorAction.STOP_WATCHING;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.assertTransitionAllowed;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.getBlockedState;
  import static org.apache.aurora.scheduler.updater.JobUpdateStateMachine.getPausedState;
  import static org.apache.aurora.scheduler.updater.OneWayJobUpdater.EvaluationResult;
  import static org.apache.aurora.scheduler.updater.OneWayJobUpdater.OneWayStatus;
  import static org.apache.aurora.scheduler.updater.OneWayJobUpdater.OneWayStatus.SUCCEEDED;
  import static org.apache.aurora.scheduler.updater.SideEffect.InstanceUpdateStatus;
  
  /**
   * Implementation of an updater that orchestrates the process of gradually updating the
   * configuration of tasks in a job.
   * <p>
   * TODO(wfarner): Consider using AbstractIdleService here.
   */
  class JobUpdateControllerImpl implements JobUpdateController {
    private static final Logger LOG = LoggerFactory.getLogger(JobUpdateControllerImpl.class);
    private static final String FATAL_ERROR_FORMAT =
        "Unexpected problem running asynchronous updater for: %s. Triggering shutdown";
    private static final String UPDATE_AUTO_PAUSED = "Update auto paused";
  
    private final UpdateFactory updateFactory;
    private final Storage storage;
    private final ScheduledExecutorService executor;
    private final StateManager stateManager;
    private final Clock clock;
    private final PulseHandler pulseHandler;
    private final Lifecycle lifecycle;
    private final TaskEventBatchWorker taskEventBatchWorker;
    private final UpdateAgentReserver updateAgentReserver;
    private final SlaKillController slaKillController;
  
    // Currently-active updaters. An active updater is one that is rolling forward or back. Paused
    // and completed updates are represented only in storage, not here.
    private final Map<IJobKey, UpdateFactory.Update> updates =
        Collections.synchronizedMap(Maps.newHashMap());
  
    // Used only for updates that have auto pause enabled. Keeps track of what instances
    // have already been seen by the updater in order to detect when a new batch is started.
    private final Map<IJobUpdateKey, Set<Integer>> instancesSeen = new ConcurrentHashMap<>();
  
    private final LoadingCache<JobUpdateStatus, AtomicLong> jobUpdateEventStats;
    private final LoadingCache<JobUpdateAction, AtomicLong> jobUpdateActionStats;
  
    @Inject
    JobUpdateControllerImpl(
        UpdateFactory updateFactory,
        Storage storage,
        ScheduledExecutorService executor,
        StateManager stateManager,
        UpdateAgentReserver updateAgentReserver,
        Clock clock,
        Lifecycle lifecycle,
        TaskEventBatchWorker taskEventBatchWorker,
        StatsProvider statsProvider,
        SlaKillController slaKillController) {
  
      this.updateFactory = requireNonNull(updateFactory);
      this.storage = requireNonNull(storage);
      this.executor = requireNonNull(executor);
      this.stateManager = requireNonNull(stateManager);
      this.clock = requireNonNull(clock);
      this.lifecycle = requireNonNull(lifecycle);
      this.taskEventBatchWorker = requireNonNull(taskEventBatchWorker);
      this.pulseHandler = new PulseHandler(clock);
      this.updateAgentReserver = requireNonNull(updateAgentReserver);
      this.slaKillController = requireNonNull(slaKillController);
  
      this.jobUpdateEventStats = CacheBuilder.newBuilder()
          .build(new CacheLoader<JobUpdateStatus, AtomicLong>() {
            @Override
            public AtomicLong load(JobUpdateStatus status) {
              return statsProvider.makeCounter(jobUpdateStatusStatName(status));
            }
          });
      Arrays.stream(JobUpdateStatus.values())
          .forEach(status -> jobUpdateEventStats.getUnchecked(status).get());
  
      this.jobUpdateActionStats = CacheBuilder.newBuilder()
          .build(new CacheLoader<JobUpdateAction, AtomicLong>() {
            @Override
            public AtomicLong load(JobUpdateAction action) {
              return statsProvider.makeCounter(jobUpdateActionStatName(action));
            }
          });
      Arrays.stream(JobUpdateAction.values())
          .forEach(action -> jobUpdateActionStats.getUnchecked(action).get());
    }
  
    @Override
    public void start(final IJobUpdate update, final AuditData auditData)
        throws UpdateStateException {
  
      requireNonNull(update);
      requireNonNull(auditData);
  
      storage.write((NoResult<UpdateStateException>) storeProvider -> {
        IJobUpdateSummary summary = update.getSummary();
        IJobUpdateInstructions instructions = update.getInstructions();
        IJobKey job = summary.getKey().getJob();
  
        // Validate the update configuration by making sure we can create an updater for it.
        updateFactory.newUpdate(update.getInstructions(), true);
  
        if (instructions.getInitialState().isEmpty() && !instructions.isSetDesiredState()) {
          throw new IllegalArgumentException("Update instruction is a no-op.");
        }
  
        List<IJobUpdateDetails> activeJobUpdates =
            storeProvider.getJobUpdateStore().fetchJobUpdates(queryActiveByJob(job));
        if (!activeJobUpdates.isEmpty()) {
          if (activeJobUpdates.size() > 1) {
            LOG.error("Multiple active updates exist for this job. {}", activeJobUpdates);
            throw new UpdateStateException(
                String.format("Multiple active updates exist for this job. %s", activeJobUpdates));
          }
  
          IJobUpdateDetails activeUpdate = activeJobUpdates.stream().findFirst().get();
          throw new UpdateInProgressException("An active update already exists for this job, "
              + "please terminate it before starting another. "
              + "Active updates are those in states " + Updates.ACTIVE_JOB_UPDATE_STATES,
              activeUpdate.getUpdate().getSummary());
        }
  
        LOG.info("Starting update for job " + job);
  
        storeProvider.getJobUpdateStore().saveJobUpdate(update);
  
        JobUpdateStatus status = ROLLING_FORWARD;
        if (isCoordinatedUpdate(instructions)) {
          status = ROLL_FORWARD_AWAITING_PULSE;
          pulseHandler.initializePulseState(update, status, 0L);
        }
  
        recordAndChangeJobUpdateStatus(
            storeProvider,
            summary.getKey(),
            addAuditData(newEvent(status), auditData));
      });
    }
  
    @Override
    public void assertNotUpdating(IJobKey job) throws JobUpdatingException {
      requireNonNull(job);
  
      if (storage.read(p -> !p.getJobUpdateStore()
          .fetchJobUpdates(queryActiveByJob(job)).isEmpty())) {
  
        throw new JobUpdatingException("Job is currently updating");
      }
    }
  
    @Override
    public void pause(final IJobUpdateKey key, AuditData auditData) throws UpdateStateException {
      requireNonNull(key);
      LOG.info("Attempting to pause update " + key);
      unscopedChangeUpdateStatus(
          key,
          Functions.compose(createAuditedEvent(auditData), GET_PAUSE_STATE));
    }
  
    @Override
    public void resume(final IJobUpdateKey key, final AuditData auditData)
        throws UpdateStateException {
  
      requireNonNull(key);
      requireNonNull(auditData);
      LOG.info("Attempting to resume update " + key);
      storage.write((NoResult<UpdateStateException>) storeProvider -> {
        Optional<IJobUpdateDetails> details = storeProvider.getJobUpdateStore().fetchJobUpdate(key);
  
        if (!details.isPresent()) {
          throw new UpdateStateException("Update does not exist: " + key);
        }
  
        IJobUpdate update = details.get().getUpdate();
        Function<JobUpdateStatus, JobUpdateStatus> stateChange =
            isCoordinatedAndPulseExpired(key, update.getInstructions())
                ? GET_BLOCKED_RESUME_STATE
                : GET_ACTIVE_RESUME_STATE;
  
        JobUpdateStatus newStatus = stateChange.apply(update.getSummary().getState().getStatus());
        changeUpdateStatus(
            storeProvider,
            update.getSummary(),
            addAuditData(newEvent(newStatus), auditData));
      });
    }
  
    @Override
    public void abort(IJobUpdateKey key, AuditData auditData) throws UpdateStateException {
      unscopedChangeUpdateStatus(
          key,
          Functions.compose(createAuditedEvent(auditData), Functions.constant(ABORTED)));
    }
  
    @Override
    public void rollback(IJobUpdateKey key, AuditData auditData) throws UpdateStateException {
      unscopedChangeUpdateStatus(
          key,
          Functions.compose(createAuditedEvent(auditData), Functions.constant(ROLLING_BACK)));
    }
  
    private static Function<JobUpdateStatus, JobUpdateEvent> createAuditedEvent(
        final AuditData auditData) {
  
      return status -> addAuditData(newEvent(status), auditData);
    }
  
    private static final Ordering<IJobUpdateEvent> CHRON_ORDERING =
        Ordering.from(Comparator.comparingLong(IJobUpdateEvent::getTimestampMs));
  
    private long inferLastPulseTimestamp(IJobUpdateEvent mostRecent) {
      // Pulse timestamps are not durably stored by design. However, on system recovery,
      // setting the timestamp of the last pulse to 0L (aka no pulse) is not correct.
      // By inspecting the job update events we can infer a reasonable time stamp to initialize to.
      // In this case, if the upgrade was not waiting for a pulse previously, we can reuse the
      // timestamp of the last event. This does reset the counter for pulses, but reflects the
      // most likely behaviour of a healthy system.
  
      long ts = 0L;
  
      if (!AWAITING_PULSE_STATES.contains(mostRecent.getStatus())) {
        ts = mostRecent.getTimestampMs();
      }
  
      return ts;
    }
  
    public boolean isAutoPauseEnabled(IJobUpdateStrategy strategy) {
      if (strategy.isSetBatchStrategy()) {
        return strategy.getBatchStrategy().isAutopauseAfterBatch();
      }
  
      if (strategy.isSetVarBatchStrategy()) {
        return strategy.getVarBatchStrategy().isAutopauseAfterBatch();
      }
  
      return false;
    }
  
    @Override
    public void systemResume() {
      storage.write((NoResult.Quiet) storeProvider -> {
        for (IJobUpdateDetails details
            : storeProvider.getJobUpdateStore().fetchJobUpdates(ACTIVE_QUERY)) {
  
          IJobUpdateSummary summary = details.getUpdate().getSummary();
          IJobUpdateInstructions instructions = details.getUpdate().getInstructions();
          IJobUpdateKey key = summary.getKey();
          JobUpdateStatus status = summary.getState().getStatus();
          // This is safe because we always write at least one job update event on job update creation
          IJobUpdateEvent latestEvent = CHRON_ORDERING.max(details.getUpdateEvents());
  
          if (isCoordinatedUpdate(instructions)) {
            LOG.info("Automatically restoring pulse state for " + key);
  
            long pulseMs = inferLastPulseTimestamp(latestEvent);
            pulseHandler.initializePulseState(details.getUpdate(), status, pulseMs);
          }
          // Since the restart causes the update to lose state, we backfill instances seen if:
          // a) The update has auto pause after batch enabled
          // b) The update is not currently paused.
          // This restores necessary state for any update that was ROLLING_FORWARD when
          // the scheduler was restarted to avoid crashing the scheduler.
          // We do not backfill when the update is in the ROLL_FORWARD_PAUSED status as the subsequent
          // resume will correctly re-initialize the seen instances state for the update.
          if (isAutoPauseEnabled(instructions.getSettings().getUpdateStrategy())
              && latestEvent.getStatus() == ROLLING_FORWARD) {
            LOG.info("Re-populating previously seen instances for " + key);
  
            instancesSeen.put(key,
                storeProvider.getJobUpdateStore()
                    .fetchJobUpdate(key)
                    .get()
                    .getInstanceEvents()
                    .stream()
                    .map(e -> e.getInstanceId())
                    .collect(Collectors.toCollection(HashSet::new)));
          }
  
          if (AUTO_RESUME_STATES.contains(status)) {
            LOG.info("Automatically resuming update " + key);
  
            try {
              changeJobUpdateStatus(storeProvider, key, newEvent(status), false);
            } catch (UpdateStateException e) {
              throw new RuntimeException(e);
            }
          }
        }
      });
    }
  
    @Override
    public JobUpdatePulseStatus pulse(final IJobUpdateKey key) throws UpdateStateException {
      final PulseState state = pulseHandler.pulseAndGet(key);
      if (state == null) {
        LOG.info("Not pulsing inactive job update: " + key);
        return JobUpdatePulseStatus.FINISHED;
      }
  
      LOG.debug(
          "Job update {} has been pulsed. Timeout of {} msec is reset.",
          key,
          state.getPulseTimeoutMs());
  
      if (JobUpdateStateMachine.isAwaitingPulse(state.getStatus())) {
        // Attempt to unblock a job update previously blocked on expired pulse.
        executor.execute(shutdownOnError(
            lifecycle,
            LOG,
            String.format(FATAL_ERROR_FORMAT, key),
            () -> {
              try {
                unscopedChangeUpdateStatus(
                    key,
                    status -> new JobUpdateEvent().setStatus(GET_UNBLOCKED_STATE.apply(status)));
              } catch (UpdateStateException e) {
                LOG.error(String.format("Error processing job update pulse for %s: %s", key, e));
              }
            }));
      }
  
      return JobUpdatePulseStatus.OK;
    }
  
    @Override
    public void instanceChangedState(final IScheduledTask updatedTask) {
      instanceChanged(
          InstanceKeys.from(
              updatedTask.getAssignedTask().getTask().getJob(),
              updatedTask.getAssignedTask().getInstanceId()),
          Optional.of(updatedTask));
    }
  
    @Override
    public void instanceDeleted(IInstanceKey instance) {
      // This is primarily used to detect when an instance was stuck in PENDING and killed, which
      // results in deletion.
      instanceChanged(instance, Optional.empty());
    }
  
    private void instanceChanged(final IInstanceKey instance, final Optional<IScheduledTask> state) {
      taskEventBatchWorker.execute(storeProvider -> {
        IJobKey job = instance.getJobKey();
        UpdateFactory.Update update = updates.get(job);
        if (update != null) {
          if (update.getUpdater().containsInstance(instance.getInstanceId())) {
            // We check to see if the state change is specified, and if it is, ensure that the new
            // state matches the current state. We do this because events are processed asynchronously
            // and it is possible for an old event trigger an action that should not be triggered
            // for the actual updated state.
            if (!state.isPresent() || isLatestState(storeProvider, state.get())) {
              LOG.info("Forwarding task change for " + InstanceKeys.toString(instance));
              try {
                evaluateUpdater(
                    storeProvider,
                    update,
                    getOnlyMatch(storeProvider.getJobUpdateStore(), queryActiveByJob(job)),
                    ImmutableMap.of(instance.getInstanceId(), state));
              } catch (UpdateStateException e) {
                throw new RuntimeException(e);
              }
            } else {
              LOG.info("Ignoring out of date task change for " + instance);
            }
          } else {
            LOG.info("Instance " + instance + " is not part of active update for "
                + JobKeys.canonicalString(job));
          }
        }
        return BatchWorker.NO_RESULT;
      });
    }
  
    /**
     * Check to see that a given {@link IScheduledTask} still exists in storage and has the same
     * status.
     */
    private boolean isLatestState(MutableStoreProvider storeProvider, IScheduledTask reportedState) {
      Optional<IScheduledTask> currentState = storeProvider
          .getTaskStore()
          .fetchTask(reportedState.getAssignedTask().getTaskId());
  
      return currentState
          .map(iScheduledTask -> iScheduledTask.getStatus() == reportedState.getStatus())
          .orElse(false);
    }
  
    private IJobUpdateSummary getOnlyMatch(JobUpdateStore store, IJobUpdateQuery query) {
      return Iterables.getOnlyElement(store.fetchJobUpdates(query)).getUpdate().getSummary();
    }
  
    @VisibleForTesting
    static IJobUpdateQuery queryActiveByJob(IJobKey job) {
      return IJobUpdateQuery.build(new JobUpdateQuery()
          .setJobKey(job.newBuilder())
          .setUpdateStatuses(Updates.ACTIVE_JOB_UPDATE_STATES));
    }
  
    /**
     * Changes the state of an update, without the 'scope' of an update ID.  This should only be used
     * when responding to outside inputs that are inherently un-scoped, such as a user action or task
     * state change.
     *
     * @param key Update identifier.
     * @param stateChange State change computation, based on the current state of the update.
     * @throws UpdateStateException If no active update exists for the provided {@code job}, or
     *                              if the proposed state transition is not allowed.
     */
    private void unscopedChangeUpdateStatus(
        final IJobUpdateKey key,
        final Function<? super JobUpdateStatus, JobUpdateEvent> stateChange)
        throws UpdateStateException {
  
      storage.write((NoResult<UpdateStateException>) storeProvider -> {
  
        Optional<IJobUpdateDetails> update = storeProvider.getJobUpdateStore().fetchJobUpdate(key);
        if (!update.isPresent()) {
          throw new UpdateStateException("Update does not exist " + key);
        }
  
        IJobUpdateSummary summary = update.get().getUpdate().getSummary();
        changeUpdateStatus(storeProvider, summary, stateChange.apply(summary.getState().getStatus()));
      });
    }
  
    private void changeUpdateStatus(
        MutableStoreProvider storeProvider,
        IJobUpdateSummary updateSummary,
        JobUpdateEvent event) throws UpdateStateException {
  
      if (updateSummary.getState().getStatus() == event.getStatus()) {
        return;
      }
  
      assertTransitionAllowed(updateSummary.getState().getStatus(), event.getStatus());
      recordAndChangeJobUpdateStatus(storeProvider, updateSummary.getKey(), event);
    }
  
    private void recordAndChangeJobUpdateStatus(
        MutableStoreProvider storeProvider,
        IJobUpdateKey key,
        JobUpdateEvent event) throws UpdateStateException {
  
      changeJobUpdateStatus(storeProvider, key, event, true);
    }
  
    private void changeJobUpdateStatus(
        MutableStoreProvider storeProvider,
        IJobUpdateKey key,
        JobUpdateEvent proposedEvent,
        boolean record) throws UpdateStateException {
  
      JobUpdateStore.Mutable updateStore = storeProvider.getJobUpdateStore();
      JobUpdateStatus status = proposedEvent.getStatus();
  
      LOG.info("Update {} is now in state {}", key, status);
      if (record) {
        updateStore.saveJobUpdateEvent(
            key,
            IJobUpdateEvent.build(proposedEvent.setTimestampMs(clock.nowMillis()).setStatus(status)));
        jobUpdateEventStats.getUnchecked(status).incrementAndGet();
      }
  
      if (JobUpdateStore.TERMINAL_STATES.contains(status)) {
        pulseHandler.remove(key);
      } else {
        pulseHandler.updatePulseStatus(key, status);
      }
  
      MonitorAction action = JobUpdateStateMachine.getActionForStatus(status);
      IJobKey job = key.getJob();
      if (action == STOP_WATCHING) {
        updates.remove(job);
      } else if (action == ROLL_FORWARD || action == ROLL_BACK) {
        if (action == ROLL_BACK) {
          updates.remove(job);
        } else {
          checkState(!updates.containsKey(job), "Updater already exists for %s", job);
        }
  
        IJobUpdate jobUpdate = updateStore.fetchJobUpdate(key).get().getUpdate();
        UpdateFactory.Update update;
        try {
          update = updateFactory.newUpdate(jobUpdate.getInstructions(), action == ROLL_FORWARD);
        } catch (RuntimeException e) {
          LOG.warn("Uncaught exception: " + e, e);
          changeJobUpdateStatus(
              storeProvider,
              key,
              newEvent(ERROR).setMessage("Internal scheduler error: " + e.getMessage()),
              true);
          return;
        }
        updates.put(job, update);
        evaluateUpdater(
            storeProvider,
            update,
            jobUpdate.getSummary(),
            ImmutableMap.of());
      }
    }
  
    private static Optional<IScheduledTask> getActiveInstance(
        TaskStore taskStore,
        IJobKey job,
        int instanceId) {
  
      return Optional.ofNullable(Iterables.getOnlyElement(
          taskStore.fetchTasks(Query.instanceScoped(job, instanceId).active()), null));
    }
  
    private static final Set<InstanceUpdateStatus> NOOP_INSTANCE_UPDATE =
        ImmutableSet.of(InstanceUpdateStatus.WORKING, InstanceUpdateStatus.SUCCEEDED);
  
    private static boolean isCoordinatedUpdate(IJobUpdateInstructions instructions) {
      return instructions.getSettings().getBlockIfNoPulsesAfterMs() > 0;
    }
  
    private boolean isCoordinatedAndPulseExpired(
        IJobUpdateKey key,
        IJobUpdateInstructions instructions) {
  
      if (isCoordinatedUpdate(instructions)) {
        PulseState pulseState = pulseHandler.get(key);
        boolean result = pulseState == null || pulseState.isBlocked(clock);
        LOG.info("Coordinated update {} pulse expired: {}", key, result);
        return result;
      } else {
        return false;
      }
    }
  
    @VisibleForTesting
    static final String PULSE_TIMEOUT_MESSAGE = "Pulses from external service have timed out.";
  
    // Determines whether it is necessary to skip evaluating a side effect if we will be pausing
    // the update after this evaluation
    private boolean skipSideEffect(
        boolean pauseAfterBatch,
        SideEffect sideEffect) {
      return pauseAfterBatch
          && Collections.disjoint(
          sideEffect.getStatusChanges(),
          InstanceUpdateStatus.TERMINAL_STATUSES);
    }
  
    private void evaluateUpdater(
        final MutableStoreProvider storeProvider,
        final UpdateFactory.Update update,
        IJobUpdateSummary summary,
        Map<Integer, Optional<IScheduledTask>> changedInstance) throws UpdateStateException {
  
      JobUpdateStatus updaterStatus = summary.getState().getStatus();
      final IJobUpdateKey key = summary.getKey();
  
      JobUpdateStore.Mutable updateStore = storeProvider.getJobUpdateStore();
  
      IJobUpdateInstructions instructions = updateStore.fetchJobUpdate(key).get()
          .getUpdate().getInstructions();
      if (isCoordinatedAndPulseExpired(key, instructions)) {
        // Move coordinated update into awaiting pulse state.
        JobUpdateStatus blockedStatus = getBlockedState(summary.getState().getStatus());
        changeUpdateStatus(
            storeProvider,
            summary,
            newEvent(blockedStatus).setMessage(PULSE_TIMEOUT_MESSAGE));
        return;
      }
  
      InstanceStateProvider<Integer, Optional<IScheduledTask>> stateProvider =
          instanceId -> getActiveInstance(storeProvider.getTaskStore(), key.getJob(), instanceId);
  
      EvaluationResult<Integer> result = update.getUpdater().evaluate(changedInstance, stateProvider);
  
      LOG.info("{} evaluation result: {}", key, result);
      final boolean autoPauseAfterCurrentBatch =
          isAutoPauseEnabled(instructions.getSettings().getUpdateStrategy())
          && maybeAutoPause(summary, result);
      if (autoPauseAfterCurrentBatch) {
        changeUpdateStatus(storeProvider,
            summary,
            newEvent(getPausedState(summary.getState().getStatus())).setMessage(UPDATE_AUTO_PAUSED));
      }
  
      for (Map.Entry<Integer, SideEffect> entry : result.getSideEffects().entrySet()) {
        // If we're pausing after processing this set of side effects, only process the side effects
        // which are in a terminal state in order to avoid starting new shards after the pause
        // has kicked in.
        if (skipSideEffect(autoPauseAfterCurrentBatch, entry.getValue())) {
          continue;
        }
  
        Iterable<InstanceUpdateStatus> statusChanges;
  
        int instanceId = entry.getKey();
        List<IJobInstanceUpdateEvent> savedEvents = updateStore.fetchJobUpdate(key).get()
            .getInstanceEvents()
            .stream()
            .filter(e -> e.getInstanceId() == instanceId)
            .collect(Collectors.toList());
  
        Set<JobUpdateAction> savedActions =
            savedEvents.stream().map(EVENT_TO_ACTION).collect(Collectors.toSet());
  
        // Don't bother persisting a sequence of status changes that represents an instance that
        // was immediately recognized as being healthy and in the desired state.
        if (entry.getValue().getStatusChanges().equals(NOOP_INSTANCE_UPDATE)
            && savedEvents.isEmpty()) {
  
          LOG.info("Suppressing no-op update for instance " + instanceId);
          statusChanges = ImmutableSet.of();
        } else {
          statusChanges = entry.getValue().getStatusChanges();
        }
  
        for (InstanceUpdateStatus statusChange : statusChanges) {
          JobUpdateAction action = STATE_MAP.get(Pair.of(statusChange, updaterStatus));
          requireNonNull(action);
  
          // A given instance update action may only be issued once during the update lifecycle.
          // Suppress duplicate events due to pause/resume operations.
          if (savedActions.contains(action)) {
            LOG.info("Suppressing duplicate update {} for instance {}.", action, instanceId);
          } else {
            IJobInstanceUpdateEvent event = IJobInstanceUpdateEvent.build(
                new JobInstanceUpdateEvent()
                    .setInstanceId(instanceId)
                    .setTimestampMs(clock.nowMillis())
                    .setAction(action));
            updateStore.saveJobInstanceUpdateEvent(summary.getKey(), event);
            jobUpdateActionStats.getUnchecked(action).incrementAndGet();
          }
        }
      }
  
      OneWayStatus status = result.getStatus();
      if (status == SUCCEEDED || status == OneWayStatus.FAILED) {
        if (SideEffect.hasActions(result.getSideEffects().values())) {
          throw new IllegalArgumentException(
              "A terminal state should not specify actions: " + result);
        }
  
        JobUpdateEvent event = new JobUpdateEvent();
        if (status == SUCCEEDED) {
          event.setStatus(update.getSuccessStatus());
        } else {
          event.setStatus(update.getFailureStatus());
          // Generate a transition message based on one (arbitrary) instance in the group that pushed
          // the update over the failure threshold (in all likelihood this group is of size 1).
          // This is done as a rough cut to aid in diagnosing a failed update, as generating a
          // complete summary would likely be of dubious value.
          for (Map.Entry<Integer, SideEffect> entry : result.getSideEffects().entrySet()) {
            Optional<Failure> failure = entry.getValue().getFailure();
            if (failure.isPresent()) {
              event.setMessage(failureMessage(entry.getKey(), failure.get()));
              break;
            }
          }
        }
  
        changeUpdateStatus(storeProvider, summary, event);
        return;
      }
  
      LOG.info("Executing side-effects for update of {}: {}",
          key,
          result.getSideEffects().entrySet());
      for (Map.Entry<Integer, SideEffect> entry : result.getSideEffects().entrySet()) {
        IInstanceKey instance = InstanceKeys.from(key.getJob(), entry.getKey());
  
        Optional<InstanceAction> action = entry.getValue().getAction();
        if (action.isPresent() && !skipSideEffect(autoPauseAfterCurrentBatch, entry.getValue())) {
          Optional<InstanceActionHandler> handler = action.get().getHandler();
          if (handler.isPresent()) {
            Optional<Amount<Long, Time>> reevaluateDelay = handler.get().getReevaluationDelay(
                instance,
                instructions,
                storeProvider,
                stateManager,
                updateAgentReserver,
                updaterStatus,
                key,
                slaKillController);
            if (reevaluateDelay.isPresent()) {
              executor.schedule(
                  getDeferredEvaluator(instance, key),
                  reevaluateDelay.get().getValue(),
                  reevaluateDelay.get().getUnit().getTimeUnit());
            }
          }
        }
      }
    }
  
    @VisibleForTesting
    static final Function<IJobInstanceUpdateEvent, JobUpdateAction> EVENT_TO_ACTION =
        IJobInstanceUpdateEvent::getAction;
  
    @VisibleForTesting
    static String failureMessage(int instanceId, Failure failure) {
      return String.format("Latest failure: instance %d %s", instanceId, failure.getReason());
    }
  
    /**
     * Associates an instance updater state change and the job's update status to an action.
     */
    private static final Map<Pair<InstanceUpdateStatus, JobUpdateStatus>, JobUpdateAction> STATE_MAP =
        ImmutableMap.<Pair<InstanceUpdateStatus, JobUpdateStatus>, JobUpdateAction>builder()
            .put(
                Pair.of(InstanceUpdateStatus.WORKING, ROLLING_FORWARD),
                JobUpdateAction.INSTANCE_UPDATING)
            .put(
                Pair.of(InstanceUpdateStatus.SUCCEEDED, ROLLING_FORWARD),
                JobUpdateAction.INSTANCE_UPDATED)
            .put(
                Pair.of(InstanceUpdateStatus.FAILED, ROLLING_FORWARD),
                JobUpdateAction.INSTANCE_UPDATE_FAILED)
            .put(
                Pair.of(InstanceUpdateStatus.WORKING, ROLLING_BACK),
                JobUpdateAction.INSTANCE_ROLLING_BACK)
            .put(
                Pair.of(InstanceUpdateStatus.SUCCEEDED, ROLLING_BACK),
                JobUpdateAction.INSTANCE_ROLLED_BACK)
            .put(
                Pair.of(InstanceUpdateStatus.FAILED, ROLLING_BACK),
                JobUpdateAction.INSTANCE_ROLLBACK_FAILED)
            .build();
  
    private static JobUpdateEvent newEvent(JobUpdateStatus status) {
      return new JobUpdateEvent().setStatus(status);
    }
  
    private static JobUpdateEvent addAuditData(JobUpdateEvent event, AuditData auditData) {
      return event.setMessage(auditData.getMessage().orElse(null))
          .setUser(auditData.getUser());
    }
  
    private Runnable getDeferredEvaluator(final IInstanceKey instance, final IJobUpdateKey key) {
      return shutdownOnError(
          lifecycle,
          LOG,
          String.format(FATAL_ERROR_FORMAT, "Key: " + key + " Instance key: " + instance),
          () -> storage.write((NoResult.Quiet) storeProvider -> {
            IJobUpdateSummary summary =
                storeProvider.getJobUpdateStore().fetchJobUpdate(key).get().getUpdate().getSummary();
            JobUpdateStatus status = summary.getState().getStatus();
            // Suppress this evaluation if the updater is not currently active.
            if (JobUpdateStateMachine.isActive(status)) {
              UpdateFactory.Update update = updates.get(instance.getJobKey());
              try {
                evaluateUpdater(
                    storeProvider,
                    update,
                    summary,
                    ImmutableMap.of(
                        instance.getInstanceId(),
                        getActiveInstance(
                            storeProvider.getTaskStore(),
                            instance.getJobKey(),
                            instance.getInstanceId())));
              } catch (UpdateStateException e) {
                LOG.error(String.format("Error running deferred evaluation for %s: %s", instance, e));
                throw new RuntimeException(e);
              }
            }
          }));
    }
  
    @VisibleForTesting
    private boolean maybeAutoPause(IJobUpdateSummary summary, EvaluationResult<Integer> result) {
      JobUpdateStatus updaterStatus = summary.getState().getStatus();
      final IJobUpdateKey key = summary.getKey();
  
      // Only apply auto-pause to an update rolling forward.
      if (updaterStatus != ROLLING_FORWARD) {
        return false;
      }
  
      if (instancesSeen.containsKey(key)) {
        Set<Integer> instancesCached = instancesSeen.get(key);
        Set<Integer> instancesBeingUpdated = result.getSideEffects().keySet();
  
        if (result.getStatus() == SUCCEEDED) {
          instancesSeen.remove(key);
          return false;
        }
  
        // If the update evaluation is dealing with new instances, that signals we are at a barrier
        // crossing.
        if (!instancesCached.containsAll(instancesBeingUpdated)) {
          instancesCached.addAll(instancesBeingUpdated);
          return true;
        }
      } else {
        instancesSeen.put(key, new HashSet<Integer>(result.getSideEffects().keySet()));
      }
      return false;
    }
  
    private static class PulseHandler {
      private final Clock clock;
  
      // TODO(maxim): expose this data via a debug endpoint AURORA-1103.
      // Currently active coordinated update pulse states. A pulse state is added when a coordinated
      // update is created and removed only when an update reaches terminal state. A PAUSED update
      // pulse state is still retained in the map and accepts pulses.
      private final Map<IJobUpdateKey, PulseState> pulseStates = Maps.newHashMap();
  
      PulseHandler(Clock clock) {
        this.clock = requireNonNull(clock);
      }
  
      synchronized void initializePulseState(IJobUpdate update, JobUpdateStatus status, long ts) {
        pulseStates.put(update.getSummary().getKey(), new PulseState(
            status,
            update.getInstructions().getSettings().getBlockIfNoPulsesAfterMs(),
            ts));
      }
  
      synchronized PulseState pulseAndGet(IJobUpdateKey key) {
        PulseState state = pulseStates.get(key);
        if (state != null) {
          state = pulseStates.put(key, new PulseState(
              state.getStatus(),
              state.getPulseTimeoutMs(),
              clock.nowMillis()));
        }
        return state;
      }
  
      synchronized void updatePulseStatus(IJobUpdateKey key, JobUpdateStatus status) {
        PulseState state = pulseStates.get(key);
        if (state != null) {
          pulseStates.put(key, new PulseState(
              status,
              state.getPulseTimeoutMs(),
              state.getLastPulseMs()));
        }
      }
  
      synchronized void remove(IJobUpdateKey key) {
        pulseStates.remove(key);
      }
  
      synchronized PulseState get(IJobUpdateKey key) {
        return pulseStates.get(key);
      }
    }
  
    private static class PulseState {
      private final JobUpdateStatus status;
      private final long pulseTimeoutMs;
      private final long lastPulseMs;
  
      PulseState(JobUpdateStatus status, long pulseTimeoutMs, long lastPulseMs) {
        this.status = requireNonNull(status);
        this.pulseTimeoutMs = pulseTimeoutMs;
        this.lastPulseMs = lastPulseMs;
      }
  
      JobUpdateStatus getStatus() {
        return status;
      }
  
      long getPulseTimeoutMs() {
        return pulseTimeoutMs;
      }
  
      long getLastPulseMs() {
        return lastPulseMs;
      }
  
      boolean isBlocked(Clock clock) {
        return clock.nowMillis() - lastPulseMs >= pulseTimeoutMs;
      }
    }
  }
  