Method instrumentation

Author: laurit

instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodInstrumentation.java

@@ -23,19 +23,19 @@
 import io.opentelemetry.instrumentation.api.incubator.semconv.util.ClassAndMethod;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.util.Collection;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.enumeration.EnumerationDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
 
-@SuppressWarnings("EnumOrdinal")
 public class MethodInstrumentation implements TypeInstrumentation {
   private final String className;
-  private final Map<String, SpanKind> methodNames;
+  private final Map<SpanKind, Collection<String>> methodNames;
 
-  public MethodInstrumentation(String className, Map<String, SpanKind> methodNames) {
+  public MethodInstrumentation(String className, Map<SpanKind, Collection<String>> methodNames) {
     this.className = className;
     this.methodNames = methodNames;
   }
@@ -63,36 +63,37 @@ public ElementMatcher<TypeDescription> typeMatcher() {
 
   @Override
   public void transform(TypeTransformer transformer) {
-    transformer.applyAdviceToMethod(
-        namedOneOf(methodNames.keySet().toArray(new String[0])).and(isMethod()),
-        mapping ->
-            mapping
-                .bind(
-                    MethodReturnType.class,
-                    (instrumentedType, instrumentedMethod, assigner, argumentHandler, sort) ->
-                        Advice.OffsetMapping.Target.ForStackManipulation.of(
-                            instrumentedMethod.getReturnType().asErasure()))
-                .bind(
-                    SpanKindOrdinal.class,
-                    (instrumentedType, instrumentedMethod, assigner, argumentHandler, sort) ->
-                        Advice.OffsetMapping.Target.ForStackManipulation.of(
-                            new EnumerationDescription.ForLoadedEnumeration(
-                                methodNames.get(instrumentedMethod.getName())))),
-        MethodInstrumentation.class.getName() + "$MethodAdvice");
+    for (Map.Entry<SpanKind, Collection<String>> entry : methodNames.entrySet()) {
+      SpanKind spanKind = entry.getKey();
+      Collection<String> names = entry.getValue();
+      transformer.applyAdviceToMethod(
+          namedOneOf(names.toArray(new String[0])).and(isMethod()),
+          mapping ->
+              mapping
+                  .bind(
+                      MethodReturnType.class,
+                      (instrumentedType, instrumentedMethod, assigner, argumentHandler, sort) ->
+                          Advice.OffsetMapping.Target.ForStackManipulation.of(
+                              instrumentedMethod.getReturnType().asErasure()))
+                  .bind(
+                      MethodSpanKind.class,
+                      new EnumerationDescription.ForLoadedEnumeration(spanKind)),
+          MethodInstrumentation.class.getName() + "$MethodAdvice");
+    }
   }
 
   // custom annotation that represents the return type of the method
   @interface MethodReturnType {}
 
   // custom annotation that represents the SpanKind of the method
-  @interface SpanKindOrdinal {}
+  @interface MethodSpanKind {}
 
   @SuppressWarnings("unused")
   public static class MethodAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(
-        @SpanKindOrdinal SpanKind spanKind,
+        @MethodSpanKind SpanKind spanKind,
         @Advice.Origin("#t") Class<?> declaringClass,
         @Advice.Origin("#m") String methodName,
         @Advice.Local("otelMethod") MethodAndType classAndMethod,

instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodInstrumentationModule.java

@@ -5,14 +5,17 @@
 
 package io.opentelemetry.javaagent.instrumentation.methods;
 
+import static java.util.Collections.emptyList;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+
 import com.google.auto.service.AutoService;
 import io.opentelemetry.api.incubator.config.DeclarativeConfigProperties;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.javaagent.bootstrap.internal.AgentInstrumentationConfig;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.tooling.config.MethodsConfigurationParser;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -35,8 +38,11 @@ private static List<TypeInstrumentation> createInstrumentations() {
         AgentInstrumentationConfig.get().getDeclarativeConfig("methods");
     List<TypeInstrumentation> list =
         methods != null ? MethodsConfig.parseDeclarativeConfig(methods) : parseConfigProperties();
+    // ensure that there is at least one instrumentation so that muzzle reference collection could
+    // work
     if (list.isEmpty()) {
-      return Collections.singletonList(new MethodInstrumentation(null, Collections.emptyMap()));
+      return singletonList(
+          new MethodInstrumentation(null, singletonMap(SpanKind.INTERNAL, emptyList())));
     }
     return list;
   }
@@ -51,11 +57,7 @@ private static List<TypeInstrumentation> parseConfigProperties() {
         .map(
             e ->
                 new MethodInstrumentation(
-                    e.getKey(),
-                    e.getValue().stream()
-                        .collect(
-                            Collectors.toMap(
-                                methodName -> methodName, ignore -> SpanKind.INTERNAL))))
+                    e.getKey(), singletonMap(SpanKind.INTERNAL, e.getValue())))
         .collect(Collectors.toList());
   }
 

instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodsConfig.java

@@ -10,7 +10,9 @@
 import io.opentelemetry.api.incubator.config.DeclarativeConfigProperties;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
-import java.util.LinkedHashMap;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.EnumMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -39,7 +41,7 @@ private static Stream<MethodInstrumentation> parseMethodInstrumentation(
       return Stream.empty();
     }
 
-    Map<String, SpanKind> methodNames = new LinkedHashMap<>();
+    Map<SpanKind, Collection<String>> methodNames = new EnumMap<>(SpanKind.class);
     for (DeclarativeConfigProperties method : config.getStructuredList("methods", emptyList())) {
       String methodName = method.getString("name");
       if (isNullOrEmpty(methodName)) {
@@ -49,7 +51,10 @@ private static Stream<MethodInstrumentation> parseMethodInstrumentation(
       }
       String spanKind = method.getString("span_kind", "INTERNAL");
       try {
-        methodNames.put(methodName, SpanKind.valueOf(spanKind.toUpperCase(Locale.ROOT)));
+        methodNames
+            .computeIfAbsent(
+                SpanKind.valueOf(spanKind.toUpperCase(Locale.ROOT)), unused -> new ArrayList<>())
+            .add(methodName);
       } catch (IllegalArgumentException e) {
         logger.log(
             Level.WARNING,