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