diff --git a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/build.gradle new file mode 100644 index 00000000000..fd09aca0fea --- /dev/null +++ b/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) +} diff --git a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbClientInstrumentation.java new file mode 100644 index 00000000000..2dd22d1ae6c --- /dev/null +++ b/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 interceptors) { + for (ExecutionInterceptor interceptor : interceptors) { + if (interceptor instanceof DynamoDbInterceptor) { + return; // list already has our interceptor, return to builder + } + } + interceptors.add(new DynamoDbInterceptor()); + } + } +} diff --git a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java new file mode 100644 index 00000000000..86cf378bc45 --- /dev/null +++ b/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 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 keys = ((UpdateItemRequest) request).key(); + DynamoDbUtil.exportTagsWithKnownKeys(span, keys); + } else if (request instanceof DeleteItemRequest) { + Map keys = ((DeleteItemRequest) request).key(); + DynamoDbUtil.exportTagsWithKnownKeys(span, keys); + } + } +} diff --git a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbUtil.java b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbUtil.java new file mode 100644 index 00000000000..021965bec22 --- /dev/null +++ b/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 keys) { + if (keys == null || keys.isEmpty()) { + return; + } + + if (keys.size() == 1) { + // Single primary key case + Map.Entry 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 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))); + } + } +} diff --git a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy new file mode 100644 index 00000000000..7e839fcdb3a --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy @@ -0,0 +1,523 @@ +import datadog.trace.agent.test.AgentTestRunner +import datadog.trace.api.DDSpanTypes +import datadog.trace.core.tagprocessor.SpanPointersProcessor +import groovy.json.JsonSlurper +import org.testcontainers.containers.GenericContainer +import org.testcontainers.utility.DockerImageName +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider +import software.amazon.awssdk.regions.Region +import software.amazon.awssdk.services.dynamodb.DynamoDbClient +import software.amazon.awssdk.services.dynamodb.model.AttributeAction +import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition +import software.amazon.awssdk.services.dynamodb.model.AttributeValue +import software.amazon.awssdk.services.dynamodb.model.AttributeValueUpdate +import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest +import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest +import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest +import software.amazon.awssdk.services.dynamodb.model.ExpectedAttributeValue +import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement +import software.amazon.awssdk.services.dynamodb.model.KeyType +import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput +import software.amazon.awssdk.services.dynamodb.model.PutItemRequest +import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType +import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest +import spock.lang.Shared +import java.time.Duration + +class DynamoDbClientTest extends AgentTestRunner { + static final LOCALSTACK = new GenericContainer(DockerImageName.parse("localstack/localstack")) + .withExposedPorts(4566) + .withEnv("SERVICES", "dynamodb") + .withReuse(true) + .withStartupTimeout(Duration.ofSeconds(120)) + + @Shared + DynamoDbClient dynamoDbClient + + @Shared + String tableNameOneKey + + @Shared + String tableNameTwoKeys + + def setupSpec() { + LOCALSTACK.start() + def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) + dynamoDbClient = DynamoDbClient.builder() + .endpointOverride(URI.create(endPoint)) + .region(Region.US_EAST_1) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) + .build() + + // Create test tables + tableNameOneKey = "dynamodb-one-key-table" + tableNameTwoKeys = "dynamodb-two-key-table" + + // Create one-key table + CreateTableRequest oneKeyTableRequest = CreateTableRequest.builder() + .tableName(tableNameOneKey) + .keySchema( + KeySchemaElement.builder() + .attributeName("id") + .keyType(KeyType.HASH) + .build() + ) + .attributeDefinitions( + AttributeDefinition.builder() + .attributeName("id") + .attributeType(ScalarAttributeType.S) + .build() + ) + .provisionedThroughput( + ProvisionedThroughput.builder() + .readCapacityUnits(5L) + .writeCapacityUnits(5L) + .build() + ) + .build() + dynamoDbClient.createTable(oneKeyTableRequest) + + // Create two-key table + List twoKeyTableSchema = new ArrayList<>() + twoKeyTableSchema.add(KeySchemaElement.builder() + .attributeName("primaryKey") + .keyType(KeyType.HASH) + .build()) + twoKeyTableSchema.add(KeySchemaElement.builder() + .attributeName("sortKey") + .keyType(KeyType.RANGE) + .build()) + + List twoKeyTableAttributeDefs = new ArrayList<>() + twoKeyTableAttributeDefs.add(AttributeDefinition.builder() + .attributeName("primaryKey") + .attributeType(ScalarAttributeType.S) + .build()) + twoKeyTableAttributeDefs.add(AttributeDefinition.builder() + .attributeName("sortKey") + .attributeType(ScalarAttributeType.S) + .build()) + + CreateTableRequest twoKeyTableRequest = CreateTableRequest.builder() + .tableName(tableNameTwoKeys) + .keySchema(twoKeyTableSchema) + .attributeDefinitions(twoKeyTableAttributeDefs) + .provisionedThroughput( + ProvisionedThroughput.builder() + .readCapacityUnits(5L) + .writeCapacityUnits(5L) + .build() + ) + .build() + dynamoDbClient.createTable(twoKeyTableRequest) + } + + def cleanupSpec() { + DeleteTableRequest deleteOneKeyRequest = DeleteTableRequest.builder().tableName(tableNameOneKey).build() + DeleteTableRequest deleteTwoKeyRequest = DeleteTableRequest.builder().tableName(tableNameTwoKeys).build() + dynamoDbClient.deleteTable(deleteOneKeyRequest) + dynamoDbClient.deleteTable(deleteTwoKeyRequest) + LOCALSTACK.stop() + } + + def "should add span pointer for updateItem operation on one-key table"() { + when: + // First, put an item + Map item = new HashMap<>() + item.put("id", AttributeValue.builder().s("test-id-1").build()) + item.put("data", AttributeValue.builder().s("initial-value").build()) + + PutItemRequest putRequest = PutItemRequest.builder() + .tableName(tableNameOneKey) + .item(item) + .build() + dynamoDbClient.putItem(putRequest) + + // Then update the item + Map key = new HashMap<>() + key.put("id", AttributeValue.builder().s("test-id-1").build()) + + Map updates = new HashMap<>() + updates.put("data", AttributeValueUpdate.builder() + .value(AttributeValue.builder().s("updated-value").build()) + .action(AttributeAction.PUT) + .build()) + updates.put("counter", AttributeValueUpdate.builder() + .value(AttributeValue.builder().n("1").build()) + .action(AttributeAction.PUT) + .build()) + + UpdateItemRequest updateRequest = UpdateItemRequest.builder() + .tableName(tableNameOneKey) + .key(key) + .attributeUpdates(updates) + .build() + dynamoDbClient.updateItem(updateRequest) + + then: + assertTraces(2) { + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.PutItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "PutItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameOneKey + tag "tablename", tableNameOneKey + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + } + } + } + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.UpdateItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "UpdateItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameOneKey + tag "tablename", tableNameOneKey + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + tag "_dd.span_links", { it != null } + // Assert the span links + def spanLinks = tags["_dd.span_links"] + assert spanLinks != null + def links = new JsonSlurper().parseText(spanLinks) + assert links.size() == 1 + def link = links[0] + assert link["attributes"] != null + assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND + assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION + // First 32 chars of SHA256("dynamodb-one-key-table|id|test-id-1||") + assert link["attributes"]["ptr.hash"] == "ca8daaa857b00545ed5186a915cf1ab5" + assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND + } + } + } + } + } + + def "should add span pointer for updateItem operation on two-key table"() { + when: + // First, put an item + Map item = new HashMap<>() + item.put("primaryKey", AttributeValue.builder().s("customer-123").build()) + item.put("sortKey", AttributeValue.builder().s("order-456").build()) + item.put("status", AttributeValue.builder().s("pending").build()) + item.put("total", AttributeValue.builder().n("99.99").build()) + + PutItemRequest putTwoKeyRequest = PutItemRequest.builder() + .tableName(tableNameTwoKeys) + .item(item) + .build() + dynamoDbClient.putItem(putTwoKeyRequest) + + // Then update the item + Map key = new HashMap<>() + key.put("primaryKey", AttributeValue.builder().s("customer-123").build()) + key.put("sortKey", AttributeValue.builder().s("order-456").build()) + + Map updates = new HashMap<>() + updates.put("status", AttributeValueUpdate.builder() + .value(AttributeValue.builder().s("shipped").build()) + .action(AttributeAction.PUT) + .build()) + updates.put("total", AttributeValueUpdate.builder() + .value(AttributeValue.builder().n("129.99").build()) + .action(AttributeAction.PUT) + .build()) + + UpdateItemRequest updateTwoKeyRequest = UpdateItemRequest.builder() + .tableName(tableNameTwoKeys) + .key(key) + .attributeUpdates(updates) + .build() + dynamoDbClient.updateItem(updateTwoKeyRequest) + + then: + assertTraces(2) { + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.PutItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "PutItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameTwoKeys + tag "tablename", tableNameTwoKeys + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + } + } + } + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.UpdateItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "UpdateItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameTwoKeys + tag "tablename", tableNameTwoKeys + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + tag "_dd.span_links", { it != null } + // Assert the span links + def spanLinks = tags["_dd.span_links"] + assert spanLinks != null + def links = new JsonSlurper().parseText(spanLinks) + assert links.size() == 1 + def link = links[0] + assert link["attributes"] != null + assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND + assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION + // First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|customer-123|sortKey|order-456") + assert link["attributes"]["ptr.hash"] == "90922c7899a82ea34406fdcdfb95161e" + assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND + } + } + } + } + } + + def "should add span pointer for deleteItem operation on one-key table"() { + when: + // First, put an item + Map item = new HashMap<>() + item.put("id", AttributeValue.builder().s("delete-test-id").build()) + item.put("data", AttributeValue.builder().s("to-be-deleted").build()) + + PutItemRequest putDelTestRequest = PutItemRequest.builder() + .tableName(tableNameOneKey) + .item(item) + .build() + dynamoDbClient.putItem(putDelTestRequest) + + // Then delete the item + Map key = new HashMap<>() + key.put("id", AttributeValue.builder().s("delete-test-id").build()) + + DeleteItemRequest deleteRequest = DeleteItemRequest.builder() + .tableName(tableNameOneKey) + .key(key) + .build() + dynamoDbClient.deleteItem(deleteRequest) + + then: + assertTraces(2) { + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.PutItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "PutItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameOneKey + tag "tablename", tableNameOneKey + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + } + } + } + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.DeleteItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "DeleteItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameOneKey + tag "tablename", tableNameOneKey + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + tag "_dd.span_links", { it != null } + // Assert the span links + def spanLinks = tags["_dd.span_links"] + assert spanLinks != null + def links = new JsonSlurper().parseText(spanLinks) + assert links.size() == 1 + def link = links[0] + assert link["attributes"] != null + assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND + assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION + // First 32 chars of SHA256("dynamodb-one-key-table|id|delete-test-id||") + assert link["attributes"]["ptr.hash"] == "65031164be5e929fddd274a02cba3f9f" + assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND + } + } + } + } + } + + def "should add span pointer for deleteItem operation on two-key table"() { + when: + // First, put an item + Map item = new HashMap<>() + item.put("primaryKey", AttributeValue.builder().s("user-789").build()) + item.put("sortKey", AttributeValue.builder().s("profile").build()) + item.put("status", AttributeValue.builder().s("active").build()) + item.put("lastLogin", AttributeValue.builder().s("2023-01-01").build()) + + PutItemRequest putProfileRequest = PutItemRequest.builder() + .tableName(tableNameTwoKeys) + .item(item) + .build() + dynamoDbClient.putItem(putProfileRequest) + + // Then delete the item + Map key = new HashMap<>() + key.put("primaryKey", AttributeValue.builder().s("user-789").build()) + key.put("sortKey", AttributeValue.builder().s("profile").build()) + + // Use expected instead of conditionExpression + Map expected = new HashMap<>() + expected.put("status", + ExpectedAttributeValue.builder() + .value(AttributeValue.builder().s("active").build()) + .build()) + + DeleteItemRequest deleteWithExpectedRequest = DeleteItemRequest.builder() + .tableName(tableNameTwoKeys) + .key(key) + .expected(expected) + .build() + dynamoDbClient.deleteItem(deleteWithExpectedRequest) + + then: + assertTraces(2) { + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.PutItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "PutItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameTwoKeys + tag "tablename", tableNameTwoKeys + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + } + } + } + trace(1) { + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "DynamoDb.DeleteItem" + spanType DDSpanTypes.HTTP_CLIENT + tags { + defaultTags() + tag "component", "java-aws-sdk" + tag "aws.operation", "DeleteItem" + tag "aws.service", "DynamoDb" + tag "aws_service", "DynamoDb" + tag "aws.agent", "java-aws-sdk" + tag "aws.table.name", tableNameTwoKeys + tag "tablename", tableNameTwoKeys + tag "http.method", "POST" + tag "http.status_code", 200 + tag "http.url", { it.startsWith("http://") } + tag "peer.hostname", { it != null } + tag "peer.port", { it instanceof Integer } + tag "span.kind", "client" + tag "aws.requestId", { it != null } + tag "_dd.span_links", { it != null } + // Assert the span links + def spanLinks = tags["_dd.span_links"] + assert spanLinks != null + def links = new JsonSlurper().parseText(spanLinks) + assert links.size() == 1 + def link = links[0] + assert link["attributes"] != null + assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.DYNAMODB_PTR_KIND + assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION + // First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|user-789|sortKey|profile") + assert link["attributes"]["ptr.hash"] == "e5ce1148208c6f88041c73ceb9bbbf3a" + assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND + } + } + } + } + } +} diff --git a/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbUtilTest.groovy b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbUtilTest.groovy new file mode 100644 index 00000000000..95692d9b705 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbUtilTest.groovy @@ -0,0 +1,130 @@ +import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags +import datadog.trace.instrumentation.aws.v2.dynamodb.DynamoDbUtil +import org.junit.Test +import software.amazon.awssdk.services.dynamodb.model.AttributeValue +import software.amazon.awssdk.core.SdkBytes + +class DynamoDbUtilTest { + static createMockSpan() { + def tags = [:] + + def mockSpan = [ + setTag: { String key, String value -> + tags[key] = value + return null + } + ] as AgentSpan + + return [span: mockSpan, tags: tags] + } + + @Test + void testExportTagsWithNullKeys() { + def mockData = createMockSpan() + def mockSpan = mockData.span + def tags = mockData.tags + + DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, null) + + assert tags.isEmpty() + } + + @Test + void testExportTagsWithEmptyKeys() { + def mockData = createMockSpan() + def mockSpan = mockData.span + def tags = mockData.tags + + DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, [:]) + + assert tags.isEmpty() + } + + @Test + void testExportTagsWithSingleStringKey() { + def mockData = createMockSpan() + def mockSpan = mockData.span + def tags = mockData.tags + + def keys = [ + "id": AttributeValue.builder().s("12345").build() + ] + + DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys) + + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "id" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "12345" + } + + @Test + void testExportTagsWithSingleNumberKey() { + def mockData = createMockSpan() + def mockSpan = mockData.span + def tags = mockData.tags + + def keys = [ + "count": AttributeValue.builder().n("42").build() + ] + + DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys) + + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "count" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "42" + } + + @Test + void testExportTagsWithSingleBinaryKey() { + def mockData = createMockSpan() + def mockSpan = mockData.span + def tags = mockData.tags + + def binaryData = "binary-data".getBytes() + def keys = [ + "data": AttributeValue.builder().b(SdkBytes.fromByteArray(binaryData)).build() + ] + + DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys) + + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "data" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "binary-data" + } + + @Test + void testExportTagsWithTwoKeys() { + def mockData = createMockSpan() + def mockSpan = mockData.span + def tags = mockData.tags + + def keys = [ + "id": AttributeValue.builder().s("12345").build(), + "name": AttributeValue.builder().s("item-name").build() + ] + + DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys) + + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "id" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "12345" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2] == "name" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2_VALUE] == "item-name" + } + + @Test + void testExportTagsWithTwoKeysSortsAlphabetically() { + def mockData = createMockSpan() + def mockSpan = mockData.span + def tags = mockData.tags + + def keys = [ + "bKey": AttributeValue.builder().s("abc").build(), + "aKey": AttributeValue.builder().s("zxy").build() + ] + + DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys) + + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "aKey" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "zxy" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2] == "bKey" + assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2_VALUE] == "abc" + } +} diff --git a/dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3ClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3ClientInstrumentation.java index 6f28c9aef98..0a675c969b4 100644 --- a/dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3ClientInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3ClientInstrumentation.java @@ -31,7 +31,7 @@ public void methodAdvice(MethodTransformer transformer) { @Override public String[] helperClassNames() { - return new String[] {packageName + ".S3Interceptor", packageName + ".TextMapInjectAdapter"}; + return new String[] {packageName + ".S3Interceptor"}; } public static class AwsS3BuilderAdvice { diff --git a/dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/TextMapInjectAdapter.java deleted file mode 100644 index cfefe319de3..00000000000 --- a/dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/TextMapInjectAdapter.java +++ /dev/null @@ -1,13 +0,0 @@ -package datadog.trace.instrumentation.aws.v2.s3; - -import datadog.context.propagation.CarrierSetter; - -public class TextMapInjectAdapter implements CarrierSetter { - - 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("\","); - } -} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java b/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java index fabf01fc7a4..ea86ae2e76f 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java @@ -4,6 +4,11 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpanContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.AWS_BUCKET_NAME; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.AWS_OBJECT_KEY; +import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.AWS_TABLE_NAME; +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 static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.S3_ETAG; import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; @@ -26,21 +31,37 @@ public class SpanPointersProcessor implements TagsPostProcessor { // The pointer direction will always be down. The serverless agent handles cases where the // direction is up. static final String DOWN_DIRECTION = "d"; + static final String DYNAMODB_PTR_KIND = "aws.dynamodb.item"; static final String S3_PTR_KIND = "aws.s3.object"; static final String LINK_KIND = "span-pointer"; @Override public Map processTags( Map unsafeTags, DDSpanContext spanContext, List spanLinks) { + AgentSpanLink s3Link = handleS3SpanPointer(unsafeTags); + if (s3Link != null) { + spanLinks.add(s3Link); + } + + AgentSpanLink dynamoDbLink = handleDynamoDbSpanPointer(unsafeTags); + if (dynamoDbLink != null) { + spanLinks.add(dynamoDbLink); + } + + return unsafeTags; + } + + private static AgentSpanLink handleS3SpanPointer(Map unsafeTags) { String eTag = asString(unsafeTags.remove(S3_ETAG)); if (eTag == null) { - return unsafeTags; + return null; } String bucket = asString(unsafeTags.get(AWS_BUCKET_NAME)); String key = asString(unsafeTags.get(AWS_OBJECT_KEY)); if (bucket == null || key == null) { - LOG.debug("Unable to calculate span pointer hash because could not find bucket or key tags."); - return unsafeTags; + // This might be from an S3 operation not supported by span pointers, so we skip without + // logging anything. + return null; } // Hash calculation rules: @@ -50,21 +71,51 @@ public Map processTags( } String[] components = new String[] {bucket, key, eTag}; try { - SpanAttributes attributes = - SpanAttributes.builder() - .put("ptr.kind", S3_PTR_KIND) - .put("ptr.dir", DOWN_DIRECTION) - .put("ptr.hash", generatePointerHash(components)) - .put("link.kind", LINK_KIND) - .build(); - - AgentSpanLink link = SpanLink.from(noopSpanContext(), DEFAULT_FLAGS, "", attributes); - spanLinks.add(link); + String hash = generatePointerHash(components); + return buildSpanPointer(hash, S3_PTR_KIND); } catch (Exception e) { LOG.debug("Failed to add span pointer: {}", e.getMessage()); + return null; } + } - return unsafeTags; + private static AgentSpanLink handleDynamoDbSpanPointer(Map unsafeTags) { + // Hash calculation rules: + // https://github.com/DataDog/dd-span-pointer-rules/blob/main/AWS/DynamoDB/Item/README.md + String tableName = asString(unsafeTags.get(AWS_TABLE_NAME)); + if (tableName == null) { + return null; + } + String primaryKey1Name = asString(unsafeTags.remove(DYNAMO_PRIMARY_KEY_1)); + String primaryKey1Value = asString(unsafeTags.remove(DYNAMO_PRIMARY_KEY_1_VALUE)); + if (primaryKey1Name == null || primaryKey1Value == null) { + // This might be from a DynamoDB operation not supported by span pointers, so we skip without + // logging anything. + return null; + } + + // If these don't exist, the user has a table with only partition key but no sort key. + // Then, we set them to empty strings when calculating the hash. + String primaryKey2Name = asString(unsafeTags.remove(DYNAMO_PRIMARY_KEY_2)); + String primaryKey2Value = asString(unsafeTags.remove(DYNAMO_PRIMARY_KEY_2_VALUE)); + if (primaryKey2Name == null) { + primaryKey2Name = ""; + } + if (primaryKey2Value == null) { + primaryKey2Value = ""; + } + + String[] components = + new String[] { + tableName, primaryKey1Name, primaryKey1Value, primaryKey2Name, primaryKey2Value + }; + try { + String hash = generatePointerHash(components); + return buildSpanPointer(hash, DYNAMODB_PTR_KIND); + } catch (Exception e) { + LOG.debug("Failed to add span pointer: {}", e.getMessage()); + return null; + } } private static String asString(Object o) { @@ -98,4 +149,16 @@ private static String generatePointerHash(String[] components) throws NoSuchAlgo byte[] truncatedHash = Arrays.copyOf(fullHash, 16); return Strings.toHexString(truncatedHash); } + + private static AgentSpanLink buildSpanPointer(String hash, String ptrKind) { + SpanAttributes attributes = + SpanAttributes.builder() + .put("ptr.kind", ptrKind) + .put("ptr.dir", DOWN_DIRECTION) + .put("ptr.hash", hash) + .put("link.kind", LINK_KIND) + .build(); + + return SpanLink.from(noopSpanContext(), DEFAULT_FLAGS, "", attributes); + } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java index bfd71a2abcf..899bc3f0b8d 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InstrumentationTags.java @@ -36,7 +36,13 @@ public class InstrumentationTags { public static final String TABLE_NAME = "tablename"; public static final String AWS_REQUEST_ID = "aws.requestId"; public static final String AWS_STORAGE_CLASS = "aws.storage.class"; + + // These are temporary keys used for span pointer hash calculation public static final String S3_ETAG = "s3.eTag"; + public static final String DYNAMO_PRIMARY_KEY_1 = "dynamodb.primary_key_1"; + public static final String DYNAMO_PRIMARY_KEY_1_VALUE = "dynamodb.primary_key_1_value"; + public static final String DYNAMO_PRIMARY_KEY_2 = "dynamodb.primary_key_2"; + public static final String DYNAMO_PRIMARY_KEY_2_VALUE = "dynamodb.primary_key_2_value"; public static final String BUCKET = "bucket"; public static final String CASSANDRA_CONTACT_POINTS = "db.cassandra.contact.points"; diff --git a/settings.gradle b/settings.gradle index 08043a76680..12e9d0ccb8a 100644 --- a/settings.gradle +++ b/settings.gradle @@ -193,6 +193,7 @@ include ':dd-java-agent:instrumentation:armeria-grpc' include ':dd-java-agent:instrumentation:armeria-jetty' include ':dd-java-agent:instrumentation:avro' include ':dd-java-agent:instrumentation:aws-common' +include ':dd-java-agent:instrumentation:aws-java-dynamodb-2.0' include ':dd-java-agent:instrumentation:aws-java-eventbridge-2.0' include ':dd-java-agent:instrumentation:aws-java-sdk-1.11.0' include ':dd-java-agent:instrumentation:aws-java-sdk-2.2'