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.mesos;
  
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.stream.Collectors;
  
  import javax.inject.Inject;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.Iterables;
  import com.google.protobuf.ByteString;
  
  import org.apache.aurora.Protobufs;
  import org.apache.aurora.codec.ThriftBinaryCodec;
  import org.apache.aurora.scheduler.base.JobKeys;
  import org.apache.aurora.scheduler.base.SchedulerException;
  import org.apache.aurora.scheduler.base.Tasks;
  import org.apache.aurora.scheduler.configuration.executor.ExecutorSettings;
  import org.apache.aurora.scheduler.resources.AcceptedOffer;
  import org.apache.aurora.scheduler.resources.ResourceBag;
  import org.apache.aurora.scheduler.resources.ResourceManager;
  import org.apache.aurora.scheduler.storage.entities.IAppcImage;
  import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
  import org.apache.aurora.scheduler.storage.entities.IDockerContainer;
  import org.apache.aurora.scheduler.storage.entities.IDockerImage;
  import org.apache.aurora.scheduler.storage.entities.IImage;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IMesosContainer;
  import org.apache.aurora.scheduler.storage.entities.IMesosFetcherURI;
  import org.apache.aurora.scheduler.storage.entities.IServerInfo;
  import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
  import org.apache.mesos.v1.Protos;
  import org.apache.mesos.v1.Protos.CommandInfo;
  import org.apache.mesos.v1.Protos.CommandInfo.URI.Builder;
  import org.apache.mesos.v1.Protos.ContainerInfo;
  import org.apache.mesos.v1.Protos.DiscoveryInfo;
  import org.apache.mesos.v1.Protos.ExecutorID;
  import org.apache.mesos.v1.Protos.ExecutorInfo;
  import org.apache.mesos.v1.Protos.Label;
  import org.apache.mesos.v1.Protos.Labels;
  import org.apache.mesos.v1.Protos.Offer;
  import org.apache.mesos.v1.Protos.Port;
  import org.apache.mesos.v1.Protos.Resource;
  import org.apache.mesos.v1.Protos.TaskID;
  import org.apache.mesos.v1.Protos.TaskInfo;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  import static org.apache.aurora.gen.apiConstants.TASK_FILESYSTEM_MOUNT_POINT;
  
  /**
   * A factory to create mesos task objects.
   */
  public interface MesosTaskFactory {
  
    /**
     * Creates a mesos task object.
     *
     * @param task Assigned task to translate into a task object.
     * @param offer Resource offer the task is being assigned to.
     * @return A new task.
     * @throws SchedulerException If the task could not be encoded.
     */
    TaskInfo createFrom(IAssignedTask task, Offer offer, boolean revocable) throws SchedulerException;
  
    // TODO(wfarner): Move this class to its own file to reduce visibility to package private.
    class MesosTaskFactoryImpl implements MesosTaskFactory {
      private static final Logger LOG = LoggerFactory.getLogger(MesosTaskFactoryImpl.class);
  
      private static final String AURORA_LABEL_PREFIX = "org.apache.aurora";
  
      @VisibleForTesting
      static final String METADATA_LABEL_PREFIX = AURORA_LABEL_PREFIX + ".metadata.";
  
      @VisibleForTesting
      static final String DEFAULT_PORT_PROTOCOL = "TCP";
  
      // N.B. We intentionally do not prefix this label. It was added when the explicit source field
      // was removed from the ExecutorInfo proto and named "source" per guidance from Mesos devs.
      @VisibleForTesting
      static final String SOURCE_LABEL = "source";
  
      @VisibleForTesting
      static final String TIER_LABEL = AURORA_LABEL_PREFIX + ".tier";
  
      private final ExecutorSettings executorSettings;
      private final IServerInfo serverInfo;
  
      @Inject
      MesosTaskFactoryImpl(ExecutorSettings executorSettings, IServerInfo serverInfo) {
        this.executorSettings = requireNonNull(executorSettings);
        this.serverInfo = requireNonNull(serverInfo);
      }
  
      @VisibleForTesting
      static ExecutorID getExecutorId(String taskId, String taskPrefix) {
        return ExecutorID.newBuilder().setValue(taskPrefix + taskId).build();
      }
  
      private static String getJobSourceName(IJobKey jobkey) {
        return String.join(".", jobkey.getRole(), jobkey.getEnvironment(), jobkey.getName());
      }
  
      private static String getJobSourceName(ITaskConfig task) {
        return getJobSourceName(task.getJob());
      }
  
      private static String getExecutorName(IAssignedTask task) {
        return task.getTask().getExecutorConfig().getName();
      }
  
      @VisibleForTesting
      static String getInstanceSourceName(ITaskConfig task, int instanceId) {
        return String.join(".", getJobSourceName(task), Integer.toString(instanceId));
      }
  
      @VisibleForTesting
      static String getInverseJobSourceName(IJobKey job) {
        return String.join(".", job.getName(), job.getEnvironment(), job.getRole());
      }
  
      private static byte[] serializeTask(IAssignedTask task) throws SchedulerException {
        try {
          return ThriftBinaryCodec.encode(task.newBuilder());
        } catch (ThriftBinaryCodec.CodingException e) {
          LOG.error("Unable to serialize task.", e);
          throw new SchedulerException("Internal error.", e);
        }
      }
  
      @Override
      public TaskInfo createFrom(IAssignedTask task, Offer offer, boolean revocable)
          throws SchedulerException {
  
        requireNonNull(task);
        requireNonNull(offer);
  
        ITaskConfig config = task.getTask();
  
        ResourceBag executorOverhead = executorSettings.getExecutorOverhead(config);
  
        AcceptedOffer acceptedOffer;
        // TODO(wfarner): Re-evaluate if/why we need to continue handling unset assignedPorts field.
        try {
          acceptedOffer = AcceptedOffer.create(
              offer,
              task,
              executorOverhead,
              revocable);
        } catch (ResourceManager.InsufficientResourcesException e) {
          throw new SchedulerException(e);
        }
        Iterable<Resource> resources = acceptedOffer.getTaskResources();
  
        if (LOG.isDebugEnabled()) {
          LOG.debug(
              "Setting task resources to {}",
              Iterables.transform(resources, Protobufs::toString));
        }
  
        TaskInfo.Builder taskBuilder = TaskInfo.newBuilder()
            .setName(JobKeys.canonicalString(Tasks.getJob(task)))
            .setTaskId(TaskID.newBuilder().setValue(task.getTaskId()))
            .setAgentId(offer.getAgentId())
            .addAllResources(resources);
  
        configureTaskLabels(config, taskBuilder);
  
        if (executorSettings.shouldPopulateDiscoverInfo()) {
          configureDiscoveryInfos(task, taskBuilder);
        }
  
        if (config.getContainer().isSetMesos()) {
          ExecutorInfo.Builder executorInfoBuilder = configureTaskForExecutor(task, acceptedOffer);
  
          Optional<ContainerInfo.Builder> containerInfoBuilder = configureTaskForImage(
              task.getTask().getContainer().getMesos(),
              getExecutorName(task));
          if (containerInfoBuilder.isPresent()) {
            executorInfoBuilder.setContainer(containerInfoBuilder.get());
          }
  
          taskBuilder.setExecutor(executorInfoBuilder.build());
        } else if (config.getContainer().isSetDocker()) {
          IDockerContainer dockerContainer = config.getContainer().getDocker();
          if (config.isSetExecutorConfig()) {
            ExecutorInfo.Builder execBuilder = configureTaskForExecutor(task, acceptedOffer)
                .setContainer(getDockerContainerInfo(
                    dockerContainer,
                    Optional.of(getExecutorName(task))));
            taskBuilder.setExecutor(execBuilder.build());
          } else {
            LOG.warn("Running Docker-based task without an executor.");
            taskBuilder.setContainer(getDockerContainerInfo(dockerContainer, Optional.empty()))
                .setCommand(CommandInfo.newBuilder().setShell(false));
          }
        } else {
          throw new SchedulerException("Task had no supported container set.");
        }
  
        if (taskBuilder.hasExecutor()) {
          taskBuilder.setData(ByteString.copyFrom(serializeTask(task)));
        }
  
        return taskBuilder.build();
      }
  
      private Optional<ContainerInfo.Builder> configureTaskForImage(
          IMesosContainer mesosContainer,
          String executorName) {
        requireNonNull(mesosContainer);
  
        if (mesosContainer.isSetImage()) {
          IImage image = mesosContainer.getImage();
  
          Protos.Image.Builder imageBuilder = Protos.Image.newBuilder();
  
          if (image.isSetAppc()) {
            IAppcImage appcImage = image.getAppc();
  
            imageBuilder.setType(Protos.Image.Type.APPC);
            imageBuilder.setAppc(Protos.Image.Appc.newBuilder()
                .setName(appcImage.getName())
                .setId(appcImage.getImageId()));
          } else if (image.isSetDocker()) {
            IDockerImage dockerImage = image.getDocker();
  
            imageBuilder.setType(Protos.Image.Type.DOCKER);
            imageBuilder.setDocker(Protos.Image.Docker.newBuilder()
                .setName(dockerImage.getName() + ":" + dockerImage.getTag()));
          } else {
            throw new SchedulerException("Task had no supported image set.");
          }
  
          ContainerInfo.MesosInfo.Builder mesosContainerBuilder =
              ContainerInfo.MesosInfo.newBuilder();
  
          Iterable<Protos.Volume> containerVolumes = Iterables.transform(mesosContainer.getVolumes(),
              input -> Protos.Volume.newBuilder()
              .setMode(Protos.Volume.Mode.valueOf(input.getMode().name()))
              .setHostPath(input.getHostPath())
              .setContainerPath(input.getContainerPath())
              .build());
  
          Protos.Volume volume = Protos.Volume.newBuilder()
              .setImage(imageBuilder)
              .setContainerPath(TASK_FILESYSTEM_MOUNT_POINT)
              .setMode(Protos.Volume.Mode.RO)
              .build();
  
          return Optional.of(ContainerInfo.newBuilder()
              .setType(ContainerInfo.Type.MESOS)
              .setMesos(mesosContainerBuilder)
              .addAllVolumes(executorSettings.getExecutorConfig(executorName).get().getVolumeMounts())
              .addAllVolumes(containerVolumes)
              .addVolumes(volume));
        }
  
        return Optional.empty();
      }
  
      private ContainerInfo getDockerContainerInfo(
          IDockerContainer config,
          Optional<String> executorName) {
  
        Iterable<Protos.Parameter> parameters = Iterables.transform(config.getParameters(),
            item -> Protos.Parameter.newBuilder().setKey(item.getName())
              .setValue(item.getValue()).build());
  
        ContainerInfo.DockerInfo.Builder dockerBuilder = ContainerInfo.DockerInfo.newBuilder()
            .setImage(config.getImage()).addAllParameters(parameters);
        return ContainerInfo.newBuilder()
            .setType(ContainerInfo.Type.DOCKER)
            .setDocker(dockerBuilder.build())
            .addAllVolumes(
                executorName.isPresent()
                    ? executorSettings.getExecutorConfig(executorName.get()).get().getVolumeMounts()
                    : ImmutableList.of())
            .build();
      }
  
      private static Protos.CommandInfo.URI toProtoURI(IMesosFetcherURI u) {
        Builder builder = Protos.CommandInfo.URI.newBuilder()
            .setValue(u.getValue())
            .setExecutable(false)
            .setExtract(u.isExtract())
            .setCache(u.isCache());
  
        if (u.isSetOutputFile()) {
          builder.setOutputFile(u.getOutputFile());
        }
  
        return builder.build();
      }
  
      @SuppressWarnings("deprecation") // we set the source field for backwards compat.
      private ExecutorInfo.Builder configureTaskForExecutor(
          IAssignedTask task,
          AcceptedOffer acceptedOffer) {
  
        String sourceName = getInstanceSourceName(task.getTask(), task.getInstanceId());
  
        ExecutorInfo.Builder builder =
            executorSettings.getExecutorConfig(getExecutorName(task)).get()
            .getExecutor()
            .toBuilder()
            .setExecutorId(getExecutorId(
                task.getTaskId(),
                executorSettings.getExecutorConfig(getExecutorName(task)).get().getTaskPrefix()))
            .setSource(sourceName)
            .setLabels(
                Labels.newBuilder().addLabels(
                    Label.newBuilder()
                        .setKey(SOURCE_LABEL)
                        .setValue(sourceName)));
  
        List<CommandInfo.URI> mesosFetcherUris = task.getTask().getMesosFetcherUris().stream()
            .map(u -> toProtoURI(u))
            .collect(Collectors.toList());
  
        builder.setCommand(builder.getCommand().toBuilder().addAllUris(mesosFetcherUris));
  
        Iterable<Resource> executorResources = acceptedOffer.getExecutorResources();
        if (LOG.isDebugEnabled()) {
          LOG.debug(
              "Setting executor resources to {}",
              Iterables.transform(executorResources, Protobufs::toString));
        }
        builder.clearResources().addAllResources(executorResources);
        return builder;
      }
  
      private void configureTaskLabels(ITaskConfig config, TaskInfo.Builder taskBuilder) {
        Labels.Builder labelsBuilder = Labels.newBuilder();
        labelsBuilder.addLabels(Label.newBuilder().setKey(TIER_LABEL).setValue(config.getTier()));
  
        config.getMetadata().stream().forEach(m -> labelsBuilder.addLabels(Label.newBuilder()
            .setKey(METADATA_LABEL_PREFIX + m.getKey())
            .setValue(m.getValue())
            .build()));
  
        taskBuilder.setLabels(labelsBuilder);
      }
  
      private void configureDiscoveryInfos(IAssignedTask task, TaskInfo.Builder taskBuilder) {
        DiscoveryInfo.Builder builder = taskBuilder.getDiscoveryBuilder();
        builder.setVisibility(DiscoveryInfo.Visibility.CLUSTER);
        builder.setName(getInverseJobSourceName(task.getTask().getJob()));
        builder.setEnvironment(task.getTask().getJob().getEnvironment());
        // A good sane choice for default location is current Aurora cluster name.
        builder.setLocation(serverInfo.getClusterName());
        for (Map.Entry<String, Integer> entry : task.getAssignedPorts().entrySet()) {
          builder.getPortsBuilder().addPorts(
              Port.newBuilder()
                  .setName(entry.getKey())
                  .setNumber(entry.getValue())
                  .setProtocol(DEFAULT_PORT_PROTOCOL)
          );
        }
      }
  
    }
  }
  