Don't worry, it's ... Probably Fine

Instrumentation with Annotations

08 Jan 2012

java

I’ve been fascinated by java.lang.instrumentation since I’d first properly seen it used; to introduce more finely grained line numbers whilst error handling. At first, I was a little taken aback; modifying bytecode at runtime, or editing classes in place? Madness! But the more I thought about it, the more interesting it became.

What I Wanted

When writing code, most of the time there are 2 main things I want to check:

  • Method Calls – what order are they being called in, if at all?
  • Local Variables – what am I returning/processing in my methods?

But writing all those print/logging statements? Effort! Out of academic interest in instrumentation and using file transformation, I set out to write a transformer that does exactly these things. Specifically, if I put an annotation on a method then it should display this behaviour at runtime when run with the transformer.

Starting Small

Logging method access was the easiest to implement first – we simply want to prepend all of the instructions with a “Entering ” output. So I wanted,

@MethodAccess
public void foo() {
    //Some stuff here
}

To output

Entering foo.
...

After digging around with the ASM libraries from objectweb, the easiest way to do this was to load the bytecode into a ClassFileTransformer, examine all methods with the @MethodAccess annotation and prepend statements to that MethodNode. It ended up looking like this:

private static InsnList localVar(MethodNode node) {
  InsnList list = new InsnList();

  for (LocalVariableNode l : new ArrayList<LocalVariableNode>(node.localVariables)) {

    if (0 == l.index) continue;

    list.add(new FieldInsnNode(Opcodes.GETSTATIC,
      "Ljava/lang/System;", "out", "Ljava/io/PrintStream;"));

    if (l.desc.matches("I")) {
      list.add(new VarInsnNode(Opcodes.ILOAD, l.index));
    } else if (l.desc.matches("D")) {
      list.add(new VarInsnNode(Opcodes.DLOAD, l.index));
    } else if (l.desc.matches("F")) {
      list.add(new VarInsnNode(Opcodes.FLOAD, l.index));
    } else if (l.desc.matches("L")) {
      list.add(new VarInsnNode(Opcodes.LLOAD, l.index));
    } else {
      list.add(new VarInsnNode(Opcodes.ALOAD, l.index));
    }

    list.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL,
      "java/io/PrintStream" , "println", "("+l.desc+")V"));
  }

  return list;
}

This is a good start, but there are 2 issues we need to be careful of; where we put these statements, and whether we’re trying to output variables that haven’t been initalised yet. Firstly, putting it as the penultimate set of instructions (before the final return) is insufficient, given that methods can exit before that with other returns or throwing exceptions.

We can fix that like so:

ListIterator<AbstractInsnNode> it = m.instructions.iterator();

  while (it.hasNext()) {
    AbstractInsnNode node = it.next();
    int opcode = node.getOpcode();

    if ((opcode >= IRETURN && opcode <= RETURN) || opcode == ATHROW) {
      InsnList newInstructions = localVar(m,
        m.instructions.indexOf(node), m.instructions);

      m.instructions.insertBefore(node.getPrevious(), newInstructions);
    }
  }

Since JVM opcodes are represented as ints by ASM, it suffices to check the opcode of each instruction in the method’s instruction list, and whether it is in the range of values that represent RETURN operations, or the single value for THROW.

Finally, we address the last issue, that of uninitialised variables. After some thought, it is sufficient to add a few lines to our localVars method; in abstract, when we deal with a “return” statement, we want to note its place in the instruction list, and any LocalVariableNodes that “start” after that point are discarded for that particular return statement. As code:

InsnList newInstructions = localVar(m, m.instructions.indexOf(node), m.instructions);
//...
private InsnList localVar(MethodNode node, int line, InsnList insnlist) {
  if (insnlist.indexOf(l.start) > line) {
    continue;
  }
  //Rest of the method
}

And there we are. Two annotations that respectively will log method use and log initialised local variables before returning.

Source code available here