Create S3 instrumentation + add span pointers (#8075)

Author: nhulston

dd-java-agent/instrumentation/aws-java-s3-2.0/build.gradle

@@ -0,0 +1,29 @@
+muzzle {
+  pass {
+    group = "software.amazon.awssdk"
+    module = "s3"
+    versions = "[2,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.
+  testRuntimeOnly project(':dd-java-agent:instrumentation:apache-httpclient-4')
+  testRuntimeOnly 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)
+}

dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3ClientInstrumentation.java

@@ -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, Instrumenter.HasMethodAdvice {
+  public S3ClientInstrumentation() {
+    super("s3", "aws-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());
+    }
+  }
+}

dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java

@@ -0,0 +1,59 @@
+package datadog.trace.instrumentation.aws.v2.s3;
+
+import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.S3_ETAG;
+
+import datadog.trace.api.Config;
+import datadog.trace.bootstrap.InstanceStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.core.interceptor.Context;
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
+import software.amazon.awssdk.services.s3.model.CopyObjectResponse;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+
+public class S3Interceptor implements ExecutionInterceptor {
+  private static final Logger log = LoggerFactory.getLogger(S3Interceptor.class);
+
+  public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
+      InstanceStore.of(ExecutionAttribute.class)
+          .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
+
+  private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws");
+
+  @Override
+  public void afterExecution(
+      Context.AfterExecution context, ExecutionAttributes executionAttributes) {
+    if (!CAN_ADD_SPAN_POINTERS) {
+      return;
+    }
+
+    AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
+    if (span == null) {
+      log.debug("Unable to find S3 request span. Not creating span pointer.");
+      return;
+    }
+    String eTag;
+    Object response = context.response();
+
+    // Get eTag for hash calculation.
+    // https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/s3/S3Client.html
+    if (response instanceof PutObjectResponse) {
+      eTag = ((PutObjectResponse) response).eTag();
+    } else if (response instanceof CopyObjectResponse) {
+      eTag = ((CopyObjectResponse) response).copyObjectResult().eTag();
+    } else if (response instanceof CompleteMultipartUploadResponse) {
+      eTag = ((CompleteMultipartUploadResponse) response).eTag();
+    } else {
+      return;
+    }
+
+    // Store eTag as tag, then calculate hash + add span pointers in SpanPointersProcessor.
+    // Bucket and key are already stored as tags in AwsSdkClientDecorator, so need to make redundant
+    // tags.
+    span.setTag(S3_ETAG, eTag);
+  }
+}

dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/TextMapInjectAdapter.java

@@ -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("\",");
+  }
+}