Fixing Bugs in Running Java Code with Dynamic Attach

Most developers know Java’s HotSwap feature as a debugging tool that is built into most JVMs. Using this feature, it is possible to change the implementation of a Java method without restarting a Java process which is typically used via an IDE while developing code. HotSwap can however be used just as much in a production environment. Doing so, it can be used to extend a live application or to fix minor bugs in a running program without an intermediate outage. In this article, I want to showcase dynamic attachment and apply a runtime code change with the attach and instrumentation APIs and introduce Byte Buddy, a library that offers APIs for making such code changes more convenient.

As an example, consider a running application that checks for a HTTP header named X-Priority to be present in a request to undergo special treatment by the server. The check is applied by the following utility class:

class HeaderUtility {

    static boolean isPriorityCall(HttpServletRequest request) {
        return request.getHeader("X-Pirority") != null;
    }

}

Did you spot the typo? Mistakes like these are all too common, especially if constant values are factored out to static fields that are reused in test code. In an unfortunate case, this mistake would only be discovered in a production setup where the header is generated by another application without the spelling error.

Fixing a mistake like the above might not be a problem. In the age of continuous delivery, redeploying a new version might be nothing more than the click of a button. In other cases, changes might not be so easy and a redeployment might be a complex procedure where downtime is unacceptable and living with the error is a better option. With HotSwap, there is however another option for applying small changes whilst avoiding an application restart.

Attach API: Infiltrating Another JVM with Dynamic Attachment

In order to change a live Java program, we first need a way to communicate with a running JVM. As the Java virtual machine implements a managed environment, there fortunately exists a standard API for doing so. The API in question is also known as the attachment API which is part of the official Java tooling. Using this API that the running JVM exposes, it is possible for a second Java process to communicate with it.

As a matter of fact, we all have already used this API: It is applied by any debugging and monitoring tool such as VisualVM or Java Mission Control. The APIs for applying such attachments are however not bundled with the standard Java APIs that we all know and use in our day-jobs. Instead, the API is bundled in a special file, the tools.jar which is only included in a JDK-bundled distribution of the virtual machine. To make things worse, the location of this JAR file is not set, it differs on VMs for Windows, Linux and especially Macintosh where the file is not only at a different location but also named classes.jar on some distributions. Finally, IBM has decided to even rename some of the classes that are contained in this JAR by moving all com.sun classes into the com.ibm namespace, adding another hassle. In Java 9, this mess was finally cleaned up where the tools.jar is replaced by the Jigsaw module jdk.attach.

After locating the API JAR (or module) we have to make it available to the attaching process. On the OpenJDK, the class used to connect to another VM is named VirtualMachine which offers an entry point to any VM that is run by the JDK or a regular HotSpot JVM on the same physical machine. After attaching to another virtual machine process via its process id, we are able to run a JAR file in a designated thread of the targeted VM:

// the following strings must be provided by us
String processId = processId();
String jarFileName = jarFileName();
VirtualMachine virtualMachine = VirtualMachine.attach(processId);
try {
    virtualMachine.loadAgent(jarFileName, "World!");
} finally {
    virtualMachine.detach();
}

Upon receiving a JAR file, the targeted virtual machine looks up the JAR’s manifest and locates the class under the Premain-Class attribute. This is very similar to how a VM executes a main method. With a Java agent, the VM with the denoted process id does however look for a method named agentmain which is then executed by the remote process in a dedicated thread:

public class HelloWorldAgent {

    public static void agentmain(String arg) {
        System.out.println("Hello, " + arg);
    }

}

Using this API, we are now able to print a Hello, World! message on any JVM as long as we know its process id. It is even possible to communicate with JVMs that are not part of a JDK distribution as long as the attaching VM is a JDK installation in order to access the tools.jar.

Instrumentation API: Changing the Target VM’s Program

So far, so good. But despite this successful communication with the targeted VM, we were not yet able to change code on the target VM and the bug prevails. For the latter change, Java agents can define a second parameter taking an instance of the Instrumentation type. Implementations of the latter interface offer access to several low-level functions, one of them being the ability to alter loaded code.

In order to fix the "X-Pirority" typo, let us first assume that we included a fixed class file for HeaderUtility, named typo.fix, in our agent’s JAR file next to the BugFixAgent we develop below. Additionally, we need to give our agent the ability to replace existing classes by adding Can-Redefine-Classes: true to its manifest file. With this preset, we can now redefine the class in question using the instrumentation API which accepts pairs of loaded classes and byte arrays for performing a class’s redefinition:

public class BugFixAgent {

    public static void agentmain(String arg, Instrumentation inst)
            throws Exception {
        // only if header utility is on the class path; otherwise,
        // a class can be found within any class loader by iterating
        // over the return value of Instrumentation::getAllLoadedClasses
        Class<?> headerUtility = Class.forName("HeaderUtility");

        // copy the contents of typo.fix into a byte array
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        try (InputStream input =
                BugFixAgent.class.getResourceAsStream("/typo.fix")) {
            byte[] buffer = new byte[1024];
            int length;
            while ((length = input.read(buffer)) != -1) {
                output.write(buffer, 0, length);
            }
        }

        // Apply the redefinition
        instrumentation.redefineClasses(
                new ClassDefinition(headerUtility, output.toByteArray()));
    }

}

After running the above code, the HeaderUtility class is redefined to represent its patched version. Any subsequent call to isPrivileged will now read the correct header. As a small caveat, the JVM will probably perform a full garbage collection upon applying the class redefinition and also needs to reoptimize any affected code. Together, this causes a small hick-up in the application’s performance. In most cases, however, this is still a much better option compared to fully restarting a process.

When applying code changes, it is important to assure that the new class defines the exact same fields, methods and modifiers as the class it replaces. Attempting a class redefinition that modifies any such property otherwise causes an UnsupportedOperationException. The HotSpot team is however attempting to resolve this limitation some day. Also, the OpenJDK-based dynamic code evolution virtual machine allows previewing this feature.

Tracking Memory Leaks with Byte Buddy

A simple bug-fixing agent like the above example is easy to implement when being familiar with the instrumentation API. With a bit more effort, it is however also possible to apply more generic code changes without manually creating patched class file but rather by rewriting existing class files while running the agent.

Byte Code Manipulation

Compiled Java code is represented as a list of byte code instructions. In this sense, a Java method is nothing more than a byte array where every byte either represents an instruction to the runtime or an argument to the most recent instruction. The mapping of any byte to its meaning is defined in the Java Virtual Machine Specification where the byte 0xB1 for example instructs the VM to return from a method with a void return type. Therefore, enhancing byte code is nothing more than extending a method’s byte array to include additional instructions that represent the additional logic we want to apply.

Of course, byte-by-byte code manipulations are cumbersome and error-prone. To avoid the manual process, a variety of libraries offer higher-level APIs that do not require immediate work with Java byte code. One such library is Byte Buddy (of which I am the author). One of its features is the ability to define template methods to be executed before and after a method’s original code.

Analyzing a Leaking Application

As an example, assume an application that has been running for a while was found to leak resources after a few weeks in production. Bugs like this one are hard to track down as it is often difficult to recreate a similar bad state within an observable testing environment. Therefore, instead of simply restarting the broken application to make the problem go away (for now), we can change the application code to trace the leak in the current process. To find out where the application is leaking handles, we want to track the life-cycle of any Closable object. Using a Java agent, we can change any construction or closing of any such object. With this information, we can hopefully identify what objects leak under which circumstances.

Manipulation with Templates

To implement this behavior, we first define templates that contain the code that we want to add to any Closeable object’s constructor or close method invocation. In order to learn where the leaking objects are created or closed, we want to print the stack trace on each such event. The template methods for this logic look like the following:

class ConstructionTemplate {

    @Advice.OnMethodExit
    static void exit(@Advice.This Object self) {
        new RuntimeException("Created: " + self).printStackTrace();
    }

}

class CloseTemplate {

    @Advice.OnMethodEnter
    static void enter(@Advice.This Object self) {
        new RuntimeException("Closed: " + self).printStackTrace();
    }

}

The template methods are annotated with OnMethodExit and OnMethodEnter to inform Byte Buddy when we want them to be called. Any of their parameters is annotated to indicate the value it represents within the redefined method. When applying the template, Byte Buddy then maps any access to a parameter to load the annotation’s value, like the method instance that is usually denoted by this.

In order to apply this change to any class implementing the Closable interface, Byte Buddy offers a domain specific language for creating a Java agent that matches types and methods and applies the above templates if they are applicable:

public class TracingAgent {

    public static void agentmain(String arg, Instrumentation inst) {
        new AgentBuilder.Default()
                // by default, JVM classes are not instrumented
                .ignore(none())
                .disableClassFormatChanges()
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .type(isSubTypeOf(Closable.class))
                .transform((builder, type, loader) -> builder
                        .visit(Advice
                                .to(ConstructionTemplate.class)
                                .on(isConstructor()))
                        .visit(Advice
                                .to(CloseTemplate.class)
                                .on(named("close").and(takesArguments(0)))
                        ))
                .installOn(inst);
    }

}

By default, Byte Buddy does not instrument classes defined in the java.* namespace which is changed by setting an inactive ignore matcher. Additionally, Byte Buddy needs to be instructed to never change the class file format as discussed before as we apply retransformation. Doing so, Byte Buddy automatically detects any existing or future type that implements the Closable interface and adds the above template code to its constructors and the close method.

Attaching the Agent

Byte Buddy also offers convenience methods for attaching an agent at runtime. In order to deal with the different tools.jar locations and VirtualMachine-types, Byte Buddy adds an abstraction for attachment where the correct setup is detected automatically. After packaging the above agent in a JAR file, it can be attached to the resource-leaking process by calling:

File jarFile = getAgent();
String processId = getProcessId();
ByteBuddyAgent.install(jarFile, processId);

Of course, the above agent will quickly create too much output for manual processing. To offer more convenience, it would make more sense to collect all created objects until they are released. By creating a Java agent, this is however as simple as writing a Java application that implements this behavior.

Summary

In this article, we looked into the attach API which makes it possible to inject a Java agent into any running JVM process in order. The agent is represented by a JAR file which contains a class with an agentmain method that a remote process executes in a dedicated thread. As one argument, this method receives an instance of the Instrumentation interface which allows for the redefinition of loaded classes. Code is redefined either by replacing an entire class files with a patched versions or by dynamically modifying the byte code of the existing class what can be simplified by using libraries such as Byte Buddy.

Frequently Asked Questions (FAQs) about Fixing Bugs in Running Java Code with Dynamic Attach