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

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

@@ -35,15 +35,15 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   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);
     }
   }
 }

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

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

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

@@ -22,7 +22,7 @@ public OpenTelemetryApiInstrumentationModule() {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
         new ContextInstrumentation(),
-        new ContextStorageInstrumentation(),
+        new ContextStorageWrappersInstrumentation(),
         new OpenTelemetryInstrumentation(),
         new SpanInstrumentation());
   }

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.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 static io.opentelemetry.context.Context getAgentContext(Context applicati
 
   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 Scope attach(Context toAttach) {
         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 Context current() {
     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 @@ private static class AgentContextWrapper implements Context {
     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;
     }

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