How to intercept Kotlin Coroutines?

Question

I'm trying to instrument Kotlin coroutines, similar to what's done here using a Javaagent. I don't want a Javaagent.

The first step is to intercept the creation, suspension and resumption of Coroutines defined in the DebugProbes. The code for that is as follows:

public class Instrumentor {
    private static final Logger LOG = LoggerFactory.getLogger(Instrumentor.class);

    public static void install() {
        TypeDescription typeDescription = TypePool.Default.ofSystemLoader()
                .describe("kotlin.coroutines.jvm.internal.DebugProbesKt")
                .resolve();
        new ByteBuddy()
                .redefine(typeDescription, ClassFileLocator.ForClassLoader.ofSystemLoader())
                .method(ElementMatchers.named("probeCoroutineCreated").and(ElementMatchers.takesArguments(1)))
                .intercept(MethodDelegation.to(CoroutineCreatedAdvice.class))
                .method(ElementMatchers.named("probeCoroutineResumed").and(ElementMatchers.takesArguments(1)))
                .intercept(MethodDelegation.to(CoroutineResumedAdvice.class))
                .method(ElementMatchers.named("probeCoroutineSuspended").and(ElementMatchers.takesArguments(1)))
                .intercept(MethodDelegation.to(CoroutineSuspendedAdvice.class))
                .make()
                .load(ClassLoader.getSystemClassLoader(), ClassLoadingStrategy.Default.INJECTION);

        DebugProbes.INSTANCE.install();
    }

    public static void uninstall() {
        DebugProbes.INSTANCE.uninstall();
    }

    public static class CoroutineCreatedAdvice {
        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static Continuation<Object> exit(@Advice.Return(readOnly = false) Continuation<Object> retVal) {
            LOG.info("Coroutine created: {}", retVal);
           
            return retVal;
        }
    }

    public static class CoroutineResumedAdvice {
        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void enter(@Advice.Argument(0) final Continuation<Object> continuation) {
            LOG.info("Coroutine resumed: {}", continuation);
        }
    }

    public static class CoroutineSuspendedAdvice {
        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void enter(@Advice.Argument(0) final Continuation<Object> continuation) {
            LOG.info("Coroutine suspended: {}", continuation);
        }
    }
}

JUnit5 test to trigger interception:

class CoroutineInstrumentationTest {
    companion object {
        @JvmStatic
        @BeforeAll
        fun beforeAll() {
            Instrumentor.install()
        }

        @JvmStatic
        @AfterAll
        fun afterAll() {
            Instrumentor.uninstall()
        }
    }

    @Test
    fun testInterception() {
        runBlocking {
            println("Test")
        }
    }
}

However, no interception happens (confirmed by the absence of log statements and by using a debugger). I'm new to Byte Buddy, so it's possible I'm missing something. Any ideas?

Kotlin v1.4.10, Kotlin Coroutines v1.3.9, Byte Buddy v1.10.17.

Answer 1

Are you sure the class is not yet loaded at this point? Try setting a breakpoint in ClassInjector.UsingReflection to see if you acutally walk through or of the injection is aborted due to a previously loaded class.

The cleaner solution would be a Java agent. You can use byte-buddy-agent to create one dynamically by ByteBuddyAgent.install() and then register an AgentBuilder on it.