src/main/java/org/apache/aurora/scheduler/resources/ResourceBag.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.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();
  }
}