Freitag, 18. Februar 2022

Using Byte Buddy for proxy creation

With the increasing adoption of Java 17 and its strict encapsulation, several unmaintained libraries that rely on internal JVM APIs have stopped working. One of these libraries is cglib, the code generation library, which allows to create and load proxy classes during the runtime of a JVM process. And while there are alternatives to cglib that support Java 17, migration is not always straight-forward. To ease such migration, this article discusses how Byte Buddy can be used for proxy creation and what concept changes need to be considered during a migration.

General concept


Other than cglib, Byte Buddy does not offer an API that is dedicated to the creation of proxies. Instead, Byte Buddy offers a generic API for defining classes. While this might feel less convenient at first, it typically aids the evolution of existing code over time since the proxy class generation can be adjusted without constraints.

With Byte Buddy’s general API, a proxy is therefore created by defining a subclass of the targeted class, where all methods are overridden. Since Java methods are dispatched virtually, these overridden methods will be invoked instead of the original methods. In essence, cglib defines a proxy just like that.

As an example, consider creating a proxy of the following Sample class:

public class Sample {
  public String hello() {
    return "Hello World!";
  }
}

This Sample class can be proxied with Byte Buddy by overriding the hello method. A simple way of implementing this override is by using a MethodDelegation. A method delegation requires a delegation target, typically a class that defines a single static method. To interact with the overridden method, the method declares parameters which are annotated with the expected behavior. As an example, consider the following delegation target which mimics the parameters of cglib’s MethodInterceptor:

public class Interceptor {
  @RuntimeType
  public static Object intercept(@This Object self, 
                                 @Origin Method method, 
                                 @AllArguments Object[] args, 
                                 @SuperMethod Method superMethod) throws Throwable {
    return superMethod.invoke(self, args);
  }
}

As the annotations’ names suggest, the method accepts the intercepted This instance, a description of the Origin method, AllArguments to the methods in form of an array, and a proxy to conduct a SuperCall to the original method implementation. With the above implementation, the interception simply invokes the original code which replicates the unproxied behavior. The method itself returns a RuntimeType as the returned value is cast to the actual return type which must be a String. If any other instance was returned, a ClassCastException would occur, just as with cglib.

With this Interceptor in place, Byte Buddy can create the proxy with only a few lines of code:

Class<?> type = new ByteBuddy()
  .subclass(Sample.class)
  .method(ElementMatchers.any()).intercept(MethodDelegation.to(Interceptor.class))
  .make()
  .load(Sample.class.getClassLoader())
  .getLoaded();

The resulting class can now be instantiated using the reflection API. By default, Byte Buddy mimics all constructors that the super class is declaring. In the above case, a default constructor will be made available as Sample also declares one.

Note that Byte Buddy always requires a specification of the methods to intercept. If multiple matchers are specified, each their delegation target would be considered in the reverse order of their specification. If all methods should be intercepted, the any-matcher captures all methods. By default, Byte Buddy does however ignore the Object::finalize method. All other Object methods like hashCode, equals or toString are proxied.

Caching proxied classes


With class creation and loading being expensive operations, cglib offers a built-in cache for its proxy classes. As key for this cache, cglib considers the shape of the proxy class and recognizes if it created a class with a compatible shape previously.

While this is convenient, this cache can quickly turn into a leaky abstraction that is sensitive to minor changes. Also, the caching mechanism is performing rather poorly due to its ambitious implementation of recognizing shapes. For this reason, Byte Buddy rather offers an explicit TypeCache and requires its user to specify a mechanism for identifying a cache key. When proxying a single class, the proxied Class typically suffices as a key:

TypeCache<Class<?>> cache = new TypeCache<>();
Class<?> type = cache.findOrInsert(Sample.class.getClassLoader(), Sample.class, () -> {
  return new ByteBuddy()
    .subclass(Sample.class)
    .method(ElementMatchers.any()).intercept(MethodDelegation.to(Interceptor.class))
    .make()
    .load(Sample.class.getClassLoader())
    .getLoaded();
});

With this cache a new proxy class is only created if no proxy class was previously stored for Sample. As an optional, additional argument, a monitor object can be provided. This monitor is then locked during class creation to avoid that the same proxy is created concurrently by different threads. This can increase contention but avoids unnecessary class generation.

If more complex caching is required, a dedicated library should of course be used instead of the cache that Byte Buddy offers.

Abstract methods and default values


Until now, we assumed that all proxied methods are implemented by the proxied class. But Byte Buddy - just as cglib - also intercepts abstract methods that do not offer a super method implementation. To support intercepting such methods, the previous interceptor must be adjusted, as it currently requires a super method proxy via its parameters. By setting a property for the SuperMethod annotation, the parameter can be considered as optional.

public class Interceptor {
  @RuntimeType
  public static Object intercept(@This Object self, 
                                 @Origin Method method, 
                                 @AllArguments Object[] args, 
                                 @SuperMethod(nullIfImpossible = true) Method superMethod,
                                 @Empty Object defaultValue) throws Throwable {
    if (superMethod == null) {
      return defaultValue;
    }
    return superMethod.invoke(self, args);
  }
}
In case of intercepting an abstract method, the proxy for the super method is set to null. Additionally, Empty injects a suitable null value for the intercepted method’s return type. For methods that return a reference type, this value will be null. For a primitive return type, the correct primitive zero is injected.

Managing instance-specific interceptor state


In the previous example, the interceptor method is static. In principle, method delegation can also delegate to an instance with a non-static method, but this would likely defeat the caching mechanism if the state would be specific for each created proxy.

cglib’s cache works around this limitation, but cannot handle several corner cases where the cache might start failing after minor changes. Byte Buddy, on the other hand, relies on the user to manage the state explicitly, typically by adding a field via the defineField step, which can then be read by the interceptor:

TypeCache<Class<?>> cache = new TypeCache<>();
Class<?> type = cache.findOrInsert(Sample.class.getClassLoader(), Sample.class, () -> {
  return new ByteBuddy()
    .subclass(Sample.class)
    .defineField(InterceptorState.class, "state", Visibility.PUBLIC)
    .method(ElementMatchers.any()).intercept(MethodDelegation.to(Interceptor.class))
    .make()
    .load(Sample.class.getClassLoader())
    .getLoaded();
});

With this changed definition, any proxy instance can contain a designated instance of InterceptorState. The value can then be set via reflection or via a method handle.

Within the interceptor, this InterceptorState is accessible via an additional parameter with the FieldValue annotation which accepts the field’s name as its property. Doing so, the generated class itself remains stateless and can remain cached.

Handling non-default constructors


Byte Buddy creates valid, verifiable Java classes. As such, any class must invoke a constructor of its super class in its own constructors. For proxies, this can be inconvenient as a class without a default constructor might not be easily constructible. Some libraries like objenesis work around this limitation, but those libraries rely on JVM-internal API and their usage should be avoided.

As mentioned before, Byte Buddy replicates all visible constructors of a proxied class by default. But this behavior can be adjusted by specifying a ConstructorStrategy as a second argument to ByteBuddy::subclass. For example, it is possible to use ConstructorStrategy.ForDefaultConstructor which creates a default constructor by invoking a super constructor with default arguments for all parameters. As an example, considering the below ConstructorSample, Byte Buddy can define a default constructor for the proxy which provides null as an argument to the proxied super class:

public class ConstructorSample {

  private final String value;

  public ConstructorSample(String value) {
    this.value = value;
  }

  public String hello() {
    return "Hello " + value;
  }
}

The dynamic type builder is now created by:

new ByteBuddy().subclass(
  ConstructorSample.class, 
  new ConstructorStrategy.ForDefaultConstructor(ElementMatchers.takesArguments(String.class)));

Note that this approach would result in the proxied method returning Hello null as a result and that this might cause an exception during a constructor’s invocation if null is not considered a valid argument.

Class loading and modules


When Byte Buddy defines a class, it does not yet consider how this class will be loaded. Without any specification, Byte Buddy loads a proxy in a dedicated class loader that is a child of the class loader that is provided to the load method. While this is often convenient, creating a class loader is however an expensive operation which should be avoided, if possible. As a cheaper alternative, proxy classes should be injected into existing class loaders; normally into the one that loaded the class that is being proxied.

With Java 9, the JVM introduced an official API for class injection via MethodHandles.Lookup, and of course Byte Buddy supports this API. If Byte Buddy is however used on Java 8 or earlier, this strategy is not yet available. Typically, users fall back to using sun.misc.Unsafe, a JVM-internal API. As Java 8 does not yet encapsulate internal API and since sun.misc.Unsafe is available on most JVM implementations, this fallback does not normally render a problem.

A caveat of using MethodHandles.Lookup is its call site sensitivity. If Java modules are used, the instance must be created and provided by the module that owns the package of the proxied class. Therefore, the instance of MethodHandles.Lookup must be provided to Byte Buddy and cannot be created from within the library which represents a module of its own.

Byte Buddy configures class loading behavior by instances of ClassLoadingStrategy which can be passed as a second argument to the load method. To support most JVMs, Byte Buddy already offers a convenience method that resolves the best available injection strategy for a given JVM via:

ClassLoadingStrategy.UsingLookup.withFallback(() -> MethodHandles.lookup());

With the above strategy, a method handle lookup is used if possible and internal API is only used as a fallback. Since the method handles lookup is resolved within a lambda, it also represents the context of the module that is using Byte Buddy, assuming that this is the right module to define the proxy class. Alternatively, this Callable has to be passed from the right place. If the module system is not used, however, the above approach is normally sufficient as all classes are likely located within the unnamed module of the same class loader.

Avoiding runtime proxies with build-time instrumentation


With a rising interest for Graal and AOT compilation of Java programs in general, the creation of runtime proxies has fallen somewhat out of fashion. Of course, when running a native program without a byte code-processing JVM, classes cannot be created during runtime. Fortunately, proxies can often be created during build time instead.

For build-time code generation, Byte Buddy offers a Maven and a Gradle plugin which allow for the application of Plugin instances that manipulate and create classes before runtime. For other build tools, Byte Buddy also offers a Plugin.Engine as part of Byte Buddy which can be invoked directly. As a matter of fact, the byte-buddy artifact even contains a manifest that allows for using the jar file as an invokable of the plugin engine.

To implement a plugin for creating proxies, the proxy creator needs to implement Byte Buddy’s Plugin and Plugin.Factory interfaces. A plugin specifies what classes to instrument and how the instrumentation should be applied. For an easy example, the following plugin creates a proxy for the Sample class and adds the name of this proxy as an assumed annotation ProxyType onto the Sample class:

public class SamplePlugin implements Plugin, Plugin.Factory {
  @Override
  public boolean matches(TypeDescription type) { 
    return type.getName().equals("pkg.Simple");
  }
  @Override
  public DynamicType.Builder<?> apply(DynamicType.Builder<?> builder, 
                                            TypeDescription typeDescription, 
                                            ClassFileLocator classFileLocator) {
    DynamicType helper = new ByteBuddy()
      .subclass(typeDescription)
      .defineField(InterceptorState.class, "state", Visibility.PUBLIC)
      .method(ElementMatchers.any()).intercept(MethodDelegation.to(Interceptor.class))
      .make();
    return builder
      .require(helper)
      .annotateType(AnnotationDescription.Builder.ofType(ProxyType.class)
        .define("value", helper.getTypeDescription().getName())
        .build());
  }
  @Override
  public void close() { }
  @Override
  public Plugin make() { return this; }
}

With the annotation in place, the runtime can now check for the existence of a build-time proxy and avoid code generation altogether in such a case:

TypeCache<Class<?>> cache = new TypeCache<>();
Class<?> type = cache.findOrInsert(Sample.class.getClassLoader(), Sample.class, () -> {
  ProxyType proxy = Sample.class.getAnnotation(ProxyType.class);
  if (proxy != null) {
    return proxy.value();
  }
  return new ByteBuddy()
    .subclass(Sample.class)
    .defineField(InterceptorState.class, "state", Visibility.PUBLIC)
    .method(ElementMatchers.any()).intercept(MethodDelegation.to(Interceptor.class))
    .make()
    .load(Sample.class.getClassLoader())
    .getLoaded();
});

An advantage of this approach is that the usage of the build-time plugin remains entirely optional. This allows for faster builds that only execute tests but do not create artifacts, and allows users that do not intend to AOT-compile their code to run their applications without an explicit build setup.

Note that a future version of Byte Buddy will likely make the use of Graal even easier by discovering and preparing runtime-generated classes when the Graal configuration agent is used. For performance reasons, using an explicit build tool is however expected to remain the most performant option. Do however note that this approach is somewhat restricted to classes of the compiled project since external dependencies are not processed by a build tool.

Inline proxy code without subclasses


With the above approach, the created proxies still require the use of reflection to create instances of the proxy. For an even more ambitious setup, Byte Buddy offers the Advice mechanism to change the code of classes directly. Advice is normally often used for the decoration of methods and a popular choice when developing Java agents. But it can also be used to emulate proxy behavior without creating a subclass.

As an example, the following advice class records the execution time of a method by declaring actions that are to be performed prior to invoking a method as well as after it. Advice offers similar annotations to MethodDelegation, be careful to not confuse those annotations as they are declared by different packages.

To emulate the previous behavior of the Interceptor, the following Decorator functions similarly to it. Note that the Decorator declares a set of proxies to recognize what instances are to be treated as proxies and which instances should function as if they were not proxied. Within the OnMethodEnter annotation, it is specified that the original code is skipped if a non-null value is returned.

public class Decorator {
  static final Set<Object> PROXIES = new HashSet<>();
  @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
  public static Object enter(
    @Advice.This Object self,
    @Advice.Origin Method method,
    @Advice.AllArguments Object[] arguments) throws Throwable {
   if (PROXIES.contains(self)) {
     return ProxyHandler.handle(self, method, arguments);
    } else {
      return null;
    }
  }
  @Advice.OnMethodExit
  public static void exit(
      @Advice.Enter Object enter,
      @Advice.Exit(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object returned) {
    if (enter != null) {
      returned = enter;
    }
  }
}

With this code, the original method can be invoked by temporarily removing the instance from the proxy set within the ProxyHandler.

Object returned;
Decorator.PROXIES.remove(self);
try {
  returned = method.invoke(self, arguments);
} finally {
  Decorator.PROXIES.add(self);
}

Note that this is a naive approach which will fail if the proxy is used concurrently. If a proxy needs to be thread-safe, it is normally required to define a thread-local set that contains temporarily disabled proxies.

Of course, it is not normally possible to apply this decoration during a JVMs runtime, but only at build-time, unless a Java agent is used. To still allow for a fallback-implementation, Byte Buddy does however allow for Advice being used as both decorator:

new ByteBuddy().redefine(Sample.class)
  .visit(Advice.to(Decorator.class).on(ElementMatchers.isMethod()))
  .make();

and as an interceptor for creating a subclass proxy:

new ByteBuddy().subclass(Sample.class)
  .method(ElementMatchers.isMethod())
  .intercept(Advice.to(Decorator.class))
  .make();

In this case, a build-time plugin can avoid a subclass creation where this is necessary. For example, it allows for proxying final classes or methods, if this should be supported. At the same time, inline proxies cannot proxy native methods.

Replacing other cglib utilities


cglib contains a row of other class generation utilities besides the Enhancer. I have previously written a summary of all of the library's capabilities where those are described.

The good news is that most of this functionality has become obsolete. Immutable beans are less useful today as it has become much more common to model immutable objects by for example records. And similarly other bean utilities have found better equivalents in today’s Java, especially since method and var handles have entered the stage. Especially cglib's FastMethod and FastClass utilities are no longer useful as reflection and method handles have passed the performance that is offered by these code generation tools.