How to Write a Memory Leak Unit Test

WARNING: This blog entry was imported from my old blog on blogs.sun.com (which used different blogging software), so formatting and links may not be correct.

Unit tests are great for ensuring that your functionality is correct. But how do you make sure you don't have leaks in your code, leaks that
eventually cause your application to crash when it runs out of memory? Unit tests don't typically don't run into memory problems directly because they tend to start up, run a little bit of code, and shut down -- getting a nice fresh memory environment on each test start.

One great way to help protect yourself from uptime problems is to write leak tests. These are unit tests where you first perform some operations, then do normal cleanup, and finally you assert that the objects you were using during the operation have been cleaned up completely. In other words, that there are no references left anywhere on the heap.

Testing for leaks isn't very hard; the basic trick is to create a weak reference to your to-be-cleaned object (by constructing a WeakReference wrapping your object), then you delete all your local references to the object, then you perform garbage collection, and finally you see whether the weak reference still contains your object. This works because the weak reference is handled specially by the garbage collection system.

It turns out things aren't quite as easy as that, because garbage collectors try to be smart, and simply calling System.gc() doesn't mean it's going to actually perform a complete and final garbage collection. This means that your unit test could incorrectly conclude you have a leak, since the reference is still held.

The NetBeans team has built a really good test infrastructure for this. When I wrote unit tests for NetBeans plugins, I could simply call their assertGC method to assert that the given weak reference should be garbage collected, and it would handle the rest. (Internally, they do things like actually allocating a bunch of memory chunks to really force the garbage collection to run.)

Here's an example:

// The following unit test fragment wants to ensure that a graphics object
// which shows the current selection in the authoring tool, is cleared up after
// the user clears the selection.
view.selectAll();
Node firstHandle = /* Node showing selection highlight, lookup code here */;
assertNotNull(firstHandle);

// Now clear view selection -- then assert that all the handles have disappeared
view.selectNone();

WeakReference<Node> ref = new WeakReference<Node>(firstHandle);
firstHandle = null;
assertGC("Selection handle leaked", ref);

When you write these tests you also have to make sure you null out any local variable references you are holding right there in the test.
The above test will pass if and only if the target object, firstHandle, is properly garbage collected.

But wait -- that's not the best part. Let's say you've written a leak test, and it fails. Now what? What do you do -- make the process pause and attach a profiler and try to hunt it down?

This is where their unit test support really shines. When the unit test discovers that the reference is still held somewhere, it uses its own heap walking library to hunt down the offending reference, and dumps it out as part of the test assertion failure message!!!

Here's the output of a real leak test failure:

Testcase: testMemoryLeak(SelectionHandlesTest):        FAILED
Selection handle leaked:
public static java.util.List javafx.scene.Scene@dc8a29-$dirtyCSSNodes->
java.util.ArrayList@4ed14172-elementData->
[Ljava.lang.Object;@3c22de9e-[41]->
javafx.scene.shape.Rectangle@612e4cd

In other words, we see that there is a static list in Scene named dirtyCSSNodes which is holding on to our target Rectangle. The syntax here is that you first see the class name for the object (where a prefix of [L means it's an array-of), then its system id, then the field name - and if it's an array, the array index.

When the unit test fails, it takes a while -- 20 seconds or so for the above test -- to actually produce the above trace. But the important part is that this is only slow when the test fails. It only has to work hard when you have a leak, and you don't want to have leaks!

I've been wanting to write leak tests for the authoring tool (which is written in JavaFX), since uptime matters a great deal in a tool which deals with potentially large objects (such as multimedia assets). And I realized that there is absolutely nothing NetBeans specific about the NetBeans leak unit test support. So I went and pulled out the relevant code into a separate library. The library basically contains two parts: a .jar file which contains the INSANE heap walking library, and a second jar which contains the assertGC() unit test assertion method and supporting infrastructure.

I have extracted this into standalone pieces (outside the NetBeans infrastructure) so you can get the bits easily - download leaktests.zip, then add the two jars in there on your test classpath and call assertGc() and assertSize() from your tests as described above. Here's the basic skeleton for all leak tests:

// (1) Create your objects
// (2) Look up the object you want to ensure gets cleaned up later:
Foo foo = // code to get foo
assertNotNull(foo);

// (3) Call your cleanup code which is supposed to free everything

// (4) Create a weak reference to your object reference, and null out
//   your reference
WeakReference<Foo> ref = new WeakReference<Foo>(foo);
foo = null;
assertGC("Foo leaked", ref);

All I've done is extract NetBeans code written by others so I've kept the licenses the same as for NetBeans. All credit goes to the author of INSANE, Petr Nejedly -- and to the authors of the memory assertion stuff in NbTestCase.java. All the source code for INSANE and NbTestCase are in the NetBeans mercurial repository.

In addition to assertGC, there is also assertSize(). This method can be used to ensure that the transitive size of an object graph is below a certain threshold! This can be good for writing tests to not only ensure that you don't have leaks, but that your data structures are of the rough expected size so you don't need excessive amounts of memory. There are more details on this on the INSANE home page.

One special note on JavaFX: The above leak isn't actually a leak; it is a deliberate optimization of the scenegraph, and the reference will be cleared up after the next scene pulse completes. Therefore, for unit leak tests, in addition to actually nulling out the weak references, you also want to run through a scene pulse as well. One really hacky, implementation-dependent and unsupported way to do that is to call scene.$scenePulseListener.pulse();. You probably want to isolate that in a utility method such that you can update it in one place when it needs to change...

Finally, note that I built this on JDK6. If there is interest perhaps we could create a wrapper project for this on Kenai or java.net, where people can also create say a Maven binary for this, a JDK 5 version (there is nothing JDK6 specific so it just needs a recompile, but I don't have JDK5 on this Snow Leopard Mac), etc. Hope you find this all as useful as I have!