add config to disable injecting datadog attribute in aws messaging services

Author: ojproductions

dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java

@@ -6,6 +6,7 @@
 import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
 
 import datadog.context.Context;
+import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
 import datadog.trace.api.datastreams.PathwayContext;
@@ -34,7 +35,8 @@ public class EventBridgeInterceptor implements ExecutionInterceptor {
 
   @Override
   public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) {
-    if (!(context.request() instanceof PutEventsRequest)) {
+    if (!(context.request() instanceof PutEventsRequest)
+        || !Config.get().isAwsInjectDatadogAttributeEnabled()) {
       return context.request();
     }
 

dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java

@@ -3,6 +3,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
 
 import datadog.context.Context;
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.InstanceStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import software.amazon.awssdk.core.SdkRequest;
@@ -23,10 +24,14 @@ public SfnInterceptor() {}
 
   @Override
   public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) {
+    SdkRequest request = context.request();
+    if (Config.get().isAwsInjectDatadogAttributeEnabled()) {
+      return request;
+    }
     try {
       return modifyRequestImpl(context, executionAttributes);
     } catch (Exception e) {
-      return context.request();
+      return request;
     }
   }
 

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

@@ -13,6 +13,7 @@
 import com.amazonaws.services.sns.model.PublishBatchRequestEntry;
 import com.amazonaws.services.sns.model.PublishRequest;
 import datadog.context.Context;
+import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
 import datadog.trace.bootstrap.ContextStore;
@@ -57,7 +58,7 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request
       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) {
+      if (messageAttributes.size() < 10 && Config.get().isAwsInjectDatadogAttributeEnabled()) {
         // Extract the topic name from the ARN for DSM
         String topicName = pRequest.getTopicArn();
         if (null == topicName) {

dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java

@@ -7,6 +7,7 @@
 import static datadog.trace.instrumentation.aws.v2.sns.TextMapInjectAdapter.SETTER;
 
 import datadog.context.Context;
+import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
 import datadog.trace.bootstrap.InstanceStore;
@@ -55,7 +56,8 @@ public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes execu
       PublishRequest request = (PublishRequest) context.request();
       // 10 messageAttributes is a limit from SQS, which is often used as a subscriber, therefore
       // the limit still applies here
-      if (request.messageAttributes().size() < 10) {
+      if (request.messageAttributes().size() < 10
+          && Config.get().isAwsInjectDatadogAttributeEnabled()) {
         // Get topic name for DSM
         String snsTopicArn = request.topicArn();
         if (null == snsTopicArn) {

dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsInterceptor.java

@@ -18,6 +18,7 @@
 import datadog.context.Context;
 import datadog.context.propagation.Propagator;
 import datadog.context.propagation.Propagators;
+import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
 import datadog.trace.bootstrap.ContextStore;
@@ -69,7 +70,8 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request
       }
     } else if (request instanceof ReceiveMessageRequest) {
       ReceiveMessageRequest rmRequest = (ReceiveMessageRequest) request;
-      if (rmRequest.getMessageAttributeNames().size() < 10
+      if (Config.get().isAwsInjectDatadogAttributeEnabled()
+          && rmRequest.getMessageAttributeNames().size() < 10
           && !rmRequest.getMessageAttributeNames().contains(DATADOG_KEY)) {
         List<String> attributeNames = new ArrayList<>(rmRequest.getMessageAttributeNames());
         attributeNames.add(DATADOG_KEY);

dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java

@@ -10,6 +10,7 @@
 import datadog.context.Context;
 import datadog.context.propagation.Propagator;
 import datadog.context.propagation.Propagators;
+import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
 import datadog.trace.bootstrap.InstanceStore;
@@ -77,7 +78,8 @@ public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes execu
     } else if (context.request() instanceof ReceiveMessageRequest) {
       ReceiveMessageRequest request = (ReceiveMessageRequest) context.request();
       if (request.messageAttributeNames().size() < 10
-          && !request.messageAttributeNames().contains(DATADOG_KEY)) {
+          && !request.messageAttributeNames().contains(DATADOG_KEY)
+          && Config.get().isAwsInjectDatadogAttributeEnabled()) {
         List<String> messageAttributeNames = new ArrayList<>(request.messageAttributeNames());
         messageAttributeNames.add(DATADOG_KEY);
         return request.toBuilder().messageAttributeNames(messageAttributeNames).build();

dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java

@@ -47,6 +47,7 @@ public final class TracerConfig {
 
   @Deprecated
   public static final String TRACE_SAMPLING_OPERATION_RULES = "trace.sampling.operation.rules";
+
   // JSON rules
   public static final String TRACE_SAMPLING_RULES = "trace.sampling.rules";
   public static final String SPAN_SAMPLING_RULES = "span.sampling.rules";
@@ -164,6 +165,8 @@ public final class TracerConfig {
   public static final String TRACE_CLOUD_PAYLOAD_TAGGING_MAX_TAGS =
       "trace.cloud.payload.tagging.max-tags";
   public static final String TRACE_SERVICE_DISCOVERY_ENABLED = "trace.service.discovery.enabled";
+  public static final String TRACE_AWS_INJECT_DATADOG_ATTRIBUTE =
+      "trace.aws.inject.datadog.attribute";
 
   private TracerConfig() {}
 }

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -623,6 +623,7 @@
 import static datadog.trace.api.config.TracerConfig.TRACE_AGENT_PORT;
 import static datadog.trace.api.config.TracerConfig.TRACE_AGENT_URL;
 import static datadog.trace.api.config.TracerConfig.TRACE_ANALYTICS_ENABLED;
+import static datadog.trace.api.config.TracerConfig.TRACE_AWS_INJECT_DATADOG_ATTRIBUTE;
 import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_BYTES;
 import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_ITEMS;
 import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_TAG_KEYS;
@@ -822,6 +823,7 @@ public static String getHostName() {
   private final String agentUnixDomainSocket;
   private final String agentNamedPipe;
   private final int agentTimeout;
+
   /** Should be set to {@code true} when running in agentless mode in a JVM without TLS */
   private final boolean forceClearTextHttpForIntakeClient;
 
@@ -1301,6 +1303,8 @@ public static String getHostName() {
   private final int tagValueUtf8CacheSize;
   private final int stackTraceLengthLimit;
 
+  private final boolean awsInjectDatadogAttributeEnabled;
+
   private final RumInjectorConfig rumInjectorConfig;
 
   private final boolean aiGuardEnabled;
@@ -1583,6 +1587,9 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
     awsServerless =
         getEnv("AWS_LAMBDA_FUNCTION_NAME") != null && !getEnv("AWS_LAMBDA_FUNCTION_NAME").isEmpty();
 
+    awsInjectDatadogAttributeEnabled =
+        configProvider.getBoolean(TRACE_AWS_INJECT_DATADOG_ATTRIBUTE, true);
+
     spanAttributeSchemaVersion = schemaVersionFromConfig();
 
     peerHostNameEnabled = configProvider.getBoolean(TRACE_PEER_HOSTNAME_ENABLED, true);
@@ -4684,6 +4691,10 @@ public int getStackTraceLengthLimit() {
     return stackTraceLengthLimit;
   }
 
+  public boolean isAwsInjectDatadogAttributeEnabled() {
+    return this.awsInjectDatadogAttributeEnabled;
+  }
+
   /** @return A map of tags to be applied only to the local application root span. */
   public TagMap getLocalRootSpanTags() {
     final Map<String, String> runtimeTags = getRuntimeTags();