Fix request context propagation for async executor tasks

agent/src/main/java/dev/aikido/agent/Wrappers.java

@@ -1,6 +1,11 @@
 package dev.aikido.agent;
 
 import dev.aikido.agent.wrappers.*;
+import dev.aikido.agent.wrappers.executor.AbstractExecutorServiceWrapper;
+import dev.aikido.agent.wrappers.executor.DelegatedExecutorServiceWrapper;
+import dev.aikido.agent.wrappers.executor.ForkJoinPoolWrapper;
+import dev.aikido.agent.wrappers.executor.ScheduledThreadPoolExecutorWrapper;
+import dev.aikido.agent.wrappers.executor.ThreadPoolExecutorWrapper;
 import dev.aikido.agent.wrappers.file.FileConstructorMultiArgumentWrapper;
 import dev.aikido.agent.wrappers.file.FileConstructorSingleArgumentWrapper;
 import dev.aikido.agent.wrappers.javalin.*;
@@ -17,6 +22,13 @@ public final class Wrappers {
     private Wrappers() {}
     public static final List<Wrapper> WRAPPERS = Arrays.asList(
             new PostgresWrapper(),
+
+            new DelegatedExecutorServiceWrapper(),
+            new ThreadPoolExecutorWrapper(),
+            new AbstractExecutorServiceWrapper(),
+            new ForkJoinPoolWrapper(),
+            new ScheduledThreadPoolExecutorWrapper(),
+
             new SpringMVCJakartaWrapper(),
             new SpringMVCJavaxWrapper(),
             new SpringWebfluxWrapper(),

agent/src/main/java/dev/aikido/agent/wrappers/executor/AbstractExecutorServiceWrapper.java

@@ -0,0 +1,52 @@
+package dev.aikido.agent.wrappers.executor;
+
+import dev.aikido.agent.wrappers.Wrapper;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.util.concurrent.AbstractExecutorService;
+import java.util.concurrent.Callable;
+
+import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+public class AbstractExecutorServiceWrapper implements Wrapper {
+    @Override
+    public String getName() {
+        return SubmitAdvice.class.getName();
+    }
+
+    @Override
+    public ElementMatcher getMatcher() {
+        return isMethod()
+            .and(named("submit"))
+            .and(
+                takesArguments(Runnable.class)
+                    .or(takesArguments(Callable.class))
+                    .or(takesArguments(Runnable.class, Object.class))
+            );
+    }
+
+    @Override
+    public ElementMatcher getTypeMatcher() {
+        return isSubTypeOf(AbstractExecutorService.class);
+    }
+
+    public static class SubmitAdvice {
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void before(
+            @Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
+        ) {
+            if (task instanceof Runnable) {
+                task = ExecutorContextPropagation.wrap((Runnable) task);
+            } else if (task instanceof Callable) {
+                task = ExecutorContextPropagation.wrap((Callable) task);
+            }
+        }
+    }
+}

agent/src/main/java/dev/aikido/agent/wrappers/executor/DelegatedExecutorServiceWrapper.java

@@ -0,0 +1,72 @@
+package dev.aikido.agent.wrappers.executor;
+
+import dev.aikido.agent.wrappers.Wrapper;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.lang.reflect.Method;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.concurrent.Callable;
+
+import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+public class DelegatedExecutorServiceWrapper implements Wrapper {
+    @Override
+    public String getName() {
+        return DelegatedExecutorAdvice.class.getName();
+    }
+
+    @Override
+    public ElementMatcher getMatcher() {
+        return isMethod()
+            .and(named("execute").or(named("submit")))
+            .and(
+                takesArguments(Runnable.class)
+                    .or(takesArguments(Callable.class))
+                    .or(takesArguments(Runnable.class, Object.class))
+            );
+    }
+
+    @Override
+    public ElementMatcher getTypeMatcher() {
+        return nameStartsWith("java.util.concurrent.Executors$");
+    }
+
+    public static class DelegatedExecutorAdvice {
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void before(
+            @Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
+        ) throws Exception {
+            if (task == null) {
+                return;
+            }
+
+            // This advice is applied to JDK classes loaded by the bootstrap classloader.
+            // Load agent_api reflectively because bootstrap classes cannot directly reference agent classes.
+            String jarFilePath = System.getProperty("AIK_agent_api_jar");
+            if (jarFilePath == null || jarFilePath.isBlank()) {
+                return;
+            }
+
+            URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) });
+            Class<?> contextPropagationClass = classLoader.loadClass(
+                "dev.aikido.agent_api.context.ContextPropagation"
+            );
+
+            if (task instanceof Runnable) {
+                Method wrapRunnable = contextPropagationClass.getMethod("wrap", Runnable.class);
+                task = wrapRunnable.invoke(null, task);
+            } else if (task instanceof Callable) {
+                Method wrapCallable = contextPropagationClass.getMethod("wrap", Callable.class);
+                task = wrapCallable.invoke(null, task);
+            }
+        }
+    }
+}

agent/src/main/java/dev/aikido/agent/wrappers/executor/ExecutorContextPropagation.java

@@ -0,0 +1,76 @@
+package dev.aikido.agent.wrappers.executor;
+
+import java.lang.reflect.Method;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.concurrent.Callable;
+
+// Bridges the executor advice (woven into java.util.concurrent classes) to ContextPropagation in
+// agent_api, which lives on a different classloader, and caches the reflected methods so the lookup
+// happens once. A missing AIK_agent_api_jar is treated as "not ready yet" (early startup) and
+// retried on the next call, rather than disabling propagation for the whole JVM; only a genuine
+// load failure once the path is set disables it.
+public final class ExecutorContextPropagation {
+    private static volatile Method wrapRunnableMethod;
+    private static volatile Method wrapCallableMethod;
+    private static volatile boolean disabled;
+
+    private ExecutorContextPropagation() {}
+
+    public static Runnable wrap(Runnable task) {
+        if (task == null) {
+            return task;
+        }
+        Method wrap = wrapRunnableMethod;
+        if (wrap == null) {
+            init();
+            wrap = wrapRunnableMethod;
+        }
+        if (wrap == null) {
+            return task;
+        }
+        try {
+            return (Runnable) wrap.invoke(null, task);
+        } catch (Throwable ignored) {
+            return task;
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    public static <T> Callable<T> wrap(Callable<T> task) {
+        if (task == null) {
+            return task;
+        }
+        Method wrap = wrapCallableMethod;
+        if (wrap == null) {
+            init();
+            wrap = wrapCallableMethod;
+        }
+        if (wrap == null) {
+            return task;
+        }
+        try {
+            return (Callable<T>) wrap.invoke(null, task);
+        } catch (Throwable ignored) {
+            return task;
+        }
+    }
+
+    private static synchronized void init() {
+        if (disabled || wrapRunnableMethod != null) {
+            return;
+        }
+        String jarFilePath = System.getProperty("AIK_agent_api_jar");
+        if (jarFilePath == null || jarFilePath.isBlank()) {
+            return; // not set yet during early startup - retry on a later call
+        }
+        try {
+            URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) });
+            Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.context.ContextPropagation");
+            wrapCallableMethod = clazz.getMethod("wrap", Callable.class);
+            wrapRunnableMethod = clazz.getMethod("wrap", Runnable.class);
+        } catch (Throwable ignored) {
+            disabled = true;
+        }
+    }
+}

agent/src/main/java/dev/aikido/agent/wrappers/executor/ForkJoinPoolWrapper.java

@@ -0,0 +1,52 @@
+package dev.aikido.agent.wrappers.executor;
+
+import dev.aikido.agent.wrappers.Wrapper;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ForkJoinPool;
+
+import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+public class ForkJoinPoolWrapper implements Wrapper {
+    @Override
+    public String getName() {
+        return ForkJoinAdvice.class.getName();
+    }
+
+    @Override
+    public ElementMatcher getMatcher() {
+        return isMethod()
+            .and(named("execute").or(named("submit")))
+            .and(
+                takesArguments(Runnable.class)
+                    .or(takesArguments(Callable.class))
+                    .or(takesArguments(Runnable.class, Object.class))
+            );
+    }
+
+    @Override
+    public ElementMatcher getTypeMatcher() {
+        return isSubTypeOf(ForkJoinPool.class);
+    }
+
+    public static class ForkJoinAdvice {
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void before(
+            @Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
+        ) {
+            if (task instanceof Runnable) {
+                task = ExecutorContextPropagation.wrap((Runnable) task);
+            } else if (task instanceof Callable) {
+                task = ExecutorContextPropagation.wrap((Callable) task);
+            }
+        }
+    }
+}

agent/src/main/java/dev/aikido/agent/wrappers/executor/ScheduledThreadPoolExecutorWrapper.java

@@ -0,0 +1,73 @@
+package dev.aikido.agent.wrappers.executor;
+
+import dev.aikido.agent.wrappers.Wrapper;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.lang.reflect.Method;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+public class ScheduledThreadPoolExecutorWrapper implements Wrapper {
+    @Override
+    public String getName() {
+        return ScheduleAdvice.class.getName();
+    }
+
+    @Override
+    public ElementMatcher getMatcher() {
+        return isMethod()
+            .and(named("schedule"))
+            .and(
+                takesArguments(Runnable.class, long.class, TimeUnit.class)
+                    .or(takesArguments(Callable.class, long.class, TimeUnit.class))
+            );
+    }
+
+    @Override
+    public ElementMatcher getTypeMatcher() {
+        return isSubTypeOf(ScheduledThreadPoolExecutor.class);
+    }
+
+    public static class ScheduleAdvice {
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void before(
+            @Advice.Argument(value = 0, readOnly = false, typing = DYNAMIC) Object task
+        ) throws Exception {
+            if (task == null) {
+                return;
+            }
+
+            // This advice is applied to JDK classes loaded by the bootstrap classloader.
+            // Load agent_api reflectively because bootstrap classes cannot directly reference agent classes.
+            String jarFilePath = System.getProperty("AIK_agent_api_jar");
+            if (jarFilePath == null || jarFilePath.isBlank()) {
+                return;
+            }
+
+            URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) });
+            Class<?> contextPropagationClass = classLoader.loadClass(
+                "dev.aikido.agent_api.context.ContextPropagation"
+            );
+
+            if (task instanceof Runnable) {
+                Method wrapRunnable = contextPropagationClass.getMethod("wrap", Runnable.class);
+                task = wrapRunnable.invoke(null, task);
+            } else if (task instanceof Callable) {
+                Method wrapCallable = contextPropagationClass.getMethod("wrap", Callable.class);
+                task = wrapCallable.invoke(null, task);
+            }
+        }
+    }
+}

agent/src/main/java/dev/aikido/agent/wrappers/executor/ThreadPoolExecutorWrapper.java

@@ -0,0 +1,42 @@
+package dev.aikido.agent.wrappers.executor;
+
+import dev.aikido.agent.wrappers.Wrapper;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.util.concurrent.ThreadPoolExecutor;
+
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+public class ThreadPoolExecutorWrapper implements Wrapper {
+    @Override
+    public String getName() {
+        return ExecuteAdvice.class.getName();
+    }
+
+    @Override
+    public ElementMatcher getMatcher() {
+        return isMethod()
+            .and(named("execute"))
+            .and(takesArguments(Runnable.class));
+    }
+
+    @Override
+    public ElementMatcher getTypeMatcher() {
+        return isSubTypeOf(ThreadPoolExecutor.class);
+    }
+
+    public static class ExecuteAdvice {
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void before(
+            @Advice.Argument(value = 0, readOnly = false) Runnable task
+        ) {
+            task = ExecutorContextPropagation.wrap(task);
+        }
+    }
+}