Embracing invokedynamic to tame class loaders in Java agents

By

Felix Barnsteiner,

Rafael Winterhalter

blog-thumb-ml-anomaly-laptop.png

One of the nicest things about Byte Buddy is that it allows you to write a Java agent without manually having to deal with byte code. To instrument a method, agent authors can simply write the code they want to inject in pure Java. This makes writing Java agents much more accessible and avoids complicated on-boarding requirements.

After the first successful experiments, agent authors often get hit by a wall of complexity that the JVM throws at them: class loaders (OSGi, oh my!), class visibility, dependence on internal APIs, class path scanners, and version conflicts to name a few.

In this article, we'll look at a relatively novel approach to break through this wall of complexity. The architecture, which is based on the invokedynamic bytecode instruction, a bytecode that is best known for leveraging Java’s lambda expressions, allows for a simple mental model when writing instrumentations. As a bonus, this also enables updating to a newer version of the agent at runtime, without having to restart the instrumented application. The Elastic APM Java agent started the migration to this invokedynamic-based architecture over a year ago and recently completed the migration.

Issues with traditional advice dispatching approaches

Let's consider a simple example of an agent that wants to measure the response time of Java servlets. In so-called advice methods, one can define code that should run before or after the actual method. It's also possible to get access to the arguments of the instrumented method.

@Advice.OnMethodEnter
public static long enter() {
    return System.nanoTime();
}

@Advice.OnMethodExit
public static void exit(
        @Advice.Argument(0) HttpServletRequest request,
        @Advice.Enter long startTime) {
    System.out.printf(
            "Request to %s took %d ns%n",
            request.getRequestURI(),
            System.nanoTime() - startTime);
}
In Byte Buddy, there are two main ways how advice gets applied to an instrumented method.


Inlined advice

By default, the enter and exit advice is copied into the target methods, as if the original author of the class had added the agent’s code into the method. The instrumented method, if written in plain Java, would then look something like this:

protected void service(HttpServletRequest req, HttpServletResponse resp) {
    long startTime = System.nanoTime();
    // original method body
    System.out.printf(
            "Request to %s took %d ns%n",
            request.getRequestURI(),
            System.nanoTime() - startTime);
}

The advantage is that the advice has access to any value or type that is normally reachable from the instrumented method. In the above example, this allows accessing javax.servlet.http.HttpServletRequest, even though the agent does not itself ship with that interface. As the agent’s code is run within the targeted method, it simply picks up the type definition that is already available to the method itself.

On the downside, the advice code is no longer executed in the context that it is defined within. As a result, you can, for example, not set a breakpoint in an advice method, because it is never actually called. Remember: the methods are merely used as a template.

But the real issue is that factoring code out of the advice methods or calling any methods that are normally reachable from advice is no longer possible. Since all code is now executed from the instrumented method, the agent might run on an entirely different class loader with no connection to the instrumented method, so even public methods might not be invokable from the instrumented code. We'll see more of this issue in the next section.

Delegated advice

For a similar but still very different approach, it is possible to instruct Byte Buddy to delegate to the advice methods instead. This can be controlled via the advice annotation attribute @Advice.OnMethodEnter(inline = false). By default, Byte Buddy will delegate to an advice method via a static method call. The instrumented method would then look like this:
protected void service(HttpServletRequest req, HttpServletResponse resp) {
    long startTime = AdviceClass.enter();
    // original method body
    AdviceClass.exit(req, startTime);
}

Similarly to before, it is up to the agent’s developer to ensure that the advice code is visible to the instrumented method. If the instrumented method does not share a class loader hierarchy with the agent’s code, this instrumentation would yield a NoClassDefFoundError upon reaching the above method. And even if the delegated advice is reachable from the agent, argument types such as HttpServletRequest might not be available to the agent’s class loader. This would then only move the error to the agent’s code upon its advice invocation.

Class loader issues

By default, agents get added to the system class loader when they are attached to the JVM and the java.lang.instrument.Instrumentation interface offers ways to add the agent to the bootstrap class loader. In theory, adding classes to the bootstrap class loader makes them visible everywhere. However, some class loaders (such as OSGi) only allow certain classes (such as java.*, com.sun.*) to be loaded from the system or bootstrap class loader. A common solution is to instrument all class loaders and explicitly redirect class loading of classes in certain packages directly to the bootstrap loader.

But adding classes to the system class loader and the bootstrap class loader also comes with downsides. The additional classes can slow down class path scanners or even cause failures that prevent the application from starting. See elastic/apm-agent-java#364 for an example. Also, it's not possible to unload classes of such a persistent class loader, which is an issue when designing an agent that wants to offer the possibility of its own removal during runtime.

Conceptually, there are only two approaches to overcoming these class loader issues where an advice class wants to invoke different methods that are normally shipped with the agents but where these methods might not be reachable. Either this code must be injected into the instrumented class' class loader such that they can be looked up directly from there. Or, a new class loader must be defined as a child of this former class loader where any additional types can now be located by implementing such a custom class loader.

For the first approach, Byte Buddy comes with utilities that allow classes to be injected into any class loader (net.bytebuddy.dynamic.loading.ClassInjector). While this seems like a straightforward fix, it comes with major drawbacks. The more flexible injectors are built on top of internal APIs such as sun.misc.Unsafe / jdk.internal.misc.Unsafe. And also safer-sounding class injector strategies like UsingReflection use clever workarounds to circumvent the safeguards that have been introduced in recent Java versions that usually disallow accessing private fields using Unsafe::putBoolean. As of today, it's a cat-and-mouse game between Oracle who restricts access to internal APIs and enforces visibility in the reflection API, and the discovery of new loopholes that can circumvent these. At the same time, the official gateway of using a method handle lookup is barely compatible with agents and its integration is an open issue (https://bugs.openjdk.java.net/browse/JDK-8200559). Therefore, it seems rather risky to build a whole agent architecture using the currently unsafe APIs that Oracle is aiming to lock down further.

With the second approach, all advice and helper classes are loaded in a child class loader. This works without relying on unsafe API because the class loader is implemented by the agent developer and because a class loader can get access to all types that are defined by its parent class loader.

Another advantage of loading helper classes in a dedicated class loader as opposed to injecting them into the instrumented class's class loader is that it is possible to unload these classes. This allows to fully detach the agent from the application and to attach a new version of the agent without leaving any trace of the previous version, also known as live-updating the agent. Byte Buddy already allows reverting all the instrumentations it has applied via re-transformation. When no other references to the agent helper class loaders are leaked, this makes all its objects, classes, and even the entire class loader eligible for garbage collection.

One complication with this approach is that the advice class is not visible to the instrumented class. The instrumented method HttpServlet::service from the previous example calls AdviceClass via a static method call. This would lead to a NoClassDefFoundError at runtime, as AdviceClass is not visible in the context of the HttpServlet::service method. That's because AdviceClass is loaded by a child class loader of the instrumented class (HttpServlet). While AdviceClass can access classes that are visible to the instrumented class, such as the HttpServletRequest argument, the reverse is not true.

Introducing an invokedynamic-based advice dispatching approach

There's another, lesser-known alternative to dispatching advice via a static method call. With net.bytebuddy.asm.Advice.WithCustomMapping::bootstrap, you can instruct Byte Buddy to insert an invokedynamic bytecode instruction into the instrumented methods. This instruction was added in Java 7 with the goal of better support for dynamic languages in the JVM, such as Groovy and JRuby.

In a nutshell, an invokedynamic invocation consists of two phases: looking up a CallSite and then invoking the MethodHandle the CallSite holds. If the same invokedynamic instruction is executed another time, CallSite from the initial lookup will be invoked.

The following example shows how an invokedynamic instruction looks like in the bytecode of a method.

// InvokeDynamic #1:exit:(Ljavax/servlet/ServletRequest;long)V</p> <p>invokedynamic #1076, 0


The lookup of the CallSite happens within a so-called bootstrap method. This method receives a couple of arguments for the lookup, such as the advice class name, method name, and the advice's MethodType that represents the arguments and return type. The following example shows how the bootstrap method is declared within the bytecode of a class.
BootstrapMethods:
  1: #1060 REF_invokeStatic java/lang/IndyBootstrapDispatcher.bootstrap:(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;[Ljava/lang/Object;)Ljava/lang/invoke/CallSite
    Method arguments:
      #1049 org.example.ServletAdvice
      #1050 1
      #12 javax/servlet/http/HttpServlet
      #1072 service
      #1075 REF_invokeVirtual javax/servlet/http/HttpServlet.service:(Ljavax/servlet/HttpServletRequest;Ljavax/servlet/HttpServletResponse;)V

The class that contains the bootstrap method (in this case java/lang/IndyBootstrapDispatcher.bootstrap); must be visible from any instrumented class. Therefore, this class needs to be added to the bootstrap class loader. To ensure compatibility with filtering class loaders, such as OSGi loaders, the class is placed into the java.lang package.

While this approach doesn't completely avoid class injection, injecting only a single class does result in a reduced surface of eternal classes that are added by the agent and reduces the need to refactor an existing agent if future releases of the JDK no longer allow for such injection.

In the Elastic APM Java agent, the bootstrap method will create a new class loader whose parent is the class loader of the instrumented class and load the advice and any amount of helpers from it. We can then load the advice class from this newly created class loader given the advice class name that is provided as an argument to the bootstrap method (Method arguments: org.example.ServletAdvice).

Using the other arguments of the bootstrap method, we can construct a MethodHandle and a CallSite that represents the advice method within the child class loader we created. For our needs, the target method is always the same. Thus, a ConstantCallSite can be returned that allows the JIT to inline the advice method.

Now that we only rely on exactly one class to be visible from the instrumented methods (java.lang.IndyBootstrapDispatcher), we can further isolate the agent by loading its classes that aren't specific to a particular library it instruments from a dedicated class loader. As described in the previous section, hiding the agent's classes from the regular class loader hierarchy avoids compatibility issues, for example with class path scanners. It also allows the agent to ship any dependencies, such as Byte Buddy or a logging library, without having to shade (aka relocate) the dependencies to the agent's namespace. This makes debugging the agent that much easier. Due to the isolated class loader, there is no concern about conflicting classes that may be present in the application's class loader hierarchy. More details on the implementation of one such isolated class loader can be found in the Elastic APM Java agent's source code for ShadedClassLoader.

The resulting class loader hierarchy looks like this:

Class loader hierarchy

Note that the agent helper class loader, which loads the advice and library-specific helper classes, has two parents: The class loader of the instrumented class (such as the class loader that servlet containers create for each web application) and the agent class loader. That allows advice and helper classes to access both types that are visible from the instrumented class' class loader and the agent class loader. While having multiple parents is not offered by the built-in class loaders, it is relatively straightforward to implement it oneself. Byte Buddy also comes with an implementation called net.bytebuddy.dynamic.loading.MultipleParentClassLoader. 

In summary, this section described how the invokedynamic instruction can be used to call an advice method that is loaded from a child class loader of the instrumented class' defining class loader. As a result, this allows the agent to hide its classes from the application while providing a way to invoke the isolated methods from the application classes it instruments. This is useful because the advice and all other classes loaded by this class loader can access the instrumented libraries' classes while the advice code is still executed as regular code. It also avoids injecting the advice and helper classes into the target class loader directly, which is only possible today by using internal APIs that Oracle is aiming to increasingly lock down.

AssignReturned

While advice that uses either inlining or delegation is implemented by the same API, and seems rather similar as a result, there are differences. Delegation advice cannot easily write values in the scope of the instrumented method. When advice is inlined, the advice method can simply assign values to annotated parameters which Byte Buddy then translates to a replacement of the represented value during the inlining process. As an example, the following inlined advice would replace the first argument of an instrumented method - here a Runnable - with a wrapper instance that also implements the Runnable interface, which reports any future invocations back to the agent:

@Advice.OnMethodEnter
public static void enter(
        @Advice.Argument(value = 0, readOnly = false) Runnable callback) {
    callback = new TracingRunnable(callback);
}

As the above code is inlined, the advice simply replaces the value that is assigned to the first argument of the instrumented method. As a result, the instrumented method is now executed as if its caller had already passed the .

When using delegation, this does not work, unfortunately. With delegation, the new value would only be assigned to the parameter of the advice method, without affecting the instrumented method’s assignment which would still carry the original runnable after the advice method was executed.

To offer such assignments when using delegating advice, Byte Buddy recently introduced the Advice.AssignReturned post-processor. Advice post processors are handlers that are invoked after an advice method was dispatched, to allow for additional operations that are independent of the advice that was applied. But most importantly, post processors always generate code that is inlined into the instrumented method, even if the advice itself is invoked via delegation. This allows for writing values in the scope of the instrumented method if these values were returned from the advice method. With post processors being an extension to the regular Advice implementation, they need to first be registered manually by calling:

Advice.withCustomBinding()
    .with(new Advice.AssignReturned.Factory());
As the name suggests, this post-processor allows an assignment of the value that is returned from an advice method to parameters of the instrumented method. To implement the above example, one would, for example, instruct the post-processor to assign the returned value to the instrumented method’s first argument as it was done before:
@Advice.OnMethodEnter(inline = false)
@Advice.AssignReturned.ToArguments(@ToArgument(0))
public static Runnable enter(@Advice.Argument(0) Runnable callback) {
    return new TracingRunnable(callback);
}

Just as in the inlined example, the instrumented method would now observe the TracingRunnable as its first argument as it was replaced by the post-processor. And besides assigning argument values, it is also possible to assign values to fields, to the method’s return value, its thrown exception and even to its this reference if the method is non-static.

In some cases, it might however be required to assign more than one value. With inlined advice, this is straightforward to implement by assigning multiple values within the advice method directly to each annotated parameter. With delegating advice, multiple assignments are however similarly easy to implement by returning an array as a return type and by specifying what index of the returned array contains what value.

To extend the hypothetical example, assuming that the instrumented method also requires an executor service as a second argument, we could enforce the usage of a freshly created cached thread pool by providing it as a second argument to an advice method’s returned array. When annotating the advice method’s assignments, every assignment now only needs to indicate what array index represents which of the assigned values.

@Advice.OnMethodEnter(inline = false)
@Advice.AssignReturned.ToArguments(
  @ToArgument(value = 0, index = 0, typing = DYNAMIC),
  @ToArgument(value = 1, index = 1, typing = DYNAMIC))
public static Runnable enter(@Advice.Argument(0) Runnable callback) {
    return new Object[] {
        new TracingRunnable(callback),
        Executors.newCachedThreadPool()
    };
}

Finally, as Object-typed arrays might contain non-assignable values, the annotations must specify that dynamic typing is to be used. Doing so, Byte Buddy attempts a type-casting of values prior to assigning. To avoid the potential ClassCastExceptions from affecting the instrumented application, the post processor can be configured to suppress these exceptions.

new Advice.AssignReturned().Factory()
    .withSuppressed(ClassCastException.class)

Failing to configure dynamic typing in cases when the array contains non-assignable values would lead to an exception during the instrumenting of a class. Aside from loss of instrumentation, the application will not be affected.

Trade-offs

One of the limitations of this architecture is that it's not possible to support Java 6 applications as it relies on the invokedynamic bytecode instruction that has been added in Java 7. As the Elastic APM Java agent never supported Java 6, this was not an issue in that case. Many other agents don't even support Java 7 anymore, whose market share is just around 1-5%, depending on what study is considered.

In addition to the requirement of Java 7+, the instrumented class has to be at bytecode level 51, meaning that it has to be compiled with a target of Java 7 or later. That's because it's not possible to use invokedynamic instructions for older class file versions. Some libraries, in particular older JDBC drivers, which an agent may want to instrument, are sometimes compiled with quite old class file versions. There's a relatively simple workaround, though. Using a ClassVisitor, we can let ASM re-write the bytecode to class file version 51 (Java 7). This has proven to be a stable and reliable way since this has been introduced in the Elastic APM Java agent. It does come with a bit of a performance penalty but we only need to do that for the relatively rare occasion where the class file version of the instrumented class is lower than 51.

Another thing to keep in mind is that early versions of Java 7 (before update 60, which released in May 2014) and Java 8 (before update 40, which released in March 2015) have bugs in their invokedynamic and MethodHandle support. For that reason, the Elastic APM Java agent disables itself if it's detected to run on these JVM versions.

Next Steps

Have a look at the docs to find out more about the Elastic APM Java Agent and how it can help you to identify and fix performance issues in your application. If you want to build your own Java agent, visit bytebuddy.net to learn more.

Sign up for Elastic Cloud free trial

Spin up a fully loaded deployment on the cloud provider you choose. As the company behind Elasticsearch, we bring our features and support to your Elastic clusters in the cloud.

Start free trial