Skip to content

Propagate request context across JDK executors - #343

Open
Mishenevd wants to merge 1 commit into
feat/context-propagation-primitivesfrom
feat/executor-context-propagation-instrumentation
Open

Propagate request context across JDK executors#343
Mishenevd wants to merge 1 commit into
feat/context-propagation-primitivesfrom
feat/executor-context-propagation-instrumentation

Conversation

@Mishenevd

@Mishenevd Mishenevd commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator

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, ScheduledThreadPoolExecutor and the Executors$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/AbstractExecutorService go through a cached reflection bridge (ExecutorContextPropagation).
  • ScheduledThreadPoolExecutor and the delegated executors load agent_api per call via a short-lived URLClassLoader. 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.

ExecutorWrapperTest covers this with 9 integration tests under the woven agent: propagation across every executor type plus CompletableFuture, 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.

Comment on lines +45 to +48
if (task instanceof Runnable) {
task = ExecutorContextPropagation.wrap((Runnable) task);
} else if (task instanceof Callable) {
task = ExecutorContextPropagation.wrap((Callable) task);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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

codecov Bot commented Aug 17, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Comment on lines +58 to +68
// 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);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.
@Mishenevd
Mishenevd force-pushed the feat/executor-context-propagation-instrumentation branch from 29b12e9 to cba33a0 Compare August 17, 2026 10:54
@Mishenevd
Mishenevd force-pushed the feat/context-propagation-primitives branch from c0d0a2a to 75593a3 Compare August 17, 2026 10:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant