Support Java agent instrumentation in native image

[ziyilin]

This PR supports Java agent instrumentation in native image.

Scope

Java agent can instrument and transform runtime Java classes in two ways:

1. premain: Java agent is loaded before main class starts. The classes are transformed at their first loading time.
2. main: Java agent is attached to a running JVM. The classes are transformed after the agent attached time.

This PR only supports the first kind of Java agent instrumentation, i.e. the premain way.

General Idea

The general idea of this PR is to compile the recorded transformed class into native image.
It is implemented in three stages:

1. Interception stage: native-image-agent records all the transformed classes, dynamic generated classes, and the target agent's premain method in a run. The result is writing to instrument-config.json file and the transformed classes are dumped to instrument-classes directory. The classes are organized in the form of module/package/class.
2. Build time stage: Instrumented classes are prepended to the beginning of imageCp, and patched to the corresponding modules as well, so they will override their original classes at build time.
3. Runtime stage: Add premain execution in JavaMainWrapper before executing the main method. The recorded agent premain method shall be executed at this moment.

Why don't instrument the classes at build time by starting native-image along with agent?
GraalVM is a Java application. Instrumenting at build time will change the behaviors of GraalVM itself, especially when the JDK classes are enhanced by the agent.

Agent Developers' Work

Due to the complexity of agent instrumentation, this PR cannot automatically do all the work. It relies on the agent developer to provide the following input:

1. JDK class transformation. The enhanced JDK classes can't be simply applied to replace the originals as application classes, because we don't want GraalVM work with unexpected JDK behaviors. The agent developers should use the APIs provided by this PR (see next section) to rewrite the enhancement for JDK classes, so that they will only take effect in native image.
2. Agent premain method implementation. This PR supports running premain method in native image to initiate the instrumentation. But the agent developer should make sure the code running in native image works well. For example, a premain method can do two things, 1) turn a global flag on, which activates all the instrumented code; and 2) do bytecode transformation. In native image only the first task should be executed, the second task is not necessary and cannot be executed.

I have prepared an adaption project for OpenTelemetry Java agent 1.32.0 at https://github.com/aliyun/alibabacloud-microservice-demo/tree/master/graalvm-native-image-demo/opentelemetry-agent-native. It has been tested with the Spring boot demo and works well.

Enhance JDK classes

This PR introduces new annotations (see Advice and Aspect for details) to modify methods with advice, i.e. actions taken before and after original method invocation. They follow idioms of byte-buddy Advice which is widely used in OpenTelemetry java agent. The agent developers can easily write JDK class enhancements that take effects in native image only without interfering with GraalVM at build time. The following actions are supported:

• Enhance methods in subclasses (including interface implementers) of the specified class/interface
• Enhance methods in a set of specified classes
• Intercept exceptions thrown from the original methods
• Rewrite original method parameters
• Reusable enhancement to multiple methods
• and more

This PR automatically generates substitution classes from classes with Aspect and Advice annotations and registers them before static analysis. If user defined JDK class enhancement conflicts with the one defined by GraalVM, i.e. two substitution methods are mapped to the same original method, the user defined one will be ignored.

Simple Demo

Here is a simple demo project.
agent-demo.zip

The following steps can show how this PR works:

1. Build this PR and assume the result is placed at $GRAALVM_HOME.

2. Unzip the file.

3. Run mvn package to build the demo project.

4. Run java -cp application/target/application-1.0-SNAPSHOT.jar org.example.Main. The result should be

   This is foo
   This is bar
   0

5. Run collectInstrument.sh to collect instrumentation configurations. Now the program is instrumented, the output should be

   before foo
   This is foo
   before foo
   This is bar
   2

6. Run build.sh to build the demo native image.

7. Run the built native image. The result should be the same as step 6.

Opentelemetry Demo

This PR can now work with OpenTelemetry agent. This demo project shows how to run the Spring boot project with OT agent in native image.

One thing should be noted is in run.sh I use -agentpath: for native-image-agent but not -agentlib, because the latter runs into a crash. In theory, these two options should be the same. I have reported it as a GraalVM JDK issue. More details can be found here (#8094).

Limitation

As native image has only one classloader at runtime, any function relies on multiple classloaders may not work as expected.

[christianwimmer]

Instrumented classes are prepended to the beginning of imageCp, so they will override their original classes at build time.

How does that work for modules? As far as I know, It is not possible with the Java module system to provide classes for one package from two different .jar files.

[ziyilin]

Instrumented classes are prepended to the beginning of imageCp, so they will override their original classes at build time.

How does that work for modules? As far as I know, It is not possible with the Java module system to provide classes for one package from two different .jar files.

Good question. Module system is not considered yet in this version of implementation. I have another version that dumps transformed classes as classdata files at pre-run time, then loads them at image build time to substitute the original ones. I think that will solve the module system problem. But I encountered some other issues with that solution, so turned to current easier solution.

I'll see how to make it work with module system.

---- Updated 2024/1/5-----
@christianwimmer This PR supports module system by patching the modules with instrumented classes now.

[NicklasWallgren]

Has there been any progress toward adding support for Java agents?

@vjovanov @christianwimmer @alina-yur

[CH3CHO]

Would the final solution support transforming JDK built-in classes, such as Thread, ThreadExecutorPoolExecutor, etc.? Because currently, we have them transformed to add some byte code for tracing and metrics collection, which won't work after switching to native image.

[ziyilin]

There will be a major update of this PR very soon. I have been working on issues that are exposed when we apply this PR with Opentelemetry agent.
At this moment, we have successfully collected opentelemetry agent based monitoring data on the native version of Spring boot, MySql, Redis and Kafka demo applications.

[ziyilin]

Would the final solution support transforming JDK built-in classes, such as Thread, ThreadExecutorPoolExecutor, etc.?

This issue will be addressed in the upcoming update of this PR. But it does not automatically apply your transformed JDK classes into native image, because GraalVM itself depends on those classes as well. You need to write static transformation class with APIs provided by this PR. The details will be published soon.