|
5 | 5 |
|
6 | 6 | package io.opentelemetry.javaagent.instrumentation.geode; |
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.implementsInterface; |
11 | 10 | import static io.opentelemetry.javaagent.instrumentation.geode.GeodeSingletons.instrumenter; |
|
19 | 18 | import io.opentelemetry.context.Scope; |
20 | 19 | import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; |
21 | 20 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
| 21 | +import javax.annotation.Nullable; |
22 | 22 | import net.bytebuddy.asm.Advice; |
23 | 23 | import net.bytebuddy.description.type.TypeDescription; |
24 | 24 | import net.bytebuddy.matcher.ElementMatcher; |
@@ -61,76 +61,77 @@ public void transform(TypeTransformer transformer) { |
61 | 61 | this.getClass().getName() + "$QueryAdvice"); |
62 | 62 | } |
63 | 63 |
|
64 | | - @SuppressWarnings("unused") |
65 | | - public static class SimpleAdvice { |
| 64 | + public static class AdviceScope { |
| 65 | + private final GeodeRequest request; |
| 66 | + private final Context context; |
| 67 | + private final Scope scope; |
66 | 68 |
|
67 | | - @Advice.OnMethodEnter(suppress = Throwable.class) |
68 | | - public static void onEnter( |
69 | | - @Advice.This Region<?, ?> region, |
70 | | - @Advice.Origin("#m") String methodName, |
71 | | - @Advice.Local("otelRequest") GeodeRequest request, |
72 | | - @Advice.Local("otelContext") Context context, |
73 | | - @Advice.Local("otelScope") Scope scope) { |
| 69 | + public AdviceScope(GeodeRequest request, Context context, Scope scope) { |
| 70 | + this.request = request; |
| 71 | + this.context = context; |
| 72 | + this.scope = scope; |
| 73 | + } |
74 | 74 |
|
75 | | - Context parentContext = currentContext(); |
76 | | - request = GeodeRequest.create(region, methodName, null); |
| 75 | + @Nullable |
| 76 | + public static AdviceScope start( |
| 77 | + Region<?, ?> region, String methodName, @Nullable String query) { |
| 78 | + Context parentContext = Context.current(); |
| 79 | + GeodeRequest request = GeodeRequest.create(region, methodName, query); |
77 | 80 | if (!instrumenter().shouldStart(parentContext, request)) { |
78 | | - return; |
| 81 | + return null; |
79 | 82 | } |
80 | 83 |
|
81 | | - context = instrumenter().start(parentContext, request); |
82 | | - scope = context.makeCurrent(); |
| 84 | + Context context = instrumenter().start(parentContext, request); |
| 85 | + return new AdviceScope(request, context, context.makeCurrent()); |
| 86 | + } |
| 87 | + |
| 88 | + public void end(@Nullable Throwable throwable) { |
| 89 | + if (scope != null) { |
| 90 | + scope.close(); |
| 91 | + } |
| 92 | + instrumenter().end(context, request, null, throwable); |
| 93 | + } |
| 94 | + } |
| 95 | + |
| 96 | + @SuppressWarnings("unused") |
| 97 | + public static class SimpleAdvice { |
| 98 | + |
| 99 | + @Nullable |
| 100 | + @Advice.OnMethodEnter(suppress = Throwable.class) |
| 101 | + public static AdviceScope onEnter( |
| 102 | + @Advice.This Region<?, ?> region, @Advice.Origin("#m") String methodName) { |
| 103 | + return AdviceScope.start(region, methodName, null); |
83 | 104 | } |
84 | 105 |
|
85 | 106 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
86 | 107 | public static void stopSpan( |
87 | | - @Advice.Thrown Throwable throwable, |
88 | | - @Advice.Local("otelRequest") GeodeRequest request, |
89 | | - @Advice.Local("otelContext") Context context, |
90 | | - @Advice.Local("otelScope") Scope scope) { |
91 | | - if (scope == null) { |
92 | | - return; |
| 108 | + @Advice.Thrown @Nullable Throwable throwable, |
| 109 | + @Advice.Enter @Nullable AdviceScope adviceScope) { |
| 110 | + if (adviceScope != null) { |
| 111 | + adviceScope.end(throwable); |
93 | 112 | } |
94 | | - |
95 | | - scope.close(); |
96 | | - instrumenter().end(context, request, null, throwable); |
97 | 113 | } |
98 | 114 | } |
99 | 115 |
|
100 | 116 | @SuppressWarnings("unused") |
101 | 117 | public static class QueryAdvice { |
102 | 118 |
|
| 119 | + @Nullable |
103 | 120 | @Advice.OnMethodEnter(suppress = Throwable.class) |
104 | | - public static void onEnter( |
| 121 | + public static AdviceScope onEnter( |
105 | 122 | @Advice.This Region<?, ?> region, |
106 | 123 | @Advice.Origin("#m") String methodName, |
107 | | - @Advice.Argument(0) String query, |
108 | | - @Advice.Local("otelRequest") GeodeRequest request, |
109 | | - @Advice.Local("otelContext") Context context, |
110 | | - @Advice.Local("otelScope") Scope scope) { |
111 | | - |
112 | | - Context parentContext = currentContext(); |
113 | | - request = GeodeRequest.create(region, methodName, query); |
114 | | - if (!instrumenter().shouldStart(parentContext, request)) { |
115 | | - return; |
116 | | - } |
117 | | - |
118 | | - context = instrumenter().start(parentContext, request); |
119 | | - scope = context.makeCurrent(); |
| 124 | + @Advice.Argument(0) String query) { |
| 125 | + return AdviceScope.start(region, methodName, query); |
120 | 126 | } |
121 | 127 |
|
122 | 128 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
123 | 129 | public static void stopSpan( |
124 | | - @Advice.Thrown Throwable throwable, |
125 | | - @Advice.Local("otelRequest") GeodeRequest request, |
126 | | - @Advice.Local("otelContext") Context context, |
127 | | - @Advice.Local("otelScope") Scope scope) { |
128 | | - if (scope == null) { |
129 | | - return; |
| 130 | + @Advice.Thrown @Nullable Throwable throwable, |
| 131 | + @Advice.Enter @Nullable AdviceScope adviceScope) { |
| 132 | + if (adviceScope != null) { |
| 133 | + adviceScope.end(throwable); |
130 | 134 | } |
131 | | - |
132 | | - scope.close(); |
133 | | - instrumenter().end(context, request, null, throwable); |
134 | 135 | } |
135 | 136 | } |
136 | 137 | } |
0 commit comments