<commit_before>package com.twitter.mesos.updater;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;
import com.twitter.common.base.ExceptionalFunction;
import com.twitter.common.util.StateMachine;
import com.twitter.mesos.gen.UpdateConfig;

import java.util.PriorityQueue;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.twitter.common.base.MorePreconditions.checkNotBlank;
import static com.twitter.mesos.updater.UpdateCommand.Type.ROLLBACK_TASK;
import static com.twitter.mesos.updater.UpdateCommand.Type.UPDATE_TASK;
import static com.twitter.mesos.updater.UpdateLogic.State.*;

/**
 * Abstract logic for the mesos updater.
 *
 * @author wfarner
 */
public class UpdateLogic {

  private static final Logger LOG = Logger.getLogger(UpdateLogic.class.getName());

  private final ExceptionalFunction<UpdateCommand, Integer, UpdateException> commandRunner;
  private final Set<Integer> newShards;
  private final UpdateConfig config;
  private final Set<Integer> oldShards;

  /**
   * Creates a new update logic.
   *
   * TODO(wfarner): During a rollback, need to get the updated set of shard IDs since the update
   *    may have added shard IDs.
   *
   * TODO(wfarner): Need to make sure to issue a kill command to shards that are being removed.
   *
   */
  public UpdateLogic(Set<Integer> oldShards, Set<Integer> newShards, UpdateConfig config,
      ExceptionalFunction<UpdateCommand, Integer, UpdateException> commandRunner) {
    this.oldShards = checkNotBlank(oldShards);
    this.newShards = checkNotBlank(newShards);
    this.config = checkNotNull(config);
    this.commandRunner = checkNotNull(commandRunner);
  }

  public static class UpdateException extends Exception {
    public UpdateException(String msg) {
      super(msg);
    }
    public UpdateException(String msg, Throwable cause) {
      super(msg, cause);
    }
  }

  /**
   * Runs the update.
   *
   * @return {@code true} if the update completed successfully or {@code false} if it failed.
   * @throws InterruptedException If interrupted while updating.
   */
  public boolean run() throws InterruptedException {
    StateMachine<State> stateMachine = createStateMachine();

    while (!TERMINAL_STATES.contains(stateMachine.getState())) {
      step(stateMachine);
    }

    return stateMachine.getState() == COMPLETE_SUCCESS;
  }

  // An iterable that will consume from a queue of shards being operated on in the current stage.
  // Be careful when reading from this, since the objects are removed.
  Iterable<Integer> consumingIds;

  int totalFailures = 0;

  private void step(StateMachine<State> stateMachine) throws InterruptedException {
    try {
      switch (stateMachine.getState()) {
        case INIT:
          LOG.info("Initiating update.");
          consumingIds = Iterables.consumingIterable(new PriorityQueue<Integer>(newShards));
          stateMachine.transition(CANARY);
          break;

        case CANARY:
          int failures = restartShards(UPDATE_TASK, config.getCanaryTaskCount(),
              config.getToleratedCanaryFailures() - totalFailures,
                config.getCanaryWatchDurationSecs());
          if (failures > config.getToleratedCanaryFailures()) {
            stateMachine.transition(PREPARE_ROLLBACK);
          } else {
            stateMachine.transition(UPDATE_BATCH);
          }
          break;

        case UPDATE_BATCH:
          if (Iterables.isEmpty(consumingIds)) {
            stateMachine.transition(COMPLETE_SUCCESS);
          } else {
            restartShards(UPDATE_TASK, config.getUpdateBatchSize(),
                config.getToleratedUpdateFailures() - totalFailures,
                config.getUpdateWatchDurationSecs());
            if (totalFailures > config.getToleratedUpdateFailures()) {
              stateMachine.transition(PREPARE_ROLLBACK);
            }
          }
          break;

        case PREPARE_ROLLBACK:
          LOG.info("Preparing for rollback.");
          consumingIds = Iterables.consumingIterable(new PriorityQueue<Integer>(
              Sets.difference(oldShards, ImmutableSet.copyOf(consumingIds))));
          stateMachine.transition(ROLLBACK_BATCH);
          break;

        case ROLLBACK_BATCH:
          if (Iterables.isEmpty(consumingIds)) {
            stateMachine.transition(COMPLETE_FAILED);
          } else {
            // TODO(wfarner): Can/should we do anything if this fails?
            restartShards(ROLLBACK_TASK, config.getUpdateBatchSize(), Integer.MAX_VALUE,
                config.getUpdateWatchDurationSecs());
          }
          break;

        default:
          throw new IllegalStateException("Unhandled state " + stateMachine.getState());
      }
    } catch (UpdateException e) {
      LOG.log(Level.SEVERE, "Update failed.", e);
    }
  }

  /**
   * Performs a restart on a batch of shards, and updates the total failure count.
   *
   * @param type The type of update to issue.
   * @param numShards Number of shards to pull from the queue.
   * @param toleratedFailures The number of tolerated failures in this batch.
   * @param updateWatchDurationSecs The number of seconds that new tasks must survive to be
   *    considered successfully updated.
   * @return The number of shards that did not survive the watch period.
   * @throws UpdateException If the restart attempt failed.
   */
  private int restartShards(UpdateCommand.Type type, int numShards, int toleratedFailures,
      int updateWatchDurationSecs) throws UpdateException {
    Preconditions.checkArgument(numShards > 0, "Restart of zero shards does not make sense.");

    Set<Integer> ids = ImmutableSet.copyOf(Iterables.limit(consumingIds, numShards));
    LOG.info(String.format("Issuing %s on shards %s", type, ids));

    int newFailures = commandRunner.apply(new UpdateCommand(type, ids,
        config.getRestartTimeoutSecs(), updateWatchDurationSecs, toleratedFailures));

    totalFailures += newFailures;
    return newFailures;
  }

  static enum State {
    INIT,
    CANARY,
    UPDATE_BATCH,
    PREPARE_ROLLBACK,
    ROLLBACK_BATCH,
    COMPLETE_SUCCESS,
    COMPLETE_FAILED
  }

  private static StateMachine<State> createStateMachine() {
    return StateMachine.<State>builder("UpdateLogic")
      .addState(INIT, CANARY, COMPLETE_FAILED)
      .addState(CANARY, UPDATE_BATCH, PREPARE_ROLLBACK, COMPLETE_FAILED)
      .addState(UPDATE_BATCH, UPDATE_BATCH, PREPARE_ROLLBACK, COMPLETE_SUCCESS, COMPLETE_FAILED)
      .addState(PREPARE_ROLLBACK, ROLLBACK_BATCH)
      .addState(ROLLBACK_BATCH, ROLLBACK_BATCH, COMPLETE_FAILED)
      .initialState(INIT)<commit_msg>Upgrade to guava_r08 + fix UpdateLogic/UpdateLogicTest - can no longer get away with Iterables.isEmpty(consumingIterable) + fix NameClassifier -> NPE in NameClassifierTest<commit_after>package com.twitter.mesos.updater;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;
import com.twitter.common.base.ExceptionalFunction;
import com.twitter.common.util.StateMachine;
import com.twitter.mesos.gen.UpdateConfig;

import java.util.Collection;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.twitter.common.base.MorePreconditions.checkNotBlank;
import static com.twitter.mesos.updater.UpdateCommand.Type.ROLLBACK_TASK;
import static com.twitter.mesos.updater.UpdateCommand.Type.UPDATE_TASK;
import static com.twitter.mesos.updater.UpdateLogic.State.CANARY;
import static com.twitter.mesos.updater.UpdateLogic.State.COMPLETE_FAILED;
import static com.twitter.mesos.updater.UpdateLogic.State.COMPLETE_SUCCESS;
import static com.twitter.mesos.updater.UpdateLogic.State.INIT;
import static com.twitter.mesos.updater.UpdateLogic.State.PREPARE_ROLLBACK;
import static com.twitter.mesos.updater.UpdateLogic.State.ROLLBACK_BATCH;
import static com.twitter.mesos.updater.UpdateLogic.State.UPDATE_BATCH;

/**
 * Abstract logic for the mesos updater.
 *
 * @author wfarner
 */
public class UpdateLogic {

  private static final Logger LOG = Logger.getLogger(UpdateLogic.class.getName());

  private final ExceptionalFunction<UpdateCommand, Integer, UpdateException> commandRunner;
  private final Set<Integer> newShards;
  private final UpdateConfig config;
  private final Set<Integer> oldShards;

  /**
   * Creates a new update logic.
   *
   * TODO(wfarner): During a rollback, need to get the updated set of shard IDs since the update
   *    may have added shard IDs.
   *
   * TODO(wfarner): Need to make sure to issue a kill command to shards that are being removed.
   *
   */
  public UpdateLogic(Set<Integer> oldShards, Set<Integer> newShards, UpdateConfig config,
      ExceptionalFunction<UpdateCommand, Integer, UpdateException> commandRunner) {
    this.oldShards = checkNotBlank(oldShards);
    this.newShards = checkNotBlank(newShards);
    this.config = checkNotNull(config);
    this.commandRunner = checkNotNull(commandRunner);
  }

  public static class UpdateException extends Exception {
    public UpdateException(String msg) {
      super(msg);
    }
    public UpdateException(String msg, Throwable cause) {
      super(msg, cause);
    }
  }

  /**
   * Runs the update.
   *
   * @return {@code true} if the update completed successfully or {@code false} if it failed.
   * @throws InterruptedException If interrupted while updating.
   */
  public boolean run() throws InterruptedException {
    StateMachine<State> stateMachine = createStateMachine();

    while (!TERMINAL_STATES.contains(stateMachine.getState())) {
      step(stateMachine);
    }

    return stateMachine.getState() == COMPLETE_SUCCESS;
  }


  private static class ConsumableSortedIds {
    private final Collection<Integer> sortedIds;
    private final Iterable<Integer> consumingIds;

    private ConsumableSortedIds(final Set<Integer> ids) {
      sortedIds = new PriorityQueue<Integer>(ids);
      consumingIds = Iterables.consumingIterable(sortedIds);
    }

    /**
     * Returns an iterable over sorted ids that will consume ids as they are iterated over.  When
     * all ids have been iterated over once, {@link #isEmpty()} will return {@code true}.
     *
     * @return an consuming {@code Iterable} over the sorted shard ids
     */
    public Iterable<Integer> get() {
      return consumingIds;
    }

    /**
     * @return {@code true} iff the Iterable returned by {@link #get()} has been fully iterated.
     */
    public boolean isEmpty() {
      return sortedIds.isEmpty();
    }
  }


  // An iterable that will consume from a queue of shards being operated on in the current stage.
  // Be careful when reading from this, since the objects are removed.
  private ConsumableSortedIds consumingIds;

  int totalFailures = 0;

  private void step(StateMachine<State> stateMachine) throws InterruptedException {
    try {
      switch (stateMachine.getState()) {
        case INIT:
          LOG.info("Initiating update.");
          consumingIds = new ConsumableSortedIds(newShards);
          stateMachine.transition(CANARY);
          break;

        case CANARY:
          int failures = restartShards(UPDATE_TASK, config.getCanaryTaskCount(),
              config.getToleratedCanaryFailures() - totalFailures,
                config.getCanaryWatchDurationSecs());
          if (failures > config.getToleratedCanaryFailures()) {
            stateMachine.transition(PREPARE_ROLLBACK);
          } else {
            stateMachine.transition(UPDATE_BATCH);
          }
          break;

        case UPDATE_BATCH:
          if (consumingIds.isEmpty()) {
            stateMachine.transition(COMPLETE_SUCCESS);
          } else {
            restartShards(UPDATE_TASK, config.getUpdateBatchSize(),
                config.getToleratedUpdateFailures() - totalFailures,
                config.getUpdateWatchDurationSecs());
            if (totalFailures > config.getToleratedUpdateFailures()) {
              stateMachine.transition(PREPARE_ROLLBACK);
            }
          }
          break;

        case PREPARE_ROLLBACK:
          LOG.info("Preparing for rollback.");
          consumingIds = new ConsumableSortedIds(Sets.difference(oldShards,
              ImmutableSet.copyOf(consumingIds.get())));
          stateMachine.transition(ROLLBACK_BATCH);
          break;

        case ROLLBACK_BATCH:
          if (consumingIds.isEmpty()) {
            stateMachine.transition(COMPLETE_FAILED);
          } else {
            // TODO(wfarner): Can/should we do anything if this fails?
            restartShards(ROLLBACK_TASK, config.getUpdateBatchSize(), Integer.MAX_VALUE,
                config.getUpdateWatchDurationSecs());
          }
          break;

        default:
          throw new IllegalStateException("Unhandled state " + stateMachine.getState());
      }
    } catch (UpdateException e) {
      LOG.log(Level.SEVERE, "Update failed.", e);
    }
  }

  /**
   * Performs a restart on a batch of shards, and updates the total failure count.
   *
   * @param type The type of update to issue.
   * @param numShards Number of shards to pull from the queue.
   * @param toleratedFailures The number of tolerated failures in this batch.
   * @param updateWatchDurationSecs The number of seconds that new tasks must survive to be
   *    considered successfully updated.
   * @return The number of shards that did not survive the watch period.
   * @throws UpdateException If the restart attempt failed.
   */
  private int restartShards(UpdateCommand.Type type, int numShards, int toleratedFailures,
      int updateWatchDurationSecs) throws UpdateException {
    Preconditions.checkArgument(numShards > 0, "Restart of zero shards does not make sense.");

    Set<Integer> ids = ImmutableSet.copyOf(Iterables.limit(consumingIds.get(), numShards));
    LOG.info(String.format("Issuing %s on shards %s", type, ids));

    int newFailures = commandRunner.apply(new UpdateCommand(type, ids,
        config.getRestartTimeoutSecs(), updateWatchDurationSecs, toleratedFailures));

    totalFailures += newFailures;
    return newFailures;
  }

  static enum State {
    INIT,
    CANARY,
    UPDATE_BATCH,
    PREPARE_ROLLBACK,
    ROLLBACK_BATCH,
    COMPLETE_SUCCESS,
    COMPLETE_FAILED
  }

  private static StateMachine<State> createStateMachine() {
    return StateMachine.<State>builder("UpdateLogic")
      .addState(INIT, CANARY, COMPLETE_FAILED)
      .addState(CANARY, UPDATE_BATCH, PREPARE_ROLLBACK, COMPLETE_FAILED)
      .addState(UPDATE_BATCH, UPDATE_BATCH, PREPARE_ROLLBACK, COMPLETE_SUCCESS, COMPLETE_FAILED)
      .addState(PREPARE_ROLLBACK, ROLLBACK_BATCH)
      .addState(ROLLBACK_BATCH, ROLLBACK_BATCH, COMPLETE_FAILED)
      .initialState(INIT)
{'fraction_non_alphanumeric': 0.08827430693725959, 'fraction_numerical': 0.0003979307600477517, 'mean_word_length': 3.374528575572962, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 40, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5758649', 'n_tokens_mistral': 4361, 'n_tokens_neox': 4276, 'n_words': 1063}