Instrument SpanKey directly (#5933)

* Instrument SpanKey directly

* feedback

* Make muzzle work

* Revert unrelated change

conventions/src/main/kotlin/io.opentelemetry.instrumentation.javaagent-shadowing.gradle.kts

@@ -33,8 +33,9 @@ tasks.withType<ShadowJar>().configureEach {
   relocate("io.opentelemetry.extension.aws", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.aws")
   relocate("io.opentelemetry.extension.kotlin", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.kotlin")
 
-  // this is for instrumentation on opentelemetry-api itself
+  // this is for instrumentation of opentelemetry-api and opentelemetry-instrumentation-api
   relocate("application.io.opentelemetry", "io.opentelemetry")
+  relocate("application.io.opentelemetry.instrumentation.api", "io.opentelemetry.instrumentation.api")
 
   // this is for instrumentation on java.util.logging (since java.util.logging itself is shaded above)
   relocate("application.java.util.logging", "java.util.logging")

gradle-plugins/src/main/kotlin/io.opentelemetry.instrumentation.muzzle-check.gradle.kts

@@ -101,8 +101,9 @@ tasks.withType<ShadowJar>().configureEach {
   relocate("io.opentelemetry.extension.aws", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.aws")
   relocate("io.opentelemetry.extension.kotlin", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.kotlin")
 
-  // this is for instrumentation on opentelemetry-api itself
+  // this is for instrumentation of opentelemetry-api and opentelemetry-instrumentation-api
   relocate("application.io.opentelemetry", "io.opentelemetry")
+  relocate("application.io.opentelemetry.instrumentation.api", "io.opentelemetry.instrumentation.api")
 
   // this is for instrumentation on java.util.logging (since java.util.logging itself is shaded above)
   relocate("application.java.util.logging", "java.util.logging")
@@ -285,7 +286,8 @@ fun addMuzzleTask(muzzleDirective: MuzzleDirective, versionArtifact: Artifact?,
       val instrumentationCL = createInstrumentationClassloader()
       val userCL = createClassLoaderForTask(config)
       withContextClassLoader(instrumentationCL) {
-        MuzzleGradlePluginUtil.assertInstrumentationMuzzled(instrumentationCL, userCL, muzzleDirective.assertPass.get())
+        MuzzleGradlePluginUtil.assertInstrumentationMuzzled(instrumentationCL, userCL,
+          muzzleDirective.excludedInstrumentationModules.get(), muzzleDirective.assertPass.get())
       }
 
       for (thread in Thread.getAllStackTraces().keys) {
@@ -348,6 +350,7 @@ fun inverseOf(muzzleDirective: MuzzleDirective, system: RepositorySystem, sessio
       versions.set(version)
       assertPass.set(!muzzleDirective.assertPass.get())
       excludedDependencies.set(muzzleDirective.excludedDependencies)
+      excludedInstrumentationModules.set(muzzleDirective.excludedInstrumentationModules)
     }
     inverseDirectives.add(inverseDirective)
   }

gradle-plugins/src/main/kotlin/io/opentelemetry/javaagent/muzzle/MuzzleDirective.kt

@@ -20,6 +20,7 @@ abstract class MuzzleDirective {
   abstract val skipVersions: SetProperty<String>
   abstract val additionalDependencies: ListProperty<Any>
   abstract val excludedDependencies: ListProperty<String>
+  abstract val excludedInstrumentationModules: ListProperty<String>
   abstract val assertPass: Property<Boolean>
   abstract val assertInverse: Property<Boolean>
   internal abstract val coreJdk: Property<Boolean> // use coreJdk() function below to enable
@@ -30,6 +31,7 @@ abstract class MuzzleDirective {
     skipVersions.convention(emptySet())
     additionalDependencies.convention(listOf())
     excludedDependencies.convention(listOf())
+    excludedInstrumentationModules.convention(listOf())
     assertPass.convention(false)
     assertInverse.convention(false)
     coreJdk.convention(false)
@@ -58,6 +60,15 @@ abstract class MuzzleDirective {
     excludedDependencies.add(excludeString!!)
   }
 
+  /**
+   * Excludes an instrumentation module from the current muzzle test.
+   *
+   * @param excludeString An instrumentation module class name to exclude
+   */
+  fun excludeInstrumentationModule(excludeString: String) {
+    excludedInstrumentationModules.add(excludeString)
+  }
+
   fun skip(vararg version: String?) {
     skipVersions.addAll(*version)
   }

gradle-plugins/src/main/kotlin/io/opentelemetry/javaagent/muzzle/matcher/MuzzleGradlePluginUtil.kt

@@ -52,7 +52,9 @@ class MuzzleGradlePluginUtil {
      * version passes different {@code userClassLoader}.
      */
     @Suppress("UNCHECKED_CAST")
-    fun assertInstrumentationMuzzled(agentClassLoader: ClassLoader, userClassLoader: ClassLoader, assertPass: Boolean) {
+    fun assertInstrumentationMuzzled(agentClassLoader: ClassLoader, userClassLoader: ClassLoader,
+                                     excludedInstrumentationModules: List<String>, assertPass: Boolean) {
+
       val matcherClass = agentClassLoader.loadClass("io.opentelemetry.javaagent.tooling.muzzle.ClassLoaderMatcher")
 
       // We cannot reference Mismatch class directly here, because we are loaded from a different
@@ -65,7 +67,9 @@ class MuzzleGradlePluginUtil {
         .invoke(null, userClassLoader, assertPass)
         as Map<String, List<Any>>
 
-      allMismatches.forEach { moduleName, mismatches ->
+      allMismatches.filterKeys {
+        !excludedInstrumentationModules.contains(it)
+      }.forEach { (moduleName, mismatches) ->
         val passed = mismatches.isEmpty()
         if (passed && !assertPass) {
           System.err.println("MUZZLE PASSED $moduleName BUT FAILURE WAS EXPECTED")

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

@@ -11,12 +11,10 @@ import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.function.Function;
 import javax.annotation.Nullable;
 
-@SuppressWarnings({"unchecked", "rawtypes"})
 final class InstrumentationApiContextBridging {
 
   static List<ContextKeyBridge<?, ?>> instrumentationApiBridges() {
@@ -48,30 +46,6 @@ final class InstrumentationApiContextBridging {
       // no old instrumentation-api on classpath
     }
 
-    List<String> spanKeyNames =
-        Arrays.asList(
-            // span kind keys
-            "KIND_SERVER_KEY",
-            "KIND_CLIENT_KEY",
-            "KIND_CONSUMER_KEY",
-            "KIND_PRODUCER_KEY",
-            // semantic convention keys
-            "HTTP_SERVER_KEY",
-            "RPC_SERVER_KEY",
-            "HTTP_CLIENT_KEY",
-            "RPC_CLIENT_KEY",
-            "DB_CLIENT_KEY",
-            "PRODUCER_KEY",
-            "CONSUMER_RECEIVE_KEY",
-            "CONSUMER_PROCESS_KEY");
-
-    for (String spanKeyName : spanKeyNames) {
-      ContextKeyBridge<?, ?> spanKeyBridge = spanKeyBridge(spanKeyName);
-      if (spanKeyBridge != null) {
-        bridges.add(spanKeyBridge);
-      }
-    }
-
     ContextKeyBridge<?, ?> httpRouteHolderBridge = httpRouteStateBridge();
     if (httpRouteHolderBridge != null) {
       bridges.add(httpRouteHolderBridge);
@@ -80,22 +54,6 @@ final class InstrumentationApiContextBridging {
     return bridges;
   }
 
-  @Nullable
-  private static ContextKeyBridge<Span, io.opentelemetry.api.trace.Span> spanKeyBridge(
-      String name) {
-    try {
-      return new ContextKeyBridge<>(
-          "application.io.opentelemetry.instrumentation.api.internal.SpanKey",
-          "io.opentelemetry.instrumentation.api.internal.SpanKey",
-          name,
-          Bridging::toApplication,
-          Bridging::toAgentOrNull);
-    } catch (Throwable e) {
-      // instrumentation-api may be absent on the classpath, just skip
-      return null;
-    }
-  }
-
   private static final Class<?> AGENT_HTTP_ROUTE_STATE;
   private static final MethodHandle AGENT_CREATE;
   private static final MethodHandle AGENT_GET_UPDATED_BY_SOURCE_ORDER;

instrumentation/opentelemetry-instrumentation-api/javaagent/build.gradle.kts

@@ -2,10 +2,27 @@ plugins {
   id("otel.javaagent-instrumentation")
 }
 
+// note that muzzle is not run against the current SNAPSHOT instrumentation-api, but this is ok
+// because the tests are run against the current SNAPSHOT instrumentation-api which will catch any
+// muzzle issues in SNAPSHOT instrumentation-api
+
+muzzle {
+  pass {
+    group.set("io.opentelemetry.instrumentation")
+    module.set("opentelemetry-instrumentation-api")
+    // currently all 1.0.0+ versions are alpha so they are all skipped
+    // if you want to test them anyways, comment out "alpha" from the exclusions in AcceptableVersions.kt
+    versions.set("[1.14.0-alpha,)")
+    assertInverse.set(true)
+    excludeInstrumentationModule("io.opentelemetry.javaagent.instrumentation.opentelemetryapi.OpenTelemetryApiInstrumentationModule")
+  }
+}
+
 dependencies {
   implementation(project(":instrumentation:opentelemetry-api:opentelemetry-api-1.0:javaagent"))
 
   compileOnly(project(":opentelemetry-api-shaded-for-instrumenting", configuration = "shadow"))
+  compileOnly(project(":opentelemetry-instrumentation-api-shaded-for-instrumenting", configuration = "shadow"))
 
   testImplementation(project(":instrumentation-api-semconv"))
   testImplementation(project(":instrumentation:opentelemetry-instrumentation-api:testing"))
@@ -25,6 +42,16 @@ testing {
 }
 
 configurations.configureEach {
+  if (name.startsWith("muzzle-Assert")) {
+    // some names also start with "muzzle-AssertFail", which is conveniently the same length
+    val ver = name.substring("muzzle-AssertPass-io.opentelemetry.instrumentation-opentelemetry-instrumentation-api-".length)
+    resolutionStrategy {
+      dependencySubstitution {
+        substitute(module("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api"))
+          .using(module("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api:$ver"))
+      }
+    }
+  }
   if (name.startsWith("testOldServerSpan")) {
     resolutionStrategy {
       dependencySubstitution {

instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/InstrumentationApiInstrumentationModule.java

@@ -6,7 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.instrumentationapi;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -23,12 +23,11 @@ public class InstrumentationApiInstrumentationModule extends InstrumentationModu
 
   @Override
   public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
-    return hasClassesNamed(
-        "application.io.opentelemetry.instrumentation.api.internal.HttpRouteState");
+    return hasClassesNamed("application.io.opentelemetry.instrumentation.api.internal.SpanKey");
   }
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new HttpRouteStateInstrumentation());
+    return asList(new HttpRouteStateInstrumentation(), new SpanKeyInstrumentation());
   }
 }

instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/SpanKeyBridging.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.instrumentationapi;
+
+import application.io.opentelemetry.instrumentation.api.internal.SpanKey;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+
+public final class SpanKeyBridging {
+
+  private static final Map<SpanKey, io.opentelemetry.instrumentation.api.internal.SpanKey>
+      agentSpanKeys = createMapping();
+
+  private static Map<SpanKey, io.opentelemetry.instrumentation.api.internal.SpanKey>
+      createMapping() {
+
+    Map<SpanKey, io.opentelemetry.instrumentation.api.internal.SpanKey> map = new HashMap<>();
+    map.put(SpanKey.KIND_SERVER, io.opentelemetry.instrumentation.api.internal.SpanKey.KIND_SERVER);
+    map.put(SpanKey.KIND_CLIENT, io.opentelemetry.instrumentation.api.internal.SpanKey.KIND_CLIENT);
+    map.put(
+        SpanKey.KIND_CONSUMER, io.opentelemetry.instrumentation.api.internal.SpanKey.KIND_CONSUMER);
+    map.put(
+        SpanKey.KIND_PRODUCER, io.opentelemetry.instrumentation.api.internal.SpanKey.KIND_PRODUCER);
+
+    map.put(SpanKey.HTTP_SERVER, io.opentelemetry.instrumentation.api.internal.SpanKey.HTTP_SERVER);
+    map.put(SpanKey.RPC_SERVER, io.opentelemetry.instrumentation.api.internal.SpanKey.RPC_SERVER);
+
+    map.put(SpanKey.HTTP_CLIENT, io.opentelemetry.instrumentation.api.internal.SpanKey.HTTP_CLIENT);
+    map.put(SpanKey.RPC_CLIENT, io.opentelemetry.instrumentation.api.internal.SpanKey.RPC_CLIENT);
+    map.put(SpanKey.DB_CLIENT, io.opentelemetry.instrumentation.api.internal.SpanKey.DB_CLIENT);
+
+    map.put(SpanKey.PRODUCER, io.opentelemetry.instrumentation.api.internal.SpanKey.PRODUCER);
+    map.put(
+        SpanKey.CONSUMER_RECEIVE,
+        io.opentelemetry.instrumentation.api.internal.SpanKey.CONSUMER_RECEIVE);
+    map.put(
+        SpanKey.CONSUMER_PROCESS,
+        io.opentelemetry.instrumentation.api.internal.SpanKey.CONSUMER_PROCESS);
+    return map;
+  }
+
+  @Nullable
+  public static io.opentelemetry.instrumentation.api.internal.SpanKey toAgentOrNull(
+      SpanKey applicationSpanKey) {
+    return agentSpanKeys.get(applicationSpanKey);
+  }
+
+  private SpanKeyBridging() {}
+}

instrumentation/opentelemetry-instrumentation-api/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/instrumentationapi/SpanKeyInstrumentation.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.instrumentationapi;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import application.io.opentelemetry.api.trace.Span;
+import application.io.opentelemetry.context.Context;
+import application.io.opentelemetry.instrumentation.api.internal.SpanKey;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage;
+import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.trace.Bridging;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+final class SpanKeyInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("application.io.opentelemetry.instrumentation.api.internal.SpanKey");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("storeInContext")
+            .and(takesArgument(0, named("application.io.opentelemetry.context.Context")))
+            .and(takesArgument(1, named("application.io.opentelemetry.api.trace.Span"))),
+        this.getClass().getName() + "$StoreInContextAdvice");
+    transformer.applyAdviceToMethod(
+        named("fromContextOrNull")
+            .and(takesArgument(0, named("application.io.opentelemetry.context.Context"))),
+        this.getClass().getName() + "$FromContextOrNullAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class StoreInContextAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class)
+    public static Object onEnter() {
+      return null;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This SpanKey applicationSpanKey,
+        @Advice.Argument(0) Context applicationContext,
+        @Advice.Argument(1) Span applicationSpan,
+        @Advice.Return(readOnly = false) Context newApplicationContext) {
+
+      io.opentelemetry.instrumentation.api.internal.SpanKey agentSpanKey =
+          SpanKeyBridging.toAgentOrNull(applicationSpanKey);
+      if (agentSpanKey == null) {
+        return;
+      }
+
+      io.opentelemetry.context.Context agentContext =
+          AgentContextStorage.getAgentContext(applicationContext);
+
+      io.opentelemetry.api.trace.Span agentSpan = Bridging.toAgentOrNull(applicationSpan);
+      if (agentSpan == null) {
+        return;
+      }
+
+      io.opentelemetry.context.Context newAgentContext =
+          agentSpanKey.storeInContext(agentContext, agentSpan);
+
+      newApplicationContext = AgentContextStorage.toApplicationContext(newAgentContext);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class FromContextOrNullAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class)
+    public static Object onEnter() {
+      return null;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This SpanKey applicationSpanKey,
+        @Advice.Argument(0) Context applicationContext,
+        @Advice.Return(readOnly = false) Span applicationSpan) {
+
+      io.opentelemetry.instrumentation.api.internal.SpanKey agentSpanKey =
+          SpanKeyBridging.toAgentOrNull(applicationSpanKey);
+      if (agentSpanKey == null) {
+        return;
+      }
+
+      io.opentelemetry.context.Context agentContext =
+          AgentContextStorage.getAgentContext(applicationContext);
+
+      io.opentelemetry.api.trace.Span agentSpan = agentSpanKey.fromContextOrNull(agentContext);
+
+      applicationSpan = agentSpan == null ? null : Bridging.toApplication(agentSpan);
+    }
+  }
+}

javaagent/build.gradle.kts

@@ -67,6 +67,7 @@ dependencies {
   baseJavaagentLibs(project(":instrumentation:opentelemetry-annotations-1.0:javaagent"))
   baseJavaagentLibs(project(":instrumentation:opentelemetry-api:opentelemetry-api-1.0:javaagent"))
   baseJavaagentLibs(project(":instrumentation:opentelemetry-api:opentelemetry-api-1.4:javaagent"))
+  baseJavaagentLibs(project(":instrumentation:opentelemetry-instrumentation-api:javaagent"))
   baseJavaagentLibs(project(":instrumentation:executors:javaagent"))
   baseJavaagentLibs(project(":instrumentation:internal:internal-class-loader:javaagent"))
   baseJavaagentLibs(project(":instrumentation:internal:internal-eclipse-osgi-3.6:javaagent"))

opentelemetry-instrumentation-api-shaded-for-instrumenting/build.gradle.kts

@@ -0,0 +1,22 @@
+plugins {
+  id("com.github.johnrengelman.shadow")
+
+  id("otel.java-conventions")
+}
+
+description = "opentelemetry-instrumentation-api shaded for internal javaagent usage"
+group = "io.opentelemetry.javaagent"
+
+dependencies {
+  implementation("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api")
+}
+
+// OpenTelemetry Instrumentation API shaded so that it can be used in instrumentation of
+// OpenTelemetry Instrumentation API itself,
+// and then its usage can be unshaded after OpenTelemetry Instrumentation API is shaded
+// (see more explanation in opentelemetry-api-1.0.gradle)
+tasks {
+  shadowJar {
+    relocate("io.opentelemetry.instrumentation.api", "application.io.opentelemetry.instrumentation.api")
+  }
+}

settings.gradle.kts

@@ -92,6 +92,7 @@ include(":muzzle")
 // agent projects
 include(":opentelemetry-api-shaded-for-instrumenting")
 include(":opentelemetry-ext-annotations-shaded-for-instrumenting")
+include(":opentelemetry-instrumentation-api-shaded-for-instrumenting")
 include(":javaagent-bootstrap")
 include(":javaagent-extension-api")
 include(":javaagent-tooling")