Java Instrumentation with JDK 1.6.X, Class Re-definition after loading

Revelation:
It was recently brought to my notice that classes can be redefined after the VM has started and the class has been loaded, thanks to a JDK 1.6.X feature introduction. In other words if we have class that is loaded whose simple method greet() returns "Hello Sanjay", it can be changed to return "Hello world". I knew that you could redefine the class during load time but that changing a class after it has loaded was news to me. We could do the same before with JPDA where classes where hot swapped but not with Instrumentation, at least that is my understanding. In a previous blog of mine, I had discussed the redefinition of classes to augment them with profiling information, the same was done during class load time using an agent and a byte code enhancement library, javaassist.

Research:
Armed with the ammunition from above, I explore the changes. With Java 6, one can load a java agent after the VM has started. This means, we have control on when the agent is loaded. Redefinition of classes is supported by the following changes to the instrumentation package in java 6 available on an instance of the java.lang.instrument.Instrumentation interface:

Instrumentation.retransformClasses(Class...)
Instrumentation.addTransformer(ClassFileTransformer, boolean)
Instrumentation.isModifiableClass(Class)
Instrumentation.isRetransformClassesSupported()

With the above, the research begs the following questions:

How do we enhance byte code of a Class?
A class byte code can be enhanced using a byte code engineering library. There are many popular libraries such as javassist, asm, bcel that can be used to alter a class.

How do we transform the class at run time with the enhanced byte code?
Consider the following example that describes the process,

---

// Obtain altered class byte code 
byte bytes[] = addProfiling(Foo.class); 

// Definition indicating class to be transformed and the new definition. 
ClassDefinition definition = new ClassDefinition(Foo.class, bytes); 

// Ask the instrumentation library to redefine the classes. 
intrumentationInstance.redefineClasses(definition); 

---

The addProfiling(Foo.class) method returns the bytecode of Foo.class.
NOTE: One thing we cannot do with redefining a class is change its "SCHEMA". In other words, we cannot add new methods/remove methods etc. When using JPDA for debugging, you might remember that athough you could change code inside an existing method, attempting to change the signature etc would not work.

How does one obtain a handle to the java.lang.instrument.Instrumentation instance to redefine the classes?

In my previous instrumentation example blog, the agent class was initialized via the following method public static void premain(String agentArgs, Instrumentation inst). With Java 6 in order to start the agent after the VM has started, the Agent class needs to define the method:

public static void agentmain(String agentArgs, Instrumentation inst)
Both the methods shown above provide an instance of the Instrumentation interface and thus a place to obtain a handle to the Instrumentation implementation.

An agent class with the following can facilitate the re-transformation of classes:

---

public class Agent {
  private static Instrumentation instrumentation;

  public static void agentmain(String agentArgs, Instrumentation inst) {
    Agent.instrumentation = inst;
  }

  public static void redefineClasses(ClassDefinition ...defs) throws Exception {
   Agent.instrumentation.redefineClasses(defs);
  }
}

---

How does one load the agent?

Apart from having a class that defines the agentmain method, the jar that will contain the agent must have certain entries in its MANIFEST file as shown below:

---

Agent-Class: com.welflex.Agent
Can-Retransform-Classes: true
Can-Redefine-Classes: true

---

The Can-Redefine-Classes attribute describes the ability of the agent to re-transform classes.

With JDK 5.X the agent had to specified in the command line via: -javaagent:jarpath in order to start the agent.With JDK 6.X the same is not required. We could load an agent at runtime via the location of the jar file. The loading of the agent is faciliated by the VirtualMachine class present in the tools.jar of the JDK. An example of how an agent is loaded is shown below:

---

// Location of agent jar
String agentPath = "/home/sanjay/agent.jar";

// Attach to VM
VirtualMachine vm = VirtualMachine.attach(processIdOfJvm);

// Load Agent
vm.loadAgent(agentPath);
vm.detach();

---

The VirtualMachine class is part of the package com.sun.tools.attach which is present in the tools.jar file located in $JDK_HOME/lib. Thus if loading of an agent post JVM start is desired, it naturally follows that the tools.jar has to be available in the class path.

An Example:
For the sake of discussion, let us consider the following requirements of our application:

• Select classes to profile. For example, display profiling information for class Foo but not for class Bar


• Profiling information should be removable. For example, if class Application was profiled in one test, then on a subsequent test, the profiling information for class Application should not exist unless requested by the test.

I have a class called Profiler. The profiler class has the following definition:

---

public class Profiler {
  public static void addClassesToProfile(Class ...classes) {
   ClassDefintions defs[] = Profiler.buildProfiledClassDefintions(classes);
   Profiler.loadAgent();
   InstrumentationDelegate.redefineClasses(classDefintions);
  }

  private static Map<class<?>, byte[]> byteHolderMap = new HashMap..;

  private static ClassDefinition[] buildProfiledClassDefinitions(Class classes) {
  //1.Store original class byte code in byteHolderMap
  //2.Enhance classes with profiling information and
  // return back ClassDefinitions with enhanced code.
  ....
  }

 public static void restoreClasses() {
   // For each entry in the byteHolderMap, redefine the class.
   ...
 }
...
}

---

When the method addClassesToProfile(Class classes) is invoked, the Profiler will create ClassDefintions that contain profiling information. After this the agent (InstrumentationDelegate) is loaded thereby obtaining a handle to the java.lang.instrument.Instrumentation implementation. The call to redefineClasses() will restore the original class schema. The method buildProfiledClassDefinitions() creates a Map of [Class, byte[]], that stores the original byte code of the class. Thus, when the restoreClasses() method is invoked, the original classes are re-instated.

Shown below is a unit test that utilizes the profiling agent:

• Runs a Test using the class without instrumentation
• Enhances the code with profiling information and runs the tests
• Restores the class under test to the one before the profiling information was added
• Re-runs the test with no profiling information.

---

public void testApplication() throws Exception {
System.out.println("Test Profiling Application Class.");

// Run test before Profiling, class will not have profiling information
Application appInit = new Application(new Description().setDescription("Pre-Profiling"));
System.out.println("Class not augmented yet, so should not see profiling information.");
assertEquals("Pre-Profiling", appInit.getDescription());

System.out.println("Adding Profiling information and running test, should see profiling.");

// Redefine and add Profiling information, should see profiling info
// printed out
Profiler.addClassesToProfile(Application.class);

Application app = new Application(new Description().setDescription("Simple App"));
app.rest();
app.rest(3000);

assertEquals("Simple App", app.getDescription());

// Remove profiling information
Profiler.restoreClasses();

// Classes must not display profiling information.
System.out.println("Classes have been restored, should not see profiling information now..");
Application afterRestore = new Application(new Description().setDescription("After Restore"));
assertEquals("After Restore", afterRestore.getDescription());
}

---

Running com.welflex.profiletarget.ApplicationTest
Test Profiling Application Class.
Class not augmented yet, so should not see profiling information.
Adding Profiling information and running test, should see profiling.
-> Enter Method:com.welflex.profiletarget.Application.rest()
<- Exit Method:com.welflex.profiletarget.Application.rest() completed in 2112 nano secs
-> Enter Method:com.welflex.profiletarget.Application.rest()
<- Exit Method:com.welflex.profiletarget.Application.rest() completed in 1000060102 nano secs
-> Enter Method:com.welflex.profiletarget.Application.rest(long)
<- Exit Method:com.welflex.profiletarget.Application.rest(long) completed in 980 nano secs
-> Enter Method:com.welflex.profiletarget.Application.rest(long)
<- Exit Method:com.welflex.profiletarget.Application.rest(long) completed in 3000029988 nano secs
-> Enter Method:com.welflex.profiletarget.Application.getDescription()
<- Exit Method:com.welflex.profiletarget.Application.getDescription() completed in 29749 nano secs
Classes have been restored, should not see profiling information now..
Tests run: 1, Failures: 0, Errors: 0, Skipped: 0, Time elapsed: 4.927 sec
Running the Example:
The example can be downloaded from HERE. The example uses the byte code engineering library javassist to augment the classes. The project itself is a Maven 2 project and requires JDK 1.6.X to execute. You might need to tweak the pom.xml file of the profiler-instrumentation project to point to your locally installed tools.jar, i.e., change the line ${java.home}/../lib/tools.jar, to point to your local tool.jar via a fully qualified path. I have had some difficulty getting this correctly working and had to use an absolute path to the tools.jar. After this, issue a "mvn install" from the root project directory to see the profiling example in action. Note also that if you import the project into Eclipse via Q4E, then when running the example, ensure that the profiler-integration project is closed otherwise the profiler-target project will pick the eclipse project instead of looking for the deployed profiler-integration jar file.

If you have difficulties running the example, please do not hesitate to contact me.

Closing Thoughts:
One can instrument classes on a running VM quite easily with the above concept.
Mr.Jack Shirazi has a nice example of hotpatching a java application
I have personally not tried attaching to a different VM but should be pretty straight forward. The example I have demonstrated is quite largely based on the JMockit testing framework that utilizes byte code enhancement. It is quite a powerful testing framework and I need to spend some time playing with it. But right now, I am off to play with my XBox :-)