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
This commit is contained in:
parent
0255e8e95f
commit
c96af0d51a
|
@ -35,15 +35,15 @@ public class ContextInstrumentation implements TypeInstrumentation {
|
|||
public void transform(TypeTransformer transformer) {
|
||||
transformer.applyAdviceToMethod(
|
||||
isMethod().and(isStatic()).and(named("root")),
|
||||
ContextInstrumentation.class.getName() + "$GetAdvice");
|
||||
ContextInstrumentation.class.getName() + "$WrapRootAdvice");
|
||||
}
|
||||
|
||||
@SuppressWarnings("unused")
|
||||
public static class GetAdvice {
|
||||
public static class WrapRootAdvice {
|
||||
|
||||
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
|
||||
public static void methodExit(@Advice.Return(readOnly = false) Context root) {
|
||||
root = AgentContextStorage.newContextWrapper(io.opentelemetry.context.Context.root(), root);
|
||||
root = AgentContextStorage.wrapRootContext(root);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -5,14 +5,15 @@
|
|||
|
||||
package io.opentelemetry.javaagent.instrumentation.opentelemetryapi;
|
||||
|
||||
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
|
||||
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
|
||||
import static net.bytebuddy.matcher.ElementMatchers.named;
|
||||
|
||||
import application.io.opentelemetry.context.ContextStorage;
|
||||
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
|
||||
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
|
||||
import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.function.Function;
|
||||
import net.bytebuddy.asm.Advice;
|
||||
import net.bytebuddy.description.type.TypeDescription;
|
||||
import net.bytebuddy.matcher.ElementMatcher;
|
||||
|
@ -23,31 +24,29 @@ import net.bytebuddy.matcher.ElementMatcher;
|
|||
* sure there is no dependency on a system property or possibility of a user overriding this since
|
||||
* it's required for instrumentation in the agent to work properly.
|
||||
*/
|
||||
public class ContextStorageInstrumentation implements TypeInstrumentation {
|
||||
public class ContextStorageWrappersInstrumentation implements TypeInstrumentation {
|
||||
|
||||
@Override
|
||||
public ElementMatcher<TypeDescription> typeMatcher() {
|
||||
return named("application.io.opentelemetry.context.LazyStorage");
|
||||
return named("application.io.opentelemetry.context.ContextStorageWrappers");
|
||||
}
|
||||
|
||||
@Override
|
||||
public void transform(TypeTransformer transformer) {
|
||||
transformer.applyAdviceToMethod(
|
||||
isMethod().and(isStatic()).and(named("get")),
|
||||
ContextStorageInstrumentation.class.getName() + "$GetAdvice");
|
||||
named("getWrappers"),
|
||||
ContextStorageWrappersInstrumentation.class.getName() + "$AddWrapperAdvice");
|
||||
}
|
||||
|
||||
@SuppressWarnings("unused")
|
||||
public static class GetAdvice {
|
||||
|
||||
@Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class)
|
||||
public static Object onEnter() {
|
||||
return null;
|
||||
}
|
||||
public static class AddWrapperAdvice {
|
||||
|
||||
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
|
||||
public static void methodExit(@Advice.Return(readOnly = false) ContextStorage storage) {
|
||||
storage = AgentContextStorage.INSTANCE;
|
||||
public static void methodExit(
|
||||
@Advice.Return(readOnly = false)
|
||||
List<Function<? super ContextStorage, ? extends ContextStorage>> wrappers) {
|
||||
wrappers = new ArrayList<>(wrappers);
|
||||
wrappers.add(AgentContextStorage.wrap());
|
||||
}
|
||||
}
|
||||
}
|
|
@ -22,7 +22,7 @@ public class OpenTelemetryApiInstrumentationModule extends InstrumentationModule
|
|||
public List<TypeInstrumentation> typeInstrumentations() {
|
||||
return asList(
|
||||
new ContextInstrumentation(),
|
||||
new ContextStorageInstrumentation(),
|
||||
new ContextStorageWrappersInstrumentation(),
|
||||
new OpenTelemetryInstrumentation(),
|
||||
new SpanInstrumentation());
|
||||
}
|
||||
|
|
|
@ -13,6 +13,9 @@ import application.io.opentelemetry.context.ContextStorage;
|
|||
import application.io.opentelemetry.context.Scope;
|
||||
import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.baggage.BaggageBridging;
|
||||
import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.trace.Bridging;
|
||||
import java.lang.invoke.MethodHandle;
|
||||
import java.lang.invoke.MethodHandles;
|
||||
import java.lang.invoke.MethodType;
|
||||
import java.lang.reflect.Field;
|
||||
import java.util.function.Function;
|
||||
import org.checkerframework.checker.nullness.qual.Nullable;
|
||||
|
@ -39,7 +42,80 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
|
|||
|
||||
private static final Logger logger = LoggerFactory.getLogger(AgentContextStorage.class);
|
||||
|
||||
public static final AgentContextStorage INSTANCE = new AgentContextStorage();
|
||||
// MethodHandle for ContextStorage.root() that was added in 1.5
|
||||
private static final MethodHandle CONTEXT_STORAGE_ROOT_HANDLE = getContextStorageRootHandle();
|
||||
|
||||
// unwrapped application root context
|
||||
private final Context applicationRoot;
|
||||
// wrapped application root context
|
||||
private final Context root;
|
||||
|
||||
private AgentContextStorage(ContextStorage delegate) {
|
||||
applicationRoot = getRootContext(delegate);
|
||||
root = getWrappedRootContext(applicationRoot);
|
||||
}
|
||||
|
||||
private static MethodHandle getContextStorageRootHandle() {
|
||||
try {
|
||||
return MethodHandles.lookup()
|
||||
.findVirtual(ContextStorage.class, "root", MethodType.methodType(Context.class));
|
||||
} catch (NoSuchMethodException | IllegalAccessException exception) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private static boolean has15Api() {
|
||||
return CONTEXT_STORAGE_ROOT_HANDLE != null;
|
||||
}
|
||||
|
||||
private static Context getRootContext(ContextStorage contextStorage) {
|
||||
if (has15Api()) {
|
||||
// when bridging to 1.5 api call ContextStorage.root()
|
||||
try {
|
||||
return (Context) CONTEXT_STORAGE_ROOT_HANDLE.invoke(contextStorage);
|
||||
} catch (Throwable throwable) {
|
||||
throw new IllegalStateException("Failed to get root context", throwable);
|
||||
}
|
||||
} else {
|
||||
return RootContextHolder.APPLICATION_ROOT;
|
||||
}
|
||||
}
|
||||
|
||||
private static Context getWrappedRootContext(Context rootContext) {
|
||||
if (has15Api()) {
|
||||
return new AgentContextWrapper(io.opentelemetry.context.Context.root(), rootContext);
|
||||
}
|
||||
return RootContextHolder.ROOT;
|
||||
}
|
||||
|
||||
public static Context wrapRootContext(Context rootContext) {
|
||||
if (has15Api()) {
|
||||
return rootContext;
|
||||
}
|
||||
return RootContextHolder.getRootContext(rootContext);
|
||||
}
|
||||
|
||||
// helper class for keeping track of root context when bridging to api earlier than 1.5
|
||||
private static class RootContextHolder {
|
||||
// unwrapped application root context
|
||||
static final Context APPLICATION_ROOT = Context.root();
|
||||
// wrapped application root context
|
||||
static final Context ROOT =
|
||||
new AgentContextWrapper(io.opentelemetry.context.Context.root(), APPLICATION_ROOT);
|
||||
|
||||
static Context getRootContext(Context rootContext) {
|
||||
// APPLICATION_ROOT is null when this method is called while the static initializer is
|
||||
// initializing the value of APPLICATION_ROOT field
|
||||
if (RootContextHolder.APPLICATION_ROOT == null) {
|
||||
return rootContext;
|
||||
}
|
||||
return RootContextHolder.ROOT;
|
||||
}
|
||||
}
|
||||
|
||||
public static Function<? super ContextStorage, ? extends ContextStorage> wrap() {
|
||||
return contextStorage -> new AgentContextStorage(contextStorage);
|
||||
}
|
||||
|
||||
public static io.opentelemetry.context.Context getAgentContext(Context applicationContext) {
|
||||
if (applicationContext instanceof AgentContextWrapper) {
|
||||
|
@ -54,6 +130,9 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
|
|||
|
||||
public static Context newContextWrapper(
|
||||
io.opentelemetry.context.Context agentContext, Context applicationContext) {
|
||||
if (applicationContext instanceof AgentContextWrapper) {
|
||||
applicationContext = ((AgentContextWrapper) applicationContext).applicationContext;
|
||||
}
|
||||
return new AgentContextWrapper(agentContext, applicationContext);
|
||||
}
|
||||
|
||||
|
@ -80,16 +159,17 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
|
|||
io.opentelemetry.context.Context.current();
|
||||
Context currentApplicationContext = currentAgentContext.get(APPLICATION_CONTEXT);
|
||||
if (currentApplicationContext == null) {
|
||||
currentApplicationContext = Context.root();
|
||||
}
|
||||
|
||||
if (currentApplicationContext == toAttach) {
|
||||
return Scope.noop();
|
||||
currentApplicationContext = applicationRoot;
|
||||
}
|
||||
|
||||
io.opentelemetry.context.Context newAgentContext;
|
||||
if (toAttach instanceof AgentContextWrapper) {
|
||||
newAgentContext = ((AgentContextWrapper) toAttach).toAgentContext();
|
||||
AgentContextWrapper wrapper = (AgentContextWrapper) toAttach;
|
||||
if (currentApplicationContext == wrapper.applicationContext
|
||||
&& currentAgentContext == wrapper.agentContext) {
|
||||
return Scope.noop();
|
||||
}
|
||||
newAgentContext = wrapper.toAgentContext();
|
||||
} else {
|
||||
newAgentContext = currentAgentContext.with(APPLICATION_CONTEXT, toAttach);
|
||||
}
|
||||
|
@ -102,9 +182,18 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
|
|||
io.opentelemetry.context.Context agentContext = io.opentelemetry.context.Context.current();
|
||||
Context applicationContext = agentContext.get(APPLICATION_CONTEXT);
|
||||
if (applicationContext == null) {
|
||||
applicationContext = Context.root();
|
||||
applicationContext = applicationRoot;
|
||||
}
|
||||
return new AgentContextWrapper(io.opentelemetry.context.Context.current(), applicationContext);
|
||||
if (applicationContext == applicationRoot
|
||||
&& agentContext == io.opentelemetry.context.Context.root()) {
|
||||
return root;
|
||||
}
|
||||
return new AgentContextWrapper(agentContext, applicationContext);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Context root() {
|
||||
return root;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
@ -121,6 +210,9 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
|
|||
final Context applicationContext;
|
||||
|
||||
AgentContextWrapper(io.opentelemetry.context.Context agentContext, Context applicationContext) {
|
||||
if (applicationContext instanceof AgentContextWrapper) {
|
||||
throw new IllegalStateException("Expected unwrapped context");
|
||||
}
|
||||
this.agentContext = agentContext;
|
||||
this.applicationContext = applicationContext;
|
||||
}
|
||||
|
|
|
@ -140,6 +140,11 @@ class ContextBridgeTest extends AgentInstrumentationSpecification {
|
|||
ref.get().getEntryValue("cat") == "yes"
|
||||
}
|
||||
|
||||
def "test empty current context is root context"() {
|
||||
expect:
|
||||
Context.current() == Context.root()
|
||||
}
|
||||
|
||||
// TODO (trask)
|
||||
// more tests are needed here, not sure how to implement, probably need to write some test
|
||||
// instrumentation to help test, similar to :testing-common:integration-tests
|
||||
|
|
Loading…
Reference in New Issue