Embracing invokedynamic to tame class loaders in Java agents

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);
}Read more
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 <code>net.bytebuddy.dynamic.loading.MultipleParentClassLoader</code>.

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:

We're hiring

Work for a global, distributed team where finding someone like you is just a Zoom meeting away. Flexible work with impact? Development opportunities from the start?