Influxdb indy support

instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbImplInstrumentation.java

@@ -15,7 +15,6 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import io.opentelemetry.context.Context;
-import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -70,89 +69,66 @@ public void transform(TypeTransformer transformer) {
   public static class InfluxDbQueryAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.Argument(0) Query query,
-        @Advice.AllArguments(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object[] arguments,
-        @Advice.FieldValue(value = "retrofit") Retrofit retrofit,
-        @Advice.Local("otelCallDepth") CallDepth callDepth,
-        @Advice.Local("otelRequest") InfluxDbRequest influxDbRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      callDepth = CallDepth.forClass(InfluxDBImpl.class);
+    @Advice.AssignReturned.ToAllArguments(index = 0, typing = Assigner.Typing.DYNAMIC)
+    public static Object[] onEnter(
+        @Advice.AllArguments(typing = Assigner.Typing.DYNAMIC) Object[] arguments,
+        @Advice.FieldValue(value = "retrofit") Retrofit retrofit) {
+      CallDepth callDepth = CallDepth.forClass(InfluxDBImpl.class);
       if (callDepth.getAndIncrement() > 0) {
-        return;
+        return null;
       }
 
+      Query query = arguments[0] instanceof Query ? (Query) arguments[0] : null;
       if (query == null) {
-        return;
+        return null;
       }
       Context parentContext = currentContext();
 
       HttpUrl httpUrl = retrofit.baseUrl();
-      influxDbRequest =
+      InfluxDbRequest influxDbRequest =
           InfluxDbRequest.create(
               httpUrl.host(), httpUrl.port(), query.getDatabase(), null, query.getCommand());
 
       if (!instrumenter().shouldStart(parentContext, influxDbRequest)) {
-        return;
+        return null;
       }
 
       // wrap callbacks so they'd run in the context of the parent span
       Object[] newArguments = new Object[arguments.length];
-      boolean hasChangedArgument = false;
       for (int i = 0; i < arguments.length; i++) {
         newArguments[i] = InfluxDbObjetWrapper.wrap(arguments[i], parentContext);
-        hasChangedArgument |= newArguments[i] != arguments[i];
-      }
-      if (hasChangedArgument) {
-        arguments = newArguments;
       }
 
-      context = instrumenter().start(parentContext, influxDbRequest);
-      scope = context.makeCurrent();
+      return new Object[] {newArguments, InfluxDbScope.start(parentContext, influxDbRequest)};
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void onExit(
-        @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelCallDepth") CallDepth callDepth,
-        @Advice.Local("otelRequest") InfluxDbRequest influxDbRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-
-      if (callDepth.decrementAndGet() > 0) {
-        return;
-      }
-
-      if (scope == null) {
+        @Advice.Thrown Throwable throwable, @Advice.Enter Object[] enterArgs) {
+      CallDepth callDepth = CallDepth.forClass(InfluxDBImpl.class);
+      if (callDepth.decrementAndGet() > 0 || enterArgs == null) {
         return;
       }
 
-      scope.close();
-
-      instrumenter().end(context, influxDbRequest, null, throwable);
+      ((InfluxDbScope) enterArgs[1]).end(throwable);
     }
   }
 
   @SuppressWarnings("unused")
   public static class InfluxDbModifyAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
+    public static InfluxDbScope onEnter(
         @Advice.Origin("#m") String methodName,
         @Advice.Argument(0) Object arg0,
-        @Advice.FieldValue(value = "retrofit") Retrofit retrofit,
-        @Advice.Local("otelCallDepth") CallDepth callDepth,
-        @Advice.Local("otelRequest") InfluxDbRequest influxDbRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      callDepth = CallDepth.forClass(InfluxDBImpl.class);
+        @Advice.FieldValue(value = "retrofit") Retrofit retrofit) {
+      CallDepth callDepth = CallDepth.forClass(InfluxDBImpl.class);
       if (callDepth.getAndIncrement() > 0) {
-        return;
+        return null;
       }
 
       if (arg0 == null) {
-        return;
+        return null;
       }
 
       Context parentContext = currentContext();
@@ -173,35 +149,25 @@ public static void onEnter(
         operation = "WRITE";
       }
 
-      influxDbRequest =
+      InfluxDbRequest influxDbRequest =
           InfluxDbRequest.create(httpUrl.host(), httpUrl.port(), database, operation, null);
 
       if (!instrumenter().shouldStart(parentContext, influxDbRequest)) {
-        return;
+        return null;
       }
 
-      context = instrumenter().start(parentContext, influxDbRequest);
-      scope = context.makeCurrent();
+      return InfluxDbScope.start(parentContext, influxDbRequest);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void onExit(
-        @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelCallDepth") CallDepth callDepth,
-        @Advice.Local("otelRequest") InfluxDbRequest influxDbRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-
-      if (callDepth.decrementAndGet() > 0) {
-        return;
-      }
-
-      if (scope == null) {
+        @Advice.Thrown Throwable throwable, @Advice.Enter InfluxDbScope scope) {
+      CallDepth callDepth = CallDepth.forClass(InfluxDBImpl.class);
+      if (callDepth.decrementAndGet() > 0 || scope == null) {
         return;
       }
-      scope.close();
 
-      instrumenter().end(context, influxDbRequest, null, throwable);
+      scope.end(throwable);
     }
   }
 }

instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbInstrumentationModule.java

@@ -23,10 +23,4 @@ public InfluxDbInstrumentationModule() {
   public List<TypeInstrumentation> typeInstrumentations() {
     return singletonList(new InfluxDbImplInstrumentation());
   }
-
-  @Override
-  public boolean isIndyModule() {
-    // Uses multiple Advice.Locals and argument assignment
-    return false;
-  }
 }

instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbScope.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.influxdb.v2_4;
+
+import static io.opentelemetry.javaagent.instrumentation.influxdb.v2_4.InfluxDbSingletons.instrumenter;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+/** Container used to carry state between enter and exit advices */
+public final class InfluxDbScope {
+  private final InfluxDbRequest influxDbRequest;
+  private final Context context;
+  private final Scope scope;
+
+  private InfluxDbScope(InfluxDbRequest influxDbRequest, Context context, Scope scope) {
+    this.influxDbRequest = influxDbRequest;
+    this.context = context;
+    this.scope = scope;
+  }
+
+  public static InfluxDbScope start(Context parentContext, InfluxDbRequest influxDbRequest) {
+    Context context = instrumenter().start(parentContext, influxDbRequest);
+    return new InfluxDbScope(influxDbRequest, context, context.makeCurrent());
+  }
+
+  public void end(Throwable throwable) {
+    if (scope == null) {
+      return;
+    }
+
+    scope.close();
+
+    instrumenter().end(context, influxDbRequest, null, throwable);
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/TypeTransformerImpl.java

@@ -9,22 +9,27 @@
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.bytebuddy.ExceptionHandlers;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
 final class TypeTransformerImpl implements TypeTransformer {
   private AgentBuilder.Identified.Extendable agentBuilder;
+  private final Advice.WithCustomMapping adviceMapping;
 
   TypeTransformerImpl(AgentBuilder.Identified.Extendable agentBuilder) {
     this.agentBuilder = agentBuilder;
+    adviceMapping =
+        Advice.withCustomMapping()
+            .with(new Advice.AssignReturned.Factory().withSuppressed(Throwable.class));
   }
 
   @Override
   public void applyAdviceToMethod(
       ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
     agentBuilder =
         agentBuilder.transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            new AgentBuilder.Transformer.ForAdvice(adviceMapping)
                 .include(
                     Utils.getBootstrapProxy(),
                     Utils.getAgentClassLoader(),

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/AdviceTransformer.java

@@ -50,6 +50,10 @@ static byte[] transform(byte[] bytes) {
       return cw.toByteArray();
     }
 
+    // Advices already using Advice.AssignReturned are assumed to be already compatible
+    // Those won't be transformed except for setting inline to false
+    boolean justDelegateAdvice = usesAssignReturned(classNode);
+
     // sort enter advice method before exit advice
     classNode.methods.sort(
         Comparator.comparingInt(
@@ -74,8 +78,12 @@ public MethodVisitor visitMethod(
               @Override
               public void visitEnd() {
                 super.visitEnd();
-
-                instrument(context, this, classVisitor);
+                if (justDelegateAdvice) {
+                  applyAdviceDelegation(
+                      context, this, classVisitor, exceptions.toArray(new String[0]));
+                } else {
+                  instrument(context, this, classVisitor);
+                }
               }
             };
           }
@@ -226,6 +234,30 @@ private static List<AdviceLocal> getLocals(MethodNode source) {
   }
 
   static final Type ADVICE_ON_METHOD_ENTER = Type.getType(Advice.OnMethodEnter.class);
+  private static final Type ADVICE_ASSIGN_RETURNED_TO_RETURNED =
+      Type.getType(Advice.AssignReturned.ToReturned.class);
+  private static final Type ADVICE_ASSIGN_RETURNED_TO_ARGUMENTS =
+      Type.getType(Advice.AssignReturned.ToArguments.class);
+  private static final Type ADVICE_ASSIGN_RETURNED_TO_FIELDS =
+      Type.getType(Advice.AssignReturned.ToFields.class);
+  private static final Type ADVICE_ASSIGN_RETURNED_TO_ALL_ARGUMENTS =
+      Type.getType(Advice.AssignReturned.ToAllArguments.class);
+
+  private static boolean usesAssignReturned(MethodNode source) {
+    return hasAnnotation(source, ADVICE_ASSIGN_RETURNED_TO_RETURNED)
+        || hasAnnotation(source, ADVICE_ASSIGN_RETURNED_TO_ARGUMENTS)
+        || hasAnnotation(source, ADVICE_ASSIGN_RETURNED_TO_FIELDS)
+        || hasAnnotation(source, ADVICE_ASSIGN_RETURNED_TO_ALL_ARGUMENTS);
+  }
+
+  private static boolean usesAssignReturned(ClassNode classNode) {
+    for (MethodNode mn : classNode.methods) {
+      if (usesAssignReturned(mn)) {
+        return true;
+      }
+    }
+    return false;
+  }
 
   private static boolean isEnterAdvice(MethodNode source) {
     return hasAnnotation(source, ADVICE_ON_METHOD_ENTER);
@@ -656,6 +688,14 @@ private static void instrument(
       }
     }
 
+    applyAdviceDelegation(context, methodNode, classVisitor, exceptionsArray);
+  }
+
+  private static void applyAdviceDelegation(
+      TransformationContext context,
+      MethodNode methodNode,
+      ClassVisitor classVisitor,
+      String[] exceptionsArray) {
     MethodVisitor mv =
         classVisitor.visitMethod(
             methodNode.access,