diff --git a/agent/src/main/java/dev/aikido/agent/Wrappers.java b/agent/src/main/java/dev/aikido/agent/Wrappers.java index a71c13b26..756094f05 100644 --- a/agent/src/main/java/dev/aikido/agent/Wrappers.java +++ b/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 WRAPPERS = Arrays.asList( new PostgresWrapper(), + + new DelegatedExecutorServiceWrapper(), + new ThreadPoolExecutorWrapper(), + new AbstractExecutorServiceWrapper(), + new ForkJoinPoolWrapper(), + new ScheduledThreadPoolExecutorWrapper(), + new SpringMVCJakartaWrapper(), new SpringMVCJavaxWrapper(), new SpringWebfluxWrapper(), diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/executor/AbstractExecutorServiceWrapper.java b/agent/src/main/java/dev/aikido/agent/wrappers/executor/AbstractExecutorServiceWrapper.java new file mode 100644 index 000000000..1f152c3f1 --- /dev/null +++ b/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); + } + } + } +} diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/executor/DelegatedExecutorServiceWrapper.java b/agent/src/main/java/dev/aikido/agent/wrappers/executor/DelegatedExecutorServiceWrapper.java new file mode 100644 index 000000000..ad5e0b7e4 --- /dev/null +++ b/agent/src/main/java/dev/aikido/agent/wrappers/executor/DelegatedExecutorServiceWrapper.java @@ -0,0 +1,78 @@ +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; + } + + // The wrapper returned by wrap() resolves through the parent (system) classloader, + // so this per-call loader can be closed once wrapping is done. + URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) }); + try { + 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); + } + } finally { + classLoader.close(); + } + } + } +} diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/executor/ExecutorContextPropagation.java b/agent/src/main/java/dev/aikido/agent/wrappers/executor/ExecutorContextPropagation.java new file mode 100644 index 000000000..11a88ccdb --- /dev/null +++ b/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 Callable wrap(Callable task) { + if (task == null) { + return task; + } + Method wrap = wrapCallableMethod; + if (wrap == null) { + init(); + wrap = wrapCallableMethod; + } + if (wrap == null) { + return task; + } + try { + return (Callable) 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; + } + } +} diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/executor/ForkJoinPoolWrapper.java b/agent/src/main/java/dev/aikido/agent/wrappers/executor/ForkJoinPoolWrapper.java new file mode 100644 index 000000000..b78f0a95d --- /dev/null +++ b/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); + } + } + } +} diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/executor/ScheduledThreadPoolExecutorWrapper.java b/agent/src/main/java/dev/aikido/agent/wrappers/executor/ScheduledThreadPoolExecutorWrapper.java new file mode 100644 index 000000000..a80e3a40f --- /dev/null +++ b/agent/src/main/java/dev/aikido/agent/wrappers/executor/ScheduledThreadPoolExecutorWrapper.java @@ -0,0 +1,79 @@ +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; + } + + // The wrapper returned by wrap() resolves through the parent (system) classloader, + // so this per-call loader can be closed once wrapping is done. + URLClassLoader classLoader = new URLClassLoader(new URL[] { new URL(jarFilePath) }); + try { + 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); + } + } finally { + classLoader.close(); + } + } + } +} diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/executor/ThreadPoolExecutorWrapper.java b/agent/src/main/java/dev/aikido/agent/wrappers/executor/ThreadPoolExecutorWrapper.java new file mode 100644 index 000000000..ece585980 --- /dev/null +++ b/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); + } + } +} diff --git a/agent_api/src/test/java/wrappers/ExecutorWrapperTest.java b/agent_api/src/test/java/wrappers/ExecutorWrapperTest.java new file mode 100644 index 000000000..f8fd71607 --- /dev/null +++ b/agent_api/src/test/java/wrappers/ExecutorWrapperTest.java @@ -0,0 +1,195 @@ +package wrappers; + +import dev.aikido.agent_api.context.Context; +import dev.aikido.agent_api.context.ContextObject; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; + +class ExecutorWrapperTest { + @AfterEach + void tearDown() { + Context.reset(); + } + + @Test + void scheduledThreadPoolExecutorSchedulePropagatesContext() throws Exception { + ContextObject ctx = new ContextObject(); + ctx.setRateLimitGroup("scheduled"); + Context.set(ctx); + ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(1); + try { + ContextObject onWorker = executor.schedule(Context::get, 0, TimeUnit.MILLISECONDS).get(5, TimeUnit.SECONDS); + assertEquals("scheduled", onWorker.getRateLimitGroup()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void delegatedSingleThreadExecutorSubmitPropagatesContext() throws Exception { + ContextObject ctx = new ContextObject(); + ctx.setRateLimitGroup("delegated"); + Context.set(ctx); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + assertEquals("delegated", executor.submit(Context::get).get(5, TimeUnit.SECONDS).getRateLimitGroup()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void threadPoolExecutorExecutePropagatesContext() throws Exception { + ContextObject ctx = new ContextObject(); + ctx.setRateLimitGroup("threadpool"); + Context.set(ctx); + ThreadPoolExecutor executor = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()); + AtomicReference onWorker = new AtomicReference<>(); + CountDownLatch done = new CountDownLatch(1); + try { + executor.execute(() -> { + onWorker.set(Context.get()); + done.countDown(); + }); + done.await(5, TimeUnit.SECONDS); + assertEquals("threadpool", onWorker.get().getRateLimitGroup()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void forkJoinPoolSubmitPropagatesContext() throws Exception { + ContextObject ctx = new ContextObject(); + ctx.setRateLimitGroup("forkjoin"); + Context.set(ctx); + ForkJoinPool executor = new ForkJoinPool(1); + try { + assertEquals("forkjoin", executor.submit(Context::get).get(5, TimeUnit.SECONDS).getRateLimitGroup()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void completableFutureSupplyAsyncPropagatesContext() throws Exception { + ContextObject ctx = new ContextObject(); + ctx.setRateLimitGroup("completablefuture"); + Context.set(ctx); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + ContextObject onWorker = CompletableFuture.supplyAsync(Context::get, executor).get(5, TimeUnit.SECONDS); + assertEquals("completablefuture", onWorker.getRateLimitGroup()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void concurrentTasksKeepTheirOwnContext() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + CyclicBarrier bothRunning = new CyclicBarrier(2); + Callable readWhileConcurrent = () -> { + bothRunning.await(5, TimeUnit.SECONDS); + return Context.get(); + }; + ContextObject first = new ContextObject(); + first.setRateLimitGroup("first"); + ContextObject second = new ContextObject(); + second.setRateLimitGroup("second"); + try { + Context.set(first); + Future a = executor.submit(readWhileConcurrent); + Context.set(second); + Future b = executor.submit(readWhileConcurrent); + assertEquals("first", a.get(5, TimeUnit.SECONDS).getRateLimitGroup()); + assertEquals("second", b.get(5, TimeUnit.SECONDS).getRateLimitGroup()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void concurrentTasksFromOneContextEachGetTheirOwnCopy() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + CyclicBarrier bothRunning = new CyclicBarrier(2); + Callable readWhileConcurrent = () -> { + bothRunning.await(5, TimeUnit.SECONDS); + return Context.get(); + }; + ContextObject ctx = new ContextObject(); + try { + Context.set(ctx); + Future a = executor.submit(readWhileConcurrent); + Future b = executor.submit(readWhileConcurrent); + ContextObject onA = a.get(5, TimeUnit.SECONDS); + ContextObject onB = b.get(5, TimeUnit.SECONDS); + assertNotSame(ctx, onA); + assertNotSame(ctx, onB); + assertNotSame(onA, onB); + } finally { + executor.shutdownNow(); + } + } + + @Test + void nestedSubmitPropagatesOuterContext() throws Exception { + ExecutorService outer = Executors.newSingleThreadExecutor(); + ExecutorService inner = Executors.newSingleThreadExecutor(); + ContextObject ctx = new ContextObject(); + ctx.setRateLimitGroup("outer"); + try { + Context.set(ctx); + ContextObject seen = outer.submit( + () -> inner.submit(Context::get).get(5, TimeUnit.SECONDS) + ).get(5, TimeUnit.SECONDS); + assertEquals("outer", seen.getRateLimitGroup()); + } finally { + outer.shutdownNow(); + inner.shutdownNow(); + } + } + + @Test + void pooledWorkerDoesNotLeakContextToNextTask() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Context.set(new ContextObject()); + executor.submit(Context::get).get(5, TimeUnit.SECONDS); + Context.reset(); + assertNull(executor.submit(Context::get).get(5, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + } + } + + @Test + void taskSubmittedWithoutContextRunsWithNone() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Context.reset(); + assertNull(executor.submit(Context::get).get(5, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + } + } +}