|
26 | 26 | import java.sql.Connection; |
27 | 27 | import java.sql.PreparedStatement; |
28 | 28 | import java.util.Locale; |
| 29 | +import javax.annotation.Nullable; |
29 | 30 | import net.bytebuddy.asm.Advice; |
30 | 31 | import net.bytebuddy.description.type.TypeDescription; |
31 | 32 | import net.bytebuddy.matcher.ElementMatcher; |
@@ -68,35 +69,50 @@ public static void addDbInfo( |
68 | 69 | @SuppressWarnings("unused") |
69 | 70 | public static class TransactionAdvice { |
70 | 71 |
|
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()); |
83 | 97 | } |
84 | 98 |
|
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); |
87 | 109 | } |
88 | 110 |
|
89 | 111 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
90 | 112 | 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); |
100 | 116 | } |
101 | 117 | } |
102 | 118 | } |
0 commit comments