io.prometheus.cloudwatch.CachingDimensionSource Maven / Gradle / Ivy

Go to download
package io.prometheus.cloudwatch;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.Expiry;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;
import java.util.stream.Collectors;

final class CachingDimensionSource implements DimensionSource {

  private final DimensionSource delegate;
  private final Cache cache;

  /**
   * Create a new DimensionSource that will cache the results from another {@link DimensionSource}
   *
   * @param source
   * @param config - config used to configure the expiry (ttl) for entries in the cache
   * @return a new CachingDimensionSource
   */
  CachingDimensionSource(DimensionSource source, DimensionCacheConfig config) {
    this.delegate = source;
    this.cache =
        Caffeine.newBuilder()
            .expireAfter(new DimensionExpiry(config.defaultExpiry, config.metricConfig))
            .build();
  }

  @Override
  public DimensionData getDimensions(MetricRule rule, List tagBasedResourceIds) {
    DimensionData cachedDimensions =
        this.cache.getIfPresent(new DimensionCacheKey(rule, tagBasedResourceIds));
    if (cachedDimensions != null) {
      return cachedDimensions;
    }
    DimensionData dimensions = delegate.getDimensions(rule, tagBasedResourceIds);
    this.cache.put(new DimensionCacheKey(rule, tagBasedResourceIds), dimensions);
    return dimensions;
  }

  static class DimensionExpiry implements Expiry {

    private final Duration defaultExpiry;
    private final Map durationMap;

    public DimensionExpiry(Duration defaultExpiry, List expiryOverrides) {
      this.defaultExpiry = defaultExpiry;
      this.durationMap =
          expiryOverrides.stream()
              .collect(Collectors.toMap(Function.identity(), dcp -> dcp.listMetricsCacheTtl));
    }

    @Override
    public long expireAfterCreate(DimensionCacheKey key, DimensionData value, long currentTime) {
      return durationMap.getOrDefault(key.rule, this.defaultExpiry).toNanos();
    }

    @Override
    public long expireAfterUpdate(
        DimensionCacheKey key, DimensionData value, long currentTime, long currentDuration) {
      return currentDuration;
    }

    @Override
    public long expireAfterRead(
        DimensionCacheKey key, DimensionData value, long currentTime, long currentDuration) {
      return currentDuration;
    }
  }

  static class DimensionCacheConfig {
    final Duration defaultExpiry;
    final List metricConfig = new ArrayList<>();

    DimensionCacheConfig(Duration defaultExpiry) {
      this.defaultExpiry = defaultExpiry;
    }

    /**
     * Add a MetricRule to be used to configure a custom TTL using the value from {@link
     * MetricRule#listMetricsCacheTtl} to override the default expiry
     *
     * @param metricRule
     * @return this
     */
    DimensionCacheConfig addOverride(MetricRule metricRule) {
      this.metricConfig.add(metricRule);
      return this;
    }
  }

  static class DimensionCacheKey {
    private final MetricRule rule;
    private final List tagBasedResourceIds;

    DimensionCacheKey(MetricRule rule, List tagBasedResourceIds) {
      this.rule = rule;
      this.tagBasedResourceIds = tagBasedResourceIds;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || getClass() != o.getClass()) return false;

      DimensionCacheKey that = (DimensionCacheKey) o;

      if (!Objects.equals(rule, that.rule)) return false;
      return Objects.equals(tagBasedResourceIds, that.tagBasedResourceIds);
    }

    @Override
    public int hashCode() {
      int result = rule != null ? rule.hashCode() : 0;
      result = 31 * result + (tagBasedResourceIds != null ? tagBasedResourceIds.hashCode() : 0);
      return result;
    }
  }
}