src/main/java/org/apache/aurora/scheduler/scheduling/TaskSchedulerImpl.java

/**
 * 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()));
      }
    }
  }
}