Allow advice to define custom mappings (#13751)

laurit · web-flow · commit 954ffad4fd6c · 2025-04-23T08:19:38.000+03:00
diff --git a/instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodInstrumentation.java b/instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodInstrumentation.java
@@ -19,7 +19,6 @@
 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.lang.reflect.Method;
 import java.util.Set;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -57,9 +56,18 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         namedOneOf(methodNames.toArray(new String[0])),
+        mapping ->
+            mapping.bind(
+                MethodReturnType.class,
+                (instrumentedType, instrumentedMethod, assigner, argumentHandler, sort) ->
+                    Advice.OffsetMapping.Target.ForStackManipulation.of(
+                        instrumentedMethod.getReturnType().asErasure())),
         MethodInstrumentation.class.getName() + "$MethodAdvice");
   }
 
+  // custom annotation that represents the return type of the method
+  @interface MethodReturnType {}
+
   @SuppressWarnings("unused")
   public static class MethodAdvice {
 
@@ -82,7 +90,7 @@ public static void onEnter(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
-        @Advice.Origin Method method,
+        @MethodReturnType Class<?> methodReturnType,
         @Advice.Local("otelMethod") ClassAndMethod classAndMethod,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope,
@@ -94,7 +102,7 @@ public static void stopSpan(
       scope.close();
 
       returnValue =
-          AsyncOperationEndSupport.create(instrumenter(), Void.class, method.getReturnType())
+          AsyncOperationEndSupport.create(instrumenter(), Void.class, methodReturnType)
               .asyncEnd(context, classAndMethod, returnValue, throwable);
     }
   }
diff --git a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/TypeTransformer.java b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/TypeTransformer.java
@@ -5,7 +5,9 @@
 
 package io.opentelemetry.javaagent.extension.instrumentation;
 
+import java.util.function.Function;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -17,12 +19,23 @@
  * will provide the implementation of all transformations described here.
  */
 public interface TypeTransformer {
+  /**
+   * Apply the advice class named {@code adviceClassName} to the instrumented type methods that
+   * match {@code methodMatcher}.
+   */
+  default void applyAdviceToMethod(
+      ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
+    applyAdviceToMethod(methodMatcher, Function.identity(), adviceClassName);
+  }
+
   /**
    * Apply the advice class named {@code adviceClassName} to the instrumented type methods that
    * match {@code methodMatcher}.
    */
   void applyAdviceToMethod(
-      ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName);
+      ElementMatcher<? super MethodDescription> methodMatcher,
+      Function<Advice.WithCustomMapping, Advice.WithCustomMapping> mappingCustomizer,
+      String adviceClassName);
 
   /**
    * Apply a custom ByteBuddy {@link AgentBuilder.Transformer} to the instrumented type. Note that
diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/TypeTransformerImpl.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/TypeTransformerImpl.java
@@ -8,6 +8,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.bytebuddy.ExceptionHandlers;
+import java.util.function.Function;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
@@ -26,10 +27,12 @@ final class TypeTransformerImpl implements TypeTransformer {
 
   @Override
   public void applyAdviceToMethod(
-      ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
+      ElementMatcher<? super MethodDescription> methodMatcher,
+      Function<Advice.WithCustomMapping, Advice.WithCustomMapping> mappingCustomizer,
+      String adviceClassName) {
     agentBuilder =
         agentBuilder.transform(
-            new AgentBuilder.Transformer.ForAdvice(adviceMapping)
+            new AgentBuilder.Transformer.ForAdvice(mappingCustomizer.apply(adviceMapping))
                 .include(
                     Utils.getBootstrapProxy(),
                     Utils.getAgentClassLoader(),
diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/AdviceTransformer.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/AdviceTransformer.java
@@ -155,8 +155,12 @@ private static class AdviceLocal {
   private static List<OutputArgument> getWritableArguments(MethodNode source) {
     List<OutputArgument> result = new ArrayList<>();
     if (source.visibleParameterAnnotations != null) {
-      int i = 0;
-      for (List<AnnotationNode> list : source.visibleParameterAnnotations) {
+      for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {
+        List<AnnotationNode> list = source.visibleParameterAnnotations[i];
+        if (list == null) {
+          continue;
+        }
+
         for (AnnotationNode annotationNode : list) {
           Type annotationType = Type.getType(annotationNode.desc);
           if (ADVICE_ARGUMENT.equals(annotationType) && isWriteable(annotationNode)) {
@@ -166,7 +170,6 @@ private static List<OutputArgument> getWritableArguments(MethodNode source) {
             }
           }
         }
-        i++;
       }
     }
 
@@ -178,15 +181,18 @@ private static List<OutputArgument> getWritableArguments(MethodNode source) {
   /** Argument annotated with {@code @Advice.Return(readOnly = false)} or {@code null}. */
   private static OutputArgument getWritableReturnValue(MethodNode source) {
     if (source.visibleParameterAnnotations != null) {
-      int i = 0;
-      for (List<AnnotationNode> list : source.visibleParameterAnnotations) {
+      for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {
+        List<AnnotationNode> list = source.visibleParameterAnnotations[i];
+        if (list == null) {
+          continue;
+        }
+
         for (AnnotationNode annotationNode : list) {
           Type annotationType = Type.getType(annotationNode.desc);
           if (ADVICE_RETURN.equals(annotationType) && isWriteable(annotationNode)) {
             return new OutputArgument(i, -1);
           }
         }
-        i++;
       }
     }
 
@@ -199,16 +205,19 @@ private static OutputArgument getWritableReturnValue(MethodNode source) {
   private static OutputArgument getEnterArgument(MethodNode source) {
     Type[] argumentTypes = Type.getArgumentTypes(source.desc);
     if (source.visibleParameterAnnotations != null) {
-      int i = 0;
-      for (List<AnnotationNode> list : source.visibleParameterAnnotations) {
+      for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {
+        List<AnnotationNode> list = source.visibleParameterAnnotations[i];
+        if (list == null) {
+          continue;
+        }
+
         for (AnnotationNode annotationNode : list) {
           Type annotationType = Type.getType(annotationNode.desc);
           if (ADVICE_ENTER.equals(annotationType)
               && argumentTypes[i].getDescriptor().length() > 1) {
             return new OutputArgument(i, -1);
           }
         }
-        i++;
       }
     }
 
@@ -221,8 +230,12 @@ private static OutputArgument getEnterArgument(MethodNode source) {
   private static List<AdviceLocal> getLocals(MethodNode source) {
     List<AdviceLocal> result = new ArrayList<>();
     if (source.visibleParameterAnnotations != null) {
-      int i = 0;
-      for (List<AnnotationNode> list : source.visibleParameterAnnotations) {
+      for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {
+        List<AnnotationNode> list = source.visibleParameterAnnotations[i];
+        if (list == null) {
+          continue;
+        }
+
         for (AnnotationNode annotationNode : list) {
           Type annotationType = Type.getType(annotationNode.desc);
           if (ADVICE_LOCAL.equals(annotationType)) {
@@ -232,7 +245,6 @@ private static List<AdviceLocal> getLocals(MethodNode source) {
             }
           }
         }
-        i++;
       }
     }
 
diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyTypeTransformerImpl.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyTypeTransformerImpl.java
@@ -12,6 +12,7 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
@@ -49,13 +50,15 @@ public IndyTypeTransformerImpl(
 
   @Override
   public void applyAdviceToMethod(
-      ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
+      ElementMatcher<? super MethodDescription> methodMatcher,
+      Function<Advice.WithCustomMapping, Advice.WithCustomMapping> mappingCustomizer,
+      String adviceClassName) {
     // default strategy used by AgentBuilder.Transformer.ForAdvice
     AgentBuilder.PoolStrategy poolStrategy = AgentBuilder.PoolStrategy.Default.FAST;
 
     agentBuilder =
         agentBuilder.transform(
-            new AgentBuilder.Transformer.ForAdvice(adviceMapping)
+            new AgentBuilder.Transformer.ForAdvice(mappingCustomizer.apply(adviceMapping))
                 .advice(methodMatcher, adviceClassName)
                 // advice transformation already performs uninlining
                 .with(
diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoader.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoader.java
@@ -25,9 +25,11 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.StringMatcher;
@@ -181,7 +183,9 @@ private static Set<String> getModuleAdviceNames(InstrumentationModule module) {
         new TypeTransformer() {
           @Override
           public void applyAdviceToMethod(
-              ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
+              ElementMatcher<? super MethodDescription> methodMatcher,
+              Function<Advice.WithCustomMapping, Advice.WithCustomMapping> mappingCustomizer,
+              String adviceClassName) {
             adviceNames.add(adviceClassName);
           }
 
diff --git a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/generation/AdviceClassNameCollector.java b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/generation/AdviceClassNameCollector.java
@@ -8,7 +8,9 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.function.Function;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -17,7 +19,9 @@ final class AdviceClassNameCollector implements TypeTransformer {
 
   @Override
   public void applyAdviceToMethod(
-      ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
+      ElementMatcher<? super MethodDescription> methodMatcher,
+      Function<Advice.WithCustomMapping, Advice.WithCustomMapping> mappingCustomizer,
+      String adviceClassName) {
     adviceClassNames.add(adviceClassName);
   }
 

Original file line number	Diff line number	Diff line change
`@@ -155,8 +155,12 @@ private static class AdviceLocal {`
`155`	`155`	`private static List<OutputArgument> getWritableArguments(MethodNode source) {`
`156`	`156`	`List<OutputArgument> result = new ArrayList<>();`
`157`	`157`	`if (source.visibleParameterAnnotations != null) {`
`158`		`- int i = 0;`
`159`		`- for (List<AnnotationNode> list : source.visibleParameterAnnotations) {`
	`158`	`+ for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {`
	`159`	`+ List<AnnotationNode> list = source.visibleParameterAnnotations[i];`
	`160`	`+ if (list == null) {`
	`161`	`+ continue;`
	`162`	`+ }`
	`163`	`+`
`160`	`164`	`for (AnnotationNode annotationNode : list) {`
`161`	`165`	`Type annotationType = Type.getType(annotationNode.desc);`
`162`	`166`	`if (ADVICE_ARGUMENT.equals(annotationType) && isWriteable(annotationNode)) {`
`@@ -166,7 +170,6 @@ private static List<OutputArgument> getWritableArguments(MethodNode source) {`
`166`	`170`	`}`
`167`	`171`	`}`
`168`	`172`	`}`
`169`		`- i++;`
`170`	`173`	`}`
`171`	`174`	`}`
`172`	`175`
`@@ -178,15 +181,18 @@ private static List<OutputArgument> getWritableArguments(MethodNode source) {`
`178`	`181`	`/** Argument annotated with {@code @Advice.Return(readOnly = false)} or {@code null}. */`
`179`	`182`	`private static OutputArgument getWritableReturnValue(MethodNode source) {`
`180`	`183`	`if (source.visibleParameterAnnotations != null) {`
`181`		`- int i = 0;`
`182`		`- for (List<AnnotationNode> list : source.visibleParameterAnnotations) {`
	`184`	`+ for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {`
	`185`	`+ List<AnnotationNode> list = source.visibleParameterAnnotations[i];`
	`186`	`+ if (list == null) {`
	`187`	`+ continue;`
	`188`	`+ }`
	`189`	`+`
`183`	`190`	`for (AnnotationNode annotationNode : list) {`
`184`	`191`	`Type annotationType = Type.getType(annotationNode.desc);`
`185`	`192`	`if (ADVICE_RETURN.equals(annotationType) && isWriteable(annotationNode)) {`
`186`	`193`	`return new OutputArgument(i, -1);`
`187`	`194`	`}`
`188`	`195`	`}`
`189`		`- i++;`
`190`	`196`	`}`
`191`	`197`	`}`
`192`	`198`
`@@ -199,16 +205,19 @@ private static OutputArgument getWritableReturnValue(MethodNode source) {`
`199`	`205`	`private static OutputArgument getEnterArgument(MethodNode source) {`
`200`	`206`	`Type[] argumentTypes = Type.getArgumentTypes(source.desc);`
`201`	`207`	`if (source.visibleParameterAnnotations != null) {`
`202`		`- int i = 0;`
`203`		`- for (List<AnnotationNode> list : source.visibleParameterAnnotations) {`
	`208`	`+ for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {`
	`209`	`+ List<AnnotationNode> list = source.visibleParameterAnnotations[i];`
	`210`	`+ if (list == null) {`
	`211`	`+ continue;`
	`212`	`+ }`
	`213`	`+`
`204`	`214`	`for (AnnotationNode annotationNode : list) {`
`205`	`215`	`Type annotationType = Type.getType(annotationNode.desc);`
`206`	`216`	`if (ADVICE_ENTER.equals(annotationType)`
`207`	`217`	`&& argumentTypes[i].getDescriptor().length() > 1) {`
`208`	`218`	`return new OutputArgument(i, -1);`
`209`	`219`	`}`
`210`	`220`	`}`
`211`		`- i++;`
`212`	`221`	`}`
`213`	`222`	`}`
`214`	`223`
`@@ -221,8 +230,12 @@ private static OutputArgument getEnterArgument(MethodNode source) {`
`221`	`230`	`private static List<AdviceLocal> getLocals(MethodNode source) {`
`222`	`231`	`List<AdviceLocal> result = new ArrayList<>();`
`223`	`232`	`if (source.visibleParameterAnnotations != null) {`
`224`		`- int i = 0;`
`225`		`- for (List<AnnotationNode> list : source.visibleParameterAnnotations) {`
	`233`	`+ for (int i = 0; i < source.visibleParameterAnnotations.length; i++) {`
	`234`	`+ List<AnnotationNode> list = source.visibleParameterAnnotations[i];`
	`235`	`+ if (list == null) {`
	`236`	`+ continue;`
	`237`	`+ }`
	`238`	`+`
`226`	`239`	`for (AnnotationNode annotationNode : list) {`
`227`	`240`	`Type annotationType = Type.getType(annotationNode.desc);`
`228`	`241`	`if (ADVICE_LOCAL.equals(annotationType)) {`
`@@ -232,7 +245,6 @@ private static List<AdviceLocal> getLocals(MethodNode source) {`
`232`	`245`	`}`
`233`	`246`	`}`
`234`	`247`	`}`
`235`		`- i++;`
`236`	`248`	`}`
`237`	`249`	`}`
`238`	`250`