javajvmbytecodeprofiler

How to combine ClassTransforms in the new Java 22 ClassFile Preview API


This is a very specific question regarding a very new feature. For context: I am writing my Bachelor Thesis about Runtime Code Generation in Java, and I'm working on a prototype for a profiler where I showcase my work. There, I use the Java ClassFile API to manipulate classes and inject bytecode, for example code to measure method calls or runtime of methods. Since everything has to happen at program runtime, i also need to update the classes at runtime, and that's done with the Instrumentation API and a corresponding agent.

It works perfectly when only one transformation in a class needs to be done, but as soon as multiple come up, I can't get them combined, only one is applied and the others are ignored.

Here is the code that generates the CodeTransform for better understanding:

private static CodeTransform generateMethodCodeTransform(int statisticsId, List<String> metricValues) {

    return new CodeTransform() {

        private int startSlot = 0;
        private int endSlot = 0;
        private int durationSlot = 0;

        @Override
        public void accept(CodeBuilder builder, CodeElement element) {

            if(metricValues.contains(ProfilerMethodMetric.RUNTIME.toString())) {

                if(element instanceof ReturnInstruction) {

                    builder.invokestatic(ClassDesc.of(System.class.getName()), "nanoTime", MethodTypeDesc.of(ConstantDescs.CD_long));
                    endSlot = builder.allocateLocal(TypeKind.LongType);
                    builder.lstore(endSlot);

                    builder.lload(endSlot);
                    builder.lload(startSlot);
                    builder.lsub();
                    durationSlot = builder.allocateLocal(TypeKind.LongType);
                    builder.lstore(durationSlot);

                    builder.invokestatic(ClassDesc.of(Profiler.class.getName()), "getProfilerStatistics", MethodTypeDesc.of(ClassDesc.of(ProfilerStatistics.class.getName())));
                    builder.ldc(statisticsId);
                    builder.invokevirtual(ClassDesc.of(ProfilerStatistics.class.getName()), "getMethodStatisticsById", MethodTypeDesc.of(ClassDesc.of(MethodStatistics.class.getName()), ConstantDescs.CD_int));
                    builder.lload(durationSlot);
                    builder.invokevirtual(ClassDesc.of(MethodStatistics.class.getName()), "trackRuntime", MethodTypeDesc.of(ConstantDescs.CD_void, ConstantDescs.CD_long));

                }

            }

            builder.with(element);

        }

        @Override
        public void atEnd(CodeBuilder builder) {

        }

        @Override
        public void atStart(CodeBuilder builder) {

            System.out.println("AT START METHOD CODE TRANSFORM");

            if(metricValues.contains(ProfilerMethodMetric.DEBUG.toString())) {
                builder.getstatic(ClassDesc.of(System.class.getName()), "out", ClassDesc.of(PrintStream.class.getName()));
                builder.ldc("DEBUG: This method is successfully injected");
                builder.invokevirtual(ClassDesc.of(PrintStream.class.getName()), "println", MethodTypeDesc.of(ConstantDescs.CD_void, ClassDesc.of(String.class.getName())));
            }

            if(metricValues.contains(ProfilerMethodMetric.CALLS.toString())) {
                builder.invokestatic(ClassDesc.of(Profiler.class.getName()), "getProfilerStatistics", MethodTypeDesc.of(ClassDesc.of(ProfilerStatistics.class.getName())));
                builder.ldc(statisticsId);
                builder.invokevirtual(ClassDesc.of(ProfilerStatistics.class.getName()), "getMethodStatisticsById", MethodTypeDesc.of(ClassDesc.of(MethodStatistics.class.getName()), ConstantDescs.CD_int));
                builder.invokevirtual(ClassDesc.of(MethodStatistics.class.getName()), "trackCall", MethodTypeDesc.of(ConstantDescs.CD_void));
            }

            if(metricValues.contains(ProfilerMethodMetric.RUNTIME.toString())) {
                builder.invokestatic(ClassDesc.of(System.class.getName()), "nanoTime", MethodTypeDesc.of(ConstantDescs.CD_long));
                startSlot = builder.allocateLocal(TypeKind.LongType);
                builder.lstore(startSlot);
            }

        }

    };

}

I then have a List of ClassTransforms where new ones are added like this:

transforms.add(ClassTransform.transformingMethodBodies(mm -> mm.equals(methodModel), codeTransform));

Here are two ways I tried getting them combined:

classfileBuffer is a byte[] containing the raw class which is then handed to the Instrumentation API to apply the changes.

ClassTransform combined = null;
for(ClassTransform t : transforms) {
    if(combined == null) {
        combined = t;
    } else {
        combined = combined.andThen(t);
    }
}

if(combined != null) {
    classfileBuffer = ClassFile.of().transform(classModel, combined);
}

This first way shows me trying to use andThen() which should, according to the docs, chain two class transforms together. But with this, nothing get's injected at the end, discarding all transforms for some reason.

for(ClassTransform t : transforms) {
    classfileBuffer = ClassFile.of().transform(classModel, t);
    classModel = ClassFile.of().parse(classfileBuffer);
}

This is my second approach. I just transform the classModel with the first transform, then parse a new classModel from the buffer, and repeat that over and over again. This made the most sense to me, but this is where the hard-to-explain and understand part comes in.

As you can see in my code on top, I put this debug print in my atStart() which is called at the start of each code transform for a method. Although in the end i have 2 ClassTransform s in my transforms list, only the first one is actually applied. No matter how many transforms I have, the print statement is only executed once as you can see here: https://i.imgur.com/1QsfmfB.png

But I don't understand why it's never called again. I also tried hashing the contents of the buffer, this loop only changes them once, and each further transformation is completely ignored, because the transform is never called.

EDIT

Here is a minimal reproductible example that still shows the same behavior:

package at.alexkiefer.rcgj.transformers;

import at.alexkiefer.rcgj.Test;

import java.io.PrintStream;
import java.lang.classfile.*;
import java.lang.constant.ClassDesc;
import java.lang.constant.ConstantDescs;
import java.lang.constant.MethodTypeDesc;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class MRE implements BytecodeTransformer {

    private static CodeTransform generateMethodCodeTransform() {

        return new CodeTransform() {

            @Override
            public void accept(CodeBuilder builder, CodeElement element) {

                builder.with(element);

            }

            @Override
            public void atEnd(CodeBuilder builder) {

            }

            @Override
            public void atStart(CodeBuilder builder) {

                System.out.println("\t\tAt start of methode code transform");

                builder.getstatic(ClassDesc.of(System.class.getName()), "out", ClassDesc.of(PrintStream.class.getName()));
                builder.ldc("DEBUG: This method is successfully injected");
                builder.invokevirtual(ClassDesc.of(PrintStream.class.getName()), "println", MethodTypeDesc.of(ConstantDescs.CD_void, ClassDesc.of(String.class.getName())));

            }

        };

    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {


        ClassModel classModel = ClassFile.of().parse(classfileBuffer);

        List<ClassTransform> transforms = new ArrayList<>();

        for(ClassElement e : classModel.elements()) {

            if(e instanceof MethodModel methodModel) {

                if(classBeingRedefined.equals(Test.class)) {
                    System.out.println("Generating class transform for method " + methodModel.methodName());
                    CodeTransform codeTransform = generateMethodCodeTransform();
                    transforms.add(ClassTransform.transformingMethodBodies(mm -> mm.equals(methodModel), codeTransform));
                }

            }

        }

        int counter = 1;
        for(ClassTransform t : transforms) {
            System.out.println("Transform " + counter++ + " of " + transforms.size() + " is happening now");
            System.out.println("\tClassfile buffer hash before transformation: " + Arrays.hashCode(classfileBuffer));
            classfileBuffer = ClassFile.of().transform(classModel, t);
            classModel = ClassFile.of().parse(classfileBuffer);
            System.out.println("\tClassfile buffer hash after transformation: " + Arrays.hashCode(classfileBuffer));
        }

        return classfileBuffer;

    }

}
public class Main {

    public static void main(String[] args) {

        Test test = new Test();
        test.doSomething();

        Profiler.inject();

        test.doSomething();
        test.doSomethingElse();

        Profiler.eject();

        //System.out.println(Profiler.getProfilerStatistics());

    }

}

This is the output:

Generating class transform for method doSomething
Generating class transform for method doSomethingElse
Generating class transform for method <init>
Transform 1 of 3 is happening now
    Classfile buffer hash before transformation: 1402060343
        At start of methode code transform
    Classfile buffer hash after transformation: -1078378321
Transform 2 of 3 is happening now
    Classfile buffer hash before transformation: -1078378321
    Classfile buffer hash after transformation: -1078378321
Transform 3 of 3 is happening now
    Classfile buffer hash before transformation: -1078378321
    Classfile buffer hash after transformation: -1078378321
DEBUG: This method is successfully injected

The test methods doSomething and doSomethingElse actually don't do anything but a loop that calculates some stuff that doesn't produce output. The debug print is injected successfully for the first method, in this case doSomething (can be verified sicne the function is called once before the profiler is injected and doesn't output anything). As you can see, the list contains three class transforms, and the first one does something, the other two don't. The classfile is not being transformed, this can be seen by looking at the hash of the bytes. Also the print in the code transform is only called once.

Also: The interface BytecodeTransformer which is implemented by this class is just a wrapper around the transform method of ClassFileTransformer from the instrumentation API, which makes no sense in this context but i have a need somewhere else. But this shouldn't be the issue here.


Solution

  • The problem is the predicate mm -> mm.equals(methodModel).

    There is no defined equality for MethodModel, so it uses the object identity.

    After the first transformer has been applied, you create a new ClassModel which will also have distinct MethodModel instances which do not match the previously captured objects. So the predicate evaluates to false for all subsequent transform operations.

    You need a predicate based on methodName() and methodType():

    if(e instanceof MethodModel methodModel) {
        if(classBeingRedefined.equals(Test.class)) {
            Utf8Entry name = methodModel.methodName(), type = methodModel.methodType();
            System.out.println("Generating class transform for method " + name);
            CodeTransform codeTransform = generateMethodCodeTransform();
            transforms.add(ClassTransform.transformingMethodBodies(
                mm -> mm.methodName().equals(name) && mm.methodType().equals(type),
                codeTransform));
        }
    }
    

    As a side note, since classBeingRedefined does not change throughout the entire method, it’s more efficient to test this first, at the beginning of the method, to take no action if it doesn’t match (which will happen more often when a ClassFileTransformer is invoked for all types).

    Further, in those cases where you don’t change anything, it’s slightly more efficient to return null, as it is faster for the caller to check for null than to check whether you altered the byte array (even if you return the original array, you might have changed something, so the caller has to check).