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.scheduling; import java.lang.annotation.Retention; import java.lang.annotation.Target; import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.atomic.AtomicLong; import java.util.function.Function; import java.util.stream.Collectors; import javax.inject.Inject; import javax.inject.Qualifier; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Joiner; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.collect.Sets; import com.google.common.eventbus.Subscribe; import org.apache.aurora.common.inject.TimedInterceptor.Timed; import org.apache.aurora.common.stats.Stats; import org.apache.aurora.scheduler.TierManager; import org.apache.aurora.scheduler.base.Query; import org.apache.aurora.scheduler.base.TaskGroupKey; import org.apache.aurora.scheduler.configuration.executor.ExecutorSettings; import org.apache.aurora.scheduler.events.PubsubEvent; import org.apache.aurora.scheduler.filter.AttributeAggregate; import org.apache.aurora.scheduler.filter.SchedulingFilter.ResourceRequest; import org.apache.aurora.scheduler.preemptor.BiCache; import org.apache.aurora.scheduler.preemptor.Preemptor; import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider; import org.apache.aurora.scheduler.storage.Storage.StoreProvider; import org.apache.aurora.scheduler.storage.entities.IAssignedTask; 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.lang.annotation.ElementType.FIELD; import static java.lang.annotation.ElementType.METHOD; import static java.lang.annotation.ElementType.PARAMETER; import static java.lang.annotation.RetentionPolicy.RUNTIME; import static java.util.Objects.requireNonNull; import static org.apache.aurora.gen.ScheduleStatus.PENDING; /** * An asynchronous task scheduler. Scheduling of tasks is performed on a delay, where each task * backs off after a failed scheduling attempt. * <p> * Pending tasks are advertised to the scheduler via internal pubsub notifications. */ public class TaskSchedulerImpl implements TaskScheduler { /** * Binding annotation for the time duration of reservations. */ @VisibleForTesting @Qualifier @Target({ FIELD, PARAMETER, METHOD }) @Retention(RUNTIME) public @interface ReservationDuration { } private static final Logger LOG = LoggerFactory .getLogger(org.apache.aurora.scheduler.scheduling.TaskSchedulerImpl.class); private final TaskAssigner assigner; private final Preemptor preemptor; private final ExecutorSettings executorSettings; private final TierManager tierManager; private final BiCache<String, TaskGroupKey> reservations; private final AtomicLong attemptsFired = Stats.exportLong("schedule_attempts_fired"); private final AtomicLong attemptsFailed = Stats.exportLong("schedule_attempts_failed"); private final AtomicLong attemptsNoMatch = Stats.exportLong("schedule_attempts_no_match"); @Inject TaskSchedulerImpl(TaskAssigner assigner, Preemptor preemptor, ExecutorSettings executorSettings, TierManager tierManager, BiCache<String, TaskGroupKey> reservations) { this.assigner = requireNonNull(assigner); this.preemptor = requireNonNull(preemptor); this.executorSettings = requireNonNull(executorSettings); this.tierManager = requireNonNull(tierManager); this.reservations = requireNonNull(reservations); } @Timed("task_schedule_attempt") public Set<String> schedule(MutableStoreProvider store, Set<String> taskIds) { try { return scheduleTasks(store, taskIds); } catch (RuntimeException e) { // We catch the generic unchecked exception here to ensure tasks are not abandoned // if there is a transient issue resulting in an unchecked exception. LOG.warn("Task scheduling unexpectedly failed, will be retried", e); attemptsFailed.incrementAndGet(); // Return empty set for all task IDs to be retried later. // It's ok if some tasks were already assigned, those will be ignored in the next round. return ImmutableSet.of(); } } private Map<String, IAssignedTask> fetchTasks(StoreProvider store, Set<String> ids) { Map<String, IAssignedTask> tasks = store.getTaskStore().fetchTasks(Query.taskScoped(ids).byStatus(PENDING)) .stream().map(IScheduledTask::getAssignedTask) .collect(Collectors.toMap(IAssignedTask::getTaskId, Function.identity())); if (ids.size() != tasks.size()) { LOG.warn("Failed to look up tasks {}", Joiner.on(", ").join(Sets.difference(ids, tasks.keySet()))); } return tasks; } private Set<String> scheduleTasks(MutableStoreProvider store, Set<String> ids) { LOG.debug("Attempting to schedule tasks {}", ids); Map<String, IAssignedTask> tasksById = fetchTasks(store, ids); if (tasksById.isEmpty()) { // None of the tasks were found in storage. This could be caused by a task group that was // killed by the user, for example. return ids; } // Prepare scheduling context for the tasks ITaskConfig task = Iterables.getOnlyElement( tasksById.values().stream().map(IAssignedTask::getTask).collect(Collectors.toSet())); AttributeAggregate aggregate = AttributeAggregate.getJobActiveState(store, task.getJob()); // Attempt to schedule using available resources. Set<String> launched = assigner.maybeAssign(store, ResourceRequest.fromTask(task, executorSettings, aggregate, tierManager), TaskGroupKey.from(task), ImmutableSet.copyOf(tasksById.values()), reservations.asMap()); attemptsFired.addAndGet(tasksById.size()); // Fall back to preemption for tasks not scheduled above. Set<String> unassigned = Sets.difference(tasksById.keySet(), launched); unassigned.forEach(taskId -> { // TODO(maxim): Now that preemption slots are searched asynchronously, consider // retrying a launch attempt within the current scheduling round IFF a reservation is // available. maybePreemptFor(tasksById.get(taskId), aggregate, store); }); attemptsNoMatch.addAndGet(unassigned.size()); // Return all successfully launched tasks as well as those weren't tried (not in PENDING). return Sets.union(launched, Sets.difference(ids, tasksById.keySet())); } private void maybePreemptFor(IAssignedTask task, AttributeAggregate jobState, MutableStoreProvider storeProvider) { TaskGroupKey groupKey = TaskGroupKey.from(task.getTask()); if (!reservations.getByValue(groupKey).isEmpty()) { return; } preemptor.attemptPreemptionFor(task, jobState, storeProvider) .ifPresent(slaveId -> reservations.put(slaveId, groupKey)); } @Subscribe public void taskChanged(final PubsubEvent.TaskStateChange stateChangeEvent) { if (Optional.of(PENDING).equals(stateChangeEvent.getOldState())) { IAssignedTask assigned = stateChangeEvent.getTask().getAssignedTask(); if (assigned.getSlaveId() != null) { reservations.remove(assigned.getSlaveId(), TaskGroupKey.from(assigned.getTask())); } } } }