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

Java tutorial

Introduction

Here is the source code for org.apache.aurora.scheduler.resources.ResourceBag.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.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.BinaryOperator;
import java.util.function.Predicate;
import java.util.stream.Stream;

import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Sets;

import static java.util.stream.Collectors.toMap;

import static org.apache.aurora.scheduler.resources.ResourceType.CPUS;
import static org.apache.aurora.scheduler.resources.ResourceType.DISK_MB;
import static org.apache.aurora.scheduler.resources.ResourceType.RAM_MB;

/**
 * A bag of unique resource values aggregated by {@link ResourceType}.
 */
public class ResourceBag {
    public static final ResourceBag EMPTY = new ResourceBag(ImmutableMap.of(CPUS, 0.0, RAM_MB, 0.0, DISK_MB, 0.0));

    public static final ResourceBag SMALL = new ResourceBag(
            ImmutableMap.of(CPUS, 1.0, RAM_MB, 1024.0, DISK_MB, 4096.0));

    public static final ResourceBag MEDIUM = new ResourceBag(
            ImmutableMap.of(CPUS, 4.0, RAM_MB, 8192.0, DISK_MB, 16384.0));

    public static final ResourceBag LARGE = new ResourceBag(
            ImmutableMap.of(CPUS, 8.0, RAM_MB, 16384.0, DISK_MB, 32768.0));

    public static final ResourceBag XLARGE = new ResourceBag(
            ImmutableMap.of(CPUS, 16.0, RAM_MB, 32768.0, DISK_MB, 65536.0));

    public static final Predicate<Map.Entry<ResourceType, Double>> IS_NEGATIVE = entry -> entry.getValue() < 0;

    public static final Predicate<Map.Entry<ResourceType, Double>> IS_POSITIVE = entry -> entry.getValue() > 0;

    public static final Predicate<Map.Entry<ResourceType, Double>> IS_MESOS_REVOCABLE = entry -> entry.getKey()
            .isMesosRevocable();

    private final Map<ResourceType, Double> resourceVectors;

    /**
     * Creates an instance of ResourceBag with given resource vectors (type -> value).
     *
     * @param resourceVectors Map of resource vectors.
     */
    ResourceBag(Map<ResourceType, Double> resourceVectors) {
        this.resourceVectors = ImmutableMap.copyOf(resourceVectors);
    }

    /**
     * Gets resource vectors in the bag.
     *
     * @return Map of resource vectors.
     */
    public Map<ResourceType, Double> getResourceVectors() {
        return resourceVectors;
    }

    /**
     * Convenience function to return a stream of resource vectors.
     *
     * @return A stream of resource vectors.
     */
    public Stream<Map.Entry<ResourceType, Double>> streamResourceVectors() {
        return resourceVectors.entrySet().stream();
    }

    /**
     * Gets the value of resource specified by {@code type} or 0.0.
     *
     * @param type Resource type to get value for.
     * @return Resource value or 0.0 if no mapping for {@code type} is found.
     */
    public Double valueOf(ResourceType type) {
        return resourceVectors.getOrDefault(type, 0.0);
    }

    /**
     * Adds this and other bag contents.
     *
     * @param other Other bag to add.
     * @return Result of addition.
     */
    public ResourceBag add(ResourceBag other) {
        return binaryOp(other, (l, r) -> l + r);
    }

    /**
     * Subtracts other bag contents from this.
     *
     * @param other Other bag to subtract.
     * @return Result of subtraction.
     */
    public ResourceBag subtract(ResourceBag other) {
        return binaryOp(other, (l, r) -> l - r);
    }

    /**
     * Divides this by other bag contents.
     * <p>
     * Note: any missing or zero resource values in {@code other} will result in
     * {@code java.lang.Double.POSITIVE_INFINITY}.
     * @param other Other bag to divide by.
     * @return Result of division.
     */
    public ResourceBag divide(ResourceBag other) {
        return binaryOp(other, (l, r) -> l / r);
    }

    /**
     * Applies {@code Math.max()} for each matching resource vector.
     *
     * @param other Other bag to compare with.
     * @return A new bag with max resource vectors.
     */
    public ResourceBag max(ResourceBag other) {
        return binaryOp(other, (l, r) -> Math.max(l, r));
    }

    /**
     * Scales each resource vector by {@code m}.
     *
     * @param m Scale factor.
     * @return Result of scale operation.
     */
    public ResourceBag scale(int m) {
        return new ResourceBag(
                resourceVectors.entrySet().stream().collect(toMap(Map.Entry::getKey, v -> v.getValue() * m)));
    }

    /**
     * Filters bag resources by {@code predicate}.
     *
     * @param predicate Predicate to filter by.
     * @return A new bag with resources filtered by {@code predicate}.
     */
    public ResourceBag filter(Predicate<Map.Entry<ResourceType, Double>> predicate) {
        return new ResourceBag(resourceVectors.entrySet().stream().filter(predicate)
                .collect(toMap(Map.Entry::getKey, Map.Entry::getValue)));
    }

    /**
     * Verifies whether another bag would be able to fit into this one.
     *
     * @param other Other bag to try and fit.
     * @return Whether or not the bag fits.
     */
    public boolean greaterThanOrEqualTo(ResourceBag other) {
        // Subtract the subject and check if any of the resources have a negative value.
        return subtract(other).filter(IS_NEGATIVE).getResourceVectors().isEmpty();
    }

    /**
     * Applies {@code operator} to this and {@code other} resource values. Any missing resource type
     * values on either side will get substituted with 0.0 before applying the operator.
     *
     * @param other Other ResourceBag.
     * @param operator Operator to apply.
     * @return Operation result.
     */
    private ResourceBag binaryOp(ResourceBag other, BinaryOperator<Double> operator) {
        ImmutableMap.Builder<ResourceType, Double> builder = ImmutableMap.builder();
        Set<ResourceType> resourceTypes = Sets.union(resourceVectors.keySet(), other.getResourceVectors().keySet());
        for (ResourceType type : resourceTypes) {
            Double left = resourceVectors.getOrDefault(type, 0.0);
            Double right = other.getResourceVectors().getOrDefault(type, 0.0);
            builder.put(type, operator.apply(left, right));
        }

        return new ResourceBag(builder.build());
    }

    @Override
    public boolean equals(Object o) {
        if (!(o instanceof ResourceBag)) {
            return false;
        }

        ResourceBag other = (ResourceBag) o;
        return Objects.equals(resourceVectors, other.resourceVectors);
    }

    @Override
    public int hashCode() {
        return Objects.hash(resourceVectors);
    }

    @Override
    public String toString() {
        return MoreObjects.toStringHelper(this).add("resourceVectors", resourceVectors).toString();
    }
}