11package datadog .trace .instrumentation .aws ;
22
33import static datadog .trace .agent .tooling .ClassLoaderMatcher .classLoaderHasClasses ;
4+ import static net .bytebuddy .matcher .ElementMatchers .hasSuperType ;
5+ import static net .bytebuddy .matcher .ElementMatchers .isAbstract ;
46import static net .bytebuddy .matcher .ElementMatchers .isPublic ;
57import static net .bytebuddy .matcher .ElementMatchers .named ;
8+ import static net .bytebuddy .matcher .ElementMatchers .not ;
69import static net .bytebuddy .matcher .ElementMatchers .takesArguments ;
710
811import com .amazonaws .client .builder .AwsClientBuilder ;
1316import datadog .trace .agent .tooling .Instrumenter ;
1417import io .opentracing .contrib .aws .TracingRequestHandler ;
1518import io .opentracing .util .GlobalTracer ;
16- import java .util .Arrays ;
19+ import java .util .ArrayList ;
1720import java .util .List ;
1821import net .bytebuddy .agent .builder .AgentBuilder ;
1922import net .bytebuddy .asm .Advice ;
@@ -25,41 +28,44 @@ public final class AWSClientInstrumentation implements Instrumenter {
2528 public AgentBuilder instrument (final AgentBuilder agentBuilder ) {
2629 return agentBuilder
2730 .type (
28- named ("com.amazonaws.client.builder.AwsSyncClientBuilder" )
29- .or (named ("com.amazonaws.client.builder.AwsAsyncClientBuilder" )),
31+ hasSuperType (named ("com.amazonaws.client.builder.AwsClientBuilder" )),
3032 classLoaderHasClasses (
31- "com.amazonaws.http.client.HttpClientFactory" ,
32- "com.amazonaws.http.apache.utils.ApacheUtils" ))
33+ // aws classes used by opentracing contrib helpers
34+ "com.amazonaws.handlers.RequestHandler2" ,
35+ "com.amazonaws.Request" ,
36+ "com.amazonaws.Response" ,
37+ "com.amazonaws.handlers.HandlerContextKey" ))
3338 .transform (
3439 new HelperInjector (
3540 "io.opentracing.contrib.aws.TracingRequestHandler" ,
3641 "io.opentracing.contrib.aws.SpanDecorator" ))
3742 .transform (
3843 DDAdvice .create ()
3944 .advice (
40- named ("build" ).and (takesArguments (0 )).and (isPublic ()),
45+ named ("build" ).and (takesArguments (0 )).and (isPublic ()). and ( not ( isAbstract ())) ,
4146 AWSClientAdvice .class .getName ()))
4247 .asDecorator ();
4348 }
4449
4550 public static class AWSClientAdvice {
46-
4751 @ Advice .OnMethodExit (onThrowable = Throwable .class , suppress = Throwable .class )
48- public static void addHandler (@ Advice .This final AwsClientBuilder builder ) {
49-
50- final RequestHandler2 handler = new TracingRequestHandler (GlobalTracer .get ());
51-
52+ public static void addHandler (@ Advice .This final AwsClientBuilder <?, ?> builder ) {
5253 List <RequestHandler2 > handlers = builder .getRequestHandlers ();
53-
54- if (handlers == null || handlers . isEmpty () ) {
55- handlers = Arrays . asList ( handler );
54+ boolean hasDDHandler = false ;
55+ if (null == handlers ) {
56+ handlers = new ArrayList < RequestHandler2 >( 1 );
5657 } else {
57- // Check if we already added the handler
58- if (!(handlers .get (0 ) instanceof TracingRequestHandler )) {
59- handlers .add (0 , handler );
58+ for (RequestHandler2 handler : handlers ) {
59+ if (handler instanceof TracingRequestHandler ) {
60+ hasDDHandler = true ;
61+ break ;
62+ }
6063 }
6164 }
62- builder .setRequestHandlers (handler );
65+ if (!hasDDHandler ) {
66+ handlers .add (new TracingRequestHandler (GlobalTracer .get ()));
67+ builder .setRequestHandlers (handlers .toArray (new RequestHandler2 [0 ]));
68+ }
6369 }
6470 }
6571}
0 commit comments