Merge pull request #214 from DataDog/tyler/aws-fixes
Change AWS instrumentation to use constructor
This commit is contained in:
commit
0382f8916e
|
@ -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]));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue