Migrate aws instrumentation to byte buddy.

dd-java-agent/dd-java-agent.gradle

@@ -21,6 +21,7 @@ dependencies {
   compile project(':dd-java-agent:tooling')
   compile project(':dd-trace-annotations')
 
+  compile project(':dd-java-agent:integrations:aws-sdk')
   compile project(':dd-java-agent:integrations:servlet-2')
   compile project(':dd-java-agent:integrations:servlet-3')
   compile project(':dd-java-agent:integrations:spring-web')

dd-java-agent/integrations/aws-sdk/aws-sdk.gradle

@@ -12,3 +12,18 @@ versionScan {
     "com.amazonaws.http.request.HttpRequestFactory": null,
   ]
 }
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+dependencies {
+  compileOnly group: 'com.amazonaws', name: 'aws-java-sdk-core', version: '1.11.119'
+
+  compile project(':dd-trace')
+  compile project(':dd-java-agent:tooling')
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+
+  compile group: 'io.opentracing.contrib', name: 'opentracing-aws-sdk', version: '0.0.2'
+  compile group: 'com.google.auto.service', name: 'auto-service', version: '1.0-rc3'
+}

dd-java-agent/integrations/aws-sdk/src/main/java/dd/inst/aws/AWSClientInstrumentation.java

@@ -0,0 +1,59 @@
+package dd.inst.aws;
+
+import static dd.trace.ClassLoaderHasClassMatcher.classLoaderHasClasses;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+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 dd.trace.Instrumenter;
+import io.opentracing.contrib.aws.TracingRequestHandler;
+import io.opentracing.util.GlobalTracer;
+import java.util.Arrays;
+import java.util.List;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class AWSClientInstrumentation implements Instrumenter {
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        .type(
+            named("com.amazonaws.client.builder.AwsSyncClientBuilder")
+                .or(named("com.amazonaws.client.builder.AwsAsyncClientBuilder")),
+            classLoaderHasClasses(
+                "com.amazonaws.http.client.HttpClientFactory",
+                "com.amazonaws.http.apache.utils.ApacheUtils"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    named("build").and(takesArguments(0)).and(isPublic()),
+                    AWSClientAdvice.class.getName()))
+        .asDecorator();
+  }
+
+  public static class AWSClientAdvice {
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void addHandler(@Advice.This final AwsClientBuilder builder) {
+
+      final RequestHandler2 handler = new TracingRequestHandler(GlobalTracer.get());
+
+      List<RequestHandler2> handlers = builder.getRequestHandlers();
+
+      if (handlers == null || handlers.isEmpty()) {
+        handlers = Arrays.asList(handler);
+      } else {
+        // Check if we already added the handler
+        if (!(handlers.get(0) instanceof TracingRequestHandler)) {
+          handlers.add(0, handler);
+        }
+      }
+      builder.setRequestHandlers(handler);
+    }
+  }
+}

dd-java-agent/integrations/helpers/src/main/java/com/datadoghq/agent/integration/AWSClientHelper.java

@@ -1,49 +0,0 @@
-package com.datadoghq.agent.integration;
-
-import com.amazonaws.client.builder.AwsClientBuilder;
-import com.amazonaws.handlers.RequestHandler2;
-import io.opentracing.contrib.aws.TracingRequestHandler;
-import java.lang.reflect.Field;
-import java.util.Arrays;
-import java.util.List;
-import org.jboss.byteman.rule.Rule;
-
-/**
- * Patch the AWS Client during the building steps. This opentracing integration is compatible with
- * the latest release of the AWS SDK
- */
-public class AWSClientHelper extends DDAgentTracingHelper<AwsClientBuilder> {
-
-  public AWSClientHelper(Rule rule) {
-    super(rule);
-  }
-
-  /**
-   * Strategy: we add a tracing handler to the client when it has just been built. We intercept the
-   * return value of the com.amazonaws.client.builder.AwsClientBuilder.build() method and add the
-   * handler.
-   *
-   * @param client The fresh AWS client instance
-   * @return The same instance, but patched
-   * @throws Exception
-   */
-  protected AwsClientBuilder doPatch(AwsClientBuilder client) throws Exception {
-
-    RequestHandler2 handler = new TracingRequestHandler(tracer);
-
-    Field field = AwsClientBuilder.class.getDeclaredField("requestHandlers");
-    field.setAccessible(true);
-    List<RequestHandler2> handlers = (List<RequestHandler2>) field.get(client);
-
-    if (handlers == null || handlers.isEmpty()) {
-      handlers = Arrays.asList(handler);
-    } else {
-      // Check if we already added the handler
-      if (!(handlers.get(0) instanceof TracingRequestHandler)) {
-        handlers.add(0, handler);
-      }
-    }
-    client.setRequestHandlers((RequestHandler2[]) handlers.toArray());
-    return client;
-  }
-}

dd-java-agent/src/main/resources/initializer-rules.btm

@@ -13,18 +13,6 @@
 #ENDRULE
 
 
-# Instrument AWS SDK client
-# ==========================
-RULE AwsClientBuilder-init
-CLASS ^com.amazonaws.client.builder.AwsClientBuilder
-METHOD <init>
-AT EXIT
-IF TRUE
-DO
-  com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
-ENDRULE
-
-
 # Instrument Cassandra client
 # ===========================
 RULE Cluster$Manager-init

dd-java-agent/src/main/resources/integration-rules.btm

@@ -21,20 +21,6 @@ DO
 ENDRULE
 
 
-
-# Instrument AWS SDK client
-# ==========================
-RULE opentracing-aws-sdk
-CLASS ^com.amazonaws.client.builder.AwsClientBuilder
-METHOD build()
-HELPER com.datadoghq.agent.integration.AWSClientHelper
-AT ENTRY
-IF TRUE
-DO
-  patch($0);
-ENDRULE
-
-
 # Instrument Cassandra client
 # ===========================
 RULE opentracing-cassandra-driver