Create DynamoDB instrumentation + add span pointers for updateItem and deleteItem (#8490)

Author: nhulston

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

@@ -0,0 +1,29 @@
+muzzle {
+  pass {
+    group = "software.amazon.awssdk"
+    module = "dynamodb"
+    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: 'dynamodb', version: '2.30.22'
+
+  // 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:dynamodb:2.30.22'
+  testImplementation 'org.testcontainers:localstack:1.20.1'
+
+  latestDepTestImplementation group: 'software.amazon.awssdk', name: 'dynamodb', version: '+'
+}
+
+tasks.withType(Test).configureEach {
+  usesService(testcontainersLimit)
+}

dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbClientInstrumentation.java

@@ -0,0 +1,48 @@
+package datadog.trace.instrumentation.aws.v2.dynamodb;
+
+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 DynamoDbClientInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  public DynamoDbClientInstrumentation() {
+    super("dynamodb", "aws-dynamodb");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("resolveExecutionInterceptors")),
+        DynamoDbClientInstrumentation.class.getName() + "$AwsDynamoDbBuilderAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".DynamoDbInterceptor", packageName + ".DynamoDbUtil"};
+  }
+
+  public static class AwsDynamoDbBuilderAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void addHandler(@Advice.Return final List<ExecutionInterceptor> interceptors) {
+      for (ExecutionInterceptor interceptor : interceptors) {
+        if (interceptor instanceof DynamoDbInterceptor) {
+          return; // list already has our interceptor, return to builder
+        }
+      }
+      interceptors.add(new DynamoDbInterceptor());
+    }
+  }
+}

dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java

@@ -0,0 +1,49 @@
+package datadog.trace.instrumentation.aws.v2.dynamodb;
+
+import datadog.trace.api.Config;
+import datadog.trace.bootstrap.InstanceStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.core.SdkRequest;
+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.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
+
+public class DynamoDbInterceptor implements ExecutionInterceptor {
+  private static final Logger log = LoggerFactory.getLogger(DynamoDbInterceptor.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 DynamoDb request span. Not creating span pointer.");
+      return;
+    }
+
+    SdkRequest request = context.request();
+    if (request instanceof UpdateItemRequest) {
+      Map<String, AttributeValue> keys = ((UpdateItemRequest) request).key();
+      DynamoDbUtil.exportTagsWithKnownKeys(span, keys);
+    } else if (request instanceof DeleteItemRequest) {
+      Map<String, AttributeValue> keys = ((DeleteItemRequest) request).key();
+      DynamoDbUtil.exportTagsWithKnownKeys(span, keys);
+    }
+  }
+}

dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbUtil.java

@@ -0,0 +1,72 @@
+package datadog.trace.instrumentation.aws.v2.dynamodb;
+
+import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_1;
+import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE;
+import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_2;
+import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_2_VALUE;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+
+public class DynamoDbUtil {
+  /**
+   * Extracts a string value from a DynamoDB AttributeValue.
+   *
+   * @param value The AttributeValue to extract from
+   * @return The extracted string value
+   */
+  private static String extractValueAsString(AttributeValue value) {
+    if (value == null) {
+      return "";
+    }
+
+    if (value.s() != null) {
+      return value.s();
+    } else if (value.n() != null) {
+      return value.n();
+    } else if (value.b() != null) {
+      // For binary values, convert the bytes back to string
+      return new String(value.b().asByteArray(), java.nio.charset.Charset.defaultCharset());
+    }
+
+    return "";
+  }
+
+  /**
+   * Gets primary key/values and exports them as temporary tags on the span so that
+   * SpanPointersProcessor.java can complete the span pointer creation.
+   *
+   * @param span The span to set the temporary tags on
+   * @param keys The primary key/values to extract from
+   */
+  static void exportTagsWithKnownKeys(AgentSpan span, Map<String, AttributeValue> keys) {
+    if (keys == null || keys.isEmpty()) {
+      return;
+    }
+
+    if (keys.size() == 1) {
+      // Single primary key case
+      Map.Entry<String, AttributeValue> entry = keys.entrySet().iterator().next();
+      span.setTag(DYNAMO_PRIMARY_KEY_1, entry.getKey());
+      span.setTag(DYNAMO_PRIMARY_KEY_1_VALUE, extractValueAsString(entry.getValue()));
+    } else {
+      // Sort keys alphabetically
+      List<String> keyNames = new ArrayList<>(keys.keySet());
+      Collections.sort(keyNames);
+
+      // First key (alphabetically)
+      String primaryKey1Name = keyNames.get(0);
+      span.setTag(DYNAMO_PRIMARY_KEY_1, primaryKey1Name);
+      span.setTag(DYNAMO_PRIMARY_KEY_1_VALUE, extractValueAsString(keys.get(primaryKey1Name)));
+
+      // Second key
+      String primaryKey2Name = keyNames.get(1);
+      span.setTag(DYNAMO_PRIMARY_KEY_2, primaryKey2Name);
+      span.setTag(DYNAMO_PRIMARY_KEY_2_VALUE, extractValueAsString(keys.get(primaryKey2Name)));
+    }
+  }
+}