remove binary datadog attribute if present in JMS SQS instrumentation

DataDog · Jul 5, 2024 · feb4c21 · feb4c21
1 parent 0327c4d
commit feb4c21
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/.../src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsJmsMessageInstrumentation.java b/.../src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsJmsMessageInstrumentation.java
@@ -7,6 +7,7 @@
 
 import com.amazon.sqs.javamessaging.message.SQSMessage;
 import com.amazonaws.services.sqs.model.Message;
+import com.amazonaws.services.sqs.model.MessageAttributeValue;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
@@ -32,6 +33,20 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static class CopyTracePropertyAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(2) Message sqsMessage) {
+      Map<String, MessageAttributeValue> messageAttributes = sqsMessage.getMessageAttributes();
+      MessageAttributeValue ddAttribute = messageAttributes.get("_datadog");
+      if (ddAttribute != null && "Binary".equals(ddAttribute.getDataType())) {
+        // binary message attributes are not supported by amazon-sqs-java-messaging-lib, and there
+        // is a chance we might introduce one, either when the message was sent from SNS or from a
+        // DD-instrumented Javascript app.
+        // When we reach this point, the value would already have been used by the aws-sqs
+        // instrumentation, so we can safely remove it.
+        messageAttributes.remove("_datadog");
+      }
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onExit(
         @Advice.Argument(2) Message sqsMessage, @Advice.FieldValue("properties") Map properties)