src/main/java/org/apache/aurora/scheduler/stats/ResourceCounter.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.stats;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

import javax.inject.Inject;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Predicate;
import com.google.common.base.Predicates;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.Iterables;

import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.resources.ResourceBag;
import org.apache.aurora.scheduler.resources.ResourceManager;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.Storage.StorageException;
import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;

import static org.apache.aurora.scheduler.quota.QuotaManager.DEDICATED;
import static org.apache.aurora.scheduler.quota.QuotaManager.NON_PROD_SHARED;
import static org.apache.aurora.scheduler.quota.QuotaManager.PROD_SHARED;
import static org.apache.aurora.scheduler.quota.QuotaManager.QUOTA_RESOURCES;

/**
 * Computes aggregate metrics about resource allocation and consumption in the scheduler.
 */
public class ResourceCounter {
  private final Storage storage;

  @Inject
  ResourceCounter(Storage storage) {
    this.storage = Objects.requireNonNull(storage);
  }

  private Iterable<ITaskConfig> getTasks(Query.Builder query) throws StorageException {
    return Iterables.transform(
        Storage.Util.fetchTasks(storage, query),
        Tasks::getConfig);
  }

  private static final Function<MetricType, Metric> TO_METRIC = Metric::new;

  /**
   * Computes totals for each of the {@link MetricType}s.
   *
   * @return aggregates for each metric type.
   * @throws StorageException if there was a problem fetching tasks from storage.
   */
  public List<Metric> computeConsumptionTotals() throws StorageException {
    List<Metric> counts = FluentIterable.from(Arrays.asList(MetricType.values()))
        .transform(TO_METRIC)
        .toList();

    for (ITaskConfig task : getTasks(Query.unscoped().active())) {
      for (Metric count : counts) {
        count.accumulate(task);
      }
    }
    return counts;
  }

  /**
   * Computes total quota allocations.
   *
   * @return Total allocated quota.
   * @throws StorageException if there was a problem fetching quotas from storage.
   */
  public Metric computeQuotaAllocationTotals() throws StorageException {
    return storage.read(storeProvider -> {
      Metric allocation = new Metric();
      storeProvider.getQuotaStore().fetchQuotas().values().forEach(allocation::accumulate);
      return allocation;
    });
  }

  /**
   * Returns quota allocations by role (as metrics).
   *
   * @return A map of role to allocated quota metrics.
   * @throws StorageException if there was a problem fetching quotas from storage.
   */
  public Map<String, Metric> computeQuotaAllocationByRole() throws StorageException {
    return storage.read(storeProvider -> {
      return storeProvider.getQuotaStore().fetchQuotas().entrySet().stream()
          .collect(Collectors.toMap(e -> e.getKey(), e -> {
            Metric allocation = new Metric();
            allocation.accumulate(e.getValue());
            return allocation;
          }));
    });
  }

  /**
   * Computes arbitrary resource aggregates based on a query, a filter, and a grouping function.
   *
   * @param query Query to select tasks for aggregation.
   * @param filter Filter to apply on query result tasks.
   * @param keyFunction Function to define aggregation groupings.
   * @param <K> Key type.
   * @return A map from the keys to their aggregates based on the tasks fetched.
   * @throws StorageException if there was a problem fetching tasks from storage.
   */
  public <K> Map<K, Metric> computeAggregates(
      Query.Builder query,
      Predicate<ITaskConfig> filter,
      Function<ITaskConfig, K> keyFunction) throws StorageException {

    LoadingCache<K, Metric> metrics = CacheBuilder.newBuilder()
        .build(new CacheLoader<K, Metric>() {
          @Override
          public Metric load(K key) {
            return new Metric();
          }
        });
    for (ITaskConfig task : Iterables.filter(getTasks(query), filter)) {
      metrics.getUnchecked(keyFunction.apply(task)).accumulate(task);
    }
    return metrics.asMap();
  }

  public enum MetricType {
    TOTAL_CONSUMED(Predicates.<ITaskConfig>alwaysTrue()),
    DEDICATED_CONSUMED(DEDICATED),
    QUOTA_CONSUMED(PROD_SHARED),
    FREE_POOL_CONSUMED(NON_PROD_SHARED);

    public final Predicate<ITaskConfig> filter;

    MetricType(Predicate<ITaskConfig> filter) {
      this.filter = filter;
    }
  }

  public static class Metric {
    public final MetricType type;
    private ResourceBag bag;

    public Metric() {
      this(MetricType.TOTAL_CONSUMED);
    }

    public Metric(MetricType type) {
      this(type, ResourceBag.EMPTY);
    }

    public Metric(Metric copy) {
      this(copy.type, copy.bag);
    }

    @VisibleForTesting
    Metric(MetricType type, ResourceBag bag) {
      this.type = type;
      this.bag = bag;
    }

    void accumulate(ITaskConfig task) {
      if (type.filter.apply(task)) {
        bag = bag.add(QUOTA_RESOURCES.apply(task));
      }
    }

    void accumulate(IResourceAggregate aggregate) {
      bag = bag.add(ResourceManager.bagFromAggregate(aggregate));
    }

    public ResourceBag getBag() {
      return bag;
    }

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

      Metric other = (Metric) o;
      return Objects.equals(other.type, this.type)
          && Objects.equals(other.bag, this.bag);
    }

    @Override
    public int hashCode() {
      return Objects.hash(type, bag);
    }
  }
}