Add instrumentation type to instrumenter customizer (#15227)

Author: laurit

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/instrumenter/InstrumenterCustomizer.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.OperationMetrics;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.api.semconv.http.HttpClientAttributesExtractor;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
 
@@ -33,6 +34,16 @@ public interface InstrumenterCustomizer {
    */
   String getInstrumentationName();
 
+  /**
+   * Tests whether given instrumenter produces telemetry of specified type. Instrumentation type is
+   * detected based on the standard {@link AttributesExtractor} implementations used by this
+   * instrumenter e.g. instrumenters that use {@link HttpClientAttributesExtractor} have {@link
+   * InstrumentationType#HTTP_CLIENT} type.
+   *
+   * @return the name of the instrumentation this customizer targets
+   */
+  boolean hasType(InstrumentationType type);
+
   /**
    * Adds a single {@link AttributesExtractor} to the instrumenter. This extractor will be used to
    * extract attributes from requests and responses during the request lifecycle.
@@ -94,4 +105,16 @@ default InstrumenterCustomizer setSpanNameExtractor(
    */
   InstrumenterCustomizer setSpanNameExtractor(
       UnaryOperator<SpanNameExtractor<?>> spanNameExtractor);
+
+  /** Types of instrumentation. */
+  enum InstrumentationType {
+    HTTP_CLIENT,
+    HTTP_SERVER,
+    DB_CLIENT,
+    RPC_CLIENT,
+    RPC_SERVER,
+    MESSAGING_PRODUCER,
+    MESSAGING_CONSUMER_RECEIVE,
+    MESSAGING_CONSUMER_PROCESS
+  }
 }

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/instrumenter/internal/InstrumenterCustomizerImpl.java

@@ -11,6 +11,9 @@
 import io.opentelemetry.instrumentation.api.instrumenter.OperationMetrics;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
 import io.opentelemetry.instrumentation.api.internal.InternalInstrumenterCustomizer;
+import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.function.UnaryOperator;
 
 /**
@@ -19,6 +22,19 @@
  */
 @SuppressWarnings({"unchecked", "rawtypes"})
 public final class InstrumenterCustomizerImpl implements InstrumenterCustomizer {
+  private static final Map<InstrumentationType, SpanKey> typeToSpanKey = new HashMap<>();
+
+  static {
+    typeToSpanKey.put(InstrumentationType.HTTP_CLIENT, SpanKey.HTTP_CLIENT);
+    typeToSpanKey.put(InstrumentationType.HTTP_SERVER, SpanKey.HTTP_SERVER);
+    typeToSpanKey.put(InstrumentationType.DB_CLIENT, SpanKey.DB_CLIENT);
+    typeToSpanKey.put(InstrumentationType.RPC_CLIENT, SpanKey.RPC_CLIENT);
+    typeToSpanKey.put(InstrumentationType.RPC_SERVER, SpanKey.RPC_SERVER);
+    typeToSpanKey.put(InstrumentationType.MESSAGING_PRODUCER, SpanKey.PRODUCER);
+    typeToSpanKey.put(InstrumentationType.MESSAGING_CONSUMER_RECEIVE, SpanKey.CONSUMER_RECEIVE);
+    typeToSpanKey.put(InstrumentationType.MESSAGING_CONSUMER_PROCESS, SpanKey.CONSUMER_RECEIVE);
+  }
+
   private final InternalInstrumenterCustomizer customizer;
 
   public InstrumenterCustomizerImpl(InternalInstrumenterCustomizer customizer) {
@@ -30,6 +46,15 @@ public String getInstrumentationName() {
     return customizer.getInstrumentationName();
   }
 
+  @Override
+  public boolean hasType(InstrumentationType type) {
+    SpanKey spanKey = typeToSpanKey.get(type);
+    if (spanKey == null) {
+      throw new IllegalArgumentException("unexpected instrumentation type: " + type);
+    }
+    return customizer.hasType(spanKey);
+  }
+
   @Override
   public InstrumenterCustomizer addAttributesExtractor(AttributesExtractor<?, ?> extractor) {
     customizer.addAttributesExtractor(extractor);

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

@@ -395,15 +395,26 @@ private void propagateOperationListenersToOnEnd() {
 
   private static <REQUEST, RESPONSE> void applyCustomizers(
       InstrumenterBuilder<REQUEST, RESPONSE> builder) {
-    for (InternalInstrumenterCustomizerProvider provider :
-        InternalInstrumenterCustomizerUtil.getInstrumenterCustomizerProviders()) {
+    List<InternalInstrumenterCustomizerProvider> customizerProviders =
+        InternalInstrumenterCustomizerUtil.getInstrumenterCustomizerProviders();
+    if (customizerProviders.isEmpty()) {
+      return;
+    }
+
+    Set<SpanKey> spanKeys = builder.getSpanKeysFromAttributesExtractors();
+    for (InternalInstrumenterCustomizerProvider provider : customizerProviders) {
       provider.customize(
           new InternalInstrumenterCustomizer<REQUEST, RESPONSE>() {
             @Override
             public String getInstrumentationName() {
               return builder.instrumentationName;
             }
 
+            @Override
+            public boolean hasType(SpanKey type) {
+              return spanKeys.contains(type);
+            }
+
             @Override
             public void addAttributesExtractor(AttributesExtractor<REQUEST, RESPONSE> extractor) {
               builder.addAttributesExtractor(extractor);

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/internal/InternalInstrumenterCustomizer.java

@@ -19,6 +19,8 @@ public interface InternalInstrumenterCustomizer<REQUEST, RESPONSE> {
 
   String getInstrumentationName();
 
+  boolean hasType(SpanKey type);
+
   void addAttributesExtractor(AttributesExtractor<REQUEST, RESPONSE> extractor);
 
   void addAttributesExtractors(

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/internal/InstrumentationCustomizerTest.java

@@ -22,6 +22,7 @@
 import io.opentelemetry.api.trace.SpanId;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.incubator.instrumenter.InstrumenterCustomizer.InstrumentationType;
 import io.opentelemetry.instrumentation.api.incubator.instrumenter.InstrumenterCustomizerProvider;
 import io.opentelemetry.instrumentation.api.incubator.instrumenter.internal.InternalInstrumenterCustomizerProviderImpl;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
@@ -110,6 +111,43 @@ void testGetInstrumentationName() {
     assertThat(customizerCalled).isTrue();
   }
 
+  @Test
+  void testHasType() {
+    AtomicBoolean customizerCalled = new AtomicBoolean();
+    setCustomizer(
+        customizer -> {
+          customizerCalled.set(true);
+          assertThat(customizer.hasType(InstrumentationType.HTTP_CLIENT)).isTrue();
+          assertThat(customizer.hasType(InstrumentationType.HTTP_SERVER)).isFalse();
+        });
+
+    class TestAttributesExtractor implements AttributesExtractor<Object, Object>, SpanKeyProvider {
+      @Override
+      public void onStart(AttributesBuilder attributes, Context parentContext, Object request) {}
+
+      @Override
+      public void onEnd(
+          AttributesBuilder attributes,
+          Context context,
+          Object request,
+          @Nullable Object response,
+          @Nullable Throwable error) {}
+
+      @Nullable
+      @Override
+      public SpanKey internalGetSpanKey() {
+        return SpanKey.HTTP_CLIENT;
+      }
+    }
+
+    Instrumenter.<Map<String, String>, Map<String, String>>builder(
+            otelTesting.getOpenTelemetry(), "test", unused -> "span")
+        .addAttributesExtractor(new TestAttributesExtractor())
+        .buildInstrumenter();
+
+    assertThat(customizerCalled).isTrue();
+  }
+
   @Test
   void testAddAttributesExtractor() {
     AtomicBoolean customizerCalled = new AtomicBoolean();