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