Implementing Aspect Oriented Programming Framework - Java Instrumentation

What is Aspect Oriented Programing: Aspect Orented Programing or AOP is a design pattern or programming practice, in which the code for cross cutting concern is separated out from the main program logic, thus effectively reducing the program complexity. Now what is a cross cutting concern? They are application non-functional requirements that apply to multiple functionalities. A simple example would be logging. Logging is done whenever some predefined things happen (for example a method is invoked). Normally, every method that is invoked would have to contain the code for logging. In AOP, we write the logging code in a separate method may be in a separate class and then configure our AOP framework to do logging whenever a method is invoked. There are so many good tutorials available on AOP, one of which is here. However, in Java SE environment, it certainly requires a change in the method's byte code made by the AOP framework. I will demonstrate how an AOP framework performs load time weaving (changes byte code during runtime).

What is Java Instrumentation: Well instrumentation is modification of java byte code for the purpose of gathering data that are to be utilized by several tools. These data can be used by profilers, loggers etc. However, certain AOP frameworks like AspectJ use the instrumentation framework to implement load time weaving of their advices to the point-cuts.

Java Agent: A java agent is a class with a premain method that gets invoked before the main method is called on the main class. Java agents are loaded using command line option in the java command. The following will give a simple example.

SimpleJavaAgentDemo.java SimpleJavaAgent.java

package com.geekyarticles.instrumentation;

public class SimpleJavaAgentDemo {

    public static void main(String[] args) {
        System.out.println("Hello! its the main method");

    }

}

package com.geekyarticles.instrumentation;

import java.lang.instrument.Instrumentation;
/**
 * Does nothing more than do a few System.out
 * @author debasish
 *
 */
public class SimpleJavaAgent {
    public static void premain(String args, Instrumentation instrumentation){
        System.out.println("Premain");
        System.out.println("Number of classes loaded: "+instrumentation.getAllLoadedClasses().length);
    }
}

To be able to run this code, you must package com.geekyarticles.instrumentation.SimpleJavaAgent in a jar file. Let's call it agent.jar. The MANIFEST.MF file must contain the line Premain-Class: com.geekyarticles.instrumentation.SimpleJavaAgent. Now you can run this command

java -javaagent:agent.jar com.geekyarticles.instrumentation.SimpleJavaAgentDemo

This will give the following output.

Premain
Number of classes loaded: 360
Hello! its the main method

Note here that the output from the premain method of the agent class comes before the main method's output.

Class File Transformation: Note above that the premain method in the agent receives an Instrumentation. Instrumentation has methods to add ClassFileTransformer objects, the transform methods of which get invoked in the order they were added. So, every transformer can do its own bit of changes every time a class is loaded or reloaded. Class reloading can be done by custom class-loaders and this feature is provided by some Java EE servers for dynamically detecting changes in the program.

For example, here I will create a ClassFileTransformer that enables logging (Now doing only System.out). I will use BCEL to achieve this. If you do not know about BCEL, you can read my earlier articles here. We will have something as below.

try {
    ByteArrayInputStream byteIn=new ByteArrayInputStream(classfileBuffer);
    ClassParser classParser=new ClassParser(byteIn, null);
    
    JavaClass originalClass=classParser.parse();
    
    ClassGen classGen=new ClassGen(originalClass);
    
    ConstantPoolGen cp=classGen.getConstantPool();

    for(Method methodclassGen.getMethods()){
        MethodGen methodGen=new MethodGen(method, classGen.getClassName(), cp);
        InstructionList il=methodGen.getInstructionList();
        
        InstructionList modiIL=new InstructionList();
        modiIL.append(new GETSTATIC(cp.addFieldref("java.lang.System", "out", "Ljava/io/PrintStream;")));
        modiIL.append(new LDC(cp.addString("Calling: "+method)));
        modiIL.append(new INVOKEVIRTUAL(cp.addMethodref("java.io.PrintStream", "println", "(Ljava/lang/String;)V")));
        
        il.insert(modiIL);
        methodGen.setMaxStack();
        methodGen.setMaxLocals();
        classGen.removeMethod(method);
        classGen.addMethod(methodGen.getMethod());
        
        ByteArrayOutputStream out=new ByteArrayOutputStream();
        classGen.update();
        classGen.getJavaClass().dump(out);
        return out.toByteArray();
            
            
    }
        
        
} catch (ClassFormatException e) {            
    e.printStackTrace();
} catch (IOException e) {
    e.printStackTrace();
} catch(Throwable t){
    t.printStackTrace();
}

This method is pretty straight forward. We pick every method and insert our logging code at the beginning of its code. However, we have to make sure that only classes in the package com.geekyarticles.instrumentation gets this code insertion, and also we want that only methods annotated with our custom annotation @Loggable gets this logging feature (this is something you always have in any AOP).

The Problem with Annotation: Unfortuanately, the current stable version of BCEL (ie 5.2) does not understand annotations. So we have to parse annotation our selves. Annotations are stored as attributes in the method. The attribute name for an annotation visible at runtime is RuntimeVisibleAnnotations. So we need to attach a reader for this kind of custom attribute. First we need to creat a class AnnotationAttribute. Then we need to attach an AttributeReader to Attribute class through a static method. The createAttribute method of this AttributeReader must return an instance of AnnotationAttribute.

Once this is done, we will get an instance of AnnotationAttribute every time we encounter an Annotation. At this moment, our annotation @Loggable does not have any attribute, so we can skip processing attributes for an annotation. In this case, only the first 6 bytes of the attribute data will be useful. The first two bytes form a pointer to the class the annotation is used in (in the constant pool). The second two bytes point to the signature of the annotation in the constant pool. The last two bytes represent the number of attributes (hence you can only have at most 65535 attributes. In our case, this is always zero. The rest of the code is pretty straight forward. I have added comments whenever necessary.

LoggingJavaAgent.java Loggable.java LoggingClassTransformer.java AnnotationAttribute.java

package com.geekyarticles.instrumentation;

import java.lang.instrument.Instrumentation;

public class LoggingJavaAgent {
    public static void premain(String args, Instrumentation instrumentation){
        instrumentation.addTransformer(new LoggingClassTransformer());
    }
}

package com.geekyarticles.instrumentation;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Loggable {

}

package com.geekyarticles.instrumentation;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Arrays;

import org.apache.bcel.classfile.Attribute;
import org.apache.bcel.classfile.AttributeReader;
import org.apache.bcel.classfile.ClassFormatException;
import org.apache.bcel.classfile.ClassParser;
import org.apache.bcel.classfile.ConstantPool;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.ClassGen;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.GETSTATIC;
import org.apache.bcel.generic.INVOKEVIRTUAL;
import org.apache.bcel.generic.InstructionList;
import org.apache.bcel.generic.LDC;
import org.apache.bcel.generic.MethodGen;

public class LoggingClassTransformer implements ClassFileTransformer{

    static{
        Attribute.addAttributeReader("RuntimeVisibleAnnotations", new AttributeReader() {
            
            
            
            @Override
            public Attribute createAttribute(int name_index, int length,
                    DataInputStream file, ConstantPool constant_pool) {
                byte [] info=null;
                try {
                    
                    System.out.println(length);

                    info=new byte[length];

                    file.readFully(info);
                    
                    
                    System.out.println(Arrays.toString(info));
                    
                } catch (Exception e) {
                    e.printStackTrace();
                }
                return new AnnotationAttribute(name_index, length, info, constant_pool);
            }
        });
    }
    
    @Override
    public byte[] transform(ClassLoader loader, String className,
            Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
            byte[] classfileBuffer) throws IllegalClassFormatException {
            if(!className.startsWith("com/geekyarticles/instrumentation/")){
                return null;
            }
        try {
            ByteArrayInputStream byteIn=new ByteArrayInputStream(classfileBuffer);
            ClassParser classParser=new ClassParser(byteIn, null);
            
            JavaClass originalClass=classParser.parse();
            
            ClassGen classGen=new ClassGen(originalClass);
            
            ConstantPoolGen cp=classGen.getConstantPool();

            for(Method methodclassGen.getMethods()){
                MethodGen methodGen=new MethodGen(method, classGen.getClassName(), cp);

                for(Attribute attrmethod.getAttributes()){
                    //It may look a little weird, but this is how you get the name of an annotation
                    if(attr instanceof AnnotationAttribute && ((AnnotationAttribute)attr).getSignature().equals("Lcom/geekyarticles/instrumentation/Loggable;")){
                        InstructionList il=methodGen.getInstructionList();
                        
                        InstructionList modiIL=new InstructionList();
                        modiIL.append(new GETSTATIC(cp.addFieldref("java.lang.System", "out", "Ljava/io/PrintStream;")));
                        modiIL.append(new LDC(cp.addString("Calling: "+method)));
                        modiIL.append(new INVOKEVIRTUAL(cp.addMethodref("java.io.PrintStream", "println", "(Ljava/lang/String;)V")));
                        
                        il.insert(modiIL);
                        methodGen.setMaxStack();
                        methodGen.setMaxLocals();
                        classGen.removeMethod(method);
                        classGen.addMethod(methodGen.getMethod());
                        
                        ByteArrayOutputStream out=new ByteArrayOutputStream();
                        classGen.update();
                        classGen.getJavaClass().dump(out);
                        return out.toByteArray();
                    }
                    
                }
                
                
                
            }
            
            
        } catch (ClassFormatException e) {            
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch(Throwable t){
            t.printStackTrace();
        }
        
        return null;
    }

}

package com.geekyarticles.instrumentation;

import java.io.DataOutputStream;
import java.io.IOException;

import org.apache.bcel.Constants;
import org.apache.bcel.classfile.Attribute;
import org.apache.bcel.classfile.ConstantPool;
import org.apache.bcel.classfile.ConstantUtf8;
import org.apache.bcel.classfile.Visitor;

public class AnnotationAttribute extends Attribute{
    private byte[] bytes;
    private String signature=null;

    public String getSignature() {
        return signature;
    }

    public void setSignature(String signature) {
        this.signature = signature;
    }

    /**
     * Create a non-standard attribute.
     *
     * @param name_index Index in constant pool
     * @param length Content length in bytes
     * @param bytes Attribute contents
     * @param constant_pool Array of constants
     */
    public AnnotationAttribute(int name_index, int length, byte[] bytes, ConstantPool constant_pool) {
        super(Constants.ATTR_UNKNOWN, name_index, length, constant_pool);
        this.bytes = bytes;
        int signatureIndex=(bytes[2]<<8)+bytes[3];

        signature=((ConstantUtf8)(constant_pool.getConstant(signatureIndex))).getBytes();
        
    }

    /**
     * Called by objects that are traversing the nodes of the tree implicitely
     * defined by the contents of a Java class. I.e., the hierarchy of methods,
     * fields, attributes, etc. spawns a tree of objects.
     *
     * @param v Visitor object
     */
    public void accept( Visitor v ) {
        
    }


    /**
     * Dump unknown bytes to file stream.
     *
     * @param file Output file stream
     * @throws IOException
     */
    //If this is missed out, the attribute will not be properly dumped
    //when we dump the class, which will create all sorts of problems
    
    public final void dump(DataOutputStream file ) throws IOException {
        super.dump(file);

        if (length > 0) {
            file.write(bytes, 0, length);
        }
    }


    /**
     * @return String representation.
     */
    public final String toString() {

        return "(Annotation attribute)";

    }


    /**
     * @return deep copy of this attribute
     */
    public Attribute copy( ConstantPool _constant_pool ) {
        AnnotationAttribute c = (AnnotationAttribute) clone();
        if (bytes != null) {
            c.bytes = new byte[bytes.length];
            System.arraycopy(bytes, 0, c.bytes, 0, bytes.length);
        }
        c.constant_pool = _constant_pool;
        return c;
    }
}

These above classes should be present in the agent jar file. The following demo class can be used as a main class for demonstration.

LoggingJavaAgentDemo.java

package com.geekyarticles.instrumentation;

public class LoggingJavaAgentDemo {

    public static void main(String[] args) {
        System.out.println("Logging java agent");
        demoMethod();
    }
    
    @Loggable
    public static void demoMethod(){
        System.out.println("Demo method");
    }

}

We are now all set to implement our own AOP framework.