Skip to content

Commit 75593a3

Browse files
MishenevdMishenevd
committed
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. Wrapping is idempotent and passes through null / already-wrapped / no-context tasks. Unit-tested in isolation without agent weaving.
1 parent 468e2c9 commit 75593a3

4 files changed

Lines changed: 313 additions & 0 deletions

File tree

‎agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingCallable.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package dev.aikido.agent_api.context;
2+
3+
import java.util.concurrent.Callable;
4+
5+
public final class ContextPropagatingCallable<T> implements Callable<T> {
6+
private final Callable<T> delegate;
7+
private final ContextObject context;
8+
9+
public ContextPropagatingCallable(Callable<T> delegate, ContextObject context) {
10+
this.delegate = delegate;
11+
this.context = context;
12+
}
13+
14+
@Override
15+
public T call() throws Exception {
16+
ContextObject previousContext = Context.get();
17+
18+
try {
19+
Context.set(context);
20+
return delegate.call();
21+
} finally {
22+
if (previousContext != null) {
23+
Context.set(previousContext);
24+
} else {
25+
Context.reset();
26+
}
27+
}
28+
}
29+
}

‎agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagatingRunnable.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package dev.aikido.agent_api.context;
2+
3+
public final class ContextPropagatingRunnable implements Runnable {
4+
private final Runnable delegate;
5+
private final ContextObject context;
6+
7+
public ContextPropagatingRunnable(Runnable delegate, ContextObject context) {
8+
this.delegate = delegate;
9+
this.context = context;
10+
}
11+
12+
@Override
13+
public void run() {
14+
ContextObject previousContext = Context.get();
15+
16+
try {
17+
Context.set(context);
18+
delegate.run();
19+
} finally {
20+
if (previousContext != null) {
21+
Context.set(previousContext);
22+
} else {
23+
Context.reset();
24+
}
25+
}
26+
}
27+
}

‎agent_api/src/main/java/dev/aikido/agent_api/context/ContextPropagation.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package dev.aikido.agent_api.context;
2+
3+
import java.util.concurrent.Callable;
4+
5+
public final class ContextPropagation {
6+
private ContextPropagation() {}
7+
8+
public static Runnable wrap(Runnable task) {
9+
if (task == null || task instanceof ContextPropagatingRunnable) {
10+
return task;
11+
}
12+
13+
ContextObject context = Context.get();
14+
if (context == null) {
15+
return task;
16+
}
17+
18+
return new ContextPropagatingRunnable(task, context);
19+
}
20+
21+
public static <T> Callable<T> wrap(Callable<T> task) {
22+
if (task == null || task instanceof ContextPropagatingCallable) {
23+
return task;
24+
}
25+
26+
ContextObject context = Context.get();
27+
if (context == null) {
28+
return task;
29+
}
30+
31+
return new ContextPropagatingCallable<>(task, context);
32+
}
33+
}

‎agent_api/src/test/java/context/ContextPropagationTest.java‎

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
package context;
2+
3+
import dev.aikido.agent_api.context.Context;
4+
import dev.aikido.agent_api.context.ContextObject;
5+
import dev.aikido.agent_api.context.ContextPropagatingCallable;
6+
import dev.aikido.agent_api.context.ContextPropagatingRunnable;
7+
import dev.aikido.agent_api.context.ContextPropagation;
8+
import org.junit.jupiter.api.AfterEach;
9+
import org.junit.jupiter.api.Assertions;
10+
import org.junit.jupiter.api.Test;
11+
12+
import java.util.concurrent.Callable;
13+
import java.util.concurrent.atomic.AtomicBoolean;
14+
import java.util.concurrent.atomic.AtomicReference;
15+
16+
class ContextPropagationTest {
17+
@AfterEach
18+
void tearDown() {
19+
Context.reset();
20+
}
21+
22+
@Test
23+
void wrapRunnableReturnsNullForNullTask() {
24+
Assertions.assertNull(ContextPropagation.wrap((Runnable) null));
25+
}
26+
27+
@Test
28+
void wrapCallableReturnsNullForNullTask() {
29+
Assertions.assertNull(ContextPropagation.wrap((Callable<Object>) null));
30+
}
31+
32+
@Test
33+
void wrapRunnableReturnsOriginalTaskWhenNoContextIsSet() {
34+
Runnable task = () -> {};
35+
36+
Runnable wrapped = ContextPropagation.wrap(task);
37+
38+
Assertions.assertSame(task, wrapped);
39+
}
40+
41+
@Test
42+
void wrapCallableReturnsOriginalTaskWhenNoContextIsSet() {
43+
Callable<String> task = () -> "ok";
44+
45+
Callable<String> wrapped = ContextPropagation.wrap(task);
46+
47+
Assertions.assertSame(task, wrapped);
48+
}
49+
50+
@Test
51+
void wrapRunnableReturnsSameTaskWhenAlreadyWrapped() {
52+
ContextObject contextObject = new ContextObject();
53+
Runnable task = new ContextPropagatingRunnable(() -> {}, contextObject);
54+
55+
Runnable wrapped = ContextPropagation.wrap(task);
56+
57+
Assertions.assertSame(task, wrapped);
58+
}
59+
60+
@Test
61+
void wrapCallableReturnsSameTaskWhenAlreadyWrapped() {
62+
ContextObject contextObject = new ContextObject();
63+
Callable<String> task = new ContextPropagatingCallable<>(() -> "ok", contextObject);
64+
65+
Callable<String> wrapped = ContextPropagation.wrap(task);
66+
67+
Assertions.assertSame(task, wrapped);
68+
}
69+
70+
@Test
71+
void wrapRunnableCapturesCurrentContext() {
72+
ContextObject requestContext = new ContextObject();
73+
Context.set(requestContext);
74+
75+
AtomicReference<ContextObject> contextDuringRun = new AtomicReference<>();
76+
Runnable wrapped = ContextPropagation.wrap(() -> contextDuringRun.set(Context.get()));
77+
78+
Context.reset();
79+
wrapped.run();
80+
81+
Assertions.assertSame(requestContext, contextDuringRun.get());
82+
Assertions.assertNull(Context.get(), "Expected worker context to be cleared after task execution");
83+
}
84+
85+
@Test
86+
void wrapCallableCapturesCurrentContext() throws Exception {
87+
ContextObject requestContext = new ContextObject();
88+
Context.set(requestContext);
89+
90+
Callable<ContextObject> wrapped = ContextPropagation.wrap(Context::get);
91+
92+
Context.reset();
93+
ContextObject contextDuringCall = wrapped.call();
94+
95+
Assertions.assertSame(requestContext, contextDuringCall);
96+
Assertions.assertNull(Context.get(), "Expected worker context to be cleared after task execution");
97+
}
98+
99+
@Test
100+
void contextPropagatingRunnableRestoresPreviousWorkerContext() {
101+
ContextObject capturedContext = new ContextObject();
102+
ContextObject previousWorkerContext = new ContextObject();
103+
104+
ContextPropagatingRunnable task = new ContextPropagatingRunnable(
105+
() -> Assertions.assertSame(capturedContext, Context.get()),
106+
capturedContext
107+
);
108+
109+
Context.set(previousWorkerContext);
110+
task.run();
111+
112+
Assertions.assertSame(previousWorkerContext, Context.get());
113+
}
114+
115+
@Test
116+
void contextPropagatingCallableRestoresPreviousWorkerContext() throws Exception {
117+
ContextObject capturedContext = new ContextObject();
118+
ContextObject previousWorkerContext = new ContextObject();
119+
120+
ContextPropagatingCallable<ContextObject> task = new ContextPropagatingCallable<>(
121+
Context::get,
122+
capturedContext
123+
);
124+
125+
Context.set(previousWorkerContext);
126+
ContextObject contextDuringCall = task.call();
127+
128+
Assertions.assertSame(capturedContext, contextDuringCall);
129+
Assertions.assertSame(previousWorkerContext, Context.get());
130+
}
131+
132+
@Test
133+
void contextPropagatingRunnableClearsContextWhenWorkerHadNoPreviousContext() {
134+
ContextObject capturedContext = new ContextObject();
135+
136+
ContextPropagatingRunnable task = new ContextPropagatingRunnable(
137+
() -> Assertions.assertSame(capturedContext, Context.get()),
138+
capturedContext
139+
);
140+
141+
Context.reset();
142+
task.run();
143+
144+
Assertions.assertNull(Context.get());
145+
}
146+
147+
@Test
148+
void contextPropagatingCallableClearsContextWhenWorkerHadNoPreviousContext() throws Exception {
149+
ContextObject capturedContext = new ContextObject();
150+
151+
ContextPropagatingCallable<ContextObject> task = new ContextPropagatingCallable<>(
152+
Context::get,
153+
capturedContext
154+
);
155+
156+
Context.reset();
157+
ContextObject contextDuringCall = task.call();
158+
159+
Assertions.assertSame(capturedContext, contextDuringCall);
160+
Assertions.assertNull(Context.get());
161+
}
162+
163+
@Test
164+
void contextPropagatingRunnableRestoresPreviousWorkerContextAfterException() {
165+
ContextObject capturedContext = new ContextObject();
166+
ContextObject previousWorkerContext = new ContextObject();
167+
168+
ContextPropagatingRunnable task = new ContextPropagatingRunnable(
169+
() -> {
170+
throw new IllegalStateException("boom");
171+
},
172+
capturedContext
173+
);
174+
175+
Context.set(previousWorkerContext);
176+
177+
Assertions.assertThrows(IllegalStateException.class, task::run);
178+
Assertions.assertSame(previousWorkerContext, Context.get());
179+
}
180+
181+
@Test
182+
void contextPropagatingCallableRestoresPreviousWorkerContextAfterException() {
183+
ContextObject capturedContext = new ContextObject();
184+
ContextObject previousWorkerContext = new ContextObject();
185+
186+
ContextPropagatingCallable<String> task = new ContextPropagatingCallable<>(
187+
() -> {
188+
throw new IllegalStateException("boom");
189+
},
190+
capturedContext
191+
);
192+
193+
Context.set(previousWorkerContext);
194+
195+
Assertions.assertThrows(IllegalStateException.class, task::call);
196+
Assertions.assertSame(previousWorkerContext, Context.get());
197+
}
198+
199+
@Test
200+
void wrappedRunnableRunsDelegate() {
201+
ContextObject requestContext = new ContextObject();
202+
Context.set(requestContext);
203+
204+
AtomicBoolean delegateCalled = new AtomicBoolean(false);
205+
Runnable wrapped = ContextPropagation.wrap(() -> delegateCalled.set(true));
206+
207+
Context.reset();
208+
wrapped.run();
209+
210+
Assertions.assertTrue(delegateCalled.get());
211+
}
212+
213+
@Test
214+
void wrappedCallableReturnsDelegateResult() throws Exception {
215+
ContextObject requestContext = new ContextObject();
216+
Context.set(requestContext);
217+
218+
Callable<String> wrapped = ContextPropagation.wrap(() -> "ok");
219+
220+
Context.reset();
221+
222+
Assertions.assertEquals("ok", wrapped.call());
223+
}
224+
}

0 commit comments

Comments
 (0)