Advanced instrumentation with ByteBuddy Agent

Reading Time: 7 minutes

The Aspect-Oriented Programming paradigm allows to achieve a great level of Separation of Concerns (SoC) in the application’s codebase. All the necessary components that do not relate to the business logic directly, like security, metrics, logging, etc., can be nicely isolated into auxiliary modules that will be wired up with the business logic indirectly, usually during the compilation or in runtime.

In the Java world, there are not so many frameworks for AOP. The notable ones are AspectJ and Spring AOP. Both balance the complexity and functionality. The process of attaching the extra modules to the business logic is called weaving. In short, AspectJ supports compile-time weaving that modifies output class files (variation – post-compile weaving is used for 3rd party classes), and load-time weaving that intercepts the classloader and modifies the classes in the process of defining them to the JVM. Alternatively, Spring AOP has a much more simple model, it allows for runtime weaving only that works just for Spring Beans defined in the Application Context. The weaving is done via proxies generated by Spring.

In addition to these two frameworks, there is one more great library that exists for pretty similar objectives, it is called Byte Buddy, developed by Rafael Winterhalter. This library is pretty well known for its highly simplified and convenient facilities that allow runtime code generation and modification. Byte Buddy makes it easy to create a new type that can redefine existing types, intercept methods, and much more to that. But in addition, it can also change the behavior of existing classes. This is done via Byte Buddy Agent that can operate with Java Instrumentation API to do necessary changes to type definitions. This may sound similar to AspectJ load-time weaving, which in fact is, but practically there is a big difference between the two. The main benefit of using Byte Buddy is simplicity. The API is very intuitive, the extra tooling is not required, and overall integration of it does not force any changes in the build process, run configuration or so.

Now it is time to stop the theory and apply the Byte Buddy in practice. A simple, yet important example could be the detection of certain methods invocations. In fact, this idea has been inspired by the BlockHound project (detector of blocking calls from non-blocking reactive threads), developed by Sergei Egorov.

Disclaimer

TL; DR;

A complete Demo project is available on GitHub:
https://github.com/alexey-anufriev/byte-buddy-agent-interceptor-demo

Byte Buddy Agent

The basic setup of Byte Buddy Agent looks very simple. It needs to be installed first, and then extended with necessary transformations. Since the agent is intended to inline some code into the target classes it is important to remember a couple of principles of this process:

1. The visitor approach must be applied to attach an interceptor to an invocation.
2. The interceptor must be pretty simple, without any class members, because the body of the intercepting method will be inlined as-is.

With these basic steps, minimal instrumentation can already work. But the proposed example requires a couple more challenges to be solved.

Challenge #1. Complex interceptor’s logic.

It may happen that the interceptor requires some additional classes to have a complete logic of interception. In the Demo, the interceptor requires two additional classes.

Now it is important to remember that Byte Buddy works on top of Boostrap Classloader, which sits at the top of the classloaders hierarchy. This means that necessary classes have to be provided to the Boostrap Classloader before they are loaded by the Application Classloader, otherwise, each classloader will load its own copy of the class, and classes compatibility will be corrupted. For example, static fields will not be working as the values will be located in different memory areas, thus interceptor will not be able to access values set by the application and vice versa.

But it is quite an important component of the Demo, since there is a global switch for the interceptor that can be controlled from the application.

To load necessary classes explicitly by the Bootstrap Classloader they must be explicitly appended to it. To do this the classes must be packed into a temporary JAR and brought to the classloader. An important note here is that classes should not be referenced by the type but rather by name before they are appended to avoid eager loading.

Challenge #2. Native methods.

At this moment interceptor works with almost all methods. Still, the problematic ones are native methods. So if there will be a need to intercept Thread.sleep then it must be done a bit differently.

Since native methods have no body, they cannot be instrumented as they don’t have a place where the code can be inlined. But instrumentation API has a possibility to overcome this limitation. In particular, there is a way to supply a prefix for native methods. This allows renaming of original methods and injecting delegate methods with original names.

For example by supplying a prefix $PREFIX$_, the original method may become:

private static native void $PREFIX$_sleep(long millis); // private to hide it

And the delegate method may look like this:

public static void sleep(long millis) {
  $PREFIX$_sleep(long millis);
}

Now there is a space for inlining the code.

According to the documentation for setNativeMethodPrefix method this renaming is safe and JVM will use the prefix for proper resolution of the original native methods.

The only tricky part here is the process of renaming and creating delegate methods. It requires an understanding of JVM bytecode and knowledge of the ASM library. In the Demo, it is pretty well documented in the following transformer: NativePrefixerTransformer.

One last important notice here is that renaming must happen before the actual instrumentation via Byte Buddy, otherwise, the methods will not be found.

Now interceptor is ready to be attached and can handle configured invocations.

P.S.

Many thanks to Rafael Winterhalter and Sergei Egorov who stayed patient no matter how annoying I was, and helped me to understand Byte Buddy much better and solve the aforementioned challenges.