org.apache.aurora.scheduler.resources.ResourceManager.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.aurora.scheduler.resources.ResourceManager.java

Source

/**
 * 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.resources;

import java.util.EnumSet;
import java.util.Map;
import java.util.Set;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import com.google.common.base.Predicate;
import com.google.common.base.Predicates;
import com.google.common.collect.Iterables;

import org.apache.aurora.gen.ResourceAggregate;
import org.apache.aurora.scheduler.TierInfo;
import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
import org.apache.aurora.scheduler.storage.entities.IResource;
import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.apache.aurora.scheduler.storage.log.ThriftBackfill;
import org.apache.mesos.v1.Protos.Resource;

import static org.apache.aurora.scheduler.resources.ResourceType.BY_MESOS_NAME;
import static org.apache.aurora.scheduler.resources.ResourceType.fromResource;
import static org.apache.mesos.v1.Protos.Offer;

/**
 * Manages resources and provides Aurora/Mesos translation.
 */
public final class ResourceManager {
    private ResourceManager() {
        // Utility class.
    }

    /**
     * TODO(maxim): reduce visibility by redirecting callers to #getRevocableOfferResources().
     */
    public static final Predicate<Resource> REVOCABLE = r -> !fromResource(r).isMesosRevocable()
            || r.hasRevocable();

    /**
     * TODO(maxim): reduce visibility by redirecting callers to #getNonRevocableOfferResources().
     */
    public static final Predicate<Resource> NON_REVOCABLE = r -> !r.hasRevocable();

    private static final Function<IResource, ResourceType> RESOURCE_TO_TYPE = r -> fromResource(r);

    private static final Function<Resource, ResourceType> MESOS_RESOURCE_TO_TYPE = r -> fromResource(r);

    private static final Function<IResource, Double> QUANTIFY_RESOURCE = r -> fromResource(r)
            .getAuroraResourceConverter().quantify(r.getRawValue());

    private static final Function<Resource, Double> QUANTIFY_MESOS_RESOURCE = r -> fromResource(r)
            .getMesosResourceConverter().quantify(r);

    private static final BinaryOperator<Double> REDUCE_VALUES = (l, r) -> l + r;

    /**
     * TODO(rdelvalle): Remove filters when arbitrary resources are fully supported (AURORA-1328).
     */
    private static final Predicate<Resource> SUPPORTED_RESOURCE = r -> BY_MESOS_NAME.containsKey(r.getName());

    /**
     * Gets offer resources matching specified {@link ResourceType}.
     *
     * @param offer Offer to get resources from.
     * @param type {@link ResourceType} to filter resources by.
     * @return Offer resources matching {@link ResourceType}.
     */
    public static Iterable<Resource> getOfferResources(Offer offer, ResourceType type) {
        return Iterables.filter(Iterables.filter(offer.getResourcesList(), SUPPORTED_RESOURCE),
                r -> fromResource(r).equals(type));
    }

    /**
     * Gets Mesos-revocable offer resources.
     *
     * @param offer Offer to get resources from.
     * @return Mesos-revocable offer resources.
     */
    public static Iterable<Resource> getRevocableOfferResources(Offer offer) {
        return Iterables.filter(offer.getResourcesList(), Predicates.and(SUPPORTED_RESOURCE, REVOCABLE));
    }

    /**
     * Gets non-Mesos-revocable offer resources.
     *
     * @param offer Offer to get resources from.
     * @return Non-Mesos-revocable offer resources.
     */
    public static Iterable<Resource> getNonRevocableOfferResources(Offer offer) {
        return Iterables.filter(offer.getResourcesList(), Predicates.and(SUPPORTED_RESOURCE, NON_REVOCABLE));
    }

    /**
     * Gets offer resources filtered by the provided {@code tierInfo} instance.
     *
     * @param offer Offer to get resources from.
     * @param tierInfo Tier info.
     * @return Offer resources filtered by {@code tierInfo}.
     */
    public static Iterable<Resource> getOfferResources(Offer offer, TierInfo tierInfo) {
        return tierInfo.isRevocable() ? getRevocableOfferResources(offer) : getNonRevocableOfferResources(offer);
    }

    /**
     * Gets offer resoruces filtered by the {@code tierInfo} and {@code type}.
     *
     * @param offer Offer to get resources from.
     * @param tierInfo Tier info.
     * @param type Resource type.
     * @return Offer resources filtered by {@code tierInfo} and {@code type}.
     */
    public static Iterable<Resource> getOfferResources(Offer offer, TierInfo tierInfo, ResourceType type) {

        return Iterables.filter(getOfferResources(offer, tierInfo), r -> fromResource(r).equals(type));
    }

    /**
     * Same as {@link #getTaskResources(ITaskConfig, ResourceType)}.
     *
     * @param task Scheduled task to get resources from.
     * @param type {@link ResourceType} to filter resources by.
     * @return Task resources matching {@link ResourceType}.
     */
    public static Iterable<IResource> getTaskResources(IScheduledTask task, ResourceType type) {
        return getTaskResources(task.getAssignedTask().getTask(), type);
    }

    /**
     * Gets task resources matching specified {@link ResourceType}.
     *
     * @param task Task config to get resources from.
     * @param type {@link ResourceType} to filter resources by.
     * @return Task resources matching {@link ResourceType}.
     */
    public static Iterable<IResource> getTaskResources(ITaskConfig task, ResourceType type) {
        return Iterables.filter(task.getResources(), r -> fromResource(r).equals(type));
    }

    /**
     * Gets task resources matching any of the specified resource types.
     *
     * @param task Task config to get resources from.
     * @param typesToMatch EnumSet of resource types.
     * @return Task resources matching any of the resource types.
     */
    public static Iterable<IResource> getTaskResources(ITaskConfig task, EnumSet<ResourceType> typesToMatch) {

        return Iterables.filter(task.getResources(), r -> typesToMatch.contains(fromResource(r)));
    }

    /**
     * Gets unique task resource types.
     *
     * @param task Task to get resource types from.
     * @return Set of {@link ResourceType} instances representing task resources.
     */
    public static Set<ResourceType> getTaskResourceTypes(IAssignedTask task) {
        Set<ResourceType> types = task.getTask().getResources().stream().map(RESOURCE_TO_TYPE)
                .collect(Collectors.toSet());
        return types.isEmpty() ? types : EnumSet.copyOf(types);
    }

    /**
     * Gets the quantity of the Mesos resource specified by {@code type}.
     *
     * @param resources Mesos resources.
     * @param type Type of resource to quantify.
     * @return Aggregate Mesos resource value.
     */
    public static Double quantityOfMesosResource(Iterable<Resource> resources, ResourceType type) {
        return StreamSupport.stream(resources.spliterator(), false).filter(r -> SUPPORTED_RESOURCE.apply(r))
                .filter(r -> fromResource(r).equals(type)).map(QUANTIFY_MESOS_RESOURCE).reduce(REDUCE_VALUES)
                .orElse(0.0);
    }

    /**
     * Gets the quantity of resource specified by {@code type}.
     *
     * @param resources Resources.
     * @param type Type of resource to quantify.
     * @return Aggregate resource value.
     */
    public static Double quantityOf(Iterable<IResource> resources, ResourceType type) {
        return quantityOf(StreamSupport.stream(resources.spliterator(), false)
                .filter(r -> fromResource(r).equals(type)).collect(Collectors.toList()));
    }

    /**
     * Gets the quantity of resources. Caller to ensure all resources are of the same type.
     *
     * @param resources Resources to sum up.
     * @return Aggregate resource value.
     */
    public static Double quantityOf(Iterable<IResource> resources) {
        return StreamSupport.stream(resources.spliterator(), false).map(QUANTIFY_RESOURCE).reduce(REDUCE_VALUES)
                .orElse(0.0);
    }

    /**
     * Creates a {@link ResourceBag} from resources.
     *
     * @param resources Resources to convert.
     * @return A {@link ResourceBag} instance.
     */
    public static ResourceBag bagFromResources(Iterable<IResource> resources) {
        return bagFromResources(resources, RESOURCE_TO_TYPE, QUANTIFY_RESOURCE);
    }

    /**
     * Creates a {@link ResourceBag} from Mesos resources.
     *
     * @param resources Mesos resources to convert.
     * @return A {@link ResourceBag} instance.
     */
    public static ResourceBag bagFromMesosResources(Iterable<Resource> resources) {
        return bagFromResources(Iterables.filter(resources, SUPPORTED_RESOURCE), MESOS_RESOURCE_TO_TYPE,
                QUANTIFY_MESOS_RESOURCE);
    }

    /**
     * Creates a {@link ResourceBag} from {@link IResourceAggregate}.
     *
     * @param aggregate {@link IResourceAggregate} to convert.
     * @return A {@link ResourceBag} instance.
     */
    public static ResourceBag bagFromAggregate(IResourceAggregate aggregate) {
        return new ResourceBag(
                aggregate.getResources().stream().collect(Collectors.toMap(RESOURCE_TO_TYPE, QUANTIFY_RESOURCE)));
    }

    /**
     * Creates a {@link IResourceAggregate} from {@link ResourceBag}.
     *
     * @param bag {@link ResourceBag} to convert.
     * @return A {@link IResourceAggregate} instance.
     */
    public static IResourceAggregate aggregateFromBag(ResourceBag bag) {
        return ThriftBackfill
                .backfillResourceAggregate(new ResourceAggregate().setResources(bag.streamResourceVectors()
                        .map(e -> IResource.newBuilder(e.getKey().getValue(),
                                e.getKey().getAuroraResourceConverter().valueOf(e.getValue())))
                        .collect(Collectors.toSet())));
    }

    private static <T> ResourceBag bagFromResources(Iterable<T> resources, Function<T, ResourceType> typeMapper,
            Function<T, Double> valueMapper) {

        return new ResourceBag(StreamSupport.stream(resources.spliterator(), false)
                .collect(Collectors.groupingBy(typeMapper)).entrySet().stream()
                .collect(Collectors.toMap(Map.Entry::getKey,
                        group -> group.getValue().stream().map(valueMapper).reduce(REDUCE_VALUES).orElse(0.0))));
    }

    /**
     * Thrown when there are insufficient resources to satisfy a request.
     */
    public static class InsufficientResourcesException extends RuntimeException {
        InsufficientResourcesException(String message) {
            super(message);
        }
    }

}