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));}");
      }
    }
  });
}

The call to the method of interest is replaced with another code, which also performs the original call via the special statement "$_ = $proceed($$);".

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 (see the code below). 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).

JavassistInjectTask.java (click it to expand):

package example;

import java.io.File;
import java.util.*;
import javassist.CannotCompileException;
import org.apache.tools.ant.*;
import org.apache.tools.ant.types.*;

/**
 * Invoke PerfmonAopInjector on a set of classes to instrument some
 * methods and calls to some methods.
 * <p>
 * The Javassist library must be on the classpath.
 */
public class JavassistInjectTask extends Task {

    /**
     * The destination of the stream. If <code>null</code>, the system
     * console is used.
     */
    private File outputFolder = new File("instrumented");

    /**
     * Stores a collection of file sets and/or file lists, used to
     * select multiple files for concatenation.
     */
    private Vector inputFilesets = new Vector();

	private int logLevel = PerfmonAopInjector.LOG_INFO;

	private List interceptedExecutions = new LinkedList();
	private List interceptedCalls = new LinkedList();

    /**
     * Sets the destination file, or uses the console if not specified.
     * @param destinationFile the destination file
     */
    public void setOutputFolder(File destinationFile) {
        this.outputFolder = destinationFile;
    }

    /**
     * Set of files to concatenate.
     * @param set the set of files
     */
    public void addFileset(FileSet set) {
        inputFilesets.addElement(set);
    }

    public void execute() throws BuildException {
    	try {
    		tryExecute();
    	} catch (BuildException e) {
    		log("TASK FAILED: " + e, Project.MSG_ERR);
    		e.printStackTrace();
    		throw e;
    	} catch (RuntimeException e) {
    		log("TASK FAILED: " + e, Project.MSG_ERR);
    		e.printStackTrace();
    		throw e;
    	}
    }

    public void tryExecute() throws BuildException {

        log("STARTING TO INJECT MONITORING...", Project.MSG_WARN);

        verifyInputs();

        // Iterate thru the sources - paths, filesets and filelists
        final List inClassFileDirs = collectTargetClassFileDirs();

        log("Loaded class dirs: " + inClassFileDirs, Project.MSG_INFO);

        final PerfmonAopInjector perfmonAopInjector = createPerfmonAopInjector();
        setInjectorClassPath(perfmonAopInjector, inClassFileDirs);

        injectMethodExecutionsMonitoring(perfmonAopInjector);

        injectMethodCallsMonitoring(perfmonAopInjector, inClassFileDirs);

        log("DONE. Classes and methods modified:\n" + perfmonAopInjector.getAllLogMessages(), Project.MSG_INFO);
    }

	private PerfmonAopInjector createPerfmonAopInjector() {
		final PerfmonAopInjector perfmonAopInjector = new PerfmonAopInjector();
        perfmonAopInjector.setLogLevel(logLevel);
        perfmonAopInjector.setOutputFolder(this.outputFolder);
        log("createPerfmonAopInjector: method executions to instrument: " + interceptedExecutions, Project.MSG_DEBUG);
        perfmonAopInjector.setInstrumentedMethodExecutions(toArray(interceptedExecutions));
        log("createPerfmonAopInjector: method calls to instrument: " + interceptedCalls, Project.MSG_DEBUG);
        perfmonAopInjector.setInstrumentedMethodCalls(toArray(interceptedCalls));
		return perfmonAopInjector;
	}

	private MonitoredMethodDescriptor[] toArray(final List monitoredMethodDescriptors) {
		return (MonitoredMethodDescriptor[]) monitoredMethodDescriptors.toArray(
				new MonitoredMethodDescriptor[monitoredMethodDescriptors.size()]);
	}

	private void injectMethodCallsMonitoring(
			final PerfmonAopInjector perfmonAopInjector,
			final List inClassFileDirs) {
		log("Going to inject monitoring of method *calls*...", Project.MSG_INFO);

		for (Iterator iterator = inClassFileDirs.iterator(); iterator.hasNext();) {
        	final ClassFileDir classFileDir = (ClassFileDir) iterator.next();
        	injectCallMonitoringForClasses(perfmonAopInjector, classFileDir.getClassFiles() );
		}
	}

	private void injectMethodExecutionsMonitoring(
			final PerfmonAopInjector perfmonAopInjector) {

		log("Going to inject monitoring into method *executions*...", Project.MSG_INFO);

		try {
			perfmonAopInjector.monitorConfiguredMethodExecutions();
		} catch (Exception e) {
			String msg = e.toString();
			if (e instanceof CannotCompileException) {
				if (e.getMessage().indexOf("no such class") >= 0) {
					msg = "A class cannot be found. Make sure that not only the classes to be " +
							"instrumented but also the classes used in the injected code are " +
							"in the class file file set. Cause: " + e;
				}
			}
			throw new BuildException("Injecting monitoring into methods failed: " + msg, e);
		}
	}

    private void injectCallMonitoringForClasses(final PerfmonAopInjector perfmonAopInjector,
    		final List classFiles) {
    	for (Iterator iterator = classFiles.iterator(); iterator.hasNext();) {
			final File classFile = (File) iterator.next();
			injectCallsMonitoringForClass(perfmonAopInjector, classFile);
		}
	}

    private void setInjectorClassPath(final PerfmonAopInjector perfmonAopInjector,
    		final List classFileDirs) {

    	for (Iterator iterator = classFileDirs.iterator(); iterator.hasNext();) {
			final ClassFileDir clasDir = (ClassFileDir) iterator.next();
	    	perfmonAopInjector.addClassDir(clasDir.getBaseDir());
		}
	}

	private void injectCallsMonitoringForClass(final PerfmonAopInjector perfmonAopInjector, final File classFile) {
		log("Processing file " + classFile.getName() + "...", Project.MSG_DEBUG);
		final String targetClass = relativeFileToClass(classFile);
		try {
			perfmonAopInjector.monitorConfiguredMethodCallsIn(targetClass);
		} catch (Exception e) {
			throw new BuildException("Injecting method calls monitoring into the class " +
					classFile + " failed.", e);
		}
    }

	private String relativeFileToClass(final File classFile) {
		final String dotSeparatedName = classFile.getPath()
			.replace(File.separatorChar, '.')
//			.replace('$', '.')
			;
		log("relativeFileToClass: Converting file '" + classFile + "' to '" + dotSeparatedName +
				"'.substring", Project.MSG_DEBUG);
		return dotSeparatedName.substring(
				0, dotSeparatedName.indexOf(".class"));
	}

	private List collectTargetClassFileDirs() {
    	final List allClassFileDirs = new Vector();
		for (Enumeration e = inputFilesets.elements(); e.hasMoreElements();) {
        	final FileSet fileSet = (FileSet) e.nextElement();
            final DirectoryScanner scanner =
                    fileSet.getDirectoryScanner(getProject());

            final File classesDir = fileSet.getDir(getProject());
			final List inClassFiles = checkAddFiles(
            		classesDir,
            		scanner.getIncludedFiles());

			allClassFileDirs.add(new ClassFileDir(classesDir, inClassFiles));
        }
		return allClassFileDirs;
	}

	private void verifyInputs() {
		if (inputFilesets.size() == 0) {
            throw new BuildException(
                "At least one class folder must be provided");
        }
		if(outputFolder == null || !outputFolder.isDirectory()) {
			throw new BuildException("The attribute outputFolder must be set to an existing folder; is: " + outputFolder);
		}
	}

    /**
     * Reset state to default.
     */
    public void reset() {
        outputFolder = null;
        inputFilesets.removeAllElements();
    }

    private List checkAddFiles(final File classDir, final String[] classesRelative) {
    	final List inClassFiles = new Vector();

        for (int i = 0; i < classesRelative.length; ++i) {
            final File file = new File(classesRelative[i]);
            if (!file.getName().endsWith(".class")) {
            	log("checkAddFiles: Ignoring non-class file " + file, Project.MSG_WARN);
            } else if (file.getName().matches(".*\\$[0-9]+.class")) {
            	log("checkAddFiles: Ignoring the anonymous inner class " + file, Project.MSG_WARN);
            } else {
            	inClassFiles.add(file);
            }
        }
        return inClassFiles;
    }

    /**
     * Set the logging level. Level should be one of
     * <ul>
     *  <li>error</li>
     *  <li>warning</li>
     *  <li>info</li>
     *  <li>debug</li>
     * </ul>
     * <p>The default is &quot;info&quot;.</p>
     * @param echoLevel the logging level
     */
    public void setLoglevel(final LogLevel echoLevel) {
        String option = echoLevel.getValue();
        if (option.equals("error")) {
            logLevel = PerfmonAopInjector.LOG_WARN;
        } else if (option.equals("warning")) {
            logLevel = PerfmonAopInjector.LOG_WARN;
        } else if (option.equals("info")) {
            logLevel = PerfmonAopInjector.LOG_INFO;
        } else {
            // must be "debug"
        	logLevel = PerfmonAopInjector.LOG_DEBUG;
        }
    }

    public void addConfiguredExecution(final MethodDescriptor methodExecution) {
    	addCallOrExecution(interceptedExecutions, methodExecution);
    }

    public void addConfiguredCall(final MethodDescriptor methodCall) {
    	addCallOrExecution(interceptedCalls, methodCall);
    }

    private void addCallOrExecution(final List destination, final MethodDescriptor method) {
    	method.validate();
    	destination.add(new MonitoredMethodDescriptor(
    			method.getType(), method.getName(), method.getMetric()));
    }

    //@Override
    public void log(String msg, int msgLevel) {
    	int actualLogLevel = msgLevel;
    	// Increase debug to info to make sure it's logged if desired level is debug
    	if (msgLevel == Project.MSG_DEBUG && this.logLevel == PerfmonAopInjector.LOG_DEBUG) {
    		actualLogLevel = Project.MSG_INFO;
    	}
    	super.log(msg, actualLogLevel);
    }

    /**
     * Represents folder containing .class files.
     */
    public static class ClassFileDir {
    	private List classFiles;
    	private File baseDir;

		public ClassFileDir(File baseDir, List classFiles) {
			this.baseDir = baseDir;
			this.classFiles = classFiles;
		}
		public List getClassFiles() {
			return classFiles;
		}
		public File getBaseDir() {
			return baseDir;
		}
		public String toString() {
			return "[dir=" + baseDir + ",classes=" + classFiles + "]";
		}
    }

    /**
     * The enumerated values for the level attribute.
     */
    public static class LogLevel extends EnumeratedAttribute {
        /**
         * @see EnumeratedAttribute#getValues
         * @return the strings allowed for the level attribute
         */
        public String[] getValues() {
            return new String[] {"error", "warning", "info", "debug"};
        }
    }

    public static class MethodDescriptor {
    	private String type;
    	private String name;
    	private String metric;
		public String getType() {
			return type;
		}
		public void setType(String type) {
			this.type = type;
		}
		public String getName() {
			return name;
		}
		public void setName(String name) {
			this.name = name;
		}
		public String getMetric() {
			return metric;
		}
		public void setMetric(String metric) {
			this.metric = metric;
		}
    	public void validate() throws BuildException {
    		String error = "";
    		if (name == null) {
    			error += "The attribute name is required for a method.";
    		}
    		if (type == null) {
    			error += " The attribute type is required for a method.";
    		}

    		if (error.length() > 0) {
    			throw new BuildException(error);
    		}
    	}
    }

}

PS: If using maven, you'll be happy to 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. Ant manual: Writing Your Own Task - incl. Supporting nested elements (you only need create<name>, add<name>,  or addConfigured<name>)
2. Rob Lybarger: Introduction to Custom Ant Tasks (2006) - the basics
3. Rob Lybarger: More on Custom Ant Tasks (2006) - about nested elements
4. Stefan Bodewig: Ant 1.6 for Task Writers (2005)

Tags: java library