|
16 | 16 | import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; |
17 | 17 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
18 | 18 | import java.lang.reflect.Method; |
| 19 | +import javax.annotation.Nullable; |
19 | 20 | import net.bytebuddy.asm.Advice; |
20 | 21 | import net.bytebuddy.description.type.TypeDescription; |
21 | 22 | import net.bytebuddy.matcher.ElementMatcher; |
@@ -53,35 +54,54 @@ public void transform(TypeTransformer transformer) { |
53 | 54 | @SuppressWarnings("unused") |
54 | 55 | public static class InvokeAndEncodeResponseAdvice { |
55 | 56 |
|
| 57 | + public static class AdviceScope { |
| 58 | + private final Context context; |
| 59 | + private final Scope scope; |
| 60 | + |
| 61 | + private AdviceScope(Context context, Scope scope) { |
| 62 | + this.context = context; |
| 63 | + this.scope = scope; |
| 64 | + } |
| 65 | + |
| 66 | + @Nullable |
| 67 | + public static AdviceScope start(Method method) { |
| 68 | + Context parentContext = Context.current(); |
| 69 | + if (!instrumenter().shouldStart(parentContext, method)) { |
| 70 | + return null; |
| 71 | + } |
| 72 | + Context context = |
| 73 | + instrumenter().start(parentContext, method).with(GwtSingletons.RPC_CONTEXT_KEY, true); |
| 74 | + return new AdviceScope(context, context.makeCurrent()); |
| 75 | + } |
| 76 | + |
| 77 | + public void end(Method method, @Nullable Throwable throwable) { |
| 78 | + scope.close(); |
| 79 | + instrumenter().end(context, method, null, throwable); |
| 80 | + } |
| 81 | + } |
| 82 | + |
| 83 | + @Nullable |
56 | 84 | @Advice.OnMethodEnter(suppress = Throwable.class) |
57 | | - public static void onEnter( |
58 | | - @Advice.Argument(1) Method method, |
59 | | - @Advice.Local("otelContext") Context context, |
60 | | - @Advice.Local("otelScope") Scope scope) { |
61 | | - context = |
62 | | - instrumenter() |
63 | | - .start(Java8BytecodeBridge.currentContext(), method) |
64 | | - .with(GwtSingletons.RPC_CONTEXT_KEY, true); |
65 | | - scope = context.makeCurrent(); |
| 85 | + public static AdviceScope onEnter(@Advice.Argument(1) Method method) { |
| 86 | + return AdviceScope.start(method); |
66 | 87 | } |
67 | 88 |
|
68 | 89 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
69 | 90 | public static void onExit( |
70 | 91 | @Advice.Argument(1) Method method, |
71 | | - @Advice.Local("otelContext") Context context, |
72 | | - @Advice.Local("otelScope") Scope scope, |
73 | | - @Advice.Thrown Throwable throwable) { |
74 | | - scope.close(); |
75 | | - |
76 | | - instrumenter().end(context, method, null, throwable); |
| 92 | + @Advice.Thrown @Nullable Throwable throwable, |
| 93 | + @Advice.Enter @Nullable AdviceScope adviceScope) { |
| 94 | + if (adviceScope != null) { |
| 95 | + adviceScope.end(method, throwable); |
| 96 | + } |
77 | 97 | } |
78 | 98 | } |
79 | 99 |
|
80 | 100 | @SuppressWarnings("unused") |
81 | 101 | public static class EncodeResponseForFailureAdvice { |
82 | 102 |
|
83 | 103 | @Advice.OnMethodEnter(suppress = Throwable.class) |
84 | | - public static void onEnter(@Advice.Argument(1) Throwable throwable) { |
| 104 | + public static void onEnter(@Advice.Argument(1) @Nullable Throwable throwable) { |
85 | 105 | if (throwable == null) { |
86 | 106 | return; |
87 | 107 | } |
|
0 commit comments