diff --git a/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbImplInstrumentation.java b/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbImplInstrumentation.java index 92456cc181da..72e5ee14bb5f 100644 --- a/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbImplInstrumentation.java +++ b/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,67 +69,48 @@ 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); } } @@ -138,21 +118,17 @@ public static void onExit( 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); } } } diff --git a/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbInstrumentationModule.java b/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbInstrumentationModule.java index 3b95cc8b1c29..3aba153fc408 100644 --- a/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbInstrumentationModule.java +++ b/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 typeInstrumentations() { return singletonList(new InfluxDbImplInstrumentation()); } - - @Override - public boolean isIndyModule() { - // Uses multiple Advice.Locals and argument assignment - return false; - } } diff --git a/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbScope.java b/instrumentation/influxdb-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbScope.java new file mode 100644 index 000000000000..6daaea80f727 --- /dev/null +++ b/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); + } +} 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 index 1f1059f7e492..1382acc66e2d 100644 --- 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 @@ -9,14 +9,19 @@ 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 @@ -24,7 +29,7 @@ public void applyAdviceToMethod( ElementMatcher methodMatcher, String adviceClassName) { agentBuilder = agentBuilder.transform( - new AgentBuilder.Transformer.ForAdvice() + new AgentBuilder.Transformer.ForAdvice(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 index 55400956d2ad..ac1f46a99e3a 100644 --- 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 @@ -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 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,