Bridge span keys defined in instrumentation api (#3911)

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 class AgentContextStorage implements ContextStorage, AutoCloseable {
             "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 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
         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 @@ public class AgentContextStorage implements ContextStorage, AutoCloseable {
       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());
+  }
+}

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

@@ -0,0 +1,85 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanKey;
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
+import java.lang.reflect.Field;
+
+public class AgentSpanTestingTracer extends BaseTracer {
+  private static final AgentSpanTestingTracer TRACER = new AgentSpanTestingTracer();
+
+  private AgentSpanTestingTracer() {
+    super(GlobalOpenTelemetry.get());
+  }
+
+  public static AgentSpanTestingTracer tracer() {
+    return TRACER;
+  }
+
+  public Context startServerSpan(String name) {
+    Context parentContext = Context.current();
+    SpanBuilder spanBuilder = spanBuilder(parentContext, name, SpanKind.SERVER);
+    return withServerSpan(parentContext, spanBuilder.startSpan());
+  }
+
+  public Context startConsumerSpan(String name) {
+    Context parentContext = Context.current();
+    SpanBuilder spanBuilder = spanBuilder(parentContext, name, SpanKind.CONSUMER);
+    return withConsumerSpan(parentContext, spanBuilder.startSpan());
+  }
+
+  public Context startClientSpan(String name) {
+    Context parentContext = Context.current();
+    SpanBuilder spanBuilder = spanBuilder(parentContext, name, SpanKind.CLIENT);
+    return withClientSpan(parentContext, spanBuilder.startSpan());
+  }
+
+  public Context startSpanWithAllKeys(String name) {
+    Context parentContext = Context.current();
+    SpanBuilder spanBuilder = spanBuilder(parentContext, name, SpanKind.INTERNAL);
+    Span span = spanBuilder.startSpan();
+    Context newContext = parentContext.with(span);
+    for (SpanKey spanKey : SpanKeyAccess.getSpanKeys()) {
+      newContext = spanKey.storeInContext(newContext, span);
+    }
+    return newContext;
+  }
+
+  @Override
+  protected String getInstrumentationName() {
+    return "agent-span-test-instrumentation";
+  }
+
+  private static class SpanKeyAccess {
+
+    public static SpanKey[] getSpanKeys() {
+      return new SpanKey[] {
+        SpanKey.SERVER,
+        SpanKey.CONSUMER,
+        getSpanKeyByName("HTTP_CLIENT"),
+        getSpanKeyByName("RPC_CLIENT"),
+        getSpanKeyByName("DB_CLIENT"),
+        getSpanKeyByName("MESSAGING_PRODUCER"),
+        SpanKey.ALL_CLIENTS,
+        SpanKey.ALL_PRODUCERS
+      };
+    }
+
+    private static SpanKey getSpanKeyByName(String name) {
+      try {
+        Field field = SpanKey.class.getDeclaredField(name);
+        field.setAccessible(true);
+        return (SpanKey) field.get(name);
+      } catch (NoSuchFieldException | IllegalAccessException exception) {
+        throw new IllegalStateException("Failed to find span key named " + name, exception);
+      }
+    }
+  }
+}

settings.gradle.kts

@@ -253,6 +253,7 @@ include(":instrumentation:okhttp:okhttp-3.0:library")
 include(":instrumentation:okhttp:okhttp-3.0:testing")
 include(":instrumentation:opentelemetry-annotations-1.0:javaagent")
 include(":instrumentation:opentelemetry-api-1.0:javaagent")
+include(":instrumentation:opentelemetry-api-1.0:testing")
 include(":instrumentation:opentelemetry-api-metrics-1.0:javaagent")
 include(":instrumentation:oshi:javaagent")
 include(":instrumentation:oshi:library")