Influxdb indy support (#11710)

Co-authored-by: Jonas Kunz <jonas.kunz@elastic.co>
Co-authored-by: Lauri Tulmin <ltulmin@splunk.com>
This commit is contained in:
César 2024-08-13 15:35:00 +02:00 committed by GitHub
parent 2af49b4c04
commit 2100a1646c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 114 additions and 70 deletions

View File

@ -15,7 +15,6 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
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 class InfluxDbImplInstrumentation implements TypeInstrumentation {
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) {
@Advice.Thrown Throwable throwable, @Advice.Enter Object[] enterArgs) {
CallDepth callDepth = CallDepth.forClass(InfluxDBImpl.class);
if (callDepth.decrementAndGet() > 0 || enterArgs == null) {
return;
}
if (scope == null) {
return;
}
scope.close();
instrumenter().end(context, influxDbRequest, null, throwable);
((InfluxDbScope) enterArgs[1]).end(throwable);
}
}
@ -138,21 +118,17 @@ public class InfluxDbImplInstrumentation implements TypeInstrumentation {
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 class InfluxDbImplInstrumentation implements TypeInstrumentation {
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) {
@Advice.Thrown Throwable throwable, @Advice.Enter InfluxDbScope scope) {
CallDepth callDepth = CallDepth.forClass(InfluxDBImpl.class);
if (callDepth.decrementAndGet() > 0 || scope == null) {
return;
}
if (scope == null) {
return;
}
scope.close();
instrumenter().end(context, influxDbRequest, null, throwable);
scope.end(throwable);
}
}
}

View File

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

View File

@ -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);
}
}

View File

@ -9,14 +9,19 @@ import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
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 @@ final class TypeTransformerImpl implements TypeTransformer {
ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
agentBuilder =
agentBuilder.transform(
new AgentBuilder.Transformer.ForAdvice()
new AgentBuilder.Transformer.ForAdvice(adviceMapping)
.include(
Utils.getBootstrapProxy(),
Utils.getAgentClassLoader(),

View File

@ -50,6 +50,10 @@ class AdviceTransformer {
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 @@ class AdviceTransformer {
@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 @@ class AdviceTransformer {
}
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 @@ class AdviceTransformer {
}
}
applyAdviceDelegation(context, methodNode, classVisitor, exceptionsArray);
}
private static void applyAdviceDelegation(
TransformationContext context,
MethodNode methodNode,
ClassVisitor classVisitor,
String[] exceptionsArray) {
MethodVisitor mv =
classVisitor.visitMethod(
methodNode.access,