shz.cache.LfuCache Maven / Gradle / Ivy

Go to download
package shz.cache;

import shz.linked.ConcurrentLDNode;

import java.lang.ref.Reference;
import java.lang.ref.ReferenceQueue;
import java.util.Comparator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.AtomicInteger;

public final class LfuCache extends LocalCache {
    private static final class Node {
        private final ConcurrentLDNode node;
        private final AtomicInteger times;
        private Reference ref;

        private Node(K key, Reference ref) {
            node = ConcurrentLDNode.of(key);
            times = new AtomicInteger();
            this.ref = ref;
        }
    }

    private static class HT {
        private final ConcurrentLDNode head;
        private final ConcurrentLDNode tail;

        private HT() {
            head = tail = ConcurrentLDNode.of(null);
            head.next(tail);
            tail.prev(head);
        }
    }

    private final int capacity;
    private final ConcurrentHashMap> cache;
    private final ConcurrentSkipListMap> times;

    public LfuCache(int capacity, RefType refType, ReferenceQueue queue) {
        super(refType, queue);
        this.capacity = capacity;
        cache = new ConcurrentHashMap<>(capacity);
        times = new ConcurrentSkipListMap<>(Comparator.comparingInt(t -> t));
    }

    public LfuCache(int capacity, RefType refType) {
        this(capacity, refType, null);
    }

    public LfuCache(int capacity) {
        this(capacity, RefType.SOFT);
    }

    @Override
    public V get(K key) {
        Node node = cache.get(key);
        if (node == null) return null;
        int useTimes = node.times.getAndIncrement();
        poll(useTimes, node.node);
        if (node.ref == null || node.ref.isEnqueued()) {
            cache.remove(key);
            return null;
        }
        offer(node.times.get(), node.node);
        return node.ref.get();
    }

    private K poll(int useTimes, ConcurrentLDNode node) {
        node.poll();
        HT ht = times.get(useTimes);
        if (ht != null) if (ht.head.next() == ht.tail) times.remove(useTimes);
        return node.val;
    }

    private void offer(int useTimes, ConcurrentLDNode node) {
        times.computeIfAbsent(useTimes, k -> new HT<>()).tail.addPrev(node);
    }

    @Override
    public void put(K key, V val) {
        Node node = cache.get(key);
        if (node == null) {
            K k;
            while (cache.size() >= capacity && (k = poll()) != null) cache.remove(k);
            node = cache.computeIfAbsent(key, t -> new Node<>(key, getReference(val)));
        }
        int useTimes = node.times.getAndIncrement();
        poll(useTimes, node.node);
        if (node.ref == null || node.ref.isEnqueued()) {
            cache.remove(key);
            return;
        }
        node.ref = getReference(val);
        offer(node.times.get(), node.node);
    }

    @Override
    public boolean containsKey(K key) {
        return cache.containsKey(key);
    }

    private K poll() {
        Integer firstKey;
        try {
            firstKey = times.firstKey();
        } catch (NoSuchElementException e) {
            return null;
        }
        return poll(firstKey, times.get(firstKey).head.next());
    }

    @Override
    public int size() {
        int size = 0;
        for (Map.Entry> kv : cache.entrySet()) {
            Node node = kv.getValue();
            if (node.ref == null || node.ref.isEnqueued()) {
                poll(node.times.get(), node.node);
                cache.remove(kv.getKey());
            } else ++size;
        }
        return size;
    }

    @Override
    public void delete(K key) {
        Node node = cache.get(key);
        if (node == null) return;
        poll(node.times.get(), node.node);
        if (node.ref != null) node.ref.clear();
        cache.remove(key);
    }

    @Override
    public void clear() {
        cache.forEach((key, node) -> {
            if (node == null) return;
            poll(node.times.get(), node.node);
            if (node.ref != null) node.ref.clear();
        });
        cache.clear();
        times.clear();
    }
}