TaskConfigManager.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.storage.db;

import java.util.Map;
import java.util.Optional;

import javax.inject.Inject;

import com.google.common.collect.Maps;

import org.apache.aurora.scheduler.storage.db.views.DBResourceAggregate;
import org.apache.aurora.scheduler.storage.db.views.DbTaskConfig;
import org.apache.aurora.scheduler.storage.entities.IAppcImage;
import org.apache.aurora.scheduler.storage.entities.IConstraint;
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.IMesosContainer;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.apache.aurora.scheduler.storage.entities.IValueConstraint;

import static java.util.Objects.requireNonNull;

class TaskConfigManager {
  private final TaskConfigMapper configMapper;
  private final JobKeyMapper jobKeyMapper;

  @Inject
  TaskConfigManager(TaskConfigMapper configMapper, JobKeyMapper jobKeyMapper) {
    this.configMapper = requireNonNull(configMapper);
    this.jobKeyMapper = requireNonNull(jobKeyMapper);
  }

  private Optional<Long> getConfigRow(ITaskConfig config) {
    // NOTE: The 'config' object passed in MUST have all version-relevant fields populated in order
    // to correctly compare with objects loaded from DB. This may not hold true if a 'config' is
    // passed from storage recovery routine during version downgrade and fields are not properly
    // backfilled. See AURORA-1603 for more details.

    // We could optimize this slightly by first comparing the un-hydrated row and breaking early.
    Map<ITaskConfig, DbTaskConfig> rowsByConfig =
        Maps.uniqueIndex(
            configMapper.selectConfigsByJob(config.getJob()),
            DbTaskConfig::toImmutable);

    return Optional.ofNullable(rowsByConfig.get(config)).map(DbTaskConfig::getRowId);
  }

  long insert(ITaskConfig config) {
    InsertResult configInsert = new InsertResult();

    // Determine whether this config is already stored.
    Optional<Long> existingRow = getConfigRow(config);
    if (existingRow.isPresent()) {
      return existingRow.get();
    }

    jobKeyMapper.merge(config.getJob());
    configMapper.insert(config, configInsert);
    for (IConstraint constraint : config.getConstraints()) {
      InsertResult constraintResult = new InsertResult();
      configMapper.insertConstraint(configInsert.getId(), constraint, constraintResult);
      switch (constraint.getConstraint().getSetField()) {
        case VALUE:
          IValueConstraint valueConstraint = constraint.getConstraint().getValue();
          InsertResult valueResult = new InsertResult();
          configMapper.insertValueConstraint(
              constraintResult.getId(),
              valueConstraint,
              valueResult);
          configMapper.insertValueConstraintValues(
              valueResult.getId(),
              valueConstraint.getValues());
          break;

        case LIMIT:
          configMapper.insertLimitConstraint(
              constraintResult.getId(),
              constraint.getConstraint().getLimit());
          break;

        default:
          throw new IllegalStateException(
              "Unhandled constraint type " + constraint.getConstraint().getSetField());
      }
    }

    if (!config.getResources().isEmpty()) {
      configMapper.insertResources(
          configInsert.getId(),
          DBResourceAggregate.pairsFromResources(config.getResources()));
    }

    if (!config.getRequestedPorts().isEmpty()) {
      configMapper.insertRequestedPorts(configInsert.getId(), config.getRequestedPorts());
    }

    if (!config.getTaskLinks().isEmpty()) {
      configMapper.insertTaskLinks(configInsert.getId(), config.getTaskLinks());
    }

    if (!config.getMetadata().isEmpty()) {
      configMapper.insertMetadata(configInsert.getId(), config.getMetadata());
    }

    if (!config.getMesosFetcherUris().isEmpty()) {
      configMapper.insertMesosFetcherUris(configInsert.getId(), config.getMesosFetcherUris());
    }

    if (config.getContainer().isSetDocker()) {
      IDockerContainer container = config.getContainer().getDocker();
      InsertResult containerInsert = new InsertResult();
      configMapper.insertContainer(configInsert.getId(), container, containerInsert);
      if (!container.getParameters().isEmpty()) {
        configMapper.insertDockerParameters(containerInsert.getId(), container.getParameters());
      }
    } else if (config.getContainer().isSetMesos()
        && config.getContainer().getMesos().isSetImage()) {

      IMesosContainer container = config.getContainer().getMesos();
      IImage image = container.getImage();

      switch (image.getSetField()) {
        case DOCKER:
          IDockerImage dockerImage = image.getDocker();
          configMapper.insertDockerImage(
              configInsert.getId(),
              dockerImage.getName(),
              dockerImage.getTag());
          break;
        case APPC:
          IAppcImage appcImage = image.getAppc();
          configMapper.insertAppcImage(
              configInsert.getId(),
              appcImage.getName(),
              appcImage.getImageId());
          break;
        default:
          throw new IllegalStateException("Unexpected image type: " + image.getSetField());
      }

      if (!container.getVolumes().isEmpty()) {
        configMapper.insertVolumes(configInsert.getId(), container.getVolumes());
      }
    }

    return configInsert.getId();
  }
}