From a311d5f57d84fcb14ea4bd808c122d5c30642f02 Mon Sep 17 00:00:00 2001 From: Tyler Benson Date: Tue, 6 Feb 2018 14:55:42 +1000 Subject: [PATCH] Change AWS instrumentation to use constructor Instead of requiring the builder. --- .../aws/AWSClientInstrumentation.java | 40 +++++++------------ .../src/test/groovy/AWSClientTest.groovy | 27 ++++++++++++- 2 files changed, 39 insertions(+), 28 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-sdk/src/main/java/datadog/trace/instrumentation/aws/AWSClientInstrumentation.java b/dd-java-agent/instrumentation/aws-sdk/src/main/java/datadog/trace/instrumentation/aws/AWSClientInstrumentation.java index 4dc112a47e..08a6443edf 100644 --- a/dd-java-agent/instrumentation/aws-sdk/src/main/java/datadog/trace/instrumentation/aws/AWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-sdk/src/main/java/datadog/trace/instrumentation/aws/AWSClientInstrumentation.java @@ -1,14 +1,11 @@ package datadog.trace.instrumentation.aws; import static datadog.trace.agent.tooling.ClassLoaderMatcher.classLoaderHasClasses; -import static net.bytebuddy.matcher.ElementMatchers.hasSuperType; +import static net.bytebuddy.matcher.ElementMatchers.declaresField; import static net.bytebuddy.matcher.ElementMatchers.isAbstract; -import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.isConstructor; import static net.bytebuddy.matcher.ElementMatchers.named; -import static net.bytebuddy.matcher.ElementMatchers.not; -import static net.bytebuddy.matcher.ElementMatchers.takesArguments; -import com.amazonaws.client.builder.AwsClientBuilder; import com.amazonaws.handlers.RequestHandler2; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.DDAdvice; @@ -16,8 +13,6 @@ import datadog.trace.agent.tooling.HelperInjector; import datadog.trace.agent.tooling.Instrumenter; import io.opentracing.contrib.aws.TracingRequestHandler; import io.opentracing.util.GlobalTracer; -import java.util.ArrayList; -import java.util.Collections; import java.util.List; import net.bytebuddy.agent.builder.AgentBuilder; import net.bytebuddy.asm.Advice; @@ -33,7 +28,10 @@ public final class AWSClientInstrumentation extends Instrumenter.Configurable { public AgentBuilder apply(final AgentBuilder agentBuilder) { return agentBuilder .type( - hasSuperType(named("com.amazonaws.client.builder.AwsClientBuilder")), + isAbstract() + .and( + named("com.amazonaws.AmazonWebServiceClient") + .and(declaresField(named("requestHandler2s")))), classLoaderHasClasses( // aws classes used by opentracing contrib helpers "com.amazonaws.handlers.RequestHandler2", @@ -44,33 +42,23 @@ public final class AWSClientInstrumentation extends Instrumenter.Configurable { new HelperInjector( "io.opentracing.contrib.aws.TracingRequestHandler", "io.opentracing.contrib.aws.SpanDecorator")) - .transform( - DDAdvice.create() - .advice( - named("build").and(takesArguments(0)).and(isPublic()).and(not(isAbstract())), - AWSClientAdvice.class.getName())) + .transform(DDAdvice.create().advice(isConstructor(), AWSClientAdvice.class.getName())) .asDecorator(); } public static class AWSClientAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void addHandler(@Advice.This final AwsClientBuilder builder) { - List handlers = builder.getRequestHandlers(); + @Advice.OnMethodExit(suppress = Throwable.class) + public static void addHandler( + @Advice.FieldValue("requestHandler2s") final List handlers) { boolean hasDDHandler = false; - if (null == handlers) { - handlers = Collections.emptyList(); - } else { - for (final RequestHandler2 handler : handlers) { - if (handler instanceof TracingRequestHandler) { - hasDDHandler = true; - break; - } + for (final RequestHandler2 handler : handlers) { + if (handler instanceof TracingRequestHandler) { + hasDDHandler = true; + break; } } if (!hasDDHandler) { - handlers = new ArrayList<>(handlers); // copy since returned list is unmodifiable handlers.add(new TracingRequestHandler(GlobalTracer.get())); - builder.setRequestHandlers(handlers.toArray(new RequestHandler2[0])); } } } diff --git a/dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy b/dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy index edf8aca263..e386920217 100644 --- a/dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy @@ -1,9 +1,11 @@ import com.amazonaws.AmazonWebServiceClient import com.amazonaws.auth.AWSStaticCredentialsProvider import com.amazonaws.auth.AnonymousAWSCredentials +import com.amazonaws.auth.BasicAWSCredentials import com.amazonaws.client.builder.AwsClientBuilder import com.amazonaws.handlers.RequestHandler2 import com.amazonaws.regions.Regions +import com.amazonaws.services.s3.AmazonS3Client import com.amazonaws.services.s3.AmazonS3ClientBuilder import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDTags @@ -16,7 +18,7 @@ import static ratpack.groovy.test.embed.GroovyEmbeddedApp.ratpack class AWSClientTest extends AgentTestRunner { - def "request handler is hooked up"() { + def "request handler is hooked up with builder"() { setup: def builder = AmazonS3ClientBuilder.standard() .withRegion(Regions.US_EAST_1) @@ -36,6 +38,27 @@ class AWSClientTest extends AgentTestRunner { false | 1 | 0 } + def "request handler is hooked up with constructor"() { + setup: + String accessKey = "asdf" + String secretKey = "qwerty" + def credentials = new BasicAWSCredentials(accessKey, secretKey) + def client = new AmazonS3Client(credentials) + if (addHandler) { + client.addRequestHandler(new RequestHandler2() {}) + } + + expect: + client.requestHandler2s != null + client.requestHandler2s.size() == size + client.requestHandler2s.get(0).getClass().getSimpleName() == "TracingRequestHandler" + + where: + addHandler | size + true | 2 + false | 1 + } + def "send request with mocked back end"() { setup: def receivedHeaders = new AtomicReference() @@ -47,7 +70,7 @@ class AWSClientTest extends AgentTestRunner { } } } - AwsClientBuilder.EndpointConfiguration endpoint = new AwsClientBuilder.EndpointConfiguration("http://localhost:$server.address.port", "us-west-2"); + AwsClientBuilder.EndpointConfiguration endpoint = new AwsClientBuilder.EndpointConfiguration("http://localhost:$server.address.port", "us-west-2") AmazonWebServiceClient client = AmazonS3ClientBuilder .standard()