Memory leaks are often tricky to debug. In complex situations, advanced tools such as memory profilers are needed. But what about unit testing potential memory leaks? Ever wondered how to write, say, a JUnit test to check for memory leaks in a given piece of your Java code?

Here’s the sketch of a simple solution used in our internal unit tests. This solution is based on weak references (java.lang.ref.WeakReference). The idea is straight-forward:

1. Keep a weak reference on the instance of the object you want to test.
2. Perform the code that is supposed to make this instance garbage collectable.
3. Request for garbage collection.
4. Check that the weak reference has been cleared. If not, it’s likely you detected a memory leak.

For convenience, a testing utility around the following lines can be written:

public class MemLeakTracker
{
  private HashMap _map = new HashMap();

  public void register(Object obj, String id)
  {
    if (_map.get(id) != null)
      throw new IllegalArgumentException(
                       "Object already stored under id " + id);
    _map.put(id, new WeakReference(obj, new ReferenceQueue()));
  }

  public boolean isGarbageCollectable(String id)
  {
    gc(); // ask for garbage collection
    WeakReference ref = (WeakReference)_map.get(id);
    if (ref == null)
      throw new RuntimeException(
                        "No object stored under id " + id);
    return ref.get() == null;
  }
  ...
}

How gc() of our utility could look like? The JDK API provides Runtime.gc(). There’s no guarantee that all garbage-collectable objects are actually garbage-collected. To maximize the chances the GC really happens, we cumulate various simple tricks:

private void gc()
{
  Runtime rt = Runtime.getRuntime();
  for (int i = 0; i < 3; i++) {
    try {
      allocateMemory((int)(2e6));
    } catch (Throwable th) {
      th.printStackTrace();
    }
    for (int j = 0; j < 3; j++)
      rt.gc();
    }
  rt.runFinalization();
  try {
    Thread.currentThread().sleep(50);
  } catch (Throwable th) { th.printStackTrace(); }
}

private void allocateMemory(int memAmount)
{
  byte[] big = new byte[memAmount];
  // Fight against clever compilers/JVMs that may not allocate
  // unless we actually use the elements of the array
  int total = 0;
  for (int i = 0; i < 10; i++) {
    // we don't touch all the elements, would take too long.
    if (i % 2 == 0)
      total += big[i];
    else
      total -= big[i];
  }
}

This detects a real memory leak for sure: if the leak exists, by no way the object will be garbage-collected. The converse is not guaranteed, since actual garbage collection is at the discretion of the JVM. So this approach never fails to detect a leak (this is the most important) but in principle it can lead to “false positive” errors (wrongly claming a memory leak exists). However, in practice, the tricks for “forcing” the GC seem good enough. In our experience, we never faced a false positive error.

Once this testing utility is written, a unit test for memory leaks looks as simple as this example:

MemLeakTracker memLeakTracker = new MemLeakTracker();

SomeClass obj = new SomeClass(...);

memLeakTracker.register(obj, "myObj");

// code that is expected to make "obj" garbage collectable, for instance:
obj.cleanup();
obj = null;

if (!memLeakTracker.isGarbageCollectable("myObj"))
  fail("myObj not gc-able!");

Being easy to write such a unit test, fixing a memory leak can easily end up with writing the corresponding unit test.

Posted by Adrian Vasiliu in Development practices. You can follow any responses to this entry through the RSS 2.0 feed. You can leave a response, or trackback from your own site.