wip

Gemfile

@@ -101,4 +101,4 @@ 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'
+# gem 'aws-sdk'

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

@@ -18,24 +18,45 @@ def add_handlers(handlers, _)
 
  # Generates Spans for all interactions with AWS
  class Handler < Seahorse::Client::Handler
+ # Some services contain trace propagation information (e.g. SQS) that affect what active trace
+ # we'll use for the AWS span.
+ # But because this information is only available after the request is made, we need to make the AWS
+ # request first, then create the trace and span with correct distributed trace parenting.
  def call(context)
- Tracing.trace(Ext::SPAN_COMMAND) do |span, trace|
- @handler.call(context).tap do
- annotate!(span, trace, ParsedContext.new(context))
- end
+ config = configuration
+
+ # Find the AWS service instrumentation
+ parsed_context = ParsedContext.new(context)
+ aws_service = parsed_context.safely(:resource).split('.')[0]
+ handler = Datadog::Tracing::Contrib::Aws::SERVICE_HANDLERS[aws_service]
+
+ # Execute handler stack, to ensure we have the response object before the trace and span are created
+ start_time = Core::Utils::Time.now.utc # Save the start time as the span creation is delayed
+ begin
+ response = @handler.call(context)
+ rescue Exception => e
+ # Catch exception to reraise it inside the trace block, to ensure the span has correct error information
+ end
+
+ Tracing.trace(Ext::SPAN_COMMAND, start_time: start_time) do |span, trace|
+ handler.before_span(config, context, response) if handler
+
+ annotate!(config, span, trace, parsed_context, aws_service)
+
+ raise e if e
  end
+
+ response
  end
 
  private
 
  # rubocop:disable Metrics/AbcSize
- def annotate!(span, trace, context)
- config = configuration
+ def annotate!(config, span, trace, context, aws_service)
  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])

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

@@ -7,17 +7,25 @@ module Aws
  module Service
  # Base class for all AWS service-specific tag handlers.
  class Base
+ def before_span(config, context, response); end
  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]
+ # Extract the `_datadog` message attribute and decode its JSON content.
+ def extract_propagation!(response, data_type)
+ messages = response.data.messages
+
+ return unless (message = messages[0])
+
+ message_attributes = message.message_attributes
 
  return unless message_attributes && (datadog = message_attributes['_datadog'])
 
- Tracing.continue_trace!(Contrib.extract(datadog))
+ if ((data = datadog[data_type]) && (parsed_data = JSON.parse(data)))
+ Tracing.continue_trace!(Contrib.extract(parsed_data))
+ end
  end
 
  def inject_propagation(trace, params, data_type)

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

@@ -12,16 +12,15 @@ module Aws
  module Service
  # SNS tag handlers.
  class SNS < Base
+ PROPAGATION_DATATYPE = 'Binary'
+
  def process(config, trace, context)
  return unless config[:propagation]
 
  case context.operation
  when :publish
- inject_propagation(trace, context.params, 'Binary')
- when :publish_batch
- context.params[:publish_batch_request_entries].each do |entry|
- inject_propagation(trace, entry, 'Binary')
- end
+ inject_propagation(trace, context.params, PROPAGATION_DATATYPE)
+ # TODO: when :publish_batch # Future support for batch publishing
  end
  end
 

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

@@ -10,28 +10,21 @@ module Aws
  module Service
  # SQS tag handlers.
  class SQS < Base
- def before_span(config, 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
+ DATATYPE = 'String'
+ def before_span(config, context, response)
+ return unless context.operation == :receive_message && config[:propagation]
+
+ # Parent the current trace based on distributed message attributes
+ extract_propagation!(response, 'string_value') if config[:parentage_style] == 'distributed'
  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
+ inject_propagation(trace, context.params, 'String')
+ # TODO when :send_message_batch # Future support for batch sending
  end
  end
 

lib/datadog/tracing/trace_digest.rb

@@ -180,6 +180,58 @@ def merge(field_value_pairs)
  }.merge!(field_value_pairs)
  )
  end
+
+ def ==(other)
+ self.class == other.class &&
+ span_id == other.span_id &&
+ span_name == other.span_name &&
+ span_resource == other.span_resource &&
+ span_service == other.span_service &&
+ span_type == other.span_type &&
+ trace_distributed_tags == other.trace_distributed_tags &&
+ trace_hostname == other.trace_hostname &&
+ trace_id == other.trace_id &&
+ trace_name == other.trace_name &&
+ trace_origin == other.trace_origin &&
+ trace_process_id == other.trace_process_id &&
+ trace_resource == other.trace_resource &&
+ trace_runtime_id == other.trace_runtime_id &&
+ trace_sampling_priority == other.trace_sampling_priority &&
+ trace_service == other.trace_service &&
+ trace_distributed_id == other.trace_distributed_id &&
+ trace_flags == other.trace_flags &&
+ trace_state == other.trace_state &&
+ trace_state_unknown_fields == other.trace_state_unknown_fields &&
+ span_remote == other.span_remote
+ end
+
+ alias eql? ==
+
+ def hash
+ [
+ self.class,
+ span_id,
+ span_name,
+ span_resource,
+ span_service,
+ span_type,
+ trace_distributed_tags,
+ trace_hostname,
+ trace_id,
+ trace_name,
+ trace_origin,
+ trace_process_id,
+ trace_resource,
+ trace_runtime_id,
+ trace_sampling_priority,
+ trace_service,
+ trace_distributed_id,
+ trace_flags,
+ trace_state,
+ trace_state_unknown_fields,
+ span_remote
+ ].hash
+ end
  end
  end
 end

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

@@ -397,8 +397,16 @@
 
  let(:responses) do
  { receive_message: {
- messages: []
- } }
+ messages: [
+ message_attributes: {
+ '_datadog' => {
+ string_value: 'String',
+ data_type: 'String'
+ }
+ }
+ ]
+ }
+ }
  end
 
  it 'generates a span' do

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

@@ -0,0 +1,165 @@
+require 'datadog/tracing/contrib/aws/parsed_context'
+
+require 'aws-sdk-sqs'
+
+RSpec.shared_examples 'injects AWS 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(Datadog::Tracing::Contrib::Aws::ParsedContext, params: params, operation: operation) }
+ 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: data_type
+ }
+ )
+ end
+ end
+
+ context 'with existing message attributes' do
+ let(:params) { { message_attributes: message_attributes } }
+ let(:message_attributes) { { 'existing' => { data_type: 'Number', string_value: 1 } } }
+
+ 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
+
+RSpec.shared_examples 'extract AWS attribute propagation' do
+ subject(:extract_propagation) { service.before_span(config, context, response) }
+
+ let(:config) { { propagation: true, parentage_style: parentage_style } }
+ let(:parentage_style) { 'distributed' }
+ let(:trace) do
+ Datadog::Tracing::TraceOperation.new(
+ id: 1,
+ parent_span_id: 2,
+ remote_parent: true,
+ trace_state: 'unrelated=state',
+ sampling_priority: 0
+ )
+ end
+ let(:context) { instance_double(Datadog::Tracing::Contrib::Aws::ParsedContext, operation: operation) }
+ let(:response) do
+ result = Aws::SQS::Types::ReceiveMessageResult.new(messages: messages)
+ Seahorse::Client::Response.new(data: result)
+ end
+ let(:messages) { [] }
+
+ before { Datadog.configure { |c| c.tracing.instrument :aws } }
+
+ context 'without message attributes' do
+ context 'without an active trace' do
+ it 'does not create trace' do
+ extract_propagation
+ expect(Datadog::Tracing.active_trace).to be_nil
+ end
+ end
+
+ context 'with an active trace' do
+ before { Datadog::Tracing.continue_trace!(trace.to_digest) }
+
+ it 'does not change active trace' do
+ extract_propagation
+ expect(Datadog::Tracing.active_trace.to_digest).to eq(trace.to_digest)
+ end
+ end
+ end
+
+ context 'with message attributes' do
+ let(:messages) { [message] }
+ let(:message) { Aws::SQS::Types::Message.new(message_attributes: message_attributes) }
+ let(:message_attributes) { { '_datadog' => attribute } }
+ let(:attribute) { Aws::SQS::Types::MessageAttributeValue.new(
+ string_value:
+ '{"traceparent":"00-00000000000000000000000000000001-0000000000000002-00",' \
+ '"tracestate":"dd=p:0000000000000002,unrelated=state"}',
+ data_type: data_type
+ ) }
+
+ context 'without an active trace' do
+ it 'creates trace' do
+ extract_propagation
+ expect(Datadog::Tracing.active_trace.to_digest).to eq(trace.to_digest)
+ end
+ end
+
+ context 'with an active trace' do
+ it 'overrides the existing trace' do
+ existing_trace = Datadog::Tracing.continue_trace!(nil)
+ expect { extract_propagation }.to(
+ change { Datadog::Tracing.active_trace.to_digest }.from(existing_trace.to_digest).to(trace.to_digest)
+ )
+ end
+ end
+
+ context 'with a local parentage style' do
+ let(:parentage_style) { 'local' }
+
+ it 'does not create a remote trace' do
+ extract_propagation
+ expect(Datadog::Tracing.active_trace).to be_nil
+ end
+ end
+
+ context 'with multiple messages' do
+ let(:messages) { [message, other_message] }
+ let(:other_message) { Aws::SQS::Types::Message.new(message_attributes: other_message_attributes) }
+ let(:other_message_attributes) { { '_datadog' => other_attribute } }
+ let(:other_attribute) { Aws::SQS::Types::MessageAttributeValue.new(
+ string_value:
+ '{"traceparent":"00-00000000000000000000000000000008-0000000000000009-00",' \
+ '"tracestate":"dd=p:0000000000000009,oops=not-this-one"}',
+ data_type: data_type
+ ) }
+
+ it 'extracts the first message attributes' do
+ extract_propagation
+ expect(Datadog::Tracing.active_trace.to_digest).to eq(trace.to_digest)
+ end
+ end
+ end
+
+ context 'disabled' do
+ let(:config) { { propagation: false } }
+
+ it 'does not add a propagation attribute' do
+ expect { extract_propagation }.to_not(change { params })
+ end
+ end
+end