Propagate request context across JDK executors - #343
Conversation
| if (task instanceof Runnable) { | ||
| task = ExecutorContextPropagation.wrap((Runnable) task); | ||
| } else if (task instanceof Callable) { | ||
| task = ExecutorContextPropagation.wrap((Callable) task); |
There was a problem hiding this comment.
🟡 Medium - Bootstrap executor advice calls an agent-only helper that JDK classes cannot resolve
The new AbstractExecutorService, ThreadPoolExecutor, and ForkJoinPool advices inject direct calls to ExecutorContextPropagation.wrap(...) into java.util.concurrent classes even though this JVM setup never injects agent helper classes into the bootstrap classloader. The project already handles bootstrap-loaded JDK wrappers via reflective loading for that reason, and these advices also suppress any linkage error, so submit/execute on core executors silently runs without propagated request context. As a result, async work scheduled onto common JDK executors loses the request metadata that Zen uses to attribute and enforce protections on downstream sinks.
Show fix
Do not reference dev.aikido.agent... helpers directly from advice woven into bootstrap-loaded JDK classes. Either move these executor wrappers to the same reflective bridge pattern already used for bootstrap JDK wrappers, or explicitly append the helper classes to the bootstrap classloader search before instrumenting java.util.concurrent so the injected calls can actually resolve.
More info - Reply on this comment to give feedback or ignore the issue.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
| // 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); |
There was a problem hiding this comment.
🟡 Medium - Delegated and scheduled executor wrappers reopen agent_api.jar on every task submission
Both bootstrap-safe executor wrappers now allocate a fresh URLClassLoader, load ContextPropagation, reflect the wrap(...) method, and close the loader for every submit/execute or schedule call. These methods are on the hot path for common asynchronous work, so the change adds repeated JAR parsing and reflective lookup overhead to every task dispatch instead of amortizing it once per JVM. Under request-driven executor usage this can materially reduce throughput and increase allocation pressure for the very async workloads this feature targets.
Show fix
Cache the reflected ContextPropagation class and wrap methods across calls instead of constructing a new URLClassLoader per task. If bootstrap isolation prevents direct helper references, initialize the reflective bridge lazily once in bootstrap-safe code and reuse the cached Method/MethodHandle objects for all subsequent submissions and schedules.
More info - Reply on this comment to give feedback or ignore the issue.
Weave ThreadPoolExecutor, ForkJoinPool, AbstractExecutorService, ScheduledThreadPoolExecutor and the Executors$Delegated* services so a task submitted from a request thread runs under that request's Context on the pool worker. Helper-backed wrappers share a cached reflection bridge (ExecutorContextPropagation); the Scheduled/Delegated wrappers load agent_api per call through a URLClassLoader that is closed once wrapping is done. Integration tests cover preservation, nested submits, concurrent isolation and pooled-worker reuse.
29b12e9 to
cba33a0
Compare
c0d0a2a to
75593a3
Compare
Second of the stack, on top of #342 — this is where the primitives get used.
We weave the JDK executor types so a task submitted from a request thread runs with that request's context on the pool worker:
ThreadPoolExecutor,ForkJoinPool,AbstractExecutorService,ScheduledThreadPoolExecutorand theExecutors$Delegated*services. In practice that covers@Async,CompletableFuture, plain thread pools and scheduled tasks — all of which used to lose the context on the hop and silently miss attacks in the async path.Two loading strategies, on purpose:
ThreadPool/ForkJoin/AbstractExecutorServicego through a cached reflection bridge (ExecutorContextPropagation).ScheduledThreadPoolExecutorand the delegated executors loadagent_apiper call via a short-livedURLClassLoader. They run early during class loading, where the cached bridge deadlocks, so the per-call loader is required here — and it's now closed after wrapping.ExecutorWrapperTestcovers this with 9 integration tests under the woven agent: propagation across every executor type plusCompletableFuture, nested submits, two concurrent tasks each keeping their own context, a pooled worker not leaking context into the next task, and the no-context passthrough.