Add basic S3 instrumentation

DataDog · nhulston · Jan 22, 2025 · Dec 10, 2024 · Dec 10, 2024 · Dec 10, 2024
commit 48d5d9fb1ceb53fc76916c9954f4730189d6917f
@@ -0,0 +1,29 @@
+muzzle {
+  pass {
+    group = "software.amazon.awssdk"
+    module = "s3"
+    versions = "[2.0,3)"
+    assertInverse = true
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test')
+
+dependencies {
+  compileOnly group: 'software.amazon.awssdk', name: 's3', version: '2.29.26'
+
+  // Include httpclient instrumentation for testing because it is a dependency for aws-sdk.
+  testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4')
+  testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-2.2')
+  testImplementation 'software.amazon.awssdk:s3:2.29.26'
+  testImplementation 'org.testcontainers:localstack:1.20.1'
+
+  latestDepTestImplementation group: 'software.amazon.awssdk', name: 's3', version: '+'
+}
+
+tasks.withType(Test).configureEach {
+  usesService(testcontainersLimit)
+}
@@ -0,0 +1,48 @@
+package datadog.trace.instrumentation.aws.v2.s3;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import java.util.List;
+import net.bytebuddy.asm.Advice;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+
+@AutoService(InstrumenterModule.class)
+public final class S3ClientInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType {
+  public S3ClientInstrumentation() {
+    super("s3");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
        isMethod().and(named("resolveExecutionInterceptors")),
+        S3ClientInstrumentation.class.getName() + "$AwsS3BuilderAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".S3Interceptor", packageName + ".TextMapInjectAdapter"};
+  }
+
+  public static class AwsS3BuilderAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void addHandler(@Advice.Return final List<ExecutionInterceptor> interceptors) {
+      for (ExecutionInterceptor interceptor : interceptors) {
+        if (interceptor instanceof S3Interceptor) {
+          return; // list already has our interceptor, return to builder
+        }
+      }
+      interceptors.add(new S3Interceptor());
+    }
+  }
+}
@@ -0,0 +1,13 @@
+package datadog.trace.instrumentation.aws.v2.s3;
+
+import software.amazon.awssdk.core.interceptor.Context;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+
+public class S3Interceptor implements ExecutionInterceptor {
+  @Override
+  public void afterExecution(
+      Context.AfterExecution context, ExecutionAttributes executionAttributes) {
+    System.out.println("S3Interceptor afterExecution()");
+  }
+}
@@ -0,0 +1,13 @@
+package datadog.trace.instrumentation.aws.v2.s3;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+
+public class TextMapInjectAdapter implements AgentPropagation.Setter<StringBuilder> {
+
+  public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter();
+
+  @Override
+  public void set(final StringBuilder builder, final String key, final String value) {
+    builder.append('"').append(key).append("\":\"").append(value).append("\",");
+  }
+}
@@ -196,6 +196,7 @@ include ':dd-java-agent:instrumentation:aws-java-sns-1.0'
 include ':dd-java-agent:instrumentation:aws-java-sns-2.0'
 include ':dd-java-agent:instrumentation:aws-java-sqs-1.0'
 include ':dd-java-agent:instrumentation:aws-java-sqs-2.0'
+include ':dd-java-agent:instrumentation:aws-java-s3-2.0'
 include ':dd-java-agent:instrumentation:aws-lambda-handler'
 include ':dd-java-agent:instrumentation:axis-2'
 include ':dd-java-agent:instrumentation:axway-api'