[Question] Add logging around native method in Java agent

[PiotrStoklosa]

Hi,

I’m working on a Java agent program and would like to add logging around certain native methods from java.lang.Thread and java.lang.Object. For example, I want to log the milliseconds argument when calling Thread.sleep(long milliseconds) or track how many times Object.notify() is called. Is there a way to achieve this kind of logging for native methods? I’ve tried several combinations of adding Advice at entry, but while it works for non-native methods, it doesn’t work for the native ones.

[raphw]

There are two approaches to this. The simplest solution would be to Member substitution where you can replace the method with a delegate.

A more complex, you can use an agent and a native method prefix. This would allow you to instrument the method as if it was written in byte code, but it is a bit more challenging to setup

[PiotrStoklosa]

Thanks @raphw,
I tried the first suggested solution. I successfully replaced a custom method (from the monitored application) with my own method. I needed to move this method to another JAR to make it visible to both the application and my Java agent, as all my Java agent classes are loaded with a custom class loader. However, when I tried to substitute Thread.sleep(long), it was not replaced. What's the difference between these two examples? I added this JAR with the replacement method to the bootstrap search (inst.appendToBootstrapClassLoaderSearch(jar);, also this class is loaded from SystemClassLoader). Here is my builder:

Method m;
try {
    m = SleepSubstitution.class.getMethod("sleep2", long.class);
} catch (NoSuchMethodException e) {
    throw new RuntimeException(e);
}
agentBuilder = agentBuilder
        .type(ElementMatchers.named("java.lang.Thread"))
        .transform((builder, typeDescription, classLoader, module, isRedefinition) ->
            builder.visit(MemberSubstitution.strict()
                    .method(ElementMatchers.named("sleep")
                            .and(ElementMatchers.takesArguments(long.class))
                            .and(ElementMatchers.returns(void.class)))
                    .replaceWith(m)
                    .on(ElementMatchers.any())));
return agentBuilder;

my custom sleep2 method signature -

public static void sleep2(long millis) throws InterruptedException

I also confirmed that the builder correctly identifies the java.lang.Thread.sleep(long) method. After installing this builder into the instrumentation, I retransformed the Thread class since it was already loaded. Should I add anything else to make it work, considering that this class belongs to the java.lang package and is already loaded (and when the method was custom from the monitored application, the replacement worked)?

Interestingly, I observed the following error, even though I am not substituting the java.lang.Thread.getName() method. This error disappears when I remove the Thread.sleep replacement code.

[Byte Buddy] ERROR java.lang.Thread [null, module java.base, Thread[main,5,main], loaded=true]
java.lang.IllegalStateException: No substitution found within public final java.lang.String java.lang.Thread.getName() of class java.lang.Thread
at net.bytebuddy.asm.MemberSubstitution$SubstitutingMethodVisitor.visitMaxs(MemberSubstitution.java:7360)
at net.bytebuddy.jar.asm.ClassReader.readCode(ClassReader.java:2664)
at net.bytebuddy.jar.asm.ClassReader.readMethod(ClassReader.java:1512)

[raphw]

You would need to add a custom .ignore(...) matcher where boot methods are ignored by default. You will also need to enable .disableClassFormatChanges() what disables some transformations that you do not use however. Also, you will need to enable a RetransformationStrategy.RETRANSFORM.

[PiotrStoklosa]

My Agent already injects advices into java.lang classes, so I already have this configuration in place:

new AgentBuilder
                .Default()
                .ignore(none())
                .disableClassFormatChanges()
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .with(AgentBuilder.RedefinitionStrategy.DiscoveryStrategy.Reiterating.INSTANCE)
                .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
                .with(AgentBuilder.TypeStrategy.Default.REBASE);

[raphw]

Of course, the exception is java.lang.String java.lang.Thread.getName(). You can set: failIfEmpty(false) on MemberSubstitution.

You are intercepting methods where no methods are substituted.

[PiotrStoklosa]

@raphw
Sorry if this is a silly question, but I'm new to class retransformation. I couldn't find the failIfEmpty(false) method on MemberSubstitution, but I did find failIfNoMatch(false). However, it didn't work either (The error disappeared, but the methods are not substituted). I don't understand how this could help or what it does to the java.lang.String java.lang.Thread.getName() method. I am replacing one specific method:

agentBuilder = agentBuilder
        .type(ElementMatchers.named("java.lang.Thread"))
        .transform((builder, typeDescription, classLoader, module, isRedefinition) ->
            builder.visit(MemberSubstitution.strict()
                    .method(ElementMatchers.named("sleep")
                            .and(ElementMatchers.takesArguments(long.class))
                            .and(ElementMatchers.returns(void.class)))
                    .replaceWith(m)
                    .on(ElementMatchers.any())));

with another specific method:

m = SleepSubstitution.class.getMethod("sleep2", long.class);

public class SleepSubstitution {
    public static void sleep2(long millis) throws InterruptedException {
        System.out.println("Intercepted sleep: " + millis);
    }
}

that has the same signature (the only difference is the native keyword). I didn't provide anything related to the java.lang.String java.lang.Thread.getName() method, so why am I getting an error related to MemberSubstitution on this method? Also, I still don't understand why the sleep method is not substituted. Could it be that native methods can't be replaced by MemberSubstitution? If yes, I'm open to the second suggested solution using the native prefix. However, I've tried this for quite some time, and I was getting an error stating that I can't add methods to already loaded classes (i.e., java.lang.Thread). Of course, I'm looking for the simplest solution here.

[raphw]

I meant this method: https://github.com/raphw/byte-buddy/blob/master/byte-buddy-dep/src/main/java/net/bytebuddy/asm/MemberSubstitution.java#L273

I am however unsure if we misunderstand each other. Substitution is replacing method calls and field access within a method. So you can replace the call to Thread.sleep here:

void myMethod() throws Exception {
  Thread.sleep(1000);
}

But you cannot instrument Thread.sleep itself which is native.

[PiotrStoklosa]

Yes, I added this:

        agentBuilder = agentBuilder
                .type(ElementMatchers.named("java.lang.Thread"))
                .transform((builder, typeDescription, classLoader, module, isRedefinition) ->
                    builder.visit(MemberSubstitution.strict()
                            .method(ElementMatchers.named("sleep")
                                    .and(ElementMatchers.takesArguments(long.class))
                                    .and(ElementMatchers.returns(void.class)))
                            .replaceWith(m)
                                    .failIfNoMatch(false)
                            .on(ElementMatchers.any())));
        return agentBuilder;

but I don't know how it could help, as I don't understand the reason why I can see this error.
I think we are talking about the same thing. As I understand it, MemberSubstitutionis a tool that enables redirecting all method calls to be substituted with calls to different method. So, all Thread.sleep(long) calls should be replaced with SleepSubstitution.sleep2(long), regardless of calling Thread.sleep() inside this new method later. I only want to achieve this behavior, so I will add the call to Thread.sleep() later; for now, I just want to see that the calls are being substituted.

[raphw]

But in your case, you are only applying the transformation to type(ElementMatchers.named("java.lang.Thread")). Is this intended?

[PiotrStoklosa]

Hmm, I'm wondering - when I want to substitute Thread.sleep(long) calls (replace every Thread.sleep(long) call made in the monitored application) with my custom method, I need to transform Thread, but now I'm questioning whether this is correct. Would you say that I should target a different type in my scenario?

[raphw]

You would need to transform any class that contains the method call, what is potentially any type.

[PiotrStoklosa]

You're right! When I changed the monitoring to any type (excluding the SleepSubstitution class to prevent recursive calls), I was able to monitor Thread.sleep(long) calls. However, I noticed that not all calls were being replaced, especially those inside lambdas. After some investigation, I found the LambdaInstrumentationStrategy and enabled it, but it didn't help. Now, my builder looks like this:

try {
    m = SleepSubstitution.class.getMethod("sleep2", long.class);
} catch (NoSuchMethodException e) {
    throw new RuntimeException(e);
}
agentBuilder = agentBuilder
        .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
        .type(ElementMatchers.not(ElementMatchers.nameContains("org.threadmonitoring.SleepSubstitution")))
        .transform((builder, typeDescription, classLoader, module, isRedefinition) ->
                builder.visit(MemberSubstitution.strict()
                        .method(ElementMatchers.named("sleep")
                                .and(ElementMatchers.takesArguments(long.class))
                                .and(ElementMatchers.returns(void.class)))
                        .replaceWith(m)
                        .failIfNoMatch(false)
                        .on(ElementMatchers.any())));

However, my lambda is still not instrumented. Here’s the monitored code:

ExecutorService executor = Executors.newFixedThreadPool(10);
executor.execute(new Runnable() {
                     @Override
                     public void run() {
                         try {
                             Thread.sleep(1000);
                         } catch (InterruptedException e) {
                             throw new RuntimeException(e);
                         }
                     }
                 }
);
executor.execute(() -> {
    try {
        System.out.println("Calling Sleep 2000");
        Thread.sleep(2000);
    } catch (InterruptedException e) {
        throw new RuntimeException(e);
    }
});
executor.shutdown();

The Thread.sleep(1000) call in the execute method is substituted (since it's in an anonymous class), but the Thread.sleep(2000) call in the second execute call is not substituted (as it's in a lambda), even though lambda is detected and transformed by the Java Agent based on the logs:

...
[Byte Buddy] DISCOVERY org.example.Main$$Lambda$ByteBuddy$6 [jdk.internal.loader.ClassLoaders$AppClassLoader@6ed3ef1, null, Thread[main,5,main], loaded=false]
[Byte Buddy] TRANSFORM org.example.Main$$Lambda$ByteBuddy$6 [jdk.internal.loader.ClassLoaders$AppClassLoader@6ed3ef1, null, Thread[main,5,main], loaded=false]
[Byte Buddy] COMPLETE org.example.Main$$Lambda$ByteBuddy$6 [jdk.internal.loader.ClassLoaders$AppClassLoader@6ed3ef1, null, Thread[main,5,main], loaded=false]
Calling Sleep 2000
...

Isn't this a bug? The Java Agent is built with JDK 11, the monitored application is running on Java 17.

[raphw]

Lambda expressions are just synthetic methods in the declaring class. Are you ignoring synthetics?

[PiotrStoklosa]

Isn't .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED) along with .ignore(none()) sufficient to ensure synthetics are not ignored?

[raphw]

The lambda strategy should not be relevant, as the code is contained elsewhere. But you are right, it should be instrumented. Could you create a simple reproducer?

[PiotrStoklosa]

Yes, I have created a simple reproducer and attached the following files: agent_and_application_jars.zip (which contains lambda_agent-1.0-SNAPSHOT.jar and monitored_application-1.0-SNAPSHOT.jar), as well as two Gradle projects (lambda_agent.zip and monitored_application.zip) that include the source code. Both the Java Agent and the monitored application are built with JDK 11.

monitored_application.zip
lambda_agent.zip
agent_and_application_jars.zip

[raphw]

You do not want to retransform, you want to decorate the class. Use

            .with(AgentBuilder.TypeStrategy.Default.DECORATE)

[PiotrStoklosa]

It works! Now the method call is replaced in lambdas. Why does the Decorate type strategy work, but Rebase doesn't?

[raphw]

Rebase pretends like it is creating a subclass but then merges the subclass into the base class.

This way, the method is not considered declared, and can therefore not be visited. With decoration, you run Byte Buddy in a different mode where this is not the case.

[PiotrStoklosa]

Sure, thank you @raphw, you helped me a lot!