Caffeine.java

/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package br.ufrgs.inf.prosoft.cache;

import com.github.benmanes.caffeine.cache.Expiry;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
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.stream.Collectors;

/**
 *
 * @author root
 * @param <K>
 * @param <V>
 */
public class Caffeine<K, V> implements Cache<K, V>, AutoCloseable {

    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 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<Optional<K>, Optional<V>>() {
                    @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);
        }
    }

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

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

    @Override
    public void put(K key, V value, long timeToLive) {
        this.keyHasTTL.put(Optional.ofNullable(key), timeToLive);
        put(key, value);
    }

    @Override
    public void put(K key, V value) {
        this.cache.put(Optional.ofNullable(key), Optional.ofNullable(value));
        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());
    }
}