diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index 05a8dbca4c8..6c8b9068364 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "com.amazonaws" module = "aws-java-sdk-sns" - versions = "[1.12.0,)" + versions = "[1.12.0,2)" assertInverse = true } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java deleted file mode 100644 index e56450d1ddd..00000000000 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java +++ /dev/null @@ -1,21 +0,0 @@ -package datadog.trace.instrumentation.aws.v1.sns; - -import com.amazonaws.services.sns.model.MessageAttributeValue; -import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; -import java.util.Map; - -public class MessageAttributeInjector - implements AgentPropagation.Setter> { - - public static final MessageAttributeInjector SETTER = new MessageAttributeInjector(); - - @Override - public void set( - final Map carrier, final String key, final String value) { - // 10 messageAttributes is a limit from SQS, which is often used as a subscriber and therefore - // still apply here - if (carrier.size() < 10 && !carrier.containsKey(key)) { - carrier.put(key, new MessageAttributeValue().withDataType("String").withStringValue(value)); - } - } -} diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java index 440986b4508..b6565469b19 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java @@ -17,10 +17,9 @@ @AutoService(InstrumenterModule.class) public final class SnsClientInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForSingleType { - private static final String INSTRUMENTATION_NAME = "aws-sdk"; public SnsClientInstrumentation() { - super(INSTRUMENTATION_NAME); + super("sns", "aws-sdk"); } @Override @@ -37,9 +36,7 @@ public void methodAdvice(MethodTransformer transformer) { @Override public String[] helperClassNames() { - return new String[] { - packageName + ".SnsInterceptor", packageName + ".MessageAttributeInjector" - }; + return new String[] {packageName + ".SnsInterceptor", packageName + ".TextMapInjectAdapter"}; } @Override diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java index a63288e4c83..270a02dda85 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java @@ -1,10 +1,11 @@ package datadog.trace.instrumentation.aws.v1.sns; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; -import static datadog.trace.instrumentation.aws.v1.sns.MessageAttributeInjector.SETTER; +import static datadog.trace.instrumentation.aws.v1.sns.TextMapInjectAdapter.SETTER; import com.amazonaws.AmazonWebServiceRequest; import com.amazonaws.handlers.RequestHandler2; +import com.amazonaws.services.sns.model.MessageAttributeValue; import com.amazonaws.services.sns.model.PublishBatchRequest; import com.amazonaws.services.sns.model.PublishBatchRequestEntry; import com.amazonaws.services.sns.model.PublishRequest; @@ -12,6 +13,9 @@ import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Map; public class SnsInterceptor extends RequestHandler2 { @@ -29,14 +33,44 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request final AgentSpan span = newSpan(request); // note: modifying message attributes has to be done before marshalling, otherwise the changes // are not reflected in the actual request (and the MD5 check on send will fail). - propagate().inject(span, pRequest.getMessageAttributes(), SETTER, TracePropagationStyle.XRAY); - + Map messageAttributes = pRequest.getMessageAttributes(); + // 10 messageAttributes is a limit from SQS, which is often used as a subscriber, therefore + // the limit still applies here + if (messageAttributes.size() < 10) { + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + messageAttributes.put( + "_datadog", + new MessageAttributeValue() + .withDataType( + "Binary") // Use Binary since SNS subscription filter policies fail silently + // with JSON strings + // https://github.com/DataDog/datadog-lambda-js/pull/269 + .withBinaryValue( + ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)))); + } } else if (request instanceof PublishBatchRequest) { PublishBatchRequest pmbRequest = (PublishBatchRequest) request; final AgentSpan span = newSpan(request); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + ByteBuffer binaryValue = + ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)); for (PublishBatchRequestEntry entry : pmbRequest.getPublishBatchRequestEntries()) { - propagate().inject(span, entry.getMessageAttributes(), SETTER, TracePropagationStyle.XRAY); + Map messageAttributes = entry.getMessageAttributes(); + if (messageAttributes.size() < 10) { + binaryValue.rewind(); + messageAttributes.put( + "_datadog", + new MessageAttributeValue().withDataType("Binary").withBinaryValue(binaryValue)); + } } } return request; diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/TextMapInjectAdapter.java new file mode 100644 index 00000000000..6f2bb8f888b --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/TextMapInjectAdapter.java @@ -0,0 +1,13 @@ +package datadog.trace.instrumentation.aws.v1.sns; + +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; + +public class TextMapInjectAdapter implements AgentPropagation.Setter { + + 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-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index 05a8281f359..c54ee6cba62 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -5,7 +5,6 @@ import com.amazonaws.services.sns.AmazonSNSClient import com.amazonaws.services.sns.AmazonSNSClientBuilder import datadog.trace.agent.test.naming.VersionedNamingTestBase import datadog.trace.agent.test.utils.TraceUtils -import datadog.trace.api.DDSpanId import datadog.trace.api.DDSpanTypes import datadog.trace.api.config.GeneralConfig import datadog.trace.bootstrap.instrumentation.api.Tags @@ -142,8 +141,14 @@ abstract class SnsClientTest extends VersionedNamingTestBase { and: messageBody["Message"] == "sometext" - messageBody["MessageAttributes"]["X-Amzn-Trace-Id"]["Value"] =~ - /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ + String base64EncodedString = messageBody["MessageAttributes"]["_datadog"]["Value"] + byte[] decodedBytes = base64EncodedString.decodeBase64() + String decodedString = new String(decodedBytes, "UTF-8") + JsonSlurper slurper = new JsonSlurper() + Map traceContextInJson = slurper.parseText(decodedString) + traceContextInJson['x-datadog-trace-id'] == sendSpan.traceId.toString() + traceContextInJson['x-datadog-parent-id'] == sendSpan.spanId.toString() + traceContextInJson['x-datadog-sampling-priority'] == "1" } } @@ -175,7 +180,7 @@ class SnsClientV1ForkedTest extends SnsClientTest { @Override String expectedOperation(String awsService, String awsOperation) { - if (awsService == "SNS"&& awsOperation == "Publish") { + if (awsService == "SNS" && awsOperation == "Publish") { return "aws.sns.send" } return "http.client.request" diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle index d9e25837058..945e4f30c2a 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "software.amazon.awssdk" module = "sns" - versions = "[2.2.0,)" + versions = "[2.2.0,3)" assertInverse = true } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java deleted file mode 100644 index ee60d19301b..00000000000 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java +++ /dev/null @@ -1,22 +0,0 @@ -package datadog.trace.instrumentation.aws.v2.sns; - -import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; -import java.util.Map; -import software.amazon.awssdk.services.sns.model.MessageAttributeValue; - -public class MessageAttributeInjector - implements AgentPropagation.Setter> { - - public static final MessageAttributeInjector SETTER = new MessageAttributeInjector(); - - @Override - public void set( - final Map carrier, final String key, final String value) { - // 10 messageAttributes is a limit from SQS, which is often used as a subscriber and therefore - // still apply here - if (carrier.size() < 10 && !carrier.containsKey(key)) { - carrier.put( - key, MessageAttributeValue.builder().dataType("String").stringValue(value).build()); - } - } -} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java index aad95b7b235..0e24d267baf 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java @@ -16,10 +16,8 @@ @AutoService(InstrumenterModule.class) public final class SnsClientInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForSingleType { - private static final String INSTRUMENTATION_NAME = "aws-sdk"; - public SnsClientInstrumentation() { - super(INSTRUMENTATION_NAME); + super("sns", "aws-sdk"); } @Override @@ -36,9 +34,7 @@ public void methodAdvice(MethodTransformer transformer) { @Override public String[] helperClassNames() { - return new String[] { - packageName + ".SnsInterceptor", packageName + ".MessageAttributeInjector" - }; + return new String[] {packageName + ".SnsInterceptor", packageName + ".TextMapInjectAdapter"}; } @Override diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index e32de673414..9c995f52c39 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -1,14 +1,16 @@ package datadog.trace.instrumentation.aws.v2.sns; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; -import static datadog.trace.instrumentation.aws.v2.sns.MessageAttributeInjector.SETTER; +import static datadog.trace.instrumentation.aws.v2.sns.TextMapInjectAdapter.SETTER; import datadog.trace.api.TracePropagationStyle; import datadog.trace.bootstrap.InstanceStore; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; import java.util.Map; +import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.SdkRequest; import software.amazon.awssdk.core.interceptor.Context; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; @@ -35,16 +37,40 @@ public SdkRequest modifyRequest( Map messageAttributes = new HashMap<>(request.messageAttributes()); final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); - propagate().inject(span, messageAttributes, SETTER, TracePropagationStyle.XRAY); + // 10 messageAttributes is a limit from SQS, which is often used as a subscriber, therefore + // the limit still applies here + if (messageAttributes.size() < 10) { + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + + messageAttributes.put( + "_datadog", // Use Binary since SNS subscription filter policies fail silently with JSON + // strings https://github.com/DataDog/datadog-lambda-js/pull/269 + MessageAttributeValue.builder() + .dataType("Binary") + .binaryValue(SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8)) + .build()); + } return request.toBuilder().messageAttributes(messageAttributes).build(); } else if (context.request() instanceof PublishBatchRequest) { PublishBatchRequest request = (PublishBatchRequest) context.request(); final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); ArrayList entries = new ArrayList<>(); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + SdkBytes binaryValue = SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8); for (PublishBatchRequestEntry entry : request.publishBatchRequestEntries()) { Map messageAttributes = new HashMap<>(entry.messageAttributes()); - propagate().inject(span, messageAttributes, SETTER, TracePropagationStyle.XRAY); + messageAttributes.put( + "_datadog", + MessageAttributeValue.builder().dataType("Binary").binaryValue(binaryValue).build()); entries.add(entry.toBuilder().messageAttributes(messageAttributes).build()); } return request.toBuilder().publishBatchRequestEntries(entries).build(); diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/TextMapInjectAdapter.java new file mode 100644 index 00000000000..cfe0368e298 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/TextMapInjectAdapter.java @@ -0,0 +1,13 @@ +package datadog.trace.instrumentation.aws.v2.sns; + +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; + +public class TextMapInjectAdapter implements AgentPropagation.Setter { + + 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-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy index 48e9921296f..f18dc377dcb 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -1,7 +1,6 @@ import datadog.trace.agent.test.naming.VersionedNamingTestBase import datadog.trace.agent.test.utils.TraceUtils -import datadog.trace.api.DDSpanId import datadog.trace.api.DDSpanTypes import datadog.trace.api.config.GeneralConfig import datadog.trace.bootstrap.instrumentation.api.Tags @@ -11,7 +10,6 @@ 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.sns.SnsClient -import software.amazon.awssdk.services.sns.model.PublishRequest import software.amazon.awssdk.services.sns.model.PublishResponse import software.amazon.awssdk.services.sqs.SqsClient import software.amazon.awssdk.services.sqs.model.QueueAttributeName @@ -123,8 +121,14 @@ abstract class SnsClientTest extends VersionedNamingTestBase { and: messageBody["Message"] == "sometext" - messageBody["MessageAttributes"]["X-Amzn-Trace-Id"]["Value"] =~ - /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ + String base64EncodedString = messageBody["MessageAttributes"]["_datadog"]["Value"] + byte[] decodedBytes = base64EncodedString.decodeBase64() + String decodedString = new String(decodedBytes, "UTF-8") + JsonSlurper slurper = new JsonSlurper() + Map traceContextInJson = slurper.parseText(decodedString) + traceContextInJson['x-datadog-trace-id'] == sendSpan.traceId.toString() + traceContextInJson['x-datadog-parent-id'] == sendSpan.spanId.toString() + traceContextInJson['x-datadog-sampling-priority'] == "1" } } @@ -156,7 +160,7 @@ class SnsClientV1ForkedTest extends SnsClientTest { @Override String expectedOperation(String awsService, String awsOperation) { - if (awsService == "SNS"&& awsOperation == "Publish") { + if (awsService == "Sns" && awsOperation == "Publish") { return "aws.sns.send" } return "http.client.request"