<commit_before>package com.twitter.mesos.scheduler.storage.mem;

import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Level;
import java.util.logging.Logger;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicate;
import com.google.common.base.Supplier;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Sets;

import org.apache.commons.lang.StringUtils;

import com.twitter.common.args.Arg;
import com.twitter.common.args.CmdLine;
import com.twitter.common.base.Closure;
import com.twitter.common.inject.TimedInterceptor.Timed;
import com.twitter.common.quantity.Amount;
import com.twitter.common.quantity.Time;
import com.twitter.common.stats.Stats;
import com.twitter.mesos.Tasks;
import com.twitter.mesos.gen.ScheduledTask;
import com.twitter.mesos.gen.TaskQuery;
import com.twitter.mesos.gen.TwitterTaskInfo;
import com.twitter.mesos.scheduler.Query;
import com.twitter.mesos.scheduler.storage.TaskStore;

import static com.google.common.base.Preconditions.checkNotNull;

/**
 * An in-memory task store.
 */
public class MemTaskStore implements TaskStore.Mutable.Transactioned {

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

  @CmdLine(name = "slow_query_log_threshold",
      help = "Log all queries that take at least this long to execute.")
  private static final Arg<Amount<Long, Time>> SLOW_QUERY_LOG_THRESHOLD =
      Arg.create(Amount.of(25L, Time.MILLISECONDS));

  /**
   * COPIER is separate from deepCopy to capture timing information.  Only the instrumented
   * {@link #deepCopy(ScheduledTask)} should interact directly with {@code COPIER}.
   */
  private static final Function<ScheduledTask, ScheduledTask> COPIER =
      Util.deepCopier();
  private final Function<ScheduledTask, ScheduledTask> deepCopy =
      new Function<ScheduledTask, ScheduledTask>() {
        @Override public ScheduledTask apply(ScheduledTask input) {
          return deepCopy(input);
        }
      };

  private final long slowQueryThresholdNanos = SLOW_QUERY_LOG_THRESHOLD.get().as(Time.NANOSECONDS);

  private final TransactionalMap<String, ScheduledTask> tasks =
      TransactionalMap.wrap(Maps.<String, ScheduledTask>newHashMap());
  private final TransactionalMap<String, Set<String>> tasksByJobKey =
      TransactionalMap.wrap(Maps.<String, Set<String>>newHashMap());

  private final AtomicLong taskQueriesById = Stats.exportLong("task_queries_by_id");
  private final AtomicLong taskQueriesByJob = Stats.exportLong("task_queries_by_job");
  private final AtomicLong taskQueriesAll = Stats.exportLong("task_queries_all");

  @Override
  public void commit() {
    tasks.commit();
    tasksByJobKey.commit();
  }

  @Override
  public void rollback() {
    tasks.rollback();
    tasksByJobKey.rollback();
  }

  @Timed("mem_storage_deep_copies")
  protected ScheduledTask deepCopy(ScheduledTask input) {
    return COPIER.apply(input);
  }

  @Timed("mem_storage_fetch_tasks")
  @Override
  public ImmutableSet<ScheduledTask> fetchTasks(TaskQuery query) {
    checkNotNull(query);

    long start = System.nanoTime();
    ImmutableSet<ScheduledTask> result = immutableMatches(query).toSet();
    long durationNanos = System.nanoTime() - start;
    Level level = (durationNanos >= slowQueryThresholdNanos) ? Level.INFO : Level.FINE;
    if (LOG.isLoggable(level)) {
      LOG.log(level, "Query took "
          + Amount.of(durationNanos, Time.NANOSECONDS).as(Time.MILLISECONDS) + " ms: " + query);
    }

    return result;
  }

  @Override
  public ImmutableSet<ScheduledTask> fetchTasks(Supplier<TaskQuery> querySupplier) {
    return fetchTasks(querySupplier.get());
  }

  @Timed("mem_storage_fetch_task_ids")
  @Override
  public Set<String> fetchTaskIds(TaskQuery query) {
    checkNotNull(query);

    return mutableMatches(query).transform(Tasks.SCHEDULED_TO_ID).toSet();
  }

  private Map<String, Collection<String>> taskIdsByJobKey(Iterable<ScheduledTask> toIndex) {
    return Multimaps.transformValues(
        Multimaps.index(toIndex, Tasks.SCHEDULED_TO_JOB_KEY),
        Tasks.SCHEDULED_TO_ID)
        .asMap();
  }

  @Override
  public Set<String> fetchTaskIds(Supplier<TaskQuery> querySupplier) {
    return fetchTaskIds(querySupplier.get());
  }

  @Timed("mem_storage_save_tasks")
  @Override
  public void saveTasks(Set<ScheduledTask> newTasks) {
    checkNotNull(newTasks);
    Preconditions.checkState(Tasks.ids(newTasks).size() == newTasks.size(),
        "Proposed new tasks would create task ID collision.");

    Set<ScheduledTask> immutable =
        FluentIterable.from(newTasks).transform(deepCopy).toSet();
    tasks.putAll(Maps.uniqueIndex(immutable, Tasks.SCHEDULED_TO_ID));

    // Update job key index.
    for (Map.Entry<String, Collection<String>> entry : taskIdsByJobKey(immutable).entrySet()) {
      ImmutableSet.Builder<String> newIds = ImmutableSet.builder();
      newIds.addAll(entry.getValue());
      Set<String> existingIds = tasksByJobKey.get(entry.getKey());
      if (existingIds != null) {
        newIds.addAll(existingIds);
      }

      // This is subtle but important: when modifying the items associated with a key, we must
      // re-insert rather than updating the contents of the set.  The transactional map has no
      // awareness of the contents of a value, so modifying the value set would break the
      // transactional guarantee.
      tasksByJobKey.put(entry.getKey(), newIds.build());
    }
  }

  @Timed("mem_storage_delete_all_tasks")
  @Override
  public void deleteAllTasks() {
    tasks.clear();
    tasksByJobKey.clear();
  }

  @Timed("mem_storage_delete_tasks")
  @Override
  public void deleteTasks(Set<String> taskIds) {
    checkNotNull(taskIds);

    // TransactionalMap does not presently support keySet modification, requiring iteration here.
    ImmutableList.Builder<ScheduledTask> gone = ImmutableList.builder();
    for (String id : taskIds) {
      ScheduledTask task = tasks.remove(id);
      if (task != null) {
        gone.add(task);
      }
    }
    for (Map.Entry<String, Collection<String>> entry : taskIdsByJobKey(gone.build()).entrySet()) {
      Set<String> existingIds = tasksByJobKey.get(entry.getKey());
      if (existingIds == null) {
        throw new IllegalStateException(
            "Index inconsistency - job for tasks not present in index: " + entry.getValue());
      }

      Set<String> newIds = Sets.difference(existingIds, ImmutableSet.copyOf(entry.getValue()));
      if (newIds.isEmpty()) {
        tasksByJobKey.remove(entry.getKey());
      } else {
        tasksByJobKey.put(entry.getKey(), newIds);
      }
    }
  }

  @Timed("mem_storage_mutate_tasks")
  @Override
  public ImmutableSet<ScheduledTask> mutateTasks(TaskQuery query, Closure<ScheduledTask> mutator) {
    checkNotNull(query);
    checkNotNull(mutator);

    ImmutableMap.Builder<String, ScheduledTask> mutated = ImmutableMap.builder();
    for (ScheduledTask original : immutableMatches(query)) {
      // Copy the object before invoking user code.  This is to support diff checking.
      ScheduledTask mutable = deepCopy.apply(original);
      mutator.execute(mutable);
      if (!original.equals(mutable)) {
        Preconditions.checkState(Tasks.id(original).equals(Tasks.id(mutable)),
            "A tasks ID may not be mutated.");

        // A diff is present - detach the mutable object from the closure's code to render
        // further mutation impossible.
        ScheduledTask updated = deepCopy.apply(mutable);
        mutated.put(Tasks.id(mutable), updated);<commit_msg> Copy the result of Sets.difference into the job index.<commit_after>package com.twitter.mesos.scheduler.storage.mem;

import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Level;
import java.util.logging.Logger;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicate;
import com.google.common.base.Supplier;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Sets;

import org.apache.commons.lang.StringUtils;

import com.twitter.common.args.Arg;
import com.twitter.common.args.CmdLine;
import com.twitter.common.base.Closure;
import com.twitter.common.inject.TimedInterceptor.Timed;
import com.twitter.common.quantity.Amount;
import com.twitter.common.quantity.Time;
import com.twitter.common.stats.Stats;
import com.twitter.mesos.Tasks;
import com.twitter.mesos.gen.ScheduledTask;
import com.twitter.mesos.gen.TaskQuery;
import com.twitter.mesos.gen.TwitterTaskInfo;
import com.twitter.mesos.scheduler.Query;
import com.twitter.mesos.scheduler.storage.TaskStore;

import static com.google.common.base.Preconditions.checkNotNull;

/**
 * An in-memory task store.
 */
public class MemTaskStore implements TaskStore.Mutable.Transactioned {

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

  @CmdLine(name = "slow_query_log_threshold",
      help = "Log all queries that take at least this long to execute.")
  private static final Arg<Amount<Long, Time>> SLOW_QUERY_LOG_THRESHOLD =
      Arg.create(Amount.of(25L, Time.MILLISECONDS));

  /**
   * COPIER is separate from deepCopy to capture timing information.  Only the instrumented
   * {@link #deepCopy(ScheduledTask)} should interact directly with {@code COPIER}.
   */
  private static final Function<ScheduledTask, ScheduledTask> COPIER =
      Util.deepCopier();
  private final Function<ScheduledTask, ScheduledTask> deepCopy =
      new Function<ScheduledTask, ScheduledTask>() {
        @Override public ScheduledTask apply(ScheduledTask input) {
          return deepCopy(input);
        }
      };

  private final long slowQueryThresholdNanos = SLOW_QUERY_LOG_THRESHOLD.get().as(Time.NANOSECONDS);

  private final TransactionalMap<String, ScheduledTask> tasks =
      TransactionalMap.wrap(Maps.<String, ScheduledTask>newHashMap());
  private final TransactionalMap<String, Set<String>> tasksByJobKey =
      TransactionalMap.wrap(Maps.<String, Set<String>>newHashMap());

  private final AtomicLong taskQueriesById = Stats.exportLong("task_queries_by_id");
  private final AtomicLong taskQueriesByJob = Stats.exportLong("task_queries_by_job");
  private final AtomicLong taskQueriesAll = Stats.exportLong("task_queries_all");

  @Override
  public void commit() {
    tasks.commit();
    tasksByJobKey.commit();
  }

  @Override
  public void rollback() {
    tasks.rollback();
    tasksByJobKey.rollback();
  }

  @Timed("mem_storage_deep_copies")
  protected ScheduledTask deepCopy(ScheduledTask input) {
    return COPIER.apply(input);
  }

  @Timed("mem_storage_fetch_tasks")
  @Override
  public ImmutableSet<ScheduledTask> fetchTasks(TaskQuery query) {
    checkNotNull(query);

    long start = System.nanoTime();
    ImmutableSet<ScheduledTask> result = immutableMatches(query).toSet();
    long durationNanos = System.nanoTime() - start;
    Level level = (durationNanos >= slowQueryThresholdNanos) ? Level.INFO : Level.FINE;
    if (LOG.isLoggable(level)) {
      LOG.log(level, "Query took "
          + Amount.of(durationNanos, Time.NANOSECONDS).as(Time.MILLISECONDS) + " ms: " + query);
    }

    return result;
  }

  @Override
  public ImmutableSet<ScheduledTask> fetchTasks(Supplier<TaskQuery> querySupplier) {
    return fetchTasks(querySupplier.get());
  }

  @Timed("mem_storage_fetch_task_ids")
  @Override
  public Set<String> fetchTaskIds(TaskQuery query) {
    checkNotNull(query);

    return mutableMatches(query).transform(Tasks.SCHEDULED_TO_ID).toSet();
  }

  private Map<String, Collection<String>> taskIdsByJobKey(Iterable<ScheduledTask> toIndex) {
    return Multimaps.transformValues(
        Multimaps.index(toIndex, Tasks.SCHEDULED_TO_JOB_KEY),
        Tasks.SCHEDULED_TO_ID)
        .asMap();
  }

  @Override
  public Set<String> fetchTaskIds(Supplier<TaskQuery> querySupplier) {
    return fetchTaskIds(querySupplier.get());
  }

  @Timed("mem_storage_save_tasks")
  @Override
  public void saveTasks(Set<ScheduledTask> newTasks) {
    checkNotNull(newTasks);
    Preconditions.checkState(Tasks.ids(newTasks).size() == newTasks.size(),
        "Proposed new tasks would create task ID collision.");

    Set<ScheduledTask> immutable =
        FluentIterable.from(newTasks).transform(deepCopy).toSet();
    tasks.putAll(Maps.uniqueIndex(immutable, Tasks.SCHEDULED_TO_ID));

    // Update job key index.
    for (Map.Entry<String, Collection<String>> entry : taskIdsByJobKey(immutable).entrySet()) {
      ImmutableSet.Builder<String> newIds = ImmutableSet.builder();
      newIds.addAll(entry.getValue());
      Set<String> existingIds = tasksByJobKey.get(entry.getKey());
      if (existingIds != null) {
        newIds.addAll(existingIds);
      }

      // This is subtle but important: when modifying the items associated with a key, we must
      // re-insert rather than updating the contents of the set.  The transactional map has no
      // awareness of the contents of a value, so modifying the value set would break the
      // transactional guarantee.
      tasksByJobKey.put(entry.getKey(), newIds.build());
    }
  }

  @Timed("mem_storage_delete_all_tasks")
  @Override
  public void deleteAllTasks() {
    tasks.clear();
    tasksByJobKey.clear();
  }

  @Timed("mem_storage_delete_tasks")
  @Override
  public void deleteTasks(Set<String> taskIds) {
    checkNotNull(taskIds);

    // TransactionalMap does not presently support keySet modification, requiring iteration here.
    ImmutableList.Builder<ScheduledTask> gone = ImmutableList.builder();
    for (String id : taskIds) {
      ScheduledTask task = tasks.remove(id);
      if (task != null) {
        gone.add(task);
      }
    }
    for (Map.Entry<String, Collection<String>> entry : taskIdsByJobKey(gone.build()).entrySet()) {
      Set<String> existingIds = tasksByJobKey.get(entry.getKey());
      if (existingIds == null) {
        throw new IllegalStateException(
            "Index inconsistency - job for tasks not present in index: " + entry.getValue());
      }

      Set<String> newIds = ImmutableSet.copyOf(
          Sets.<String>difference(existingIds, ImmutableSet.copyOf(entry.getValue())));
      if (newIds.isEmpty()) {
        tasksByJobKey.remove(entry.getKey());
      } else {
        tasksByJobKey.put(entry.getKey(), newIds);
      }
    }
  }

  @Timed("mem_storage_mutate_tasks")
  @Override
  public ImmutableSet<ScheduledTask> mutateTasks(TaskQuery query, Closure<ScheduledTask> mutator) {
    checkNotNull(query);
    checkNotNull(mutator);

    ImmutableMap.Builder<String, ScheduledTask> mutated = ImmutableMap.builder();
    for (ScheduledTask original : immutableMatches(query)) {
      // Copy the object before invoking user code.  This is to support diff checking.
      ScheduledTask mutable = deepCopy.apply(original);
      mutator.execute(mutable);
      if (!original.equals(mutable)) {
        Preconditions.checkState(Tasks.id(original).equals(Tasks.id(mutable)),
            "A tasks ID may not be mutated.");

        // A diff is present - detach the mutable object from the closure's code to render
        // further mutation impossible.
        ScheduledTask updated = deepCopy.apply(mutable);
        mutated.put(Tasks.id(mutable), updated);
{'fraction_non_alphanumeric': 0.1115076383671425, 'fraction_numerical': 0.00025043826696719256, 'mean_word_length': 4.116271620755926, 'pattern_counts': {'":': 0, '<': 76, '<?xml version=': 0, '>': 78, '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': '7249990', 'n_tokens_mistral': 4884, 'n_tokens_neox': 4786, 'n_words': 1001}