Skip to content

Commit

Permalink
Influxdb indy support (#11710)
Browse files Browse the repository at this point in the history 
Co-authored-by: Jonas Kunz <[email protected]>
Co-authored-by: Lauri Tulmin <[email protected]>
  • Loading branch information
@LikeTheSalad @JonasKunz @laurit
3 people authored Aug 13, 2024
1 parent 2af49b4 commit 2100a16
Show file tree
Hide file tree
Showing 5 changed files with 114 additions and 70 deletions.
88 changes: 27 additions & 61 deletions ...io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbImplInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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);
}
}
}
6 changes: 0 additions & 6 deletions .../opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbInstrumentationModule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
39 changes: 39 additions & 0 deletions ...src/main/java/io/opentelemetry/javaagent/instrumentation/influxdb/v2_4/InfluxDbScope.java
View file
Original file line number Diff line number Diff line change
@@ -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);
}
}
7 changes: 6 additions & 1 deletion ...src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/TypeTransformerImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
44 changes: 42 additions & 2 deletions .../main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/AdviceTransformer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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);
}
}
};
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 2100a16

Please sign in to comment.