Optimize codes

Author: steverao

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

@@ -302,9 +302,8 @@ private Instrumenter<REQUEST, RESPONSE> buildInstrumenter(
     List<InstrumentationCustomizer> customizers =
         INSTRUMENTATION_CUSTOMIZER_MAP.entrySet().stream()
             .filter(entry -> entry.getKey().test(instrumentationName))
-            .map(Map.Entry::getValue)
-            .findFirst()
-            .orElse(null);
+            .flatMap(entry -> entry.getValue().stream())
+            .collect(Collectors.toList());
     if (customizers != null) {
       for (InstrumentationCustomizer customizer : customizers) {
         if (customizer.getContextCustomizer() != null) {
@@ -313,6 +312,9 @@ private Instrumenter<REQUEST, RESPONSE> buildInstrumenter(
         if (customizer.getAttributesExtractor() != null) {
           addAttributesExtractor(customizer.getAttributesExtractor());
         }
+        if (customizer.getAttributesExtractors() != null) {
+          addAttributesExtractors(customizer.getAttributesExtractors());
+        }
         if (customizer.getOperationMetrics() != null) {
           addOperationMetrics(customizer.getOperationMetrics());
         }

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

@@ -9,7 +9,9 @@
 import io.opentelemetry.instrumentation.api.instrumenter.ContextCustomizer;
 import io.opentelemetry.instrumentation.api.instrumenter.OperationMetrics;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import java.util.List;
 import java.util.function.Predicate;
+import javax.annotation.Nullable;
 
 /**
  * A service provider interface (SPI) for providing customizations for instrumentation, including
@@ -55,6 +57,23 @@ default <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> getAttributes
     return null;
   }
 
+  /**
+   * Returns a list of {@link AttributesExtractor}s that will extract attributes from requests and
+   * responses during the instrumentation process.
+   *
+   * <p>This allows providing multiple extractors for a single instrumentation. The default
+   * implementation returns {@code null} for backward compatibility.
+   *
+   * @param <REQUEST> the type of request object used by the instrumented library
+   * @param <RESPONSE> the type of response object used by the instrumented library
+   * @return a list of attributes extractors, or null if not applicable
+   */
+  @Nullable
+  default <REQUEST, RESPONSE>
+      List<AttributesExtractor<REQUEST, RESPONSE>> getAttributesExtractors() {
+    return null;
+  }
+
   /**
    * Returns a new instance of a {@link ContextCustomizer} that will customize the tracing context
    * during request processing.

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

@@ -11,6 +11,8 @@
 import io.opentelemetry.instrumentation.api.instrumenter.ContextCustomizer;
 import io.opentelemetry.instrumentation.api.instrumenter.OperationMetrics;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import java.util.Collections;
+import java.util.List;
 import java.util.function.Predicate;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -50,6 +52,13 @@ AttributesExtractor<REQUEST, RESPONSE> getAttributesExtractor() {
             return (AttributesExtractor<REQUEST, RESPONSE>) attributesExtractor;
           }
 
+          @Override
+          public <REQUEST, RESPONSE>
+              List<AttributesExtractor<REQUEST, RESPONSE>> getAttributesExtractors() {
+            return Collections.singletonList(
+                (AttributesExtractor<REQUEST, RESPONSE>) attributesExtractor);
+          }
+
           @Override
           public <REQUEST> ContextCustomizer<REQUEST> getContextCustomizer() {
             return (ContextCustomizer<REQUEST>) contextCustomizer;
@@ -82,6 +91,12 @@ void testGetAttributesExtractor() {
     assertThat(extractor).isSameAs(attributesExtractor);
   }
 
+  @Test
+  void testGetAttributesExtractors() {
+    List<AttributesExtractor<Object, Object>> extractors = customizer.getAttributesExtractors();
+    assertThat(extractors).containsExactly(attributesExtractor);
+  }
+
   @Test
   void testGetContextCustomizer() {
     assertThat(customizer.getContextCustomizer()).isSameAs(contextCustomizer);

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/AgentInitializer.java

@@ -6,15 +6,11 @@
 package io.opentelemetry.javaagent.bootstrap;
 
 import io.opentelemetry.instrumentation.api.internal.ConfigPropertiesUtil;
-import io.opentelemetry.instrumentation.api.internal.ServiceLoaderUtil;
 import java.io.File;
 import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Constructor;
 import java.security.PrivilegedAction;
 import java.security.PrivilegedExceptionAction;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.ServiceLoader;
 import javax.annotation.Nullable;
 
 /**
@@ -66,7 +62,6 @@ public Void run() {
           public Void run() throws Exception {
             agentClassLoader = createAgentClassLoader("inst", javaagentFile);
             agentStarter = createAgentStarter(agentClassLoader, inst, javaagentFile);
-            setLoaderFunction();
             if (!fromPremain || !delayAgentStart()) {
               agentStarter.start();
               agentStarted = true;
@@ -217,19 +212,6 @@ private static AgentStarter createAgentStarter(
         constructor.newInstance(instrumentation, javaagentFile, isSecurityManagerSupportEnabled);
   }
 
-  /** Sets a custom loader function for the ServiceLoaderUtil to use the agentClassLoader. */
-  private static void setLoaderFunction() {
-    ServiceLoaderUtil.setLoaderFunction(
-        clazz -> {
-          List<Object> instances = new ArrayList<>();
-          ServiceLoader<?> serviceLoader = ServiceLoader.load(clazz, agentClassLoader);
-          for (Object instance : serviceLoader) {
-            instances.add(instance);
-          }
-          return instances;
-        });
-  }
-
   private AgentInitializer() {}
 
   @SuppressWarnings("SystemOut")

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentStarterImpl.java

@@ -6,6 +6,7 @@
 package io.opentelemetry.javaagent.tooling;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.internal.ServiceLoaderUtil;
 import io.opentelemetry.instrumentation.api.internal.cache.weaklockfree.WeakConcurrentMapCleaner;
 import io.opentelemetry.javaagent.bootstrap.AgentInitializer;
 import io.opentelemetry.javaagent.bootstrap.AgentStarter;
@@ -16,6 +17,8 @@
 import java.lang.instrument.Instrumentation;
 import java.lang.instrument.UnmodifiableClassException;
 import java.security.ProtectionDomain;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.ServiceLoader;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
@@ -72,6 +75,7 @@ public void start() {
 
     EarlyInitAgentConfig earlyConfig = EarlyInitAgentConfig.create();
     extensionClassLoader = createExtensionClassLoader(getClass().getClassLoader(), earlyConfig);
+    setLoaderFunction();
 
     String loggerImplementationName = earlyConfig.getString("otel.javaagent.logging");
     // default to the built-in stderr slf4j-simple logger
@@ -246,4 +250,17 @@ public void visitMethodInsn(
       return hookInserted ? cw.toByteArray() : null;
     }
   }
+
+  /** Sets a custom loader function for the ServiceLoaderUtil to use the agentClassLoader. */
+  private void setLoaderFunction() {
+    ServiceLoaderUtil.setLoaderFunction(
+        clazz -> {
+          List<Object> instances = new ArrayList<>();
+          ServiceLoader<?> serviceLoader = ServiceLoader.load(clazz, extensionClassLoader);
+          for (Object instance : serviceLoader) {
+            instances.add(instance);
+          }
+          return instances;
+        });
+  }
 }