Change the way how SpanKeys are extracted in instrumentation-api (#5546)

* Change the way how SpanKeys are extracted in instrumentation-api

* fix tests

* code review comments

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

@@ -21,10 +21,13 @@ import io.opentelemetry.instrumentation.api.instrumenter.http.HttpClientAttribut
 import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.rpc.RpcAttributesExtractor;
 import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import io.opentelemetry.instrumentation.api.internal.SpanKeyProvider;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.Nullable;
 
 /**
@@ -260,7 +263,7 @@ public final class InstrumenterBuilder<REQUEST, RESPONSE> {
   }
 
   SpanSuppressionStrategy getSpanSuppressionStrategy() {
-    Set<SpanKey> spanKeys = SpanKeyExtractor.determineSpanKeys(attributesExtractors);
+    Set<SpanKey> spanKeys = getSpanKeysFromAttributesExtractors();
     if (enableSpanSuppressionByType) {
       return SpanSuppressionStrategy.from(spanKeys);
     }
@@ -268,6 +271,18 @@ public final class InstrumenterBuilder<REQUEST, RESPONSE> {
     return SpanSuppressionStrategy.suppressNestedClients(spanKeys);
   }
 
+  private Set<SpanKey> getSpanKeysFromAttributesExtractors() {
+    return attributesExtractors.stream()
+        .filter(SpanKeyProvider.class::isInstance)
+        .map(SpanKeyProvider.class::cast)
+        .flatMap(
+            provider -> {
+              SpanKey spanKey = provider.internalGetSpanKey();
+              return spanKey == null ? Stream.of() : Stream.of(spanKey);
+            })
+        .collect(Collectors.toSet());
+  }
+
   private interface InstrumenterConstructor<RQ, RS> {
     Instrumenter<RQ, RS> create(InstrumenterBuilder<RQ, RS> builder);
 

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

@@ -1,55 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package io.opentelemetry.instrumentation.api.instrumenter;
-
-import io.opentelemetry.instrumentation.api.instrumenter.db.DbClientAttributesExtractor;
-import io.opentelemetry.instrumentation.api.instrumenter.http.HttpClientAttributesExtractor;
-import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingAttributesExtractor;
-import io.opentelemetry.instrumentation.api.instrumenter.rpc.RpcAttributesExtractor;
-import io.opentelemetry.instrumentation.api.internal.SpanKey;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-
-final class SpanKeyExtractor {
-
-  /**
-   * Automatically determines {@link SpanKey}s that should be applied to the newly constructed
-   * {@link Instrumenter} based on the {@link AttributesExtractor}s configured.
-   */
-  static Set<SpanKey> determineSpanKeys(
-      List<? extends AttributesExtractor<?, ?>> attributesExtractors) {
-    Set<SpanKey> spanKeys = new HashSet<>();
-    for (AttributesExtractor<?, ?> attributeExtractor : attributesExtractors) {
-      if (attributeExtractor instanceof HttpClientAttributesExtractor) {
-        spanKeys.add(SpanKey.HTTP_CLIENT);
-      } else if (attributeExtractor instanceof RpcAttributesExtractor) {
-        spanKeys.add(SpanKey.RPC_CLIENT);
-      } else if (attributeExtractor instanceof DbClientAttributesExtractor) {
-        spanKeys.add(SpanKey.DB_CLIENT);
-      } else if (attributeExtractor instanceof MessagingAttributesExtractor) {
-        spanKeys.add(
-            determineMessagingSpanKey((MessagingAttributesExtractor<?, ?>) attributeExtractor));
-      }
-    }
-    return spanKeys;
-  }
-
-  private static SpanKey determineMessagingSpanKey(
-      MessagingAttributesExtractor<?, ?> messagingAttributesExtractor) {
-    switch (messagingAttributesExtractor.operation()) {
-      case SEND:
-        return SpanKey.PRODUCER;
-      case RECEIVE:
-        return SpanKey.CONSUMER_RECEIVE;
-      case PROCESS:
-        return SpanKey.CONSUMER_PROCESS;
-    }
-    throw new IllegalStateException("Can't possibly happen");
-  }
-
-  private SpanKeyExtractor() {}
-}

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/db/DbClientCommonAttributesExtractor.java

@@ -7,13 +7,16 @@ package io.opentelemetry.instrumentation.api.instrumenter.db;
 
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.annotations.UnstableApi;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import io.opentelemetry.instrumentation.api.internal.SpanKeyProvider;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import javax.annotation.Nullable;
 
 abstract class DbClientCommonAttributesExtractor<
         REQUEST, RESPONSE, GETTER extends DbClientCommonAttributesGetter<REQUEST>>
-    implements AttributesExtractor<REQUEST, RESPONSE> {
+    implements AttributesExtractor<REQUEST, RESPONSE>, SpanKeyProvider {
 
   final GETTER getter;
 
@@ -36,4 +39,14 @@ abstract class DbClientCommonAttributesExtractor<
       REQUEST request,
       @Nullable RESPONSE response,
       @Nullable Throwable error) {}
+
+  /**
+   * This method is internal and is hence not for public use. Its API is unstable and can change at
+   * any time.
+   */
+  @UnstableApi
+  @Override
+  public SpanKey internalGetSpanKey() {
+    return SpanKey.DB_CLIENT;
+  }
 }

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/http/HttpClientAttributesExtractor.java

@@ -7,6 +7,9 @@ package io.opentelemetry.instrumentation.api.instrumenter.http;
 
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.annotations.UnstableApi;
+import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import io.opentelemetry.instrumentation.api.internal.SpanKeyProvider;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.util.List;
 import javax.annotation.Nullable;
@@ -22,7 +25,8 @@ import javax.annotation.Nullable;
  */
 public final class HttpClientAttributesExtractor<REQUEST, RESPONSE>
     extends HttpCommonAttributesExtractor<
-        REQUEST, RESPONSE, HttpClientAttributesGetter<REQUEST, RESPONSE>> {
+        REQUEST, RESPONSE, HttpClientAttributesGetter<REQUEST, RESPONSE>>
+    implements SpanKeyProvider {
 
   /** Creates the HTTP client attributes extractor with default configuration. */
   public static <REQUEST, RESPONSE> HttpClientAttributesExtractor<REQUEST, RESPONSE> create(
@@ -62,4 +66,14 @@ public final class HttpClientAttributesExtractor<REQUEST, RESPONSE>
     super.onEnd(attributes, context, request, response, error);
     set(attributes, SemanticAttributes.HTTP_FLAVOR, getter.flavor(request, response));
   }
+
+  /**
+   * This method is internal and is hence not for public use. Its API is unstable and can change at
+   * any time.
+   */
+  @UnstableApi
+  @Override
+  public SpanKey internalGetSpanKey() {
+    return SpanKey.HTTP_CLIENT;
+  }
 }

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/messaging/MessagingAttributesExtractor.java

@@ -7,7 +7,10 @@ package io.opentelemetry.instrumentation.api.instrumenter.messaging;
 
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.annotations.UnstableApi;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import io.opentelemetry.instrumentation.api.internal.SpanKeyProvider;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import javax.annotation.Nullable;
 
@@ -21,7 +24,7 @@ import javax.annotation.Nullable;
  * best compliance with the OpenTelemetry specification.
  */
 public abstract class MessagingAttributesExtractor<REQUEST, RESPONSE>
-    implements AttributesExtractor<REQUEST, RESPONSE> {
+    implements AttributesExtractor<REQUEST, RESPONSE>, SpanKeyProvider {
   public static final String TEMP_DESTINATION_NAME = "(temporary)";
 
   @Override
@@ -96,4 +99,22 @@ public abstract class MessagingAttributesExtractor<REQUEST, RESPONSE>
 
   @Nullable
   protected abstract String messageId(REQUEST request, @Nullable RESPONSE response);
+
+  /**
+   * This method is internal and is hence not for public use. Its API is unstable and can change at
+   * any time.
+   */
+  @UnstableApi
+  @Override
+  public SpanKey internalGetSpanKey() {
+    switch (operation()) {
+      case SEND:
+        return SpanKey.PRODUCER;
+      case RECEIVE:
+        return SpanKey.CONSUMER_RECEIVE;
+      case PROCESS:
+        return SpanKey.CONSUMER_PROCESS;
+    }
+    throw new IllegalStateException("Can't possibly happen");
+  }
 }

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/rpc/RpcAttributesExtractor.java

@@ -7,7 +7,10 @@ package io.opentelemetry.instrumentation.api.instrumenter.rpc;
 
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.annotations.UnstableApi;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import io.opentelemetry.instrumentation.api.internal.SpanKeyProvider;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import javax.annotation.Nullable;
 
@@ -21,7 +24,7 @@ import javax.annotation.Nullable;
  * specification.
  */
 public abstract class RpcAttributesExtractor<REQUEST, RESPONSE>
-    implements AttributesExtractor<REQUEST, RESPONSE> {
+    implements AttributesExtractor<REQUEST, RESPONSE>, SpanKeyProvider {
 
   @Override
   public final void onStart(AttributesBuilder attributes, Context parentContext, REQUEST request) {
@@ -48,4 +51,14 @@ public abstract class RpcAttributesExtractor<REQUEST, RESPONSE>
 
   @Nullable
   protected abstract String method(REQUEST request);
+
+  /**
+   * This method is internal and is hence not for public use. Its API is unstable and can change at
+   * any time.
+   */
+  @UnstableApi
+  @Override
+  public SpanKey internalGetSpanKey() {
+    return SpanKey.RPC_CLIENT;
+  }
 }

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.internal;
+
+import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import javax.annotation.Nullable;
+
+/**
+ * Returns the {@link SpanKey} associated with the {@link AttributesExtractor} that implements this
+ * interface.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public interface SpanKeyProvider {
+
+  @Nullable
+  SpanKey internalGetSpanKey();
+}

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterTest.java

@@ -560,6 +560,8 @@ class InstrumenterTest {
 
   @Test
   void clientNestedSuppressed_whenSameInstrumentationType() {
+    when(mockDbAttributes.internalGetSpanKey()).thenCallRealMethod();
+
     Instrumenter<Map<String, String>, Map<String, String>> instrumenterOuter =
         getInstrumenterWithType(true, mockDbAttributes);
     Instrumenter<Map<String, String>, Map<String, String>> instrumenterInner =
@@ -616,6 +618,8 @@ class InstrumenterTest {
 
   @Test
   void instrumentationTypeDetected_http() {
+    when(mockHttpClientAttributes.internalGetSpanKey()).thenCallRealMethod();
+
     Instrumenter<Map<String, String>, Map<String, String>> instrumenter =
         getInstrumenterWithType(true, mockHttpClientAttributes, new AttributesExtractor1());
 
@@ -627,6 +631,8 @@ class InstrumenterTest {
 
   @Test
   void instrumentationTypeDetected_db() {
+    when(mockDbAttributes.internalGetSpanKey()).thenCallRealMethod();
+
     Instrumenter<Map<String, String>, Map<String, String>> instrumenter =
         getInstrumenterWithType(true, mockDbAttributes, new AttributesExtractor2());
 
@@ -638,6 +644,8 @@ class InstrumenterTest {
 
   @Test
   void instrumentationTypeDetected_rpc() {
+    when(mockRpcAttributes.internalGetSpanKey()).thenCallRealMethod();
+
     Instrumenter<Map<String, String>, Map<String, String>> instrumenter =
         getInstrumenterWithType(true, mockRpcAttributes);
 
@@ -650,6 +658,7 @@ class InstrumenterTest {
   @Test
   void instrumentationTypeDetected_producer() {
     when(mockMessagingAttributes.operation()).thenReturn(MessageOperation.SEND);
+    when(mockMessagingAttributes.internalGetSpanKey()).thenCallRealMethod();
 
     Instrumenter<Map<String, String>, Map<String, String>> instrumenter =
         getInstrumenterWithType(true, mockMessagingAttributes);
@@ -663,6 +672,7 @@ class InstrumenterTest {
   @Test
   void instrumentationTypeDetected_mix() {
     when(mockMessagingAttributes.operation()).thenReturn(MessageOperation.SEND);
+    when(mockMessagingAttributes.internalGetSpanKey()).thenCallRealMethod();
 
     Instrumenter<Map<String, String>, Map<String, String>> instrumenter =
         getInstrumenterWithType(

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/SpanKeyExtractorTest.java

@@ -1,74 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package io.opentelemetry.instrumentation.api.instrumenter;
-
-import static java.util.Collections.singleton;
-import static java.util.Collections.singletonList;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
-import io.opentelemetry.instrumentation.api.instrumenter.db.DbClientAttributesExtractor;
-import io.opentelemetry.instrumentation.api.instrumenter.http.HttpClientAttributesExtractor;
-import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessageOperation;
-import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingAttributesExtractor;
-import io.opentelemetry.instrumentation.api.instrumenter.rpc.RpcAttributesExtractor;
-import io.opentelemetry.instrumentation.api.internal.SpanKey;
-import java.util.Set;
-import java.util.stream.Stream;
-import org.junit.jupiter.api.extension.ExtensionContext;
-import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.Arguments;
-import org.junit.jupiter.params.provider.ArgumentsProvider;
-import org.junit.jupiter.params.provider.ArgumentsSource;
-
-class SpanKeyExtractorTest {
-
-  @ParameterizedTest
-  @ArgumentsSource(ClientSpanKeys.class)
-  void shouldDetermineKeysForClientAttributesExtractors(
-      AttributesExtractor<?, ?> attributesExtractor, SpanKey expectedSpanKey) {
-
-    Set<SpanKey> spanKeys = SpanKeyExtractor.determineSpanKeys(singletonList(attributesExtractor));
-    assertEquals(singleton(expectedSpanKey), spanKeys);
-  }
-
-  static final class ClientSpanKeys implements ArgumentsProvider {
-
-    @Override
-    public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
-      return Stream.of(
-          Arguments.of(mock(DbClientAttributesExtractor.class), SpanKey.DB_CLIENT),
-          Arguments.of(mock(HttpClientAttributesExtractor.class), SpanKey.HTTP_CLIENT),
-          Arguments.of(mock(RpcAttributesExtractor.class), SpanKey.RPC_CLIENT));
-    }
-  }
-
-  @ParameterizedTest
-  @ArgumentsSource(MessagingSpanKeys.class)
-  void shouldDetermineKeysForMessagingAttributesExtractor(
-      MessageOperation operation, SpanKey expectedSpanKey) {
-
-    MessagingAttributesExtractor<?, ?> attributesExtractor =
-        mock(MessagingAttributesExtractor.class);
-    when(attributesExtractor.operation()).thenReturn(operation);
-
-    Set<SpanKey> spanKeys = SpanKeyExtractor.determineSpanKeys(singletonList(attributesExtractor));
-
-    assertEquals(singleton(expectedSpanKey), spanKeys);
-  }
-
-  static final class MessagingSpanKeys implements ArgumentsProvider {
-
-    @Override
-    public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
-      return Stream.of(
-          Arguments.of(MessageOperation.PROCESS, SpanKey.CONSUMER_PROCESS),
-          Arguments.of(MessageOperation.RECEIVE, SpanKey.CONSUMER_RECEIVE),
-          Arguments.of(MessageOperation.SEND, SpanKey.PRODUCER));
-    }
-  }
-}