Letting the Garbage Collector Do Callbacks

Garbage collection in Java is great, but in some cases you want to listen when an object is garbage collected. In my case I had the following scenario:

1. I have a throw away key (a composition of objects) and a value (statistics) in some map.
2. When one of the components of the key is garbage collected, the key and value should be removed from the map.

The java.util.WeakHashMap was the first thing that came to mind, but there are some problems:

1. It is not thread safe. And wrapping it in a synchronized block is not acceptable for scalability reasons.
2. The key (a WeakReference containing the real key) is removed when the real key has been garbage collected. But in my case this isn’t going to work. I don’t want my throw-away key to be removed when it is garbage collected because it would be garbage collected immediately since nobody but the WeakReference is holding a reference to it.

So I needed to go a step deeper. How can I listen to the garbage collection on an object? This can be done using a java.lang.ref.Reference (java.lang.ref.WeakReference for example) and a java.lang.ref.ReferenceQueue. When the object inside the WeakReference is garbage collected, the WeakReference is put on a ReferenceQueue. This makes it possible to do some cleanup by taking items from this ReferenceQueue. This is the first big step to listen to garbage collection.

The other two constraints are now easy to solve:

1. thread safety: use a ConcurrentHashMap implementation
2. listen to an arbitrary object instead of the key: wrap the object inside a WeakReference and store the key & map inside. If you create a thread that consumes these references from the ReferenceQueue, this thread is now able to remove the key (and value) from the map

This weekend I created a proof of concept implementation and it can be found bellow.

public class GarbageCollectingConcurrentMap<K, V> {

    private final static ReferenceQueue referenceQueue = new ReferenceQueue();

    static {
        new CleanupThread().start();
    }

    private final ConcurrentMap<K, GarbageReference<K, V>> map = new ConcurrentHashMap<K, GarbageReference<K, V>>();

    public void clear() {
        map.clear();
    }

    public V get(K key) {
        GarbageReference<K, V> ref = map.get(key);
        return ref == null ? null : ref.value;
    }

    public Object getGarbageObject(K key){
        GarbageReference<K,V> ref=map.get(key);
        return ref == null ? null : ref.get();
    }

    public Collection<K> keySet() {
        return map.keySet();
    }

    public void put(K key, V value, Object garbageObject) {
        if (key == null || value == null || garbageObject == null) throw new NullPointerException();
        if (key == garbageObject)
            throw new IllegalArgumentException("key can't be equal to garbageObject for gc to work");
        if (value == garbageObject)
            throw new IllegalArgumentException("value can't be equal to garbageObject for gc to work");

        GarbageReference reference = new GarbageReference(garbageObject, key, value, map);
        map.put(key, reference);
    }

    static class GarbageReference<K, V> extends WeakReference {
        final K key;
        final V value;
        final ConcurrentMap<K, V> map;

        GarbageReference(Object referent, K key, V value, ConcurrentMap<K, V> map) {
            super(referent, referenceQueue);
            this.key = key;
            this.value = value;
            this.map = map;
        }
    }

    static class CleanupThread extends Thread {
        CleanupThread() {
            setPriority(Thread.MAX_PRIORITY);
            setName("GarbageCollectingConcurrentMap-cleanupthread");
            setDaemon(true);
        }

        public void run() {
            while (true) {
                try {
                    GarbageReference ref = (GarbageReference) referenceQueue.remove();
                    while (true) {
                        ref.map.remove(ref.key);
                        ref = (GarbageReference) referenceQueue.remove();
                    }
                } catch (InterruptedException e) {
                    //ignore
                }
            }
        }
    }
}

I am now able to run my concurrency detector on large projects like JBoss instead of running out of memory.