From c96af0d51a16c5ee878a3c3ac566e15b4ee59949 Mon Sep 17 00:00:00 2001 From: Lauri Tulmin Date: Fri, 20 Aug 2021 23:51:29 +0300 Subject: [PATCH] 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 --- .../ContextInstrumentation.java | 6 +- ...ontextStorageWrappersInstrumentation.java} | 27 +++-- ...OpenTelemetryApiInstrumentationModule.java | 2 +- .../context/AgentContextStorage.java | 110 ++++++++++++++++-- .../src/test/groovy/ContextBridgeTest.groovy | 5 + 5 files changed, 123 insertions(+), 27 deletions(-) rename instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/{ContextStorageInstrumentation.java => ContextStorageWrappersInstrumentation.java} (65%) diff --git a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextInstrumentation.java b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextInstrumentation.java index a54f684b3e..c4beb21b9d 100644 --- a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextInstrumentation.java +++ b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextInstrumentation.java @@ -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); } } } diff --git a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageInstrumentation.java b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageWrappersInstrumentation.java similarity index 65% rename from instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageInstrumentation.java rename to instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageWrappersInstrumentation.java index ae8bdfbd80..280cc192df 100644 --- a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageInstrumentation.java +++ b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/ContextStorageWrappersInstrumentation.java @@ -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 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> wrappers) { + wrappers = new ArrayList<>(wrappers); + wrappers.add(AgentContextStorage.wrap()); } } } diff --git a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/OpenTelemetryApiInstrumentationModule.java b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/OpenTelemetryApiInstrumentationModule.java index 8ee47b61cb..10e370c77f 100644 --- a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/OpenTelemetryApiInstrumentationModule.java +++ b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/OpenTelemetryApiInstrumentationModule.java @@ -22,7 +22,7 @@ public class OpenTelemetryApiInstrumentationModule extends InstrumentationModule public List typeInstrumentations() { return asList( new ContextInstrumentation(), - new ContextStorageInstrumentation(), + new ContextStorageWrappersInstrumentation(), new OpenTelemetryInstrumentation(), new SpanInstrumentation()); } diff --git a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java index d0bac8a57e..696d660619 100644 --- a/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java +++ b/instrumentation/opentelemetry-api-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/context/AgentContextStorage.java @@ -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 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; } diff --git a/instrumentation/opentelemetry-api-1.0/javaagent/src/test/groovy/ContextBridgeTest.groovy b/instrumentation/opentelemetry-api-1.0/javaagent/src/test/groovy/ContextBridgeTest.groovy index a609bfd04d..375517a9bc 100644 --- a/instrumentation/opentelemetry-api-1.0/javaagent/src/test/groovy/ContextBridgeTest.groovy +++ b/instrumentation/opentelemetry-api-1.0/javaagent/src/test/groovy/ContextBridgeTest.groovy @@ -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