The WeakReference class, monitoring memory leak and garbage collection in a Java application

 Below is a Stack implementation that uses an internal resizeable array structure. 

public class MyStack<T> implements Stack<T> {

    private static final int CAPACITY = 100;
    private Object[] array;

    private int pos = 0;
    public MyStack() {
        this.array = new Object[CAPACITY];
    }

    @Override
    public void push(T item) {

        if (pos >= array.length / 2) {
            Object[] newArray = new Object[pos * 2];
            System.arraycopy(array, 0, newArray, 0, array.length);
            array = newArray;
        }

        array[pos++] = item;
    }

    @Override
    public T pop() {
        if (isEmpty()) {
            throw new RuntimeException("empty stack");
        }

        @SuppressWarnings("unchecked")
        T item = (T) array[pos - 1];        
        pos -= 1;

        return item;
    }

    @Override
    @SuppressWarnings("unchecked")
    public T peek() {
        if (isEmpty()) {
            throw new RuntimeException("empty stack");
        }

        return (T) array[pos - 1];
    }

    @Override
    public boolean isEmpty() {
        return pos == 0;
    }

    @Override
    public int size() {
        return pos;
    }
}

There is a memory leak here because in the pop method, even though the position of the array element decreases, the deleted element is not set to null. This means the deleted object (which can be large) will be still referenced by the internal array and thus it will not be collected by the garbage collector. 

Below is a unit test that catches this problem:

class StackMemoryLeakTest {

    static class LargeObject {
        byte[] payload = new byte[10_000_000]; // 10MB
    }

    @Test
    void testStackLeak() throws InterruptedException {
        Stack<LargeObject> stack = new MyStack<>();
        List<WeakReference<LargeObject>> refs = new ArrayList<>();
        
        for (int i = 0; i < 50; i++) {
            LargeObject obj = new LargeObject();
            refs.add(new WeakReference<>(obj));
            stack.push(obj);
            Thread.sleep(1000);
        }

        for (int i = 0; i < 50; i++) {
            stack.pop();
            Thread.sleep(1000);
            System.gc();
        }

        Thread.sleep(10000); 

        System.gc();

        Thread.sleep(30000); 

        int collected = 0;
        for (WeakReference<LargeObject> ref : refs) {
            if (ref.get() == null) {
                collected++;
            }
        }

        System.out.println("Collected objects: " + collected + " / " + refs.size());

        // Expect at least 80% of them to be collected
        assertTrue(collected >= 40, "Too many objects were retained — possible memory leak");
    }
}

Currently this test fails. Note that I've put there Thread.sleep statements so that it's easier to use a profiler to monitor the memory usage during the test run. We'll come to that a bit later. 

The line that we need to add to the pop method should set the unused object reference to null. Here is the corrected version of the pop method:

@Override
public T pop() {
    if (isEmpty()) {
        throw new RuntimeException("empty stack");
    }

    @SuppressWarnings("unchecked")
    T item = (T) array[pos - 1];
    array[pos - 1] = null; // set the unused reference to null
    pos -= 1;

    return item;
}

In order to monitor the memory usage and see the memory leak, I will use the VisualVM tool:
https://visualvm.github.io/download.html

Here is a screenshot from the tool that shows the memory leak:

As we see, the used heap keeps increasing and the garbage collector does not collect the large objects that we store in the stack.

And here is the screenshot after we add the fix to the stack implementation:

What is WeakReference class and why did we use it?

A WeakReference is a special reference type in Java that does not prevent the referenced object from being garbage collected.

If we used normal references like this in the unit test:

List<LargeObject> list = new ArrayList<>();

then the garbage collector would never collect the objects because they are still referenced strongly. That's why we use WeakReferences so that the garbage collector can collect them, and we still hold these references to verify if they are collected.