Merge pull request #214 from DataDog/tyler/aws-fixes

Change AWS instrumentation to use constructor
This commit is contained in:
Andrew Kent 2018-02-06 09:24:13 -08:00 committed by GitHub
commit 0382f8916e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 39 additions and 28 deletions

View File

@ -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<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;
}
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]));
}
}
}

View File

@ -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<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()