Add SNS/SQS trace propagation

Gemfile

@@ -93,3 +93,5 @@ end
 #
 # TODO: Remove this once the issue is resolved: https://github.com/ffi/ffi/issues/1107
 gem 'ffi', '~> 1.16.3', require: false
+
+gem 'aws-sdk'

lib/datadog/tracing/contrib.rb

@@ -3,6 +3,7 @@
 require_relative '../tracing'
 require_relative 'contrib/registry'
 require_relative 'contrib/extensions'
+require_relative 'contrib/component'
 
 module Datadog
   module Tracing
@@ -26,6 +27,40 @@ module Contrib
       # after the tracer has complete initialization. Use `Datadog::Tracing::Contrib::REGISTRY` instead
       # of `Datadog.registry` when you code might be called during tracer initialization.
       REGISTRY = Registry.new
+
+      def self.inject(digest, data)
+        raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+        @propagation.inject!(digest, data)
+      end
+
+      def self.extract(data)
+        raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+        @propagation.extract(data)
+      end
+
+      Contrib::Component.register('_contrib') do |config|
+        tracing = config.tracing
+        tracing.propagation_style # TODO: do we still need this?
+
+        @propagation = Datadog::Tracing::Distributed::Propagation.new(
+          propagation_styles: {
+            Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_B3_MULTI_HEADER =>
+              Tracing::Distributed::B3Multi.new(fetcher: Tracing::Distributed::Fetcher),
+            Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_B3_SINGLE_HEADER =>
+              Tracing::Distributed::B3Single.new(fetcher: Tracing::Distributed::Fetcher),
+            Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_DATADOG =>
+              Tracing::Distributed::Datadog.new(fetcher: Tracing::Distributed::Fetcher),
+            Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_TRACE_CONTEXT =>
+              Tracing::Distributed::TraceContext.new(fetcher: Tracing::Distributed::Fetcher),
+            Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_NONE => Tracing::Distributed::None.new
+          },
+          propagation_style_inject: tracing.propagation_style_inject,
+          propagation_style_extract: tracing.propagation_style_extract,
+          propagation_extract_first: tracing.propagation_extract_first
+        )
+      end
     end
   end
 end

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

@@ -45,7 +45,24 @@ class Settings < Contrib::Configuration::Settings
               o.type :string, nilable: true
               o.env Ext::ENV_PEER_SERVICE
             end
-          end
+
+            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.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
+            end
         end
       end
     end

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

@@ -19,32 +19,34 @@ def add_handlers(handlers, _)
         # Generates Spans for all interactions with AWS
         class Handler < Seahorse::Client::Handler
           def call(context)
-            Tracing.trace(Ext::SPAN_COMMAND) do |span|
+            Tracing.trace(Ext::SPAN_COMMAND) do |span, trace|
               @handler.call(context).tap do
-                annotate!(span, ParsedContext.new(context))
+                annotate!(span, trace, ParsedContext.new(context))
               end
             end
           end
 
           private
 
           # rubocop:disable Metrics/AbcSize
-          def annotate!(span, context)
-            span.service = configuration[:service_name]
+          def annotate!(span, trace, context)
+            config = configuration
+            span.service = config[:service_name]
             span.type = Tracing::Metadata::Ext::HTTP::TYPE_OUTBOUND
             span.name = Ext::SPAN_COMMAND
             span.resource = context.safely(:resource)
             aws_service = span.resource.split('.')[0]
             span.set_tag(Ext::TAG_AWS_SERVICE, aws_service)
             params = context.safely(:params)
             if (handler = Datadog::Tracing::Contrib::Aws::SERVICE_HANDLERS[aws_service])
+              handler.process(config, trace, context)
               handler.add_tags(span, params)
             end
 
-            if configuration[:peer_service]
+            if config[:peer_service]
               span.set_tag(
                 Tracing::Metadata::Ext::TAG_PEER_SERVICE,
-                configuration[:peer_service]
+                config[:peer_service]
               )
             end
 
@@ -61,8 +63,8 @@ def annotate!(span, context)
             span.set_tag(Tracing::Metadata::Ext::TAG_PEER_HOSTNAME, context.safely(:host))
 
             # Set analytics sample rate
-            if Contrib::Analytics.enabled?(configuration[:analytics_enabled])
-              Contrib::Analytics.set_sample_rate(span, configuration[:analytics_sample_rate])
+            if Contrib::Analytics.enabled?(config[:analytics_enabled])
+              Contrib::Analytics.set_sample_rate(span, config[:analytics_sample_rate])
             end
             Contrib::Analytics.set_measured(span)
 

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

@@ -7,7 +7,28 @@ module Aws
         module Service
           # Base class for all AWS service-specific tag handlers.
           class Base
+            def process(config, trace, context); end
             def add_tags(span, params); end
+
+            MESSAGE_ATTRIBUTES_LIMIT = 10 # Can't set more than 10 message attributes
+
+            def extract_propagation(context)
+              message_attributes = context.params[:message_attributes]
+
+              return unless message_attributes && (datadog = message_attributes['_datadog'])
+
+              Tracing.continue_trace!(Contrib.extract(datadog))
+            end
+
+            def inject_propagation(trace, context, data_type)
+              message_attributes = (context.params[:message_attributes] ||= {})
+              return if message_attributes.size >= MESSAGE_ATTRIBUTES_LIMIT
+
+              data = {}
+              if Datadog::Tracing::Contrib.inject(trace.to_digest, data)
+                message_attributes['_datadog'] = { :data_type => data_type, :binary_value => data.to_json }
+              end
+            end
           end
         end
       end

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

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'json'
+
 require_relative './base'
 require_relative '../ext'
 
@@ -10,6 +12,23 @@ module Aws
         module Service
           # SNS tag handlers.
           class SNS < Base
+            def process(config, trace, context)
+              inject_propagation(trace, context, 'Binary') if config[:propagation]
+
+              return unless config[:propagation]
+
+              case context.operation
+              when :publish
+                inject_propagation(trace, context, 'Binary')
+              when :publish_batch
+                if config[:batch_propagation]
+                  inject_propagation(trace, context, 'Binary')
+                else
+                  inject_propagation(trace, context, 'Binary')
+                end
+              end
+            end
+
             def add_tags(span, params)
               topic_arn = params[:topic_arn]
               topic_name = params[:name]

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

@@ -10,6 +10,27 @@ module Aws
         module Service
           # SQS tag handlers.
           class SQS < Base
+            def before_span(config, context)
+              if config[:propagation] && context.operation == :receive_message
+                extract_propagation(context)
+              end
+            end
+
+            def process(config, trace, context)
+              return unless config[:propagation]
+
+              case context.operation
+              when :send_message
+                inject_propagation(trace, context, 'String')
+              when :send_message_batch
+                if config[:batch_propagation]
+                  inject_propagation(trace, context, 'String')
+                else
+                  inject_propagation(trace, context, 'String')
+                end
+              end
+            end
+
             def add_tags(span, params)
               queue_url = params[:queue_url]
               queue_name = params[:queue_name]

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

@@ -379,6 +379,37 @@
           .to eq('sqs.us-stubbed-1.amazonaws.com')
       end
     end
+
+    describe '#receive_message' do
+      subject!(:receive_message) do
+        client.receive_message(
+          {
+            queue_url: 'https://sqs.us-stubbed-1.amazonaws.com/123456789012/MyQueueName',
+            attribute_names: ["All"],
+            max_number_of_messages: 1,
+            visibility_timeout: 1,
+            wait_time_seconds: 1,
+            receive_request_attempt_id: "my_receive_request_attempt_1",
+          }
+        )
+      end
+
+      let(:responses) do
+        { receive_message: {
+          messages: []
+        } }
+      end
+
+      it 'generates a span' do
+        expect(span.name).to eq('aws.command')
+        expect(span.service).to eq('aws')
+        expect(span.type).to eq('http')
+        expect(span.resource).to eq('sqs.receive_message')
+
+        expect(span.get_tag('aws.agent')).to eq('aws-sdk-ruby')
+        expect(span.get_tag('aws.operation')).to eq('receive_message')
+      end
+    end
   end
 
   context 'with an SNS client' do
@@ -389,7 +420,14 @@
         client.publish(
           {
             topic_arn: 'arn:aws:sns:us-west-2:123456789012:my-topic-name',
-            message: 'Hello, world!'
+            message: 'Hello, world!',
+            message_attributes: {
+              'String' => {
+                data_type: 'String', # required
+                string_value: 'String',
+                binary_value: 'data',
+              },
+            },
           }
         )
       end

spec/datadog/tracing/contrib/aws/service/sns_spec.rb

@@ -37,4 +37,66 @@
       expect(span).to have_received(:set_tag).with(Datadog::Tracing::Contrib::Aws::Ext::TAG_TOPIC_NAME, nil)
     end
   end
+
+  shared_examples 'injects attribute propagation' do
+    subject(:inject_propagation) { service.process(config, trace, context) }
+
+    let(:config) { { propagation: true } }
+    let(:trace) { Datadog::Tracing::TraceOperation.new(id: trace_id, parent_span_id: span_id) }
+    let(:context) { instance_double('Context', params: params) }
+    let(:params) { {} }
+    let(:trace_id) { 1 }
+    let(:span_id) { 2 }
+
+    before { Datadog.configure { |c| c.tracing.instrument :aws } }
+
+    context 'without preexisting message attributes' do
+      it 'adds a propagation attribute' do
+        inject_propagation
+        expect(params[:message_attributes]).to eq(
+          '_datadog' => {
+            binary_value:
+              '{"x-datadog-trace-id":"1","x-datadog-parent-id":"2",' \
+              '"traceparent":"00-00000000000000000000000000000001-0000000000000002-00",' \
+              '"tracestate":"dd=p:0000000000000002"}',
+            data_type: 'Binary'
+          }
+        )
+      end
+    end
+
+    context 'with existing message attributes' do
+      let(:params) { { message_attributes: message_attributes } }
+      let(:message_attributes) { { 'existing' => { data_type: 'String', string_value: 'value' } } }
+
+      it 'adds a propagation attribute' do
+        expect { inject_propagation }.to change { message_attributes.keys }.from(['existing']).to(['existing', '_datadog'])
+      end
+    end
+
+    context 'with 10 message attributes already set' do
+      let(:params) { { message_attributes: message_attributes } }
+      let(:message_attributes) do
+        Array.new(10) do |i|
+          ["attr#{i}", { data_type: 'Number', string_value: i }]
+        end.to_h
+      end
+
+      it 'does not add a propagation attribute' do
+        expect { inject_propagation }.to_not(change { params })
+      end
+    end
+
+    context 'disabled' do
+      let(:config) { { propagation: false } }
+
+      it 'does not add a propagation attribute' do
+        expect { inject_propagation }.to_not(change { params })
+      end
+    end
+  end
+
+  it_behaves_like 'injects attribute propagation' do
+    let(:service) { sns }
+  end
 end