Como implementar cache de objetos utilizando um HashMap com SoftReference?

2013/09/11

A API do Java nos fornece o WeakHashMap que é uma implementação de HashMap que utiliza WeakReference.

Eventualmente queremos implementar caches sem manter referências fortes aos objetos em cache e a própria documentação do WeakHashMap diz que ele não é legal para esse tipo de situação.

Para isso, implementei o SoftHashMap, um HashMap que utiliza SoftReference. Essa implementação pode ser útil para implementarmos cache de objetos sem manter referências fortes no HashMap.

Segue o código:

import java.lang.ref.*;
import java.util.*;

/**
 * Implementação de um HashMap que usa SoftReference's.<BR>
 * Funciona como um cache, sem manter referência forte para as chaves e valores.
 * 
 * @param <K> Tipo de dado da chave.
 * @param <V> Tipo de dado do valor.
 * 
 * @author Ricardo Artur Staroski
 */
final class SoftHashMap<K, V> extends AbstractMap<K, V> {

    /**
     * Especialização de uma SoftReference que contém não somente o valor mas também a chave.<BR>
     * Isso torna fácil a localização da Entry no HashMap depois que foi garbage collected.
     * 
     * @param <X> Tipo de dado da chave.
     * @param <Y> TIpo de dado do valor.
     */
    private static class SoftValue<X, Y> extends SoftReference<Y> {

        private final X key;

        private SoftValue(X key, Y value, ReferenceQueue<Y> queue) {
            super(value, queue);
            this.key = key;
        }
    }

    //  HashMap interno que mantém as SoftReference's
    private final Map<K, SoftReference<V>> hash = new HashMap<K, SoftReference<V>>();

    // A quantidade de referencias fortes a manter internamente
    private final int STRONG_SIZE;

    // A FIFO com as strong references
    private final LinkedList<V> strongCache = new LinkedList<V>();

    // Reference queue para as SoftReference's limpas pelo garbage collector
    private final ReferenceQueue<V> queue = new ReferenceQueue<V>();

    public SoftHashMap() {
        this(100);
    }

    public SoftHashMap(int strongSize) {
        STRONG_SIZE = strongSize;
    }

    @Override
    public void clear() {
        strongCache.clear();
        processQueue(); // joga fora as referencias que já foram garbage collected
        hash.clear();
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        // não, não e não!!!
        // o SoftHashMap serve justamente para não manter referência forte das chaves e valores
        // então não vem com essa de querer obter o Set com as Entry's
        throw new UnsupportedOperationException("entrySet()");
    }

    @Override
    public V get(Object key) {
        V value = null;
        // obter a SoftReference a partir da chave
        SoftReference<V> softRef = hash.get(key);
        if (softRef != null) {
            // a partir da SoftReference pegamos o valor,
            // que pode ser null se ele não estava no map,
            // ou se ele foi removido no método processQueue() 
            value = softRef.get();
            if (value == null) {
                // se o valor foi garbage collected, remove a Entry inteira do HashMap
                hash.remove(key);
            } else {
                // agora a gente adiciona este objeto ao começo da ReferenceQueue.
                // Uma referencia pode aparecer mais de uma vez, porque os lookups da FIFO são lentos,
                // a gente não quer varrer ela toda vez para remover duplicatas.
                strongCache.addFirst(value);
                if (strongCache.size() > STRONG_SIZE) {
                    // remove a última Entry se a lista estiver maior que STRONG_SIZE
                    strongCache.removeLast();
                }
            }
        }
        return value;
    }

    /**
     * Armazena o par chave-valor usando um objeto SoftValue.
     */
    @Override
    public V put(K key, V value) {
        processQueue(); // joga fora as referencias que já foram garbage collected
        SoftReference<V> reference = hash.put(key, new SoftValue<K, V>(key, value, queue));
        return reference != null ? reference.get() : null;
    }

    @Override
    public V remove(Object key) {
        processQueue(); // joga fora as referencias que já foram garbage collected
        SoftReference<V> reference = hash.remove(key);
        return reference != null ? reference.get() : null;
    }

    @Override
    public int size() {
        processQueue(); // joga fora as referencias que já foram garbage collected
        return hash.size();
    }

    /**
     * Varre a ReferenceQueue e remove os objetos SoftValue que já foram garbage collected.
     */
    private void processQueue() {
        // parece confuso, mas quando se declara uma ReferenceQueue<V>,
        // o método poll() returna um Reference<? extends V>.
        // no nosso caso essa Reference vai ser uma instância de SoftValue.
        SoftValue<K, V> sv;
        while ((sv = (SoftValue<K, V>) queue.poll()) != null) {
            hash.remove(sv.key);
        }
    }
}