Skip to content

Commit c96af0d

Browse files
lauritlaurit
authored
Make empty agent bridged context equal root context (#3869)
* Make empty agent bridged context equal root context * use ContextStorageWrappers * Use method handle to call ContextStorage.root() * add comment back * Add missing imort for javadoc generation
1 parent 0255e8e commit c96af0d

File tree

5 files changed

+123
-27
lines changed

5 files changed

+123
-27
lines changed

‎instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextInstrumentation.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -35,15 +35,15 @@ public ElementMatcher<TypeDescription> typeMatcher() {
3535
public void transform(TypeTransformer transformer) {
3636
transformer.applyAdviceToMethod(
3737
isMethod().and(isStatic()).and(named("root")),
38-
ContextInstrumentation.class.getName() + "$GetAdvice");
38+
ContextInstrumentation.class.getName() + "$WrapRootAdvice");
3939
}
4040

4141
@SuppressWarnings("unused")
42-
public static class GetAdvice {
42+
public static class WrapRootAdvice {
4343

4444
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
4545
public static void methodExit(@Advice.Return(readOnly = false) Context root) {
46-
root = AgentContextStorage.newContextWrapper(io.opentelemetry.context.Context.root(), root);
46+
root = AgentContextStorage.wrapRootContext(root);
4747
}
4848
}
4949
}

‎instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageInstrumentation.java ‎instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageWrappersInstrumentation.java

+13-14
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,15 @@
55

66
package io.opentelemetry.javaagent.instrumentation.opentelemetryapi;
77

8-
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
9-
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
108
import static net.bytebuddy.matcher.ElementMatchers.named;
119

1210
import application.io.opentelemetry.context.ContextStorage;
1311
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
1412
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
1513
import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage;
14+
import java.util.ArrayList;
15+
import java.util.List;
16+
import java.util.function.Function;
1617
import net.bytebuddy.asm.Advice;
1718
import net.bytebuddy.description.type.TypeDescription;
1819
import net.bytebuddy.matcher.ElementMatcher;
@@ -23,31 +24,29 @@
2324
* sure there is no dependency on a system property or possibility of a user overriding this since
2425
* it's required for instrumentation in the agent to work properly.
2526
*/
26-
public class ContextStorageInstrumentation implements TypeInstrumentation {
27+
public class ContextStorageWrappersInstrumentation implements TypeInstrumentation {
2728

2829
@Override
2930
public ElementMatcher<TypeDescription> typeMatcher() {
30-
return named("application.io.opentelemetry.context.LazyStorage");
31+
return named("application.io.opentelemetry.context.ContextStorageWrappers");
3132
}
3233

3334
@Override
3435
public void transform(TypeTransformer transformer) {
3536
transformer.applyAdviceToMethod(
36-
isMethod().and(isStatic()).and(named("get")),
37-
ContextStorageInstrumentation.class.getName() + "$GetAdvice");
37+
named("getWrappers"),
38+
ContextStorageWrappersInstrumentation.class.getName() + "$AddWrapperAdvice");
3839
}
3940

4041
@SuppressWarnings("unused")
41-
public static class GetAdvice {
42-
43-
@Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class)
44-
public static Object onEnter() {
45-
return null;
46-
}
42+
public static class AddWrapperAdvice {
4743

4844
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
49-
public static void methodExit(@Advice.Return(readOnly = false) ContextStorage storage) {
50-
storage = AgentContextStorage.INSTANCE;
45+
public static void methodExit(
46+
@Advice.Return(readOnly = false)
47+
List<Function<? super ContextStorage, ? extends ContextStorage>> wrappers) {
48+
wrappers = new ArrayList<>(wrappers);
49+
wrappers.add(AgentContextStorage.wrap());
5150
}
5251
}
5352
}

‎instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/OpenTelemetryApiInstrumentationModule.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ public OpenTelemetryApiInstrumentationModule() {
2222
public List<TypeInstrumentation> typeInstrumentations() {
2323
return asList(
2424
new ContextInstrumentation(),
25-
new ContextStorageInstrumentation(),
25+
new ContextStorageWrappersInstrumentation(),
2626
new OpenTelemetryInstrumentation(),
2727
new SpanInstrumentation());
2828
}

‎instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java

+101-9
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@
1313
import application.io.opentelemetry.context.Scope;
1414
import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.baggage.BaggageBridging;
1515
import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.trace.Bridging;
16+
import java.lang.invoke.MethodHandle;
17+
import java.lang.invoke.MethodHandles;
18+
import java.lang.invoke.MethodType;
1619
import java.lang.reflect.Field;
1720
import java.util.function.Function;
1821
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -39,7 +42,80 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
3942

4043
private static final Logger logger = LoggerFactory.getLogger(AgentContextStorage.class);
4144

42-
public static final AgentContextStorage INSTANCE = new AgentContextStorage();
45+
// MethodHandle for ContextStorage.root() that was added in 1.5
46+
private static final MethodHandle CONTEXT_STORAGE_ROOT_HANDLE = getContextStorageRootHandle();
47+
48+
// unwrapped application root context
49+
private final Context applicationRoot;
50+
// wrapped application root context
51+
private final Context root;
52+
53+
private AgentContextStorage(ContextStorage delegate) {
54+
applicationRoot = getRootContext(delegate);
55+
root = getWrappedRootContext(applicationRoot);
56+
}
57+
58+
private static MethodHandle getContextStorageRootHandle() {
59+
try {
60+
return MethodHandles.lookup()
61+
.findVirtual(ContextStorage.class, "root", MethodType.methodType(Context.class));
62+
} catch (NoSuchMethodException | IllegalAccessException exception) {
63+
return null;
64+
}
65+
}
66+
67+
private static boolean has15Api() {
68+
return CONTEXT_STORAGE_ROOT_HANDLE != null;
69+
}
70+
71+
private static Context getRootContext(ContextStorage contextStorage) {
72+
if (has15Api()) {
73+
// when bridging to 1.5 api call ContextStorage.root()
74+
try {
75+
return (Context) CONTEXT_STORAGE_ROOT_HANDLE.invoke(contextStorage);
76+
} catch (Throwable throwable) {
77+
throw new IllegalStateException("Failed to get root context", throwable);
78+
}
79+
} else {
80+
return RootContextHolder.APPLICATION_ROOT;
81+
}
82+
}
83+
84+
private static Context getWrappedRootContext(Context rootContext) {
85+
if (has15Api()) {
86+
return new AgentContextWrapper(io.opentelemetry.context.Context.root(), rootContext);
87+
}
88+
return RootContextHolder.ROOT;
89+
}
90+
91+
public static Context wrapRootContext(Context rootContext) {
92+
if (has15Api()) {
93+
return rootContext;
94+
}
95+
return RootContextHolder.getRootContext(rootContext);
96+
}
97+
98+
// helper class for keeping track of root context when bridging to api earlier than 1.5
99+
private static class RootContextHolder {
100+
// unwrapped application root context
101+
static final Context APPLICATION_ROOT = Context.root();
102+
// wrapped application root context
103+
static final Context ROOT =
104+
new AgentContextWrapper(io.opentelemetry.context.Context.root(), APPLICATION_ROOT);
105+
106+
static Context getRootContext(Context rootContext) {
107+
// APPLICATION_ROOT is null when this method is called while the static initializer is
108+
// initializing the value of APPLICATION_ROOT field
109+
if (RootContextHolder.APPLICATION_ROOT == null) {
110+
return rootContext;
111+
}
112+
return RootContextHolder.ROOT;
113+
}
114+
}
115+
116+
public static Function<? super ContextStorage, ? extends ContextStorage> wrap() {
117+
return contextStorage -> new AgentContextStorage(contextStorage);
118+
}
43119

44120
public static io.opentelemetry.context.Context getAgentContext(Context applicationContext) {
45121
if (applicationContext instanceof AgentContextWrapper) {
@@ -54,6 +130,9 @@ public static io.opentelemetry.context.Context getAgentContext(Context applicati
54130

55131
public static Context newContextWrapper(
56132
io.opentelemetry.context.Context agentContext, Context applicationContext) {
133+
if (applicationContext instanceof AgentContextWrapper) {
134+
applicationContext = ((AgentContextWrapper) applicationContext).applicationContext;
135+
}
57136
return new AgentContextWrapper(agentContext, applicationContext);
58137
}
59138

@@ -80,16 +159,17 @@ public Scope attach(Context toAttach) {
80159
io.opentelemetry.context.Context.current();
81160
Context currentApplicationContext = currentAgentContext.get(APPLICATION_CONTEXT);
82161
if (currentApplicationContext == null) {
83-
currentApplicationContext = Context.root();
84-
}
85-
86-
if (currentApplicationContext == toAttach) {
87-
return Scope.noop();
162+
currentApplicationContext = applicationRoot;
88163
}
89164

90165
io.opentelemetry.context.Context newAgentContext;
91166
if (toAttach instanceof AgentContextWrapper) {
92-
newAgentContext = ((AgentContextWrapper) toAttach).toAgentContext();
167+
AgentContextWrapper wrapper = (AgentContextWrapper) toAttach;
168+
if (currentApplicationContext == wrapper.applicationContext
169+
&& currentAgentContext == wrapper.agentContext) {
170+
return Scope.noop();
171+
}
172+
newAgentContext = wrapper.toAgentContext();
93173
} else {
94174
newAgentContext = currentAgentContext.with(APPLICATION_CONTEXT, toAttach);
95175
}
@@ -102,9 +182,18 @@ public Context current() {
102182
io.opentelemetry.context.Context agentContext = io.opentelemetry.context.Context.current();
103183
Context applicationContext = agentContext.get(APPLICATION_CONTEXT);
104184
if (applicationContext == null) {
105-
applicationContext = Context.root();
185+
applicationContext = applicationRoot;
106186
}
107-
return new AgentContextWrapper(io.opentelemetry.context.Context.current(), applicationContext);
187+
if (applicationContext == applicationRoot
188+
&& agentContext == io.opentelemetry.context.Context.root()) {
189+
return root;
190+
}
191+
return new AgentContextWrapper(agentContext, applicationContext);
192+
}
193+
194+
@Override
195+
public Context root() {
196+
return root;
108197
}
109198

110199
@Override
@@ -121,6 +210,9 @@ private static class AgentContextWrapper implements Context {
121210
final Context applicationContext;
122211

123212
AgentContextWrapper(io.opentelemetry.context.Context agentContext, Context applicationContext) {
213+
if (applicationContext instanceof AgentContextWrapper) {
214+
throw new IllegalStateException("Expected unwrapped context");
215+
}
124216
this.agentContext = agentContext;
125217
this.applicationContext = applicationContext;
126218
}

‎instrumentation/opentelemetry-api-1.0/javaagent/src/test/groovy/ContextBridgeTest.groovy

+5
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,11 @@ class ContextBridgeTest extends AgentInstrumentationSpecification {
140140
ref.get().getEntryValue("cat") == "yes"
141141
}
142142

143+
def "test empty current context is root context"() {
144+
expect:
145+
Context.current() == Context.root()
146+
}
147+
143148
// TODO (trask)
144149
// more tests are needed here, not sure how to implement, probably need to write some test
145150
// instrumentation to help test, similar to :testing-common:integration-tests

0 commit comments

Comments
 (0)