wip

lib/datadog/tracing/contrib/aws/configuration/settings.rb

@@ -46,22 +46,27 @@ class Settings < Contrib::Configuration::Settings
               o.env Ext::ENV_PEER_SERVICE
             end
 
+            # Enables distributed trace propagation for SNS and SQS messages.
+            # @default `DD_TRACE_AWS_PROPAGATION_ENABLED` environment variable, otherwise `false`
+            # @return [Boolean]
             option :propagation do |o|
               o.type :bool
-              # TODO: Add env var for this
-              # TODO: Add env var for this
-              # TODO: Add env var for this
-              # o.env Ext::ENV_PEER_SERVICE
+              o.env Ext::ENV_PROPAGATION_ENABLED
               o.default false
             end
 
-            option :batch_propagation do |o|
-              o.type :bool
-              # TODO: Add env var for this
-              # TODO: Add env var for this
-              # TODO: Add env var for this
-              # o.env Ext::ENV_PEER_SERVICE
-              o.default false
+            # Controls whether the local trace is parented to the SNS message consumed.
+            # Possible values are:
+            # `local`: The local active trace is used; SNS has no effect on trace parentage.
+            # `distributed`: The local active trace becomes a child of the propagation context from the SNS message.
+            #
+            # This option is always disable (the equivalent to`local`) if `propagation` is disabled.
+            # @default `DD_TRACE_AWS_TRACE_PARENTAGE_STYLE` environment variable, otherwise `distributed`
+            # @return [String]
+            option :parentage_style do |o|
+              o.type :string
+              o.env Ext::ENV_TRACE_PARENTAGE_STYLE
+              o.default 'distributed'
             end
         end
       end

lib/datadog/tracing/contrib/aws/ext.rb

@@ -13,6 +13,9 @@ module Ext
           # @!visibility private
           ENV_ANALYTICS_ENABLED = 'DD_TRACE_AWS_ANALYTICS_ENABLED'
           ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_AWS_ANALYTICS_SAMPLE_RATE'
+          ENV_PROPAGATION_ENABLED = 'DD_TRACE_AWS_PROPAGATION_ENABLED'
+          ENV_TRACE_PARENTAGE_STYLE = 'DD_TRACE_AWS_TRACE_PARENTAGE_STYLE'
+
           DEFAULT_PEER_SERVICE_NAME = 'aws'
           SPAN_COMMAND = 'aws.command'
           TAG_AGENT = 'aws.agent'

lib/datadog/tracing/contrib/aws/service/base.rb

@@ -12,7 +12,7 @@ def add_tags(span, params); end
 
             MESSAGE_ATTRIBUTES_LIMIT = 10 # Can't set more than 10 message attributes
 
-            def extract_propagation(context)
+            def extract_propagation!(context)
               message_attributes = context.params[:message_attributes]
 
               return unless message_attributes && (datadog = message_attributes['_datadog'])

lib/datadog/tracing/contrib/aws/service/sqs.rb

@@ -11,8 +11,12 @@ module Service
           # SQS tag handlers.
           class SQS < Base
             def before_span(config, context)
-              if config[:propagation] && context.operation == :receive_message
-                extract_propagation(context)
+              # DEV: Because we only support tracing propagation today, having separate `propagation and `propagation_style`
+              # options seems redundant. But when the DSM propagation is introduced, it's possible for `propagation` to be
+              # enable and `propagation_style` to disable, while DSM propagation is still enabled, as its data is not
+              # directly related to tracing parentage.
+              if config[:propagation] && config[:parentage_style] == 'distributed' && context.operation == :receive_message
+                extract_propagation!(context)
               end
             end
 

sig/datadog/tracing/contrib/aws/ext.rbs

@@ -6,6 +6,7 @@ module Datadog
           ENV_ENABLED: "DD_TRACE_AWS_ENABLED"
 
           ENV_PEER_SERVICE: "DD_TRACE_AWS_PEER_SERVICE"
+          ENV_PROPAGATION_ENABLED: string
           ENV_SERVICE_NAME: "DD_TRACE_AWS_SERVICE_NAME"
 
           ENV_ANALYTICS_ENABLED: "DD_TRACE_AWS_ANALYTICS_ENABLED"
@@ -14,6 +15,7 @@ module Datadog
 
           DEFAULT_PEER_SERVICE_NAME: "aws"
 
+          ENV_TRACE_PARENTAGE_STYLE: string
           PEER_SERVICE_SOURCES: Array[String]
 
           SPAN_COMMAND: "aws.command"

sig/datadog/tracing/contrib/aws/service/base.rbs

@@ -4,6 +4,8 @@ module Datadog
       module Aws
         module Service
           class Base
+            MESSAGE_ATTRIBUTES_LIMIT: int
+
             def add_tags: (untyped span, untyped params) -> nil
           end
         end

spec/datadog/tracing/contrib/aws/instrumentation_spec.rb

@@ -255,6 +255,7 @@
     end
 
     describe '#send_message_batch' do
+      # TODO: SHOULD TAG ALL MESSAGES
       subject!(:send_message_batch) do
         client.send_message_batch(
           {
@@ -473,6 +474,11 @@
       end
     end
 
+    describe '#publish_batch' do
+      # TODO: Implement this test
+      # Should tag all messages, not just one.
+    end
+
     describe '#create_topic' do
       subject!(:create_topic) do
         client.create_topic(