Manipulate bytecode with Javassist

on waitingforcode.com

Manipulate bytecode with Javassist

Javassist is a very useful library to manipulate bytecode in Java. In intuitive way it allows to access class members and change their behaviour - as well in source-code level as in bytecode one.

This article starts by the presentation of Javassist library. It explores several, but not all, features of this library. The second part shows how Javassist could be implemented in real code. It shows how Javassist can be used to add some time execution stats for methods call.

Javassist features

First of all, the mysterious name Javassist means Java programming assistant. The classes can be edited with Java source code or directly with bytecode. We can now explore main features which can be used to manipulate Java source code.

Java classes are represented as the instances of javassist.CtClass. Thanks to its accessor and mutator methods, we can easily make some programming operations: add new field, define new constructor or remove methods and fields. Once modified, the instance of CtClass must be saved through the call of writeFile() method.

We already know how to modify a class, but we don't know how to retrieve this class. It's quite easy because we only need to make two steps: get the instance of javassist.ClassPool. This container exposes several getters which helps to find desired class. But ClassPool it's not a simple read-only container. It contains also some methods to manipulate Java classes (through makeClass()) and packages (with makePackage()). One very important thing to take into account when working with ClassPool is the memory. To prevent some issues, it's advised to detach (detach()) each modified CtClass or replace ClassPool objects with new ones to make them eligible for garbage collection.

It was all about source-code manipulation. Another objects are destined to bytecode manipulation. CtClass is replaced by javassist.bytecode.ClassFile. It can be accessed from CtClass.getClassFile() or constructed directly from an I/O stream. ClassFile can accept bytecode which may be generated with javassist.bytecode.Bytecode class.

Javassist example

In our example we'll decorate each method longer than 10 lines with a code measuring its execution. To simplify reading, we'll have this simple data holder storing time execution of each decorated method:

public class StatsHolder {

  private static final Map<String, List<Long>> DATA = new HashMap<>();

  public static void notifyStats(String className, String methodName, long startTime, long endTime) {
    String key = getKey(className, methodName);
    List<Long> stats = DATA.get(key);
    if (stats == null) {
        stats = new ArrayList<>();
        DATA.put(key, stats);
    }
    stats.add(endTime - startTime);
  }

  public static List<Long> getStats(String className, String methodName) {
    return DATA.get(getKey(className, methodName));
  }

  private static String getKey(String className, String methodName) {
    return className+"_"+methodName;
  }

}

It's time to discover MeasurementDecorator:

public final class MeasurementDecorator {

  private static final String MEASURE_START_OUTPUT = "start = System.currentTimeMillis();";
  private static final String MEASURE_END_OUTPUT = "end = System.currentTimeMillis();";
  private static final Consumer<CtMethod> CONSUMER = new MethodConsumer();
  private static final ClassPool CLASS_POOL = ClassPool.getDefault();

  private MeasurementDecorator() {
    // prevents init
  }

  public static CtClass decorateClass(String className) throws NotFoundException, CannotCompileException,IOException {
    CtClass classToChange = CLASS_POOL.get(className);
    Stream.of(classToChange.getDeclaredMethods()).forEach(CONSUMER);
    classToChange.writeFile();
    classToChange.detach();
    return classToChange;
  }

  private static final class MethodConsumer implements Consumer<CtMethod> {

    @Override
    public void accept(CtMethod method) {
      try {
        // a must-to-do, otherwise start and end variables will remain unknown
        method.addLocalVariable("start", CtClass.longType);
        method.addLocalVariable("end", CtClass.longType);

        String methodName = method.getName();
        String className = method.getDeclaringClass().getName();

        // decorate the class at the begin and at the end with the call to static notifyStats() method
        // also defines the values for previously declared variables
        method.insertBefore(MEASURE_START_OUTPUT);
        method.insertAfter(MEASURE_END_OUTPUT);
        method.insertAfter("com.waitingforcode.StatsHolder.notifyStats(\""+className+"\", \""+methodName+"\", start, end);");
      } catch (CannotCompileException e) {
        throw new RuntimeException("An error occurred on decorating method "+method.getName(), e);
      }
    } 
  }

}

The decoration looks almost exactly as the normal code writting. A little bit different is only the necessity to declare local variables through dedicated method. In tests we'll see the difference between decorated method with and without class loader notification:

public class MeasurementDecoratorTest {

  private static final String SAMPLE_CLASS_NAME = "com.waitingforcode.SamplePrinter";
  private static final String OTHER_CLASS_NAME = "com.waitingforcode.OtherPrinter";

  @Test
  public void should_decorate_class() throws NotFoundException, IOException, CannotCompileException, IllegalAccessException, InstantiationException, InterruptedException {
    MeasurementDecorator.decorateClass(SAMPLE_CLASS_NAME).toClass();
    // It works because class loader was requested to load new representation of SamplePrinter
    SamplePrinter printer = new SamplePrinter();
    printer.printLine();
    printer.printLine();

    List<Long> stats = StatsHolder.getStats(SAMPLE_CLASS_NAME, "printLine");

    assertThat(stats).hasSize(2);
  }

  @Test
  public void should_not_decorate_class_after_loading() throws InterruptedException, IOException, CannotCompileException, NotFoundException, IllegalAccessException, InstantiationException {
    MeasurementDecorator.decorateClass(OTHER_CLASS_NAME);
    // It won't work because class loader doesn't know about new Class representation for OtherPrinter
    OtherPrinter printer = new OtherPrinter();
    printer.printLine();
    printer.printLine();

    List<Long> stats = StatsHolder.getStats(OTHER_CLASS_NAME, "printLine");

    assertThat(stats).isNull();
  }

}

As you can deduce thanks to comments, toClass() method notifies class loader about .class file changes. Another Javassist use could be with Java agents. You can take a quick read on the articles about Java agents and Java instrumentation package to learn more about it. The second article contains even an example of Java agent using Javassist to transform a class before loading.

We can see that Javassist brings some ways easy to implement in class redefinition. The knowledge of plain bytecode is not mandatory. Instead, we can use source Java code and leave Javassist handle bytecode under-the-hood. We can observe that in the second part of the article which contains an example of class decoration with time execution tool.

Share on: