change to the same instrumentation as other dd-trace libraries

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
   }
 }

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

dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java

@@ -1,17 +1,21 @@
 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;
 import datadog.trace.api.TracePropagationStyle;
 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<String, MessageAttributeValue> 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<String, MessageAttributeValue> messageAttributes = entry.getMessageAttributes();
+        if (messageAttributes.size() < 10) {
+          binaryValue.rewind();
+          messageAttributes.put(
+              "_datadog",
+              new MessageAttributeValue().withDataType("Binary").withBinaryValue(binaryValue));
+        }
       }
     }
     return request;

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

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"

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
   }
 }

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

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<String, MessageAttributeValue> 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<PublishBatchRequestEntry> 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<String, MessageAttributeValue> 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();

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

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"