feat(tracing): propagate W3C trace context across AWS messaging

lib/instana/instrumentation/active_job.rb

@@ -70,25 +70,10 @@ def self.extract_context(job)
       end
 
       def self.parse_w3c_traceparent(traceparent, tracestate)
-        m = traceparent.match(/[0-9a-f]{2}-(?<trace>[0-9a-f]{32})-(?<parent>[0-9a-f]{16})-(?<flags>[0-9a-f]{2})/)
-        return nil unless m
+        ctx = ::Instana::Util.parse_w3c_trace_context(traceparent, tracestate)
+        return nil unless ctx
 
-        w3c_sampled = (m['flags'].to_i(16) & 0x01) == 1
-        ts_match = tracestate&.match(/in=(?<t>[0-9a-f]+);(?<s>[0-9a-f]+)/)
-
-        if ts_match
-          ::Instana::SpanContext.new(
-            trace_id: ::Instana::Util.header_to_id(ts_match['t']),
-            span_id: ::Instana::Util.header_to_id(ts_match['s']),
-            level: w3c_sampled ? 1 : 0
-          )
-        else
-          ::Instana::SpanContext.new(
-            trace_id: ::Instana::Util.header_to_id(m['trace'][16..]),
-            span_id: ::Instana::Util.header_to_id(m['parent']),
-            level: w3c_sampled ? 1 : 0
-          )
-        end
+        ::Instana::SpanContext.new(trace_id: ctx[:trace_id], span_id: ctx[:span_id], level: ctx[:level])
       end
     end
   end

lib/instana/instrumentation/aws_sdk_lambda.rb

@@ -10,18 +10,23 @@ def call(context)
           return @handler.call(context) unless [:invoke_async, :invoke].include?(context.operation_name)
 
           if ::Instana.tracer.tracing? && context.operation_name == :invoke
-            span_context  = ::Instana.tracer.context
-            trace_payload = {
-              'X-INSTANA-T' => span_context.trace_id,
-              'X-INSTANA-S' => span_context.span_id,
-              'X-INSTANA-L' => span_context.level.to_s
-            }
-            merged = if context.params[:client_context]
-                       JSON.parse(Base64.strict_decode64(context.params[:client_context])).merge(trace_payload)
-                     else
-                       trace_payload
-                     end
-            context.params[:client_context] = Base64.strict_encode64(JSON.dump(merged))
+            span_context = ::Instana.tracer.context
+            if span_context
+              trace_payload = {
+                'X-INSTANA-T' => span_context.trace_id,
+                'X-INSTANA-S' => span_context.span_id,
+                'X-INSTANA-L' => span_context.level.to_s,
+                'traceparent' => span_context.trace_parent_header,
+                'tracestate' => span_context.trace_state_header
+              }.reject { |_, v| v.nil? || v.empty? }
+
+              merged = if context.params[:client_context]
+                         JSON.parse(Base64.strict_decode64(context.params[:client_context])).merge(trace_payload)
+                       else
+                         trace_payload
+                       end
+              context.params[:client_context] = Base64.strict_encode64(JSON.dump(merged))
+            end
           end
 
           tags = {

lib/instana/instrumentation/aws_sdk_sns.rb

@@ -14,22 +14,41 @@ def call(context)
           }.reject { |_, v| v.nil? }
 
           if context.operation_name == :publish
-            ::Instana.tracer.in_span(:sns, attributes: {sns: sns_tags}) do |_span|
+            ::Instana.tracer.in_span(:sns, attributes: { sns: sns_tags }) do |_span|
               if ::Instana.tracer.tracing?
-                span_context = ::Instana.tracer.context
-                context.params[:message_attributes] ||= {}
-                context.params[:message_attributes].merge!(
-                  'X_INSTANA_T' => { data_type: 'String', string_value: span_context.trace_id_header },
-                  'X_INSTANA_S' => { data_type: 'String', string_value: span_context.span_id_header },
-                  'X_INSTANA_L' => { data_type: 'String', string_value: span_context.level.to_s }
-                )
+                inject_trace_attributes(context.params)
+              end
+              @handler.call(context)
+            end
+          elsif context.operation_name == :publish_batch
+            ::Instana.tracer.in_span(:sns, attributes: { sns: sns_tags }) do |_span|
+              if ::Instana.tracer.tracing?
+                context.params[:publish_batch_request_entries]&.each do |entry|
+                  inject_trace_attributes(entry)
+                end
               end
               @handler.call(context)
             end
           else
             @handler.call(context)
           end
         end
+
+        private
+
+        def inject_trace_attributes(params)
+          span_context = ::Instana.tracer.context
+          return unless span_context
+
+          params[:message_attributes] ||= {}
+          available_slots = 10 - params[:message_attributes].size
+
+          extra = ::Instana::Util.aws_message_attributes_w3c(span_context)
+          extra.merge!(::Instana::Util.aws_message_attributes_instana(span_context)) if available_slots >= 5
+
+          extra.reject! { |_, v| v[:string_value].nil? || v[:string_value].empty? }
+          params[:message_attributes].merge!(extra)
+        end
       end
 
       def add_handlers(handlers, _config)

lib/instana/instrumentation/aws_sdk_sqs.rb

@@ -36,11 +36,16 @@ def call(context)
 
         def inject_instana_headers(span, params)
           params[:message_attributes] ||= {}
-          params[:message_attributes].merge!({
-                                               "X_INSTANA_T" => {data_type: 'String', string_value: span.context.trace_id},
-                                               "X_INSTANA_S" => {data_type: 'String', string_value: span.context.span_id},
-                                               "X_INSTANA_L" => {data_type: 'String', string_value: span.context.level.to_s}
-                                             })
+          span_context = span.context
+          return unless span_context
+
+          available_slots = 10 - params[:message_attributes].size
+
+          extra = ::Instana::Util.aws_message_attributes_w3c(span_context)
+          extra.merge!(::Instana::Util.aws_message_attributes_instana(span_context)) if available_slots >= 5
+
+          extra.reject! { |_, v| v[:string_value].nil? || v[:string_value].empty? }
+          params[:message_attributes].merge!(extra)
         end
 
         def tags_for(operation_name, params)

lib/instana/instrumentation/bunny.rb

@@ -134,21 +134,7 @@ def extract_instana_context(headers)
       end
 
       def extract_w3c_context(headers)
-        m = headers['traceparent'].match(
-          /[0-9a-f][0-9a-e]-(?<trace>[0-9a-f]{32})-(?<parent>[0-9a-f]{16})-(?<flags>[0-9a-f]{2})/
-        )
-        return {} unless m
-
-        level = (m['flags'].to_i(16) & 0x01) == 1 ? 1 : 0
-        ts_match = headers['tracestate']&.match(/in=(?<t>[0-9a-f]+);(?<s>[0-9a-f]+)/)
-
-        if ts_match
-          { trace_id: ts_match['t'], span_id: ts_match['s'], level: level }
-        else
-          { trace_id: ::Instana::Util.header_to_id(m['trace'][16..]),
-            span_id: ::Instana::Util.header_to_id(m['parent']),
-            level: level }
-        end
+        ::Instana::Util.parse_w3c_trace_context(headers['traceparent'], headers['tracestate']) || {}
       end
 
       def log_error(error)

lib/instana/instrumentation/grpc.rb

@@ -107,28 +107,13 @@ def extract_instana_grpc_context(metadata)
       end
 
       def extract_w3c_grpc_context(metadata)
-        m = metadata['traceparent'].match(
-          /[0-9a-f][0-9a-e]-(?<trace>[0-9a-f]{32})-(?<parent>[0-9a-f]{16})-(?<flags>[0-9a-f]{2})/
-        )
-        return nil unless m
-
-        level    = (m['flags'].to_i(16) & 0x01) == 1 ? 1 : 0
-        ts_match = metadata['tracestate']&.match(/in=(?<t>[0-9a-f]+);(?<s>[0-9a-f]+)/)
-        baggage  = {
-          external_trace_id: m['trace'],
-          external_trace_flags: m['flags'],
-          external_state: metadata['tracestate']
-        }
-
-        if ts_match
-          SpanContext.new(trace_id: ::Instana::Util.header_to_id(ts_match['t']),
-                          span_id: ::Instana::Util.header_to_id(ts_match['s']),
-                          level: level, baggage: baggage)
-        else
-          SpanContext.new(trace_id: ::Instana::Util.header_to_id(m['trace'][16..]),
-                          span_id: ::Instana::Util.header_to_id(m['parent']),
-                          level: level, baggage: baggage)
-        end
+        ctx = ::Instana::Util.parse_w3c_trace_context(metadata['traceparent'], metadata['tracestate'])
+        return nil unless ctx
+
+        SpanContext.new(trace_id: ctx[:trace_id],
+                        span_id: ctx[:span_id],
+                        level: ctx[:level],
+                        baggage: ctx.slice(:external_trace_id, :external_trace_flags, :external_state))
       end
     end
   end

lib/instana/instrumentation/shoryuken.rb

@@ -15,24 +15,45 @@ def call(_worker_instance, _queue, sqs_message, _body, &block)
         }
 
         context = incomming_context_from(sqs_message.message_attributes)
-        instana_context = Instana::SpanContext.new(trace_id: context[:trace_id], span_id: context[:span_id], level: context[:level])
-        Trace.with_span(OpenTelemetry::Trace.non_recording_span(instana_context)) do
-          ::Instana.tracer.in_span(:sqs, attributes: {sqs: sqs_tags}, &block)
+        if context[:trace_id]
+          instana_context = Instana::SpanContext.new(
+            trace_id: context[:trace_id],
+            span_id: context[:span_id],
+            level: context[:level] || 1
+          )
+          Trace.with_span(OpenTelemetry::Trace.non_recording_span(instana_context)) do
+            ::Instana.tracer.in_span(:sqs, attributes: { sqs: sqs_tags }, &block)
+          end
+        else
+          ::Instana.tracer.in_span(:sqs, attributes: { sqs: sqs_tags }, &block)
         end
       end
 
       private
 
       def incomming_context_from(attributes)
-        trace_id = read_message_header(attributes, 'X_INSTANA_T')
-        span_id = read_message_header(attributes, 'X_INSTANA_S')
-        level = read_message_header(attributes, 'X_INSTANA_L')
-
-        {
-          trace_id: trace_id,
-          span_id: span_id,
-          level: level
-        }.reject { |_, v| v.nil? }
+        return {} unless attributes
+
+        if read_message_header(attributes, 'X_INSTANA_T')
+          trace_id = read_message_header(attributes, 'X_INSTANA_T')
+          span_id = read_message_header(attributes, 'X_INSTANA_S')
+          level = read_message_header(attributes, 'X_INSTANA_L')
+
+          {
+            trace_id: trace_id,
+            span_id: span_id,
+            level: level ? level.to_i : 1
+          }.reject { |_, v| v.nil? }
+        else
+          w3c_context_from(attributes)
+        end
+      end
+
+      def w3c_context_from(attributes)
+        ::Instana::Util.parse_w3c_trace_context(
+          read_message_header(attributes, 'traceparent'),
+          read_message_header(attributes, 'tracestate')
+        ) || {}
       end
 
       def read_message_header(attributes, key)

lib/instana/util.rb

@@ -2,7 +2,7 @@
 # (c) Copyright Instana Inc. 2016
 
 module Instana
-  module Util
+  module Util # rubocop:disable Metrics/ModuleLength
     class << self
       ID_RANGE = -2**63..2**63-1
       # Retrieves and returns the source code for any ruby
@@ -171,6 +171,44 @@ def header_to_id(given)
         given
       end
 
+      # Parse W3C +traceparent+ / +tracestate+ values using the OpenTelemetry
+      # trace-context propagator.
+      #
+      # When the tracestate carries an Instana entry (+in=<trace>;<span>+) it
+      # takes precedence over the ids in the traceparent.
+      #
+      # @param traceparent [String] the traceparent header value
+      # @param tracestate [String, nil] the tracestate header value
+      #
+      # @return [Hash, nil] +nil+ when the traceparent is missing or invalid,
+      #   otherwise a hash with:
+      #   - +:trace_id+, +:span_id+, +:level+ for continuing the trace
+      #   - +:external_trace_id+, +:external_trace_flags+, +:external_state+
+      #     for re-emitting the W3C context downstream
+      #
+      def parse_w3c_trace_context(traceparent, tracestate = nil)
+        return nil unless traceparent.is_a?(String)
+
+        carrier = { 'traceparent' => traceparent, 'tracestate' => tracestate.is_a?(String) ? tracestate : nil }
+        extracted = OpenTelemetry::Trace::Propagation::TraceContext.text_map_propagator.extract(
+          carrier, context: OpenTelemetry::Context.empty
+        )
+        span_context = OpenTelemetry::Trace.current_span(extracted).context
+        return nil unless span_context.valid?
+
+        sampled = span_context.trace_flags.sampled?
+        instana_state = span_context.tracestate.value('in')&.match(/\A(?<trace>[0-9a-f]+);(?<span>[0-9a-f]+)\z/)
+
+        {
+          trace_id: header_to_id(instana_state ? instana_state[:trace] : span_context.hex_trace_id[16..]),
+          span_id: header_to_id(instana_state ? instana_state[:span] : span_context.hex_span_id),
+          level: sampled ? 1 : 0,
+          external_trace_id: span_context.hex_trace_id,
+          external_trace_flags: sampled ? '01' : '00',
+          external_state: tracestate
+        }
+      end
+
       def timeout_timestamp
         Process.clock_gettime(Process::CLOCK_MONOTONIC)
       end
@@ -229,6 +267,35 @@ def http_reason_phrase(status_code)
 
         klass.name.sub('Net::HTTP', '').gsub(/([A-Z])/, ' \1').strip
       end
+
+      # Build the W3C trace-context message attributes for an AWS message
+      # (SQS or SNS).  Returns a hash suitable for merging into
+      # +params[:message_attributes]+.
+      #
+      # @param span_context [Instana::SpanContext]
+      # @return [Hash]
+      #
+      def aws_message_attributes_w3c(span_context)
+        {
+          'traceparent' => { data_type: 'String', string_value: span_context.trace_parent_header },
+          'tracestate' => { data_type: 'String', string_value: span_context.trace_state_header }
+        }
+      end
+
+      # Build the Instana proprietary message attributes for an AWS message
+      # (SQS or SNS).  Returns a hash suitable for merging into
+      # +params[:message_attributes]+.
+      #
+      # @param span_context [Instana::SpanContext]
+      # @return [Hash]
+      #
+      def aws_message_attributes_instana(span_context)
+        {
+          'X_INSTANA_T' => { data_type: 'String', string_value: span_context.trace_id },
+          'X_INSTANA_S' => { data_type: 'String', string_value: span_context.span_id },
+          'X_INSTANA_L' => { data_type: 'String', string_value: span_context.level.to_s }
+        }
+      end
     end
   end
 end