|
26 | 26 | import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; |
27 | 27 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
28 | 28 | import java.lang.reflect.Method; |
| 29 | +import javax.annotation.Nullable; |
29 | 30 | import javax.ws.rs.Path; |
30 | 31 | import net.bytebuddy.asm.Advice; |
31 | 32 | import net.bytebuddy.description.type.TypeDescription; |
@@ -67,52 +68,62 @@ public void transform(TypeTransformer transformer) { |
67 | 68 | @SuppressWarnings("unused") |
68 | 69 | public static class JaxRsAnnotationsAdvice { |
69 | 70 |
|
70 | | - @Advice.OnMethodEnter(suppress = Throwable.class) |
71 | | - public static void nameSpan( |
72 | | - @Advice.This Object target, |
73 | | - @Advice.Origin Method method, |
74 | | - @Advice.Local("otelCallDepth") CallDepth callDepth, |
75 | | - @Advice.Local("otelHandlerData") HandlerData handlerData, |
76 | | - @Advice.Local("otelContext") Context context, |
77 | | - @Advice.Local("otelScope") Scope scope) { |
78 | | - callDepth = CallDepth.forClass(Path.class); |
79 | | - if (callDepth.getAndIncrement() > 0) { |
80 | | - return; |
| 71 | + public static class AdviceScope { |
| 72 | + private HandlerData handlerData; |
| 73 | + private final CallDepth callDepth; |
| 74 | + private Context context; |
| 75 | + private Scope scope; |
| 76 | + |
| 77 | + public AdviceScope(CallDepth callDepth) { |
| 78 | + this.callDepth = callDepth; |
81 | 79 | } |
82 | 80 |
|
83 | | - Context parentContext = Java8BytecodeBridge.currentContext(); |
84 | | - handlerData = new HandlerData(target.getClass(), method); |
| 81 | + public AdviceScope enter(Class<?> type, Method method) { |
| 82 | + if (callDepth.getAndIncrement() > 0) { |
| 83 | + return this; |
| 84 | + } |
| 85 | + |
| 86 | + Context parentContext = Java8BytecodeBridge.currentContext(); |
| 87 | + handlerData = new HandlerData(type, method); |
| 88 | + |
| 89 | + HttpServerRoute.update( |
| 90 | + parentContext, |
| 91 | + HttpServerRouteSource.CONTROLLER, |
| 92 | + JaxrsServerSpanNaming.SERVER_SPAN_NAME, |
| 93 | + handlerData); |
| 94 | + |
| 95 | + if (!instrumenter().shouldStart(parentContext, handlerData)) { |
| 96 | + return this; |
| 97 | + } |
| 98 | + |
| 99 | + context = instrumenter().start(parentContext, handlerData); |
| 100 | + scope = context.makeCurrent(); |
| 101 | + return this; |
| 102 | + } |
85 | 103 |
|
86 | | - HttpServerRoute.update( |
87 | | - parentContext, |
88 | | - HttpServerRouteSource.CONTROLLER, |
89 | | - JaxrsServerSpanNaming.SERVER_SPAN_NAME, |
90 | | - handlerData); |
| 104 | + public void exit(@Nullable Throwable throwable) { |
| 105 | + if (callDepth.decrementAndGet() > 0) { |
| 106 | + return; |
| 107 | + } |
91 | 108 |
|
92 | | - if (!instrumenter().shouldStart(parentContext, handlerData)) { |
93 | | - return; |
| 109 | + if (scope == null) { |
| 110 | + return; |
| 111 | + } |
| 112 | + scope.close(); |
| 113 | + instrumenter().end(context, handlerData, null, throwable); |
94 | 114 | } |
| 115 | + } |
95 | 116 |
|
96 | | - context = instrumenter().start(parentContext, handlerData); |
97 | | - scope = context.makeCurrent(); |
| 117 | + @Advice.OnMethodEnter(suppress = Throwable.class) |
| 118 | + public static AdviceScope nameSpan(@Advice.This Object target, @Advice.Origin Method method) { |
| 119 | + AdviceScope adviceScope = new AdviceScope(CallDepth.forClass(Path.class)); |
| 120 | + return adviceScope.enter(target.getClass(), method); |
98 | 121 | } |
99 | 122 |
|
100 | 123 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
101 | 124 | public static void stopSpan( |
102 | | - @Advice.Thrown Throwable throwable, |
103 | | - @Advice.Local("otelCallDepth") CallDepth callDepth, |
104 | | - @Advice.Local("otelHandlerData") HandlerData handlerData, |
105 | | - @Advice.Local("otelContext") Context context, |
106 | | - @Advice.Local("otelScope") Scope scope) { |
107 | | - if (callDepth.decrementAndGet() > 0) { |
108 | | - return; |
109 | | - } |
110 | | - |
111 | | - if (scope == null) { |
112 | | - return; |
113 | | - } |
114 | | - scope.close(); |
115 | | - instrumenter().end(context, handlerData, null, throwable); |
| 125 | + @Advice.Thrown @Nullable Throwable throwable, @Advice.Enter AdviceScope adviceScope) { |
| 126 | + adviceScope.exit(throwable); |
116 | 127 | } |
117 | 128 | } |
118 | 129 | } |
0 commit comments