Skip to content

Commit cafb0e4

Browse files
committed
migrate last bit to indy
1 parent d520a28 commit cafb0e4

File tree

1 file changed

+39
-23
lines changed

1 file changed

+39
-23
lines changed

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/ConnectionInstrumentation.java

Lines changed: 39 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.sql.Connection;
2727
import java.sql.PreparedStatement;
2828
import java.util.Locale;
29+
import javax.annotation.Nullable;
2930
import net.bytebuddy.asm.Advice;
3031
import net.bytebuddy.description.type.TypeDescription;
3132
import net.bytebuddy.matcher.ElementMatcher;
@@ -68,35 +69,50 @@ public static void addDbInfo(
6869
@SuppressWarnings("unused")
6970
public static class TransactionAdvice {
7071

71-
@Advice.OnMethodEnter(suppress = Throwable.class)
72-
public static void onEnter(
73-
@Advice.This Connection connection,
74-
@Advice.Origin("#m") String methodName,
75-
@Advice.Local("otelContext") Context context,
76-
@Advice.Local("otelScope") Scope scope) {
77-
Context parentContext = currentContext();
78-
DbRequest request =
79-
DbRequest.createTransaction(connection, methodName.toUpperCase(Locale.ROOT));
80-
81-
if (request == null || !transactionInstrumenter().shouldStart(parentContext, request)) {
82-
return;
72+
public static final class AdviceScope {
73+
private final DbRequest request;
74+
private final Context context;
75+
private final Scope scope;
76+
77+
private AdviceScope(DbRequest request, Context context, Scope scope) {
78+
this.request = request;
79+
this.context = context;
80+
this.scope = scope;
81+
}
82+
83+
@Nullable
84+
public static AdviceScope start(Connection connection, String methodName) {
85+
DbRequest request =
86+
DbRequest.createTransaction(connection, methodName.toUpperCase(Locale.ROOT));
87+
if (request == null) {
88+
return null;
89+
}
90+
Context parentContext = currentContext();
91+
if (!transactionInstrumenter().shouldStart(parentContext, request)) {
92+
return null;
93+
}
94+
95+
Context context = transactionInstrumenter().start(parentContext, request);
96+
return new AdviceScope(request, context, context.makeCurrent());
8397
}
8498

85-
context = transactionInstrumenter().start(parentContext, request);
86-
scope = context.makeCurrent();
99+
public void end(@Nullable Throwable throwable) {
100+
scope.close();
101+
transactionInstrumenter().end(context, request, null, throwable);
102+
}
103+
}
104+
105+
@Advice.OnMethodEnter(suppress = Throwable.class)
106+
public static AdviceScope onEnter(
107+
@Advice.This Connection connection, @Advice.Origin("#m") String methodName) {
108+
return AdviceScope.start(connection, methodName);
87109
}
88110

89111
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
90112
public static void stopSpan(
91-
@Advice.Thrown Throwable throwable,
92-
@Advice.Local("otelRequest") DbRequest request,
93-
@Advice.Local("otelContext") Context context,
94-
@Advice.Local("otelScope") Scope scope) {
95-
if (scope == null) {
96-
return;
97-
}
98-
scope.close();
99-
transactionInstrumenter().end(context, request, null, throwable);
113+
@Advice.Thrown @Nullable Throwable throwable,
114+
@Advice.Enter @Nullable AdviceScope adviceScope) {
115+
if (adviceScope != null) adviceScope.end(throwable);
100116
}
101117
}
102118
}

0 commit comments

Comments
 (0)