Caffeine.java

package br.ufrgs.inf.prosoft.cache;

import br.ufrgs.inf.prosoft.cache.tools.CSV;
import com.github.benmanes.caffeine.cache.Expiry;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

public class Caffeine<K, V> implements Cache<K, V>, AutoCloseable {

  private static final Logger LOGGER = Logger.getLogger(Caffeine.class.getName());

  private static final boolean CACHE_EVALUATE_PERFORMANCE = System.getenv("CACHE_EVENTS") == null || !System.getenv("CACHE_EVENTS").equals("false");
  private static final boolean CACHE_REGISTER_SIZE = System.getenv("CACHE_REGISTER_SIZE") != null && System.getenv("CACHE_REGISTER_SIZE").equals("true");
  private static final boolean CACHE_REFRESH_TTL = System.getenv("CACHE_REFRESH_TTL") != null && System.getenv("CACHE_REFRESH_TTL").equals("true");
  private static final Map<String, Long> METHOD_HAS_TTL = new HashMap<>();
  private static String CACHE_METHODS_TTL = System.getenv("CACHE_METHODS_TTL") != null && !System.getenv("CACHE_METHODS_TTL").isEmpty() && !System.getenv("CACHE_METHODS_TTL").equals("null")
    ? System.getenv("CACHE_METHODS_TTL") : null;
  private final com.github.benmanes.caffeine.cache.Cache<Optional<K>, Optional<V>> cache;
  private final ConcurrentHashMap<Optional<K>, Long> keyHasTTL;
  private final CachePerformance cachePerformance;
  private final ScheduledExecutorService sizeMonitor;

  public Caffeine() {
    this(new CachePerformance());
  }

  public Caffeine(long ttl) {
    this(new CachePerformance(), ttl);
  }

  public Caffeine(String name) {
    this(new CachePerformance(name));
  }

  public Caffeine(String name, long ttl) {
    this(new CachePerformance(name), ttl);
  }

  public Caffeine(String name, Long ttl, long size) {
    this(new CachePerformance(name), ttl, size);
  }

  public Caffeine(CachePerformance cachingPerformance) {
    this(cachingPerformance, null);
  }

  public Caffeine(CachePerformance cachingPerformance, Long defaultTTL) {
    this(cachingPerformance, defaultTTL, null);
  }

  public Caffeine(CachePerformance cachingPerformance, Long defaultTTL, Long size) {
    this.cachePerformance = cachingPerformance;
    this.keyHasTTL = new ConcurrentHashMap<>();
    com.github.benmanes.caffeine.cache.Caffeine<Optional<K>, Optional<V>> builder = com.github.benmanes.caffeine.cache.Caffeine.newBuilder()
      .expireAfter(new Expiry<>() {
        @Override
        public long expireAfterCreate(Optional<K> key, Optional<V> value, long currentTime) {
          Long customTTL = keyHasTTL.remove(key);
          if (customTTL != null) return customTTL * 1000000;
          if (defaultTTL == null) return Long.MAX_VALUE;
          long ttlNanoseconds = defaultTTL * 1000000;
          if (ttlNanoseconds < 0) return Long.MAX_VALUE;
          return ttlNanoseconds;
        }

        @Override
        public long expireAfterUpdate(Optional<K> key, Optional<V> value, long currentTime, long currentDuration) {
          return expireAfterCreate(key, value, currentTime);
        }

        @Override
        public long expireAfterRead(Optional<K> key, Optional<V> value, long currentTime, long currentDuration) {
          if (CACHE_REFRESH_TTL) return expireAfterCreate(key, value, currentTime);
          return currentDuration;
        }
      });
    if (size != null) builder.maximumSize(size);
    builder.removalListener((k, v, cause) -> registerEvent(EventType.INVALIDATION, v.orElse(null)));
    this.cache = builder.build();

    this.sizeMonitor = Executors.newScheduledThreadPool(0);
    if (Caffeine.CACHE_EVALUATE_PERFORMANCE) this.sizeMonitor.scheduleAtFixedRate(() ->
      this.cachePerformance.registerEvent(EventType.POPULATION, String.valueOf(size())), 0, 5, TimeUnit.SECONDS);
  }

  private synchronized static void loadTTLsForMethods() {
    if (CACHE_METHODS_TTL == null) return;
    CSV methodsHasTTL = CSV.read(CACHE_METHODS_TTL);
    methodsHasTTL.stream().forEach(it -> METHOD_HAS_TTL.put(it.get("method"), Long.valueOf(it.get("ttl"))));
    LOGGER.log(Level.INFO, "TTLs for methods file loaded");
    CACHE_METHODS_TTL = null;
  }

  @Override
  public void close() {
    try {
      Thread.sleep(1000);
    } catch (InterruptedException ignored) {
    }
    this.sizeMonitor.shutdown();
  }

  public CachePerformance getCachePerformance() {
    return this.cachePerformance;
  }

  private long getTTLForMethod() {
    loadTTLsForMethods();
    if (METHOD_HAS_TTL.isEmpty()) return 0;
    StackTraceElement[] traceElements = Thread.currentThread().getStackTrace();
    for (int i = 3; i < traceElements.length; i++) {
      String methodName = traceElements[i].getClassName() + "." + traceElements[i].getMethodName();
      if (METHOD_HAS_TTL.containsKey(methodName)) return METHOD_HAS_TTL.get(methodName);
    }
    return 0;
  }

  @Override
  public void put(K key, V value, long timeToLive) {
    long ttl = getTTLForMethod();
    if (ttl > 0) timeToLive = ttl;
    if (timeToLive > 0) this.keyHasTTL.put(Optional.ofNullable(key), timeToLive);
    put(key, value);
  }

  @Override
  public void put(K key, V value) {
    if (!this.keyHasTTL.containsKey(Optional.ofNullable(key))) {
      long ttl = getTTLForMethod();
      if (ttl > 0) this.keyHasTTL.put(Optional.ofNullable(key), ttl);
    }
    this.cache.put(Optional.ofNullable(key), Optional.ofNullable(value));
    if (this.cache.getIfPresent(Optional.ofNullable(key)) != null) registerEvent(EventType.REPLACE, value);
    else registerEvent(EventType.ADDITION, value);
  }

  @Override
  public V get(K key) throws KeyNotFoundException {
    Optional<V> get = this.cache.getIfPresent(Optional.ofNullable(key));
    if (get == null) {
      registerEvent(EventType.MISS, key);
      throw new KeyNotFoundException();
    }
    registerEvent(EventType.HIT, get);
    return get.orElse(null);
  }

  @Override
  public void invalidate(K key) {
    this.cache.invalidate(Optional.ofNullable(key));
    registerEvent(EventType.INVALIDATION, null);
  }

  public boolean containsKey(K key) {
    Optional<V> get = this.cache.getIfPresent(Optional.ofNullable(key));
    return get != null;
  }

  public Set<Map.Entry<K, V>> entrySet() {
    Set<Map.Entry<K, V>> entrySet = this.cache.asMap().entrySet().stream().collect(
      Collectors.toMap(entry -> entry.getKey().orElse(null), entry -> entry.getValue().orElse(null))
    ).entrySet();
    if (CACHE_EVALUATE_PERFORMANCE) entrySet.forEach(entry -> registerEvent(EventType.HIT, entry.getValue()));
    return entrySet;
  }

  public void forEach(BiConsumer<K, V> consumer) {
    entrySet().forEach(entry -> consumer.accept(entry.getKey(), entry.getValue()));
  }

  public int size() {
    return this.cache.asMap().size();
  }

  public boolean isEmpty() {
    return this.cache.asMap().isEmpty();
  }

  private String getIdentifier(Object object) {
    return object != null ? String.valueOf(object.hashCode()) : "null";
  }

  private void registerEvent(EventType eventType, Object object) {
    if (!CACHE_EVALUATE_PERFORMANCE) return;
    String identifier = getIdentifier(object);
    if (CACHE_REGISTER_SIZE) this.cachePerformance.registerEvent(eventType, identifier, CachePerformance.calculateObjectSize(object));
    else this.cachePerformance.registerEvent(eventType, identifier);
    if (eventType.equals(EventType.ADDITION)) this.cachePerformance.registerSize(size());
  }

  public List<V> values() {
    return this.cache.asMap().values().stream()
      .map(value -> value.orElse(null))
      .collect(Collectors.toList());
  }
}