Collecting additional metrics with JMH

[Doug Simon]

For Graal[1], I want to use JMH to micro benchmark core parts of the compiler such as how long a certain compiler transformation takes to run. To do so, I want to use a fixture that builds or clones a compiler graph prior to invoking the benchmarking method. The goal is for the benchmarking method to only contain the application of the compiler transformation and not include the cost of building/cloning the input graph.

@State(Scope.Thread)
public class GraphState {
    private StructuredGraph originalGraph;
    public StructuredGraph graph;
    public GraphState(StructuredGraph graph) {
        this.originalGraph = graph;
    }

    @Setup(Level.Invocation)
    public void beforeInvocation() {
        graph = originalGraph.copy();
    }
}

This works.

However, I would also like to profile the allocation done *only* in the @GMB method. This will allow us to easily see one possible cause for a regression in performance. I wanted to do this using com.sun.management.ThreadMXBean.getThreadAllocatedBytes(long) by modifying GraphState as follows:

    public static final AtomicLong totalAlloc = new AtomicLong();
    public static final AtomicInteger totalOps = new AtomicInteger();
    static {
        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                System.out.println("TOTAL_ALLOCATED_BYTES: " + NumberFormat.getInstance().format(totalAlloc.get()));
                System.out.println("AVERAGE_ALLOCATED_BYTES_PER_OP: " + NumberFormat.getInstance().format(totalAlloc.get() / totalOps.get()));
            }
        });
    }

    public StructuredGraph originalGraph;
    public StructuredGraph graph;
    public long allocBefore;

    @Setup(Level.Invocation)
    public void beforeInvocation() {
        graph = originalGraph.copy();
        allocBefore = GraalUtil.getCurrentThreadAllocatedBytes();
    }

    @TearDown(Level.Invocation)
    public void afterInvocation() {
        long delta = GraalUtil.getCurrentThreadAllocatedBytes() - allocBefore;
        totalAlloc.addAndGet(delta);
        totalOps.addAndGet(1);
    }

The use of statics and shutdown hooks for collecting and printing the allocation profiles feels very hacky to me. However, I can’t see any other way to do it. Is there a more elegant way to do this or am I simply trying to squeeze a square peg through a round hole?

-Doug

[1] https://wiki.openjdk.java.net/display/Graal/Main