|
20 | 20 | import io.opentelemetry.context.Scope; |
21 | 21 | import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; |
22 | 22 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
| 23 | +import javax.annotation.Nullable; |
23 | 24 | import net.bytebuddy.asm.Advice; |
| 25 | +import net.bytebuddy.asm.Advice.AssignReturned; |
| 26 | +import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument; |
24 | 27 | import net.bytebuddy.description.type.TypeDescription; |
25 | 28 | import net.bytebuddy.matcher.ElementMatcher; |
26 | 29 | import org.apache.http.HttpHost; |
27 | 30 | import org.apache.http.HttpRequest; |
| 31 | +import org.apache.http.HttpResponse; |
28 | 32 | import org.apache.http.client.ResponseHandler; |
29 | 33 | import org.apache.http.client.methods.HttpUriRequest; |
30 | 34 |
|
@@ -123,171 +127,145 @@ public void transform(TypeTransformer transformer) { |
123 | 127 | this.getClass().getName() + "$RequestWithHandlerAdvice"); |
124 | 128 | } |
125 | 129 |
|
126 | | - @SuppressWarnings("unused") |
127 | | - public static class UriRequestAdvice { |
| 130 | + public static class AdviceScope { |
| 131 | + private final ApacheHttpClientRequest otelRequest; |
| 132 | + private final Context parentContext; |
| 133 | + private final Context context; |
| 134 | + private final Scope scope; |
| 135 | + |
| 136 | + private AdviceScope( |
| 137 | + ApacheHttpClientRequest otelRequest, Context parentContext, Context context, Scope scope) { |
| 138 | + this.otelRequest = otelRequest; |
| 139 | + this.parentContext = parentContext; |
| 140 | + this.context = context; |
| 141 | + this.scope = scope; |
| 142 | + } |
128 | 143 |
|
129 | | - @Advice.OnMethodEnter(suppress = Throwable.class) |
130 | | - public static void methodEnter( |
131 | | - @Advice.Argument(0) HttpUriRequest request, |
132 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
133 | | - @Advice.Local("otelContext") Context context, |
134 | | - @Advice.Local("otelScope") Scope scope) { |
| 144 | + @Nullable |
| 145 | + public static AdviceScope start(ApacheHttpClientRequest otelRequest) { |
135 | 146 | Context parentContext = currentContext(); |
| 147 | + if (!instrumenter().shouldStart(parentContext, otelRequest)) { |
| 148 | + return null; |
| 149 | + } |
| 150 | + Context context = instrumenter().start(parentContext, otelRequest); |
| 151 | + return new AdviceScope(otelRequest, parentContext, context, context.makeCurrent()); |
| 152 | + } |
136 | 153 |
|
137 | | - otelRequest = new ApacheHttpClientRequest(request); |
| 154 | + public <T> ResponseHandler<T> wrapHandler(ResponseHandler<T> handler) { |
| 155 | + return new WrappingStatusSettingResponseHandler<>( |
| 156 | + context, parentContext, otelRequest, handler); |
| 157 | + } |
138 | 158 |
|
139 | | - if (!instrumenter().shouldStart(parentContext, otelRequest)) { |
140 | | - return; |
| 159 | + public void end(@Nullable Object result, @Nullable Throwable throwable) { |
| 160 | + scope.close(); |
| 161 | + if (throwable != null) { |
| 162 | + instrumenter().end(context, otelRequest, null, throwable); |
| 163 | + } else if (result instanceof HttpResponse) { |
| 164 | + instrumenter().end(context, otelRequest, (HttpResponse) result, null); |
141 | 165 | } |
| 166 | + // ended in WrappingStatusSettingResponseHandler |
| 167 | + } |
| 168 | + } |
| 169 | + |
| 170 | + @SuppressWarnings("unused") |
| 171 | + public static class UriRequestAdvice { |
142 | 172 |
|
143 | | - context = instrumenter().start(parentContext, otelRequest); |
144 | | - scope = context.makeCurrent(); |
| 173 | + @Advice.OnMethodEnter(suppress = Throwable.class) |
| 174 | + public static AdviceScope methodEnter(@Advice.Argument(0) HttpUriRequest request) { |
| 175 | + return AdviceScope.start(new ApacheHttpClientRequest(request)); |
145 | 176 | } |
146 | 177 |
|
147 | 178 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
148 | 179 | public static void methodExit( |
149 | 180 | @Advice.Argument(0) HttpUriRequest request, |
150 | | - @Advice.Return Object result, |
151 | | - @Advice.Thrown Throwable throwable, |
152 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
153 | | - @Advice.Local("otelContext") Context context, |
154 | | - @Advice.Local("otelScope") Scope scope) { |
155 | | - if (scope == null) { |
156 | | - return; |
157 | | - } |
| 181 | + @Advice.Return @Nullable Object result, |
| 182 | + @Advice.Thrown @Nullable Throwable throwable, |
| 183 | + @Advice.Enter @Nullable AdviceScope adviceScope) { |
158 | 184 |
|
159 | | - scope.close(); |
160 | | - ApacheHttpClientHelper.doMethodExit(context, otelRequest, result, throwable); |
| 185 | + if (adviceScope != null) { |
| 186 | + adviceScope.end(result, throwable); |
| 187 | + } |
161 | 188 | } |
162 | 189 | } |
163 | 190 |
|
164 | 191 | @SuppressWarnings("unused") |
165 | 192 | public static class UriRequestWithHandlerAdvice { |
166 | 193 |
|
| 194 | + @AssignReturned.ToArguments(@ToArgument(value = 1, index = 1)) |
167 | 195 | @Advice.OnMethodEnter(suppress = Throwable.class) |
168 | | - public static void methodEnter( |
| 196 | + public static Object[] methodEnter( |
169 | 197 | @Advice.Argument(0) HttpUriRequest request, |
170 | | - @Advice.Argument(value = 1, readOnly = false) ResponseHandler<?> handler, |
171 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
172 | | - @Advice.Local("otelContext") Context context, |
173 | | - @Advice.Local("otelScope") Scope scope) { |
174 | | - Context parentContext = currentContext(); |
175 | | - |
176 | | - otelRequest = new ApacheHttpClientRequest(request); |
177 | | - |
178 | | - if (!instrumenter().shouldStart(parentContext, otelRequest)) { |
179 | | - return; |
180 | | - } |
181 | | - |
182 | | - context = instrumenter().start(parentContext, otelRequest); |
183 | | - scope = context.makeCurrent(); |
| 198 | + @Advice.Argument(1) ResponseHandler<?> handler) { |
184 | 199 |
|
| 200 | + AdviceScope adviceScope = AdviceScope.start(new ApacheHttpClientRequest(request)); |
185 | 201 | // Wrap the handler so we capture the status code |
186 | | - if (handler != null) { |
187 | | - handler = |
188 | | - new WrappingStatusSettingResponseHandler<>( |
189 | | - context, parentContext, otelRequest, handler); |
190 | | - } |
| 202 | + return new Object[] { |
| 203 | + adviceScope, adviceScope == null ? handler : adviceScope.wrapHandler(handler) |
| 204 | + }; |
191 | 205 | } |
192 | 206 |
|
193 | 207 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
194 | 208 | public static void methodExit( |
195 | | - @Advice.Argument(0) HttpUriRequest request, |
196 | 209 | @Advice.Return Object result, |
197 | 210 | @Advice.Thrown Throwable throwable, |
198 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
199 | | - @Advice.Local("otelContext") Context context, |
200 | | - @Advice.Local("otelScope") Scope scope) { |
201 | | - if (scope == null) { |
202 | | - return; |
203 | | - } |
| 211 | + @Advice.Enter Object[] enterResult) { |
204 | 212 |
|
205 | | - scope.close(); |
206 | | - ApacheHttpClientHelper.doMethodExit(context, otelRequest, result, throwable); |
| 213 | + AdviceScope adviceScope = (AdviceScope) enterResult[0]; |
| 214 | + if (adviceScope != null) { |
| 215 | + adviceScope.end(result, throwable); |
| 216 | + } |
207 | 217 | } |
208 | 218 | } |
209 | 219 |
|
210 | 220 | @SuppressWarnings("unused") |
211 | 221 | public static class RequestAdvice { |
212 | 222 |
|
213 | 223 | @Advice.OnMethodEnter(suppress = Throwable.class) |
214 | | - public static void methodEnter( |
215 | | - @Advice.Argument(0) HttpHost host, |
216 | | - @Advice.Argument(1) HttpRequest request, |
217 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
218 | | - @Advice.Local("otelContext") Context context, |
219 | | - @Advice.Local("otelScope") Scope scope) { |
220 | | - Context parentContext = currentContext(); |
221 | | - |
222 | | - otelRequest = new ApacheHttpClientRequest(host, request); |
223 | | - |
224 | | - if (!instrumenter().shouldStart(parentContext, otelRequest)) { |
225 | | - return; |
226 | | - } |
227 | | - |
228 | | - context = instrumenter().start(parentContext, otelRequest); |
229 | | - scope = context.makeCurrent(); |
| 224 | + public static AdviceScope methodEnter( |
| 225 | + @Advice.Argument(0) HttpHost host, @Advice.Argument(1) HttpRequest request) { |
| 226 | + return AdviceScope.start(new ApacheHttpClientRequest(host, request)); |
230 | 227 | } |
231 | 228 |
|
232 | 229 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
233 | 230 | public static void methodExit( |
234 | 231 | @Advice.Return Object result, |
235 | 232 | @Advice.Thrown Throwable throwable, |
236 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
237 | | - @Advice.Local("otelContext") Context context, |
238 | | - @Advice.Local("otelScope") Scope scope) { |
239 | | - if (scope == null) { |
240 | | - return; |
241 | | - } |
| 233 | + @Advice.Enter AdviceScope adviceScope) { |
242 | 234 |
|
243 | | - scope.close(); |
244 | | - ApacheHttpClientHelper.doMethodExit(context, otelRequest, result, throwable); |
| 235 | + if (adviceScope != null) { |
| 236 | + adviceScope.end(result, throwable); |
| 237 | + } |
245 | 238 | } |
246 | 239 | } |
247 | 240 |
|
248 | 241 | @SuppressWarnings("unused") |
249 | 242 | public static class RequestWithHandlerAdvice { |
250 | 243 |
|
| 244 | + @AssignReturned.ToArguments(@ToArgument(value = 2, index = 1)) |
251 | 245 | @Advice.OnMethodEnter(suppress = Throwable.class) |
252 | | - public static void methodEnter( |
| 246 | + public static Object[] methodEnter( |
253 | 247 | @Advice.Argument(0) HttpHost host, |
254 | 248 | @Advice.Argument(1) HttpRequest request, |
255 | | - @Advice.Argument(value = 2, readOnly = false) ResponseHandler<?> handler, |
256 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
257 | | - @Advice.Local("otelContext") Context context, |
258 | | - @Advice.Local("otelScope") Scope scope) { |
259 | | - Context parentContext = currentContext(); |
260 | | - |
261 | | - otelRequest = new ApacheHttpClientRequest(host, request); |
262 | | - |
263 | | - if (!instrumenter().shouldStart(parentContext, otelRequest)) { |
264 | | - return; |
265 | | - } |
266 | | - |
267 | | - context = instrumenter().start(parentContext, otelRequest); |
268 | | - scope = context.makeCurrent(); |
269 | | - |
270 | | - // Wrap the handler so we capture the status code |
271 | | - if (handler != null) { |
272 | | - handler = |
273 | | - new WrappingStatusSettingResponseHandler<>( |
274 | | - context, parentContext, otelRequest, handler); |
275 | | - } |
| 249 | + @Advice.Argument(2) ResponseHandler<?> handler) { |
| 250 | + |
| 251 | + AdviceScope adviceScope = AdviceScope.start(new ApacheHttpClientRequest(host, request)); |
| 252 | + return new Object[] { |
| 253 | + adviceScope, |
| 254 | + // Wrap the handler so we capture the status code |
| 255 | + adviceScope == null ? handler : adviceScope.wrapHandler(handler) |
| 256 | + }; |
276 | 257 | } |
277 | 258 |
|
278 | 259 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
279 | 260 | public static void methodExit( |
280 | 261 | @Advice.Return Object result, |
281 | 262 | @Advice.Thrown Throwable throwable, |
282 | | - @Advice.Local("otelRequest") ApacheHttpClientRequest otelRequest, |
283 | | - @Advice.Local("otelContext") Context context, |
284 | | - @Advice.Local("otelScope") Scope scope) { |
285 | | - if (scope == null) { |
286 | | - return; |
287 | | - } |
| 263 | + @Advice.Enter Object[] enterResult) { |
288 | 264 |
|
289 | | - scope.close(); |
290 | | - ApacheHttpClientHelper.doMethodExit(context, otelRequest, result, throwable); |
| 265 | + AdviceScope adviceScope = (AdviceScope) enterResult[0]; |
| 266 | + if (adviceScope != null) { |
| 267 | + adviceScope.end(result, throwable); |
| 268 | + } |
291 | 269 | } |
292 | 270 | } |
293 | 271 | } |
0 commit comments