content: /**
   * 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.util.Collection;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.atomic.AtomicLong;
  
  import javax.inject.Inject;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.Iterables;
  import com.google.common.collect.Maps;
  
  import org.apache.aurora.common.stats.StatsProvider;
  import org.apache.aurora.scheduler.base.InstanceKeys;
  import org.apache.aurora.scheduler.base.TaskGroupKey;
  import org.apache.aurora.scheduler.filter.SchedulingFilter.ResourceRequest;
  import org.apache.aurora.scheduler.mesos.MesosTaskFactory;
  import org.apache.aurora.scheduler.offers.HostOffer;
  import org.apache.aurora.scheduler.offers.OfferManager;
  import org.apache.aurora.scheduler.offers.OfferManager.LaunchException;
  import org.apache.aurora.scheduler.resources.ResourceManager;
  import org.apache.aurora.scheduler.resources.ResourceType;
  import org.apache.aurora.scheduler.state.StateManager;
  import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
  import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
  import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
  import org.apache.aurora.scheduler.updater.UpdateAgentReserver;
  import org.apache.mesos.v1.Protos;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  import static org.apache.aurora.common.inject.TimedInterceptor.Timed;
  import static org.apache.aurora.gen.ScheduleStatus.ASSIGNED;
  import static org.apache.aurora.gen.ScheduleStatus.LOST;
  
  public class TaskAssignerImpl implements TaskAssigner {
    private static final Logger LOG = LoggerFactory.getLogger(TaskAssignerImpl.class);
  
    @VisibleForTesting
    static final Optional<String> LAUNCH_FAILED_MSG =
        Optional.of("Unknown exception attempting to schedule task.");
    @VisibleForTesting
    static final String ASSIGNER_LAUNCH_FAILURES = "assigner_launch_failures";
  
    private final AtomicLong launchFailures;
  
    private final StateManager stateManager;
    private final MesosTaskFactory taskFactory;
    private final OfferManager offerManager;
    private final UpdateAgentReserver updateAgentReserver;
  
    @Inject
    public TaskAssignerImpl(
        StateManager stateManager,
        MesosTaskFactory taskFactory,
        OfferManager offerManager,
        UpdateAgentReserver updateAgentReserver,
        StatsProvider statsProvider) {
  
      this.stateManager = requireNonNull(stateManager);
      this.taskFactory = requireNonNull(taskFactory);
      this.offerManager = requireNonNull(offerManager);
      this.launchFailures = statsProvider.makeCounter(ASSIGNER_LAUNCH_FAILURES);
      this.updateAgentReserver = requireNonNull(updateAgentReserver);
    }
  
    @VisibleForTesting
    IAssignedTask mapAndAssignResources(Protos.Offer offer, IAssignedTask task) {
      IAssignedTask assigned = task;
      for (ResourceType type : ResourceManager.getTaskResourceTypes(assigned)) {
        if (type.getMapper().isPresent()) {
          assigned = type.getMapper().get().mapAndAssign(offer, assigned);
        }
      }
      return assigned;
    }
  
    private Protos.TaskInfo assign(
        MutableStoreProvider storeProvider,
        Protos.Offer offer,
        String taskId,
        boolean revocable) {
  
      String host = offer.getHostname();
      IAssignedTask assigned = stateManager.assignTask(
          storeProvider,
          taskId,
          host,
          offer.getAgentId(),
          task -> mapAndAssignResources(offer, task));
      LOG.info(
          "Offer on agent {} (id {}) is being assigned task for {}.",
          host, offer.getAgentId().getValue(), taskId);
      return taskFactory.createFrom(assigned, offer, revocable);
    }
  
    private void launchUsingOffer(
        MutableStoreProvider stores,
        ResourceRequest resourceRequest,
        IAssignedTask task,
        HostOffer offer) throws LaunchException {
  
      String taskId = task.getTaskId();
      Protos.TaskInfo taskInfo =
          assign(stores, offer.getOffer(), taskId, resourceRequest.isRevocable());
      resourceRequest.getJobState().updateAttributeAggregate(offer.getAttributes());
      try {
        offerManager.launchTask(offer.getOffer().getId(), taskInfo);
      } catch (LaunchException e) {
        LOG.warn("Failed to launch task.", e);
        launchFailures.incrementAndGet();
  
        // The attempt to schedule the task failed, so we need to backpedal on the assignment.
        // It is in the LOST state and a new task will move to PENDING to replace it.
        // Should the state change fail due to storage issues, that's okay.  The task will
        // time out in the ASSIGNED state and be moved to LOST.
        stateManager.changeState(stores, taskId, Optional.of(ASSIGNED), LOST, LAUNCH_FAILED_MSG);
        throw e;
      }
    }
  
    private static final class ReservationStatus {
      final boolean taskReserving;
      final Optional<HostOffer> offer;
  
      private ReservationStatus(boolean taskReserving, Optional<HostOffer> offer) {
        this.taskReserving = taskReserving;
        this.offer = requireNonNull(offer);
      }
  
      static final ReservationStatus NOT_RESERVING = new ReservationStatus(false, Optional.empty());
      static final ReservationStatus NOT_READY = new ReservationStatus(true, Optional.empty());
  
      static ReservationStatus ready(HostOffer offer) {
        return new ReservationStatus(true, Optional.of(offer));
      }
  
      boolean isTaskReserving() {
        return taskReserving;
      }
  
      Optional<HostOffer> getOffer() {
        return offer;
      }
    }
  
    private ReservationStatus getReservation(IAssignedTask task, ResourceRequest resourceRequest) {
  
      IInstanceKey key = InstanceKeys.from(task.getTask().getJob(), task.getInstanceId());
      Optional<String> agentId = updateAgentReserver.getAgent(key);
      if (!agentId.isPresent()) {
        return ReservationStatus.NOT_RESERVING;
      }
      Optional<HostOffer> offer = offerManager.getMatching(
          Protos.AgentID.newBuilder().setValue(agentId.get()).build(),
          resourceRequest);
      if (offer.isPresent()) {
        LOG.info("Used update reservation for {} on {}", key, agentId.get());
        updateAgentReserver.release(agentId.get(), key);
        return ReservationStatus.ready(offer.get());
      } else {
        LOG.info(
            "Tried to reuse offer on {} for {}, but was not ready yet.",
            agentId.get(),
            key);
        return ReservationStatus.NOT_READY;
      }
    }
  
    /**
     * Determine whether or not the offer is reserved for a different task via preemption or
     * update affinity.
     */
    private boolean isAgentReserved(
        HostOffer offer,
        TaskGroupKey groupKey,
        Map<String, TaskGroupKey> preemptionReservations) {
  
      String agentId = offer.getOffer().getAgentId().getValue();
      boolean reservedForPreemption = Optional.ofNullable(preemptionReservations.get(agentId))
          .map(group -> !group.equals(groupKey))
          .orElse(false);
  
      return reservedForPreemption || updateAgentReserver.isReserved(agentId);
    }
  
    private static class SchedulingMatch {
      final IAssignedTask task;
      final HostOffer offer;
  
      SchedulingMatch(IAssignedTask task, HostOffer offer) {
        this.task = requireNonNull(task);
        this.offer = requireNonNull(offer);
      }
    }
  
    private Collection<SchedulingMatch> findMatches(
        ResourceRequest resourceRequest,
        TaskGroupKey groupKey,
        Set<IAssignedTask> tasks,
        Map<String, TaskGroupKey> preemptionReservations) {
  
      // Avoid matching multiple tasks against any offer.
      Map<String, SchedulingMatch> matchesByOffer = Maps.newHashMap();
  
      tasks.forEach(task -> {
        ReservationStatus reservation = getReservation(task, resourceRequest);
        Optional<HostOffer> chosenOffer;
        if (reservation.isTaskReserving()) {
          // Use the reserved offer, which may not currently exist.
          chosenOffer = reservation.getOffer();
        } else {
          // Get all offers that will satisfy the given ResourceRequest and that are not reserved
          // for updates or preemption.
          Iterable<HostOffer> matchingOffers = Iterables.filter(
              offerManager.getAllMatching(groupKey, resourceRequest),
              o -> !matchesByOffer.containsKey(o.getOffer().getId().getValue())
                  && !isAgentReserved(o, groupKey, preemptionReservations));
  
          chosenOffer = Optional.ofNullable(Iterables.getFirst(matchingOffers, null));
        }
  
        chosenOffer.ifPresent(hostOffer -> matchesByOffer.put(
            hostOffer.getOffer().getId().getValue(),
            new SchedulingMatch(task, hostOffer)));
      });
  
      return matchesByOffer.values();
    }
  
    @Timed("assigner_maybe_assign")
    @Override
    public Set<String> maybeAssign(
        MutableStoreProvider storeProvider,
        ResourceRequest resourceRequest,
        TaskGroupKey groupKey,
        Set<IAssignedTask> tasks,
        Map<String, TaskGroupKey> reservations) {
  
      ImmutableSet.Builder<String> assigned = ImmutableSet.builder();
  
      for (SchedulingMatch match : findMatches(resourceRequest, groupKey, tasks, reservations)) {
        try {
          launchUsingOffer(storeProvider, resourceRequest, match.task, match.offer);
          assigned.add(match.task.getTaskId());
        } catch (LaunchException e) {
          // Any launch exception causes the scheduling round to terminate for this TaskGroup.
          break;
        }
      }
  
      return assigned.build();
    }
  }
  