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:
Lauri Tulmin 2021-08-20 23:51:29 +03:00 committed by GitHub
parent 0255e8e95f
commit c96af0d51a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 123 additions and 27 deletions

View File

@ -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);
}
}
}

View File

@ -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());
}
}
}

View File

@ -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());
}

View File

@ -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;
}

View File

@ -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