Change AWS instrumentation to use constructor
Instead of requiring the builder.
MODIFY
dd-java-agent/instrumentation/aws-sdk/src/main/java/datadog/trace/instrumentation/aws/AWSClientInstrumentation.java
dd-java-agent/instrumentation/aws-sdk/src/main/java/datadog/trace/instrumentation/aws/AWSClientInstrumentation.java
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<RequestHandler2> handlers = builder.getRequestHandlers();
+ @Advice.OnMethodExit(suppress = Throwable.class)
+ public static void addHandler(
+ @Advice.FieldValue("requestHandler2s") final List<RequestHandler2> handlers) {
boolean hasDDHandler = false;
- if (null == handlers) {
- handlers = Collections.emptyList();
- } else {
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]));
}
}
}

MODIFY
dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy
dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy
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<Headers>()
@@ -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()

Change AWS instrumentation to use constructor

Instead of requiring the builder.