|
5 | 5 |
|
6 | 6 | package io.opentelemetry.javaagent.instrumentation.methods; |
7 | 7 |
|
8 | | -import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext; |
9 | 8 | import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed; |
10 | 9 | import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType; |
11 | 10 | import static io.opentelemetry.javaagent.instrumentation.methods.MethodSingletons.getBootstrapLoader; |
|
25 | 24 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
26 | 25 | import java.util.Collection; |
27 | 26 | import java.util.Map; |
| 27 | +import javax.annotation.Nullable; |
28 | 28 | import net.bytebuddy.asm.Advice; |
| 29 | +import net.bytebuddy.asm.Advice.AssignReturned; |
29 | 30 | import net.bytebuddy.description.enumeration.EnumerationDescription; |
30 | 31 | import net.bytebuddy.description.type.TypeDescription; |
31 | 32 | import net.bytebuddy.implementation.bytecode.assign.Assigner; |
@@ -91,42 +92,60 @@ public void transform(TypeTransformer transformer) { |
91 | 92 | @SuppressWarnings("unused") |
92 | 93 | public static class MethodAdvice { |
93 | 94 |
|
| 95 | + public static class AdviceScope { |
| 96 | + private final MethodAndType classAndMethod; |
| 97 | + private final Context context; |
| 98 | + private final Scope scope; |
| 99 | + |
| 100 | + private AdviceScope(MethodAndType classAndMethod, Context context, Scope scope) { |
| 101 | + this.classAndMethod = classAndMethod; |
| 102 | + this.context = context; |
| 103 | + this.scope = scope; |
| 104 | + } |
| 105 | + |
| 106 | + @Nullable |
| 107 | + public static AdviceScope start( |
| 108 | + SpanKind spanKind, Class<?> declaringClass, String methodName) { |
| 109 | + Context parentContext = Context.current(); |
| 110 | + MethodAndType methodAndType = |
| 111 | + MethodAndType.create(ClassAndMethod.create(declaringClass, methodName), spanKind); |
| 112 | + |
| 113 | + if (!instrumenter().shouldStart(parentContext, methodAndType)) { |
| 114 | + return null; |
| 115 | + } |
| 116 | + Context context = instrumenter().start(parentContext, methodAndType); |
| 117 | + return new AdviceScope(methodAndType, context, context.makeCurrent()); |
| 118 | + } |
| 119 | + |
| 120 | + public Object end( |
| 121 | + Class<?> methodReturnType, Object returnValue, @Nullable Throwable throwable) { |
| 122 | + scope.close(); |
| 123 | + return AsyncOperationEndSupport.create(instrumenter(), Void.class, methodReturnType) |
| 124 | + .asyncEnd(context, classAndMethod, returnValue, throwable); |
| 125 | + } |
| 126 | + } |
| 127 | + |
| 128 | + @Nullable |
94 | 129 | @Advice.OnMethodEnter(suppress = Throwable.class) |
95 | | - public static void onEnter( |
| 130 | + public static AdviceScope onEnter( |
96 | 131 | @MethodSpanKind SpanKind spanKind, |
97 | 132 | @Advice.Origin("#t") Class<?> declaringClass, |
98 | | - @Advice.Origin("#m") String methodName, |
99 | | - @Advice.Local("otelMethod") MethodAndType classAndMethod, |
100 | | - @Advice.Local("otelContext") Context context, |
101 | | - @Advice.Local("otelScope") Scope scope) { |
102 | | - Context parentContext = currentContext(); |
103 | | - classAndMethod = |
104 | | - MethodAndType.create(ClassAndMethod.create(declaringClass, methodName), spanKind); |
105 | | - |
106 | | - if (!instrumenter().shouldStart(parentContext, classAndMethod)) { |
107 | | - return; |
108 | | - } |
109 | | - |
110 | | - context = instrumenter().start(parentContext, classAndMethod); |
111 | | - scope = context.makeCurrent(); |
| 133 | + @Advice.Origin("#m") String methodName) { |
| 134 | + return AdviceScope.start(spanKind, declaringClass, methodName); |
112 | 135 | } |
113 | 136 |
|
| 137 | + @AssignReturned.ToReturned |
114 | 138 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
115 | | - public static void stopSpan( |
| 139 | + public static Object stopSpan( |
116 | 140 | @MethodReturnType Class<?> methodReturnType, |
117 | | - @Advice.Local("otelMethod") MethodAndType classAndMethod, |
118 | | - @Advice.Local("otelContext") Context context, |
119 | | - @Advice.Local("otelScope") Scope scope, |
120 | | - @Advice.Return(typing = Assigner.Typing.DYNAMIC, readOnly = false) Object returnValue, |
121 | | - @Advice.Thrown Throwable throwable) { |
122 | | - if (scope == null) { |
123 | | - return; |
124 | | - } |
125 | | - scope.close(); |
| 141 | + @Advice.Return(typing = Assigner.Typing.DYNAMIC) Object returnValue, |
| 142 | + @Advice.Thrown @Nullable Throwable throwable, |
| 143 | + @Advice.Enter @Nullable AdviceScope adviceScope) { |
126 | 144 |
|
127 | | - returnValue = |
128 | | - AsyncOperationEndSupport.create(instrumenter(), Void.class, methodReturnType) |
129 | | - .asyncEnd(context, classAndMethod, returnValue, throwable); |
| 145 | + if (adviceScope != null) { |
| 146 | + return adviceScope.end(methodReturnType, returnValue, throwable); |
| 147 | + } |
| 148 | + return returnValue; |
130 | 149 | } |
131 | 150 | } |
132 | 151 | } |
0 commit comments