Java tutorial
/** * 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.storage.mem; import java.lang.annotation.ElementType; import java.lang.annotation.Retention; import java.lang.annotation.RetentionPolicy; import java.lang.annotation.Target; import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.atomic.AtomicLong; import javax.inject.Inject; import javax.inject.Qualifier; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.Functions; import com.google.common.base.Optional; import com.google.common.base.Preconditions; import com.google.common.base.Predicate; import com.google.common.base.Predicates; import com.google.common.base.Supplier; import com.google.common.collect.FluentIterable; import com.google.common.collect.HashMultimap; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.collect.Maps; import com.google.common.collect.Multimap; import com.google.common.collect.Multimaps; import org.apache.aurora.common.base.MorePreconditions; import org.apache.aurora.common.inject.TimedInterceptor.Timed; 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.gen.ScheduledTask; import org.apache.aurora.gen.TaskConfig; import org.apache.aurora.scheduler.base.JobKeys; import org.apache.aurora.scheduler.base.Query; import org.apache.aurora.scheduler.base.Tasks; import org.apache.aurora.scheduler.storage.TaskStore; import org.apache.aurora.scheduler.storage.entities.IJobKey; import org.apache.aurora.scheduler.storage.entities.IScheduledTask; import org.apache.aurora.scheduler.storage.entities.ITaskConfig; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import static java.util.Objects.requireNonNull; /** * An in-memory task store. */ class MemTaskStore implements TaskStore.Mutable { private static final Logger LOG = LoggerFactory.getLogger(MemTaskStore.class); /** * When true, enable snapshot deflation. */ @Retention(RetentionPolicy.RUNTIME) @Target({ ElementType.METHOD, ElementType.PARAMETER }) @Qualifier public @interface SlowQueryThreshold { } private final long slowQueryThresholdNanos; private static final Function<Query.Builder, Optional<Set<IJobKey>>> QUERY_TO_JOB_KEY = JobKeys::from; private static final Function<Query.Builder, Optional<Set<String>>> QUERY_TO_SLAVE_HOST = query -> query.get() .getSlaveHosts().isEmpty() ? Optional.absent() : Optional.of(query.get().getSlaveHosts()); // Since this class operates under the API and umbrella of {@link Storage}, it is expected to be // thread-safe but not necessarily strongly-consistent unless the externally-controlled storage // lock is secured. To adhere to that, these data structures are individually thread-safe, but // we don't lock across them because of the relaxed consistency guarantees. // Note that this behavior makes it possible to receive query results that are not sane, // specifically when a secondary key value is changed. In other words, we currently don't always // support the invariant that a query by slave host yields a result with all tasks matching that // slave host. This is deemed acceptable due to the fact that secondary key values are rarely // mutated in practice, and mutated in ways that are not impacted by this behavior. private final Map<String, Task> tasks = Maps.newConcurrentMap(); private final SecondaryIndex<IJobKey> jobIndex; private final List<SecondaryIndex<?>> secondaryIndices; // An interner is used here to collapse equivalent TaskConfig instances into canonical instances. // Ideally this would fall out of the object hierarchy (TaskConfig being associated with the job // rather than the task), but we intuit this detail here for performance reasons. private final Interner<TaskConfig, String> configInterner = new Interner<>(); private final AtomicLong taskQueriesById; private final AtomicLong taskQueriesAll; @Inject MemTaskStore(StatsProvider statsProvider, @SlowQueryThreshold Amount<Long, Time> slowQueryThreshold) { jobIndex = new SecondaryIndex<>(Tasks::getJob, QUERY_TO_JOB_KEY, statsProvider, "job"); secondaryIndices = ImmutableList.of(jobIndex, new SecondaryIndex<>(Tasks::scheduledToSlaveHost, QUERY_TO_SLAVE_HOST, statsProvider, "host")); slowQueryThresholdNanos = slowQueryThreshold.as(Time.NANOSECONDS); taskQueriesById = statsProvider.makeCounter("task_queries_by_id"); taskQueriesAll = statsProvider.makeCounter("task_queries_all"); } @Timed("mem_storage_fetch_task") @Override public Optional<IScheduledTask> fetchTask(String taskId) { requireNonNull(taskId); return Optional.fromNullable(tasks.get(taskId)).transform(t -> t.storedTask); } @Timed("mem_storage_fetch_tasks") @Override public ImmutableSet<IScheduledTask> fetchTasks(Query.Builder query) { requireNonNull(query); long start = System.nanoTime(); ImmutableSet<IScheduledTask> result = matches(query).toSet(); long durationNanos = System.nanoTime() - start; boolean infoLevel = durationNanos >= slowQueryThresholdNanos; long time = Amount.of(durationNanos, Time.NANOSECONDS).as(Time.MILLISECONDS); String message = "Query took {} ms: {}"; if (infoLevel) { LOG.info(message, time, query.get()); } else if (LOG.isDebugEnabled()) { LOG.debug(message, time, query.get()); } return result; } @Timed("mem_storage_get_job_keys") @Override public Set<IJobKey> getJobKeys() { return jobIndex.keySet(); } private final Function<IScheduledTask, Task> toTask = task -> new Task(task, configInterner); @Timed("mem_storage_save_tasks") @Override public void saveTasks(Set<IScheduledTask> newTasks) { requireNonNull(newTasks); Preconditions.checkState(Tasks.ids(newTasks).size() == newTasks.size(), "Proposed new tasks would create task ID collision."); Iterable<Task> canonicalized = Iterables.transform(newTasks, toTask); tasks.putAll(Maps.uniqueIndex(canonicalized, TO_ID)); for (SecondaryIndex<?> index : secondaryIndices) { index.insert(Iterables.transform(canonicalized, TO_SCHEDULED)); } } @Timed("mem_storage_delete_all_tasks") @Override public void deleteAllTasks() { tasks.clear(); for (SecondaryIndex<?> index : secondaryIndices) { index.clear(); } configInterner.clear(); } @Timed("mem_storage_delete_tasks") @Override public void deleteTasks(Set<String> taskIds) { requireNonNull(taskIds); for (String id : taskIds) { Task removed = tasks.remove(id); if (removed != null) { for (SecondaryIndex<?> index : secondaryIndices) { index.remove(removed.storedTask); } configInterner.removeAssociation(removed.storedTask.getAssignedTask().getTask().newBuilder(), id); } } } @Timed("mem_storage_mutate_task") @Override public Optional<IScheduledTask> mutateTask(String taskId, Function<IScheduledTask, IScheduledTask> mutator) { return fetchTask(taskId).transform(original -> { IScheduledTask maybeMutated = mutator.apply(original); requireNonNull(maybeMutated); if (!original.equals(maybeMutated)) { Preconditions.checkState(Tasks.id(original).equals(Tasks.id(maybeMutated)), "A task's ID may not be mutated."); tasks.put(Tasks.id(maybeMutated), toTask.apply(maybeMutated)); for (SecondaryIndex<?> index : secondaryIndices) { index.replace(original, maybeMutated); } } return maybeMutated; }); } @Timed("mem_storage_unsafe_modify_in_place") @Override public boolean unsafeModifyInPlace(String taskId, ITaskConfig taskConfiguration) { MorePreconditions.checkNotBlank(taskId); requireNonNull(taskConfiguration); Task stored = tasks.get(taskId); if (stored == null) { return false; } else { ScheduledTask updated = stored.storedTask.newBuilder(); updated.getAssignedTask().setTask(taskConfiguration.newBuilder()); tasks.put(taskId, toTask.apply(IScheduledTask.build(updated))); return true; } } private static Predicate<Task> queryFilter(Query.Builder query) { return Predicates.compose(Util.queryFilter(query), new Function<Task, IScheduledTask>() { @Override public IScheduledTask apply(Task canonicalTask) { return canonicalTask.storedTask; } }); } private Iterable<Task> fromIdIndex(Iterable<String> taskIds) { return FluentIterable.from(taskIds).transform(Functions.forMap(tasks, null)).filter(Predicates.notNull()) .toList(); } private FluentIterable<IScheduledTask> matches(Query.Builder query) { // Apply the query against the working set. Optional<? extends Iterable<Task>> from = Optional.absent(); if (query.get().getTaskIds().isEmpty()) { for (SecondaryIndex<?> index : secondaryIndices) { Optional<Iterable<String>> indexMatch = index.getMatches(query); if (indexMatch.isPresent()) { // Note: we could leverage multiple indexes here if the query applies to them, by // choosing to intersect the results. Given current indexes and query profile, this is // unlikely to offer much improvement, though. from = Optional.of(fromIdIndex(indexMatch.get())); break; } } // No indices match, fall back to a full scan. if (!from.isPresent()) { taskQueriesAll.incrementAndGet(); from = Optional.of(tasks.values()); } } else { taskQueriesById.incrementAndGet(); from = Optional.of(fromIdIndex(query.get().getTaskIds())); } return FluentIterable.from(from.get()).filter(queryFilter(query)).transform(TO_SCHEDULED); } private static final Function<Task, IScheduledTask> TO_SCHEDULED = task -> task.storedTask; private static final Function<Task, String> TO_ID = Functions.compose(Tasks::id, TO_SCHEDULED); private static class Task { private final IScheduledTask storedTask; Task(IScheduledTask storedTask, Interner<TaskConfig, String> interner) { interner.removeAssociation(storedTask.getAssignedTask().getTask().newBuilder(), Tasks.id(storedTask)); TaskConfig canonical = interner.addAssociation(storedTask.getAssignedTask().getTask().newBuilder(), Tasks.id(storedTask)); ScheduledTask builder = storedTask.newBuilder(); builder.getAssignedTask().setTask(canonical); this.storedTask = IScheduledTask.build(builder); } @Override public boolean equals(Object o) { if (!(o instanceof Task)) { return false; } Task other = (Task) o; return storedTask.equals(other.storedTask); } @Override public int hashCode() { return storedTask.hashCode(); } } @VisibleForTesting static String getIndexSizeStatName(String name) { return "task_store_index_" + name + "_items"; } /** * A non-unique secondary index on the task store. Maps a custom key type to a set of task IDs. * * @param <K> Key type. */ private static class SecondaryIndex<K> { private final Multimap<K, String> index = Multimaps.synchronizedSetMultimap(HashMultimap.create()); private final Function<IScheduledTask, K> indexer; private final Function<Query.Builder, Optional<Set<K>>> queryExtractor; private final AtomicLong hitCount; /** * Creates a secondary index that will extract keys from tasks using the provided indexer. * * @param indexer Indexing function. * @param queryExtractor Function to extract the keys relevant to a query. * @param statsProvider Stats system to export metrics to. * @param name Name to use in stats keys. */ SecondaryIndex(Function<IScheduledTask, K> indexer, Function<Query.Builder, Optional<Set<K>>> queryExtractor, StatsProvider statsProvider, String name) { this.indexer = indexer; this.queryExtractor = queryExtractor; this.hitCount = statsProvider.makeCounter("task_queries_by_" + name); statsProvider.makeGauge(getIndexSizeStatName(name), new Supplier<Number>() { @Override public Number get() { return index.size(); } }); } Set<K> keySet() { return ImmutableSet.copyOf(index.keySet()); } void insert(Iterable<IScheduledTask> tasks) { for (IScheduledTask task : tasks) { insert(task); } } void insert(IScheduledTask task) { K key = indexer.apply(task); if (key != null) { index.put(key, Tasks.id(task)); } } void clear() { index.clear(); } void remove(IScheduledTask task) { K key = indexer.apply(task); if (key != null) { index.remove(key, Tasks.id(task)); } } void replace(IScheduledTask old, IScheduledTask replacement) { synchronized (index) { remove(old); insert(replacement); } } private final Function<Set<K>, Iterable<String>> lookup = new Function<Set<K>, Iterable<String>>() { @Override public Iterable<String> apply(Set<K> keys) { hitCount.incrementAndGet(); ImmutableSet.Builder<String> builder = ImmutableSet.builder(); synchronized (index) { for (K key : keys) { builder.addAll(index.get(key)); } } return builder.build(); } }; Optional<Iterable<String>> getMatches(Query.Builder query) { return queryExtractor.apply(query).transform(lookup); } } }