Implementing Build-time Bytecode Instrumentation With Javassist

If you need to modify the code in class files at the (post-)build time without adding any third-party dependencies, for example to inject cross-cutting concerns such as logging, and you don’t wan’t to deal with the low-level byte code details, Javassist is the right tool for you. I’ve already blogged about “Injecting better logging into a binary .class using Javassist” and today I shall elaborate on the instrumentation capabilities of Javassist and its integration into the build process using a custom Ant task.

Terminology

• Instrumentation – adding code to existing .class files
• Weaving – instrumentation of physical files, i.e. applying advices to class files
• Advice – the code that is “injected” to a class file; usually we distinguish a “before”, “after”, and ‘around” advice based on how it applies to a method
• Pointcut – specifies where to apply an advice (e.g. a fully qualified class + method name or a pattern the AOP tool understands)
• Injection – the “logical” act of adding code to an existing class by an external tool
• AOP – aspect oriented programming

Javassist versus AspectJ

Why should you use Javassit over a classical AOP tool like AspectJ? Well, normally you wouldn’t because AspectJ is easier to use, less error-prone, and much more powerful. But there are cases when you cannot use it, for example you need to modify bytecode but cannot afford to add any external dependencies. Consider the following when deciding between them:

Javassist:

• Only basic (but often sufficient) instrumentation capabilities
• Build-time only – modifies .class files
• The modified code has no additional dependencies (except those you add), i.e. you don’t need the javassist.jar at the run-time
• Easy to use but not as easy as AspectJ; the code to be injected is handled over as a string, which is compiled to bytecode by Javassist

AspectJ:

• Very powerful
• Both build-time and load-time (when class gets loaded by the JVM) weaving (instrumentation) supported
• The modified code depends on the AspectJ runtime library (advices extend its base class, special objects used to provide access to the runtime information such as method parameters)
• It’s use is no different from normal Java programming, especially if you use the annotation-based syntax (@Pointcut, @Around etc.). Advices are compiled before use and thus checked by the compiler

Classical bytecode manipulation library:

• Too low-level, you need to define and add bytecode instructions, while Javassist permits you to add pieces of Java code

Instrumenting with Javassist

About some of the basic changes you can do with Javassist. This by no means an exhaustive list.

Declaring a local variable for passing data from a before to an after advice

If you need to pass some data from a before advice to an after advice, you cannot create a new local variable in the code passed to Javassist (e.g. “int myVar = 5;”). Instead of that, you must declare it via CtMethod.addLocalVariable(String name, CtClass type) and then you can use is in the code, both in before and after advices of the method.

Example:

final CtMethod method = ...;
method.addLocalVariable("startMs", CtClass.longType);
 method.insertBefore("startMs = System.currentTimeMillis();");
 method.insertAfter("{final long endMs = System.currentTimeMillis();" +
"System.out.println(\"Executed in ms: \" + (endMs-startMs));}");

Instrumenting a method execution

Adding a code at the very beginning or very end of a method:

// Advice my.example.TargetClass.myMethod(..) with a before and after advices
final ClassPool pool = ClassPool.getDefault();
final CtClass compiledClass = pool.get("my.example.TargetClass");
final CtMethod method = compiledClass.getDeclaredMethod("myMethod");

 method.addLocalVariable("startMs", CtClass.longType);
 method.insertBefore("startMs = System.currentTimeMillis();");
 method.insertAfter("{final long endMs = System.currentTimeMillis();" +
   "System.out.println(\"Executed in ms: \" + (endMs-startMs));}");

compiledClass.writeFile("/tmp/modifiedClassesFolder");
// Enjoy the new /tmp/modifiedClassesFolder/my/example/TargetClass.class

There is also CtMethod.insertAfter(String code, boolean asFinally) – JavaDoc: if asFinally “is true then the inserted bytecode is executed not only when the control normally returns but also when an exception is thrown. If this parameter is true, the inserted code cannot access local variables.”

Notice that you always pass the code as either a single statement, as in “System.out.println(\”Hi from injected!\”);” or as a block of statements, enclosed by “{” and “}”.

Instrumenting a method call

Sometimes you cannot modify a method itself, for example because it’s a system class. In that case you can instrument all calls to that method, that appear in your code. For that you need a custom ExprEditor subclass, which is a Visitor whose methods are called for individual statements (such as method calls, or instantiation with a new) in a method. You would then invoke it on all classes/methods that may call the method of interest.

In the following example, we add performance monitoring to all calls to javax.naming.NamingEnumeration.next():

final CtClass compiledClass = pool.get("my.example.TargetClass");
final CtMethod[] targetMethods = compiledClass.getDeclaredMethods();
for (int i = 0; i < targetMethods.length; i++) {
  targetMethods[i].instrument(new ExprEditor() {
    public void edit(final MethodCall m) throws CannotCompileException {
      if ("javax.naming.NamingEnumeration".equals(m.getClassName()) && "next".equals(m.getMethodName())) {
        m.replace("{long startMs = System.currentTimeMillis(); " +
            "$_ = $proceed($$); " +
            "long endMs = System.currentTimeMillis();" +
            "System.out.println(\"Executed in ms: \" + (endMs-startMs));}");
      }
    }
  });
}

Beware: What matters is the declared type on which the method is invoked, which can be an interface, as in this example, the actual implementation isn’t important. This is opposite to the method execution instrumentation, where you always instrument a concrete type.

The problem with instrumenting calls is that you need to know all the classes that (may) include them and thus need to be processed. There is no official way of listing all classes [perhaps matching a pattern] that are visible to the JVM, though ther’re are some workarounds (accessing the Sun’s ClassLoader.classes private property). The best way is thus – aside of listing them manually – to add the folder or JAR with classes to Javassist ClassPool’s internal classpath (see below) and then scan the folder/JAR for all .class files, converting their names into class names. Something like:

// Groovy code; the method instrumentCallsIn would perform the code above:
pool.appendClassPath("/path/to/a/folder");
new File("/path/to/a/folder").eachFileRecurse(FileType.FILES) {
 file -> instrumentCallsIn( pool.get(file.getAbsolutePath().replace("\.class$","").replace('/','.')) );}

Javassist and class-path configuration

You certainly wonder how does Javassist find the classes to modify. Javassist is actually extremely flexible in this regard. You obtain a class by calling

private final ClassPool pool = ClassPool.getDefault();
...
final CtClass targetClass = pool.get("target.class.ClassName");

The ClassPool can search a number of places, that are added to its internal class path via the simple call

/* ClassPath newCP = */ pool.appendClassPath("/path/to/a/folder/OR/jar/OR/(jarFolder/*)");

The supported class path sources are clear from the available implementations of  ClassPath: there is a ByteArrayClassPath, ClassClassPath, DirClassPath, JarClassPath, JarDirClassPath (used if the path ends with “/*”), LoaderClassPath, URLClassPath.

The important thing is that the class to be modified  or any class used in the code that you inject into it doesn’t need to be on the JVM classpath, it only needs to be on the pool’s class path.

Implementing mini-AOP with Javassist and Ant using a custom task

This part briefly describes how to instrument classes with Javassist via a custom Ant task, which can be easily integrated into a build process.

The corresponding part of the build.xml is:

<target name="declareCustomTasks" depends="compile">
   <mkdir dir="${antbuild.dir}"/>

   <!-- Javac classpath contains javassist.jar, ant.jar -->
   <javac srcdir="${antsrc.dir}" destdir="${antbuild.dir}" encoding="${encoding}" source="1.4" classpathref="monitoringInjectorTask.classpath" debug="true" />

 <taskdef name="javassistinject" classname="example.JavassistInjectTask"
   classpathref="monitoringInjectorTask.classpath" loaderref="javassistinject"/>
 <typedef name="call" classname="example.JavassistInjectTask$MethodDescriptor"
   classpathref="monitoringInjectorTask.classpath" loaderref="javassistinject"/>
 <typedef name="execution" classname="example.JavassistInjectTask$MethodDescriptor"
   classpathref="monitoringInjectorTask.classpath" loaderref="javassistinject"/>
 </target>

 <target name="injectMonitoring" depends="compile,declareCustomTasks" description="Process the compiled classes and inject calls to the performance monitoring API to some of them (currently hardcoded in PerfmonAopInjector)">

  <javassistinject outputFolder="${classes.dir}" logLevel="info">
    <fileset dir="${classes.dir}" includes="**/*.class">
    <!-- method executions to inject with performance monitoring -->
    <execution name="someSlowMethod" type="my.MyClass" />
    <!-- method calls to inject with performance monitoring -->
    <call name="search" type="javax.naming.directory.InitialDirContext" metric="ldap" />
    <call name="next" type="javax.naming.NamingEnumeration" metric="ldap" />
    <call name="hasMore" type="javax.naming.NamingEnumeration" metric="ldap" />
  </javassistinject>

 </target>

Noteworthy:

• I’ve implemented a simple custom Ant task with the class example.JavassistInjectTask, extending org.apache.tools.ant.Task. It has setters for attributes and nested elements and uses the custom class PerfmonAopInjector (not shown) to perform the actual instrumentation via Javassist API. Attributes/nested elements:
  • setLoglevel(EchoLevel level) – see the EchoTask
  • setOutputFolder(File out)
  • addConfiguredCall(MethodDescriptor call)
  • addConfiguredExecution(MethodDescriptor exec)
  • addFileset(FileSet fs) – use fs.getDirectoryScanner(super.getProject()).getIncludedFiles() to get the names of the files under the dir
• MethodDescriptor is a POJO with a no-arg public constructor and setters for its attributes (name, type, metric), which is introduced to Ant via <typedef> and its instances are passed to the JavassistInjectTask by Ant using its addConfigured<name>, where the name equlas the element’s name, i.e. the name specified in the typedef
• PerfmonAopInjector is another POJO that uses Javassist to inject execution time logging to method executions and calls as shown in the previous section, applying it to the classes/methods supplied by the JavassistInjectTask based on its <call .. /> and <execution … /> configuration
• The fileset element is used both to tell Javassist in what directory it should look for classes and to find out the classes that may contain calls that should be instrumented (listing all the .class files and converting their names to class names)
• All the typedefs use the same ClassLoader instance so that the classes can see each other, this is ensured by loaderref="javassistinject" (its value is a custom identifier, same for all three)
• The monitoringInjectorTask.classpath contains javassist.jar, ant.jar, JavassistInjectTask, PerfmonAopInjector and their helper classes
• The classes.dir contains all the classes that may need to be instrumented and the classes used in the injected code, it’s added to the Javassist’s internal classpath via ClassPool.appendClassPath(“/absolute/apth/to/the/classes.dir”)

Notice that System.out|err.println called by any referenced class are automatically  intercepted by Ant and changed into Task.log(String msg, Project.MSG_INFO) and will be thus included in Ant’s output (unless -quiet).

PS: If using maven, you’ll be happy yo know that Javassist is in a Maven repository (well, at least it has a pom.xml, so I suppose so).

Ant custom task resources

1. Rob Lybarger: Introduction to Custom Ant Tasks (2006) – the basics
2. Rob Lybarger: More on Custom Ant Tasks (2006) – about nested elements
3. Ant manual: Writing Your Own Task
4. Stefan Bodewig: Ant 1.6 for Task Writers (2005)

From http://theholyjava.wordpress.com/2010/06/25/implementing-build-time-instrumentation-with-javassist/