Bridge span keys defined in instrumentation api

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/SpanKey.java

@@ -18,7 +18,6 @@ public final class SpanKey {
   private static final ContextKey<Span> CONSUMER_KEY =
       ContextKey.named("opentelemetry-traces-span-key-consumer");
 
-  // TODO bridge these constants in AgentContextStorage
   private static final ContextKey<Span> HTTP_KEY =
       ContextKey.named("opentelemetry-traces-span-key-http");
   private static final ContextKey<Span> RPC_KEY =

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

@@ -28,8 +28,12 @@ dependencies {
   // using OpenTelemetry SDK to make sure that instrumentation doesn't cause
   // OpenTelemetrySdk.getTracerProvider() to throw ClassCastException
   testImplementation("io.opentelemetry:opentelemetry-sdk")
+  testImplementation(project(":instrumentation-api"))
 
   // @WithSpan annotation is used to generate spans in ContextBridgeTest
   testImplementation("io.opentelemetry:opentelemetry-extension-annotations")
   testInstrumentation(project(":instrumentation:opentelemetry-annotations-1.0:javaagent"))
+
+  testImplementation(project(":instrumentation:opentelemetry-api-1.0:testing"))
+  testInstrumentation(project(":instrumentation:opentelemetry-api-1.0:testing"))
 }

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

@@ -151,8 +151,26 @@ public static Context newContextWrapper(
             "io.opentelemetry.api.baggage.BaggageContextKey",
             BaggageBridging::toApplication,
             BaggageBridging::toAgent),
+        bridgeSpanKey("SERVER_KEY"),
+        bridgeSpanKey("CONSUMER_KEY"),
+        bridgeSpanKey("HTTP_KEY"),
+        bridgeSpanKey("RPC_KEY"),
+        bridgeSpanKey("DB_KEY"),
+        bridgeSpanKey("MESSAGING_KEY"),
+        bridgeSpanKey("CLIENT_KEY"),
+        bridgeSpanKey("PRODUCER_KEY"),
       };
 
+  private static ContextKeyBridge<Span, io.opentelemetry.api.trace.Span> bridgeSpanKey(
+      String name) {
+    return new ContextKeyBridge<>(
+        "application.io.opentelemetry.instrumentation.api.instrumenter.SpanKey",
+        "io.opentelemetry.instrumentation.api.instrumenter.SpanKey",
+        name,
+        Bridging::toApplication,
+        Bridging::toAgentOrNull);
+  }
+
   @Override
   public Scope attach(Context toAttach) {
     io.opentelemetry.context.Context currentAgentContext =
@@ -266,13 +284,22 @@ static class ContextKeyBridge<APPLICATION, AGENT> {
         String agentKeyHolderClassName,
         Function<AGENT, APPLICATION> toApplication,
         Function<APPLICATION, AGENT> toAgent) {
+      this(applicationKeyHolderClassName, agentKeyHolderClassName, "KEY", toApplication, toAgent);
+    }
+
+    ContextKeyBridge(
+        String applicationKeyHolderClassName,
+        String agentKeyHolderClassName,
+        String fieldName,
+        Function<AGENT, APPLICATION> toApplication,
+        Function<APPLICATION, AGENT> toAgent) {
       this.toApplication = toApplication;
       this.toAgent = toAgent;
 
       ContextKey<APPLICATION> applicationContextKey;
       try {
         Class<?> applicationKeyHolderClass = Class.forName(applicationKeyHolderClassName);
-        Field applicationContextKeyField = applicationKeyHolderClass.getDeclaredField("KEY");
+        Field applicationContextKeyField = applicationKeyHolderClass.getDeclaredField(fieldName);
         applicationContextKeyField.setAccessible(true);
         applicationContextKey = (ContextKey<APPLICATION>) applicationContextKeyField.get(null);
       } catch (Throwable t) {
@@ -283,7 +310,7 @@ static class ContextKeyBridge<APPLICATION, AGENT> {
       io.opentelemetry.context.ContextKey<AGENT> agentContextKey;
       try {
         Class<?> agentKeyHolderClass = Class.forName(agentKeyHolderClassName);
-        Field agentContextKeyField = agentKeyHolderClass.getDeclaredField("KEY");
+        Field agentContextKeyField = agentKeyHolderClass.getDeclaredField(fieldName);
         agentContextKeyField.setAccessible(true);
         agentContextKey =
             (io.opentelemetry.context.ContextKey<AGENT>) agentContextKeyField.get(null);

instrumentation/opentelemetry-api-1.0/javaagent/src/test/groovy/ContextBridgeTest.groovy

@@ -9,6 +9,10 @@ import io.opentelemetry.api.trace.Span
 import io.opentelemetry.context.Context
 import io.opentelemetry.context.ContextKey
 import io.opentelemetry.extension.annotations.WithSpan
+import io.opentelemetry.instrumentation.api.instrumenter.SpanKey
+import io.opentelemetry.instrumentation.api.tracer.ClientSpan
+import io.opentelemetry.instrumentation.api.tracer.ConsumerSpan
+import io.opentelemetry.instrumentation.api.tracer.ServerSpan
 import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.Executors
@@ -145,6 +149,64 @@ class ContextBridgeTest extends AgentInstrumentationSpecification {
     Context.current() == Context.root()
   }
 
+  def "test server span bridge"() {
+    expect:
+    AgentSpanTesting.runWithServerSpan("server") {
+      assert Span.current() != null
+      assert ServerSpan.fromContextOrNull(Context.current()) != null
+      runWithSpan("internal") {
+        assert ServerSpan.fromContextOrNull(Context.current()) != null
+      }
+    }
+  }
+
+  def "test consumer span bridge"() {
+    expect:
+    AgentSpanTesting.runWithConsumerSpan("consumer") {
+      assert Span.current() != null
+      assert ConsumerSpan.fromContextOrNull(Context.current()) != null
+      runWithSpan("internal") {
+        assert ConsumerSpan.fromContextOrNull(Context.current()) != null
+      }
+    }
+  }
+
+  def "test client span bridge"() {
+    expect:
+    AgentSpanTesting.runWithClientSpan("client") {
+      assert Span.current() != null
+      assert ClientSpan.fromContextOrNull(Context.current()) != null
+      runWithSpan("internal") {
+        assert ClientSpan.fromContextOrNull(Context.current()) != null
+      }
+    }
+  }
+
+  def "test span key bridge"() {
+    expect:
+    AgentSpanTesting.runWithAllSpanKeys("parent") {
+      assert Span.current() != null
+      def spanKeys = [
+        SpanKey.SERVER,
+        SpanKey.CONSUMER,
+        SpanKey.HTTP_CLIENT,
+        SpanKey.RPC_CLIENT,
+        SpanKey.DB_CLIENT,
+        SpanKey.MESSAGING_PRODUCER,
+        SpanKey.ALL_CLIENTS,
+        SpanKey.ALL_PRODUCERS
+      ]
+      spanKeys.each { spanKey ->
+        assert spanKey.fromContextOrNull(Context.current()) != null
+      }
+      runWithSpan("internal") {
+        spanKeys.each { spanKey ->
+          assert spanKey.fromContextOrNull(Context.current()) != null
+        }
+      }
+    }
+  }
+
   // 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

instrumentation/opentelemetry-api-1.0/testing/build.gradle.kts

@@ -0,0 +1,3 @@
+plugins {
+  id("otel.javaagent-testing")
+}

instrumentation/opentelemetry-api-1.0/testing/src/main/java/AgentSpanTesting.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+public class AgentSpanTesting {
+
+  /**
+   * Runs the provided {@code runnable} inside the scope of an SERVER span with name {@code
+   * spanName}.
+   */
+  public static void runWithServerSpan(String spanName, Runnable runnable) {
+    runnable.run();
+  }
+
+  /**
+   * Runs the provided {@code runnable} inside the scope of an CONSUMER span with name {@code
+   * spanName}.
+   */
+  public static void runWithConsumerSpan(String spanName, Runnable runnable) {
+    runnable.run();
+  }
+
+  /**
+   * Runs the provided {@code runnable} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  public static void runWithClientSpan(String spanName, Runnable runnable) {
+    runnable.run();
+  }
+
+  /**
+   * Runs the provided {@code runnable} inside the scope of an INTERNAL span with name {@code
+   * spanName}. Span is added into context under all possible keys from {@link
+   * io.opentelemetry.instrumentation.api.instrumenter.SpanKey}
+   */
+  public static void runWithAllSpanKeys(String spanName, Runnable runnable) {
+    runnable.run();
+  }
+}

instrumentation/opentelemetry-api-1.0/testing/src/main/java/AgentSpanTestingInstrumentation.java

@@ -0,0 +1,138 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class AgentSpanTestingInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("AgentSpanTesting");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("runWithServerSpan"), this.getClass().getName() + "$RunWithServerSpanAdvice");
+    transformer.applyAdviceToMethod(
+        named("runWithConsumerSpan"), this.getClass().getName() + "$RunWithConsumerSpanAdvice");
+    transformer.applyAdviceToMethod(
+        named("runWithClientSpan"), this.getClass().getName() + "$RunWithClientSpanAdvice");
+    transformer.applyAdviceToMethod(
+        named("runWithAllSpanKeys"), this.getClass().getName() + "$RunWithAllSpanKeysAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class RunWithServerSpanAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) String spanName,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      context = AgentSpanTestingTracer.tracer().startServerSpan(spanName);
+      scope = context.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      scope.close();
+      if (throwable != null) {
+        AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
+      } else {
+        AgentSpanTestingTracer.tracer().end(context);
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class RunWithConsumerSpanAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) String spanName,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      context = AgentSpanTestingTracer.tracer().startConsumerSpan(spanName);
+      scope = context.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      scope.close();
+      if (throwable != null) {
+        AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
+      } else {
+        AgentSpanTestingTracer.tracer().end(context);
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class RunWithClientSpanAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) String spanName,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      context = AgentSpanTestingTracer.tracer().startClientSpan(spanName);
+      scope = context.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      scope.close();
+      if (throwable != null) {
+        AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
+      } else {
+        AgentSpanTestingTracer.tracer().end(context);
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class RunWithAllSpanKeysAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) String spanName,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      context = AgentSpanTestingTracer.tracer().startSpanWithAllKeys(spanName);
+      scope = context.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      scope.close();
+      if (throwable != null) {
+        AgentSpanTestingTracer.tracer().endExceptionally(context, throwable);
+      } else {
+        AgentSpanTestingTracer.tracer().end(context);
+      }
+    }
+  }
+}

instrumentation/opentelemetry-api-1.0/testing/src/main/java/AgentSpanTestingInstrumentationModule.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import static java.util.Collections.singletonList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class AgentSpanTestingInstrumentationModule extends InstrumentationModule {
+  public AgentSpanTestingInstrumentationModule() {
+    super(AgentSpanTestingInstrumentationModule.class.getName());
+  }
+
+  @Override
+  public boolean isHelperClass(String className) {
+    return className.startsWith("AgentSpanTestingTracer");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new AgentSpanTestingInstrumentation());
+  }
+}