From 776ac966f867159f3d7314ed978e8a5f685d7a06 Mon Sep 17 00:00:00 2001 From: Mishenevd Date: Fri, 14 Aug 2026 16:40:58 +0200 Subject: [PATCH] Add context propagation primitives for async task execution Introduce ContextPropagatingRunnable/Callable and the ContextPropagation factory that captures the current request Context at wrap time and restores it around task execution, restoring the worker's previous context afterwards. Each task gets its own copy of the context (ContextObject.copyForPropagation) so parallel workers sharing one request never race on its mutable state. Wrapping is idempotent and passes through null / already-wrapped / no-context tasks. Unit-tested in isolation without agent weaving. --- .../dev/aikido/agent_api/context/Context.java | 7 + .../agent_api/context/ContextObject.java | 30 +- .../context/ContextPropagatingCallable.java | 24 ++ .../context/ContextPropagatingRunnable.java | 22 ++ .../agent_api/context/ContextPropagation.java | 33 ++ .../agent_api/storage/RedirectNode.java | 11 + .../java/context/ContextPropagationTest.java | 328 ++++++++++++++++++ 7 files changed, 454 insertions(+), 1 deletion(-) create mode 100644 agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java create mode 100644 agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java create mode 100644 agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java create mode 100644 agent_api/src/test/java/context/ContextPropagationTest.java diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java b/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java index 0389cc903..0382f89b4 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java @@ -15,4 +15,11 @@ public static void set(ContextObject contextObject) { public static void reset() { threadLocalContext.remove(); } + public static void restore(ContextObject previous) { + if (previous != null) { + set(previous); + } else { + reset(); + } + } } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java index 339358384..41e505bd6 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java @@ -6,7 +6,7 @@ import java.util.*; -public class ContextObject { +public class ContextObject implements Cloneable { protected String method; protected String source; protected String url; @@ -25,6 +25,34 @@ public class ContextObject { protected transient Map> cache = new HashMap<>(); protected transient Optional forcedProtectionOff = Optional.empty(); + // Async tasks get their own copy so parallel workers sharing one request's + // context never race on its mutable working state (cache, redirect nodes). + public ContextObject copyForPropagation() { + try { + ContextObject copy = (ContextObject) super.clone(); + copy.cache = copyCache(this.cache); + copy.redirectStartNodes = copyRedirectChains(this.redirectStartNodes); + return copy; + } catch (CloneNotSupportedException e) { + return this; + } + } + + private static Map> copyCache(Map> cache) { + Map> copy = new HashMap<>(); + cache.forEach((key, strings) -> copy.put(key, new HashMap<>(strings))); + return copy; + } + + private static ArrayList copyRedirectChains(ArrayList nodes) { + if (nodes == null) { + return null; + } + ArrayList copy = new ArrayList<>(nodes.size()); + nodes.forEach(starter -> copy.add(starter.copyChain())); + return copy; + } + public boolean middlewareExecuted() {return executedMiddleware; } public void setExecutedMiddleware(boolean value) { executedMiddleware = value; } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java new file mode 100644 index 000000000..f7f8b65b3 --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java @@ -0,0 +1,24 @@ +package dev.aikido.agent_api.context; + +import java.util.concurrent.Callable; + +public final class ContextPropagatingCallable implements Callable { + private final Callable delegate; + private final ContextObject context; + + public ContextPropagatingCallable(Callable delegate, ContextObject context) { + this.delegate = delegate; + this.context = context; + } + + @Override + public T call() throws Exception { + ContextObject previous = Context.get(); + try { + Context.set(context); + return delegate.call(); + } finally { + Context.restore(previous); + } + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java new file mode 100644 index 000000000..2c81e42ef --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java @@ -0,0 +1,22 @@ +package dev.aikido.agent_api.context; + +public final class ContextPropagatingRunnable implements Runnable { + private final Runnable delegate; + private final ContextObject context; + + public ContextPropagatingRunnable(Runnable delegate, ContextObject context) { + this.delegate = delegate; + this.context = context; + } + + @Override + public void run() { + ContextObject previous = Context.get(); + try { + Context.set(context); + delegate.run(); + } finally { + Context.restore(previous); + } + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java new file mode 100644 index 000000000..fb63a3c1b --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java @@ -0,0 +1,33 @@ +package dev.aikido.agent_api.context; + +import java.util.concurrent.Callable; + +public final class ContextPropagation { + private ContextPropagation() {} + + public static Runnable wrap(Runnable task) { + if (task == null || task instanceof ContextPropagatingRunnable) { + return task; + } + + ContextObject context = Context.get(); + if (context == null) { + return task; + } + + return new ContextPropagatingRunnable(task, context.copyForPropagation()); + } + + public static Callable wrap(Callable task) { + if (task == null || task instanceof ContextPropagatingCallable) { + return task; + } + + ContextObject context = Context.get(); + if (context == null) { + return task; + } + + return new ContextPropagatingCallable<>(task, context.copyForPropagation()); + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java index 690ee8753..a7c372bb0 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/RedirectNode.java @@ -36,6 +36,17 @@ public void setChild(RedirectNode child) { this.child = child; } + public RedirectNode copyChain() { + RedirectNode copy = new RedirectNode(this.url); + RedirectNode source = this.child; + RedirectNode tail = copy; + while (source != null) { + tail = new RedirectNode(tail, source.url); + source = source.child; + } + return copy; + } + @Override public boolean equals(Object obj) { if (obj == null || getClass() != obj.getClass()) { diff --git a/agent_api/src/test/java/context/ContextPropagationTest.java b/agent_api/src/test/java/context/ContextPropagationTest.java new file mode 100644 index 000000000..9c0e9ea0f --- /dev/null +++ b/agent_api/src/test/java/context/ContextPropagationTest.java @@ -0,0 +1,328 @@ +package context; + +import dev.aikido.agent_api.context.Context; +import dev.aikido.agent_api.context.ContextObject; +import dev.aikido.agent_api.context.ContextPropagatingCallable; +import dev.aikido.agent_api.context.ContextPropagatingRunnable; +import dev.aikido.agent_api.context.ContextPropagation; +import dev.aikido.agent_api.storage.RedirectNode; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.net.URL; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +class ContextPropagationTest { + @AfterEach + void tearDown() { + Context.reset(); + } + + @Test + void wrapRunnableReturnsNullForNullTask() { + Assertions.assertNull(ContextPropagation.wrap((Runnable) null)); + } + + @Test + void wrapCallableReturnsNullForNullTask() { + Assertions.assertNull(ContextPropagation.wrap((Callable) null)); + } + + @Test + void wrapRunnableReturnsOriginalTaskWhenNoContextIsSet() { + Runnable task = () -> {}; + + Runnable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapCallableReturnsOriginalTaskWhenNoContextIsSet() { + Callable task = () -> "ok"; + + Callable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapRunnableReturnsSameTaskWhenAlreadyWrapped() { + ContextObject contextObject = new ContextObject(); + Runnable task = new ContextPropagatingRunnable(() -> {}, contextObject); + + Runnable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapCallableReturnsSameTaskWhenAlreadyWrapped() { + ContextObject contextObject = new ContextObject(); + Callable task = new ContextPropagatingCallable<>(() -> "ok", contextObject); + + Callable wrapped = ContextPropagation.wrap(task); + + Assertions.assertSame(task, wrapped); + } + + @Test + void wrapRunnableCapturesCurrentContext() { + ContextObject requestContext = new ContextObject(); + requestContext.setRateLimitGroup("req"); + Context.set(requestContext); + + AtomicReference contextDuringRun = new AtomicReference<>(); + Runnable wrapped = ContextPropagation.wrap(() -> contextDuringRun.set(Context.get())); + + Context.reset(); + wrapped.run(); + + Assertions.assertEquals("req", contextDuringRun.get().getRateLimitGroup()); + Assertions.assertNull(Context.get(), "Expected worker context to be cleared after task execution"); + } + + @Test + void wrapCallableCapturesCurrentContext() throws Exception { + ContextObject requestContext = new ContextObject(); + requestContext.setRateLimitGroup("req"); + Context.set(requestContext); + + Callable wrapped = ContextPropagation.wrap(Context::get); + + Context.reset(); + ContextObject contextDuringCall = wrapped.call(); + + Assertions.assertEquals("req", contextDuringCall.getRateLimitGroup()); + Assertions.assertNull(Context.get(), "Expected worker context to be cleared after task execution"); + } + + @Test + void wrapPropagatesAnIsolatedSnapshotNotTheLiveContext() { + ContextObject requestContext = new ContextObject(); + Context.set(requestContext); + + AtomicReference contextDuringRun = new AtomicReference<>(); + Runnable wrapped = ContextPropagation.wrap(() -> contextDuringRun.set(Context.get())); + + Context.reset(); + wrapped.run(); + + ContextObject propagated = contextDuringRun.get(); + Assertions.assertNotSame(requestContext, propagated); + Assertions.assertNotSame(requestContext.getCache(), propagated.getCache()); + } + + @Test + void snapshotDeepCopiesRedirectChainNodes() throws Exception { + RedirectableContext requestContext = new RedirectableContext(); + RedirectNode starter = new RedirectNode(new URL("http://origin")); + new RedirectNode(starter, new URL("http://dest")); + requestContext.addRedirectNode(starter); + + ContextObject copy = requestContext.copyForPropagation(); + + RedirectNode originalStarter = requestContext.getRedirectStartNodes().get(0); + RedirectNode copiedStarter = copy.getRedirectStartNodes().get(0); + + Assertions.assertNotSame(originalStarter, copiedStarter); + Assertions.assertNotSame(originalStarter.getChild(), copiedStarter.getChild()); + Assertions.assertEquals("http://dest", copiedStarter.getChild().getUrl().toString()); + } + + @Test + void snapshotDeepCopiesCacheInnerMaps() { + ContextObject requestContext = new ContextObject(); + Map inner = new HashMap<>(); + inner.put("k", "v"); + requestContext.getCache().put("body", inner); + + ContextObject copy = requestContext.copyForPropagation(); + + Assertions.assertNotSame(requestContext.getCache().get("body"), copy.getCache().get("body")); + Assertions.assertEquals("v", copy.getCache().get("body").get("k")); + } + + @Test + void snapshotCacheMutationDoesNotLeakToOriginal() { + ContextObject requestContext = new ContextObject(); + Map inner = new HashMap<>(); + inner.put("k", "v"); + requestContext.getCache().put("body", inner); + + ContextObject copy = requestContext.copyForPropagation(); + copy.getCache().get("body").put("injected", "x"); + copy.getCache().put("query", new HashMap<>()); + + Assertions.assertFalse(requestContext.getCache().get("body").containsKey("injected")); + Assertions.assertFalse(requestContext.getCache().containsKey("query")); + } + + @Test + void snapshotRedirectChainMutationDoesNotLeakToOriginal() throws Exception { + RedirectableContext requestContext = new RedirectableContext(); + RedirectNode starter = new RedirectNode(new URL("http://origin")); + new RedirectNode(starter, new URL("http://dest")); + requestContext.addRedirectNode(starter); + + ContextObject copy = requestContext.copyForPropagation(); + RedirectNode copiedDest = copy.getRedirectStartNodes().get(0).getChild(); + new RedirectNode(copiedDest, new URL("http://extra")); + + RedirectNode originalDest = requestContext.getRedirectStartNodes().get(0).getChild(); + Assertions.assertNull(originalDest.getChild()); + } + + @Test + void copyChainPreservesMultiNodeChain() throws Exception { + RedirectNode a = new RedirectNode(new URL("http://a")); + RedirectNode b = new RedirectNode(a, new URL("http://b")); + RedirectNode c = new RedirectNode(b, new URL("http://c")); + + RedirectNode copy = a.copyChain(); + + Assertions.assertEquals("http://a", copy.getUrl().toString()); + Assertions.assertEquals("http://b", copy.getChild().getUrl().toString()); + Assertions.assertEquals("http://c", copy.getChild().getChild().getUrl().toString()); + Assertions.assertNotSame(b, copy.getChild()); + Assertions.assertNotSame(c, copy.getChild().getChild()); + } + + private static class RedirectableContext extends ContextObject { + RedirectableContext() { + this.redirectStartNodes = new ArrayList<>(); + } + } + + @Test + void contextPropagatingRunnableRestoresPreviousWorkerContext() { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingRunnable task = new ContextPropagatingRunnable( + () -> Assertions.assertSame(capturedContext, Context.get()), + capturedContext + ); + + Context.set(previousWorkerContext); + task.run(); + + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void contextPropagatingCallableRestoresPreviousWorkerContext() throws Exception { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingCallable task = new ContextPropagatingCallable<>( + Context::get, + capturedContext + ); + + Context.set(previousWorkerContext); + ContextObject contextDuringCall = task.call(); + + Assertions.assertSame(capturedContext, contextDuringCall); + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void contextPropagatingRunnableClearsContextWhenWorkerHadNoPreviousContext() { + ContextObject capturedContext = new ContextObject(); + + ContextPropagatingRunnable task = new ContextPropagatingRunnable( + () -> Assertions.assertSame(capturedContext, Context.get()), + capturedContext + ); + + Context.reset(); + task.run(); + + Assertions.assertNull(Context.get()); + } + + @Test + void contextPropagatingCallableClearsContextWhenWorkerHadNoPreviousContext() throws Exception { + ContextObject capturedContext = new ContextObject(); + + ContextPropagatingCallable task = new ContextPropagatingCallable<>( + Context::get, + capturedContext + ); + + Context.reset(); + ContextObject contextDuringCall = task.call(); + + Assertions.assertSame(capturedContext, contextDuringCall); + Assertions.assertNull(Context.get()); + } + + @Test + void contextPropagatingRunnableRestoresPreviousWorkerContextAfterException() { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingRunnable task = new ContextPropagatingRunnable( + () -> { + throw new IllegalStateException("boom"); + }, + capturedContext + ); + + Context.set(previousWorkerContext); + + Assertions.assertThrows(IllegalStateException.class, task::run); + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void contextPropagatingCallableRestoresPreviousWorkerContextAfterException() { + ContextObject capturedContext = new ContextObject(); + ContextObject previousWorkerContext = new ContextObject(); + + ContextPropagatingCallable task = new ContextPropagatingCallable<>( + () -> { + throw new IllegalStateException("boom"); + }, + capturedContext + ); + + Context.set(previousWorkerContext); + + Assertions.assertThrows(IllegalStateException.class, task::call); + Assertions.assertSame(previousWorkerContext, Context.get()); + } + + @Test + void wrappedRunnableRunsDelegate() { + ContextObject requestContext = new ContextObject(); + Context.set(requestContext); + + AtomicBoolean delegateCalled = new AtomicBoolean(false); + Runnable wrapped = ContextPropagation.wrap(() -> delegateCalled.set(true)); + + Context.reset(); + wrapped.run(); + + Assertions.assertTrue(delegateCalled.get()); + } + + @Test + void wrappedCallableReturnsDelegateResult() throws Exception { + ContextObject requestContext = new ContextObject(); + Context.set(requestContext); + + Callable wrapped = ContextPropagation.wrap(() -> "ok"); + + Context.reset(); + + Assertions.assertEquals("ok", wrapped.call()); + } +}