wip

Gemfile

@@ -100,5 +100,3 @@ 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.rb

@@ -2,6 +2,7 @@
 
 require_relative 'core'
 require_relative 'tracing/pipeline'
+require_relative 'tracing/distributed'
 
 module Datadog
   # Datadog APM tracing public API.

lib/datadog/tracing/contrib.rb

@@ -3,7 +3,6 @@
 require_relative '../tracing'
 require_relative 'contrib/registry'
 require_relative 'contrib/extensions'
-require_relative 'contrib/component'
 
 module Datadog
   module Tracing
@@ -27,40 +26,6 @@ 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/instrumentation.rb

@@ -34,8 +34,9 @@ def call(context)
             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
+            rescue Exception => e # rubocop:disable Lint/RescueException
               # Catch exception to reraise it inside the trace block, to ensure the span has correct error information
+              # This matches the behavior of {Datadog::Tracing::SpanOperation#measure}
             end
 
             Tracing.trace(Ext::SPAN_COMMAND, start_time: start_time) do |span, trace|

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

@@ -24,7 +24,7 @@ def extract_propagation!(response, data_type)
               return unless message_attributes && (datadog = message_attributes['_datadog'])
 
               if (data = datadog[data_type]) && (parsed_data = JSON.parse(data))
-                Tracing.continue_trace!(Contrib.extract(parsed_data))
+                Tracing.continue_trace!(Distributed.extract(parsed_data))
               end
             end
 
@@ -33,7 +33,7 @@ def inject_propagation(trace, params, data_type)
               return if message_attributes.size >= MESSAGE_ATTRIBUTES_LIMIT
 
               data = {}
-              if Datadog::Tracing::Contrib.inject(trace.to_digest, data)
+              if Distributed.inject(trace.to_digest, data)
                 message_attributes['_datadog'] = { :data_type => data_type, :binary_value => data.to_json }
               end
             end

lib/datadog/tracing/distributed.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require_relative 'distributed/b3_multi'
+require_relative 'distributed/b3_single'
+require_relative 'distributed/datadog'
+require_relative 'distributed/none'
+require_relative 'distributed/propagation'
+require_relative 'distributed/trace_context'
+require_relative 'contrib/component'
+
+module Datadog
+  module Tracing
+    # Namespace for distributed tracing propagation and correlation
+    module Distributed
+      module_function
+
+      # Inject distributed headers into the given request
+      # @param digest [Datadog::Tracing::TraceDigest] the trace to inject
+      # @param data [Hash] the request to inject
+      def inject(digest, data)
+        raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+        @propagation.inject!(digest, data)
+      end
+
+      # Extract distributed headers from the given request
+      # @param data [Hash] the request to extract from
+      # @return [Datadog::Tracing::TraceDigest,nil] the extracted trace digest or nil if none was found
+      def extract(data)
+        raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+        @propagation.extract(data)
+      end
+
+      Contrib::Component.register('distributed') do |config|
+        tracing = config.tracing
+        # DEV: evaluate propagation_style in case it overrides propagation_style_extract & propagation_extract_first
+        tracing.propagation_style
+
+        @propagation = Propagation.new(
+          propagation_styles: {
+            Configuration::Ext::Distributed::PROPAGATION_STYLE_B3_MULTI_HEADER =>
+              B3Multi.new(fetcher: Fetcher),
+            Configuration::Ext::Distributed::PROPAGATION_STYLE_B3_SINGLE_HEADER =>
+              B3Single.new(fetcher: Fetcher),
+            Configuration::Ext::Distributed::PROPAGATION_STYLE_DATADOG =>
+              Datadog.new(fetcher: Fetcher),
+            Configuration::Ext::Distributed::PROPAGATION_STYLE_TRACE_CONTEXT =>
+              TraceContext.new(fetcher: Fetcher),
+            Configuration::Ext::Distributed::PROPAGATION_STYLE_NONE => 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/trace_digest.rb

@@ -181,6 +181,7 @@ def merge(field_value_pairs)
         )
       end
 
+      # rubocop:disable Metrics/AbcSize,Metrics/PerceivedComplexity,Metrics/CyclomaticComplexity
       def ==(other)
         self.class == other.class &&
           span_id == other.span_id &&
@@ -204,6 +205,7 @@ def ==(other)
           trace_state_unknown_fields == other.trace_state_unknown_fields &&
           span_remote == other.span_remote
       end
+      # rubocop:enable Metrics/AbcSize,Metrics/PerceivedComplexity,Metrics/CyclomaticComplexity
 
       alias eql? ==
 

sig/datadog/tracing/distributed.rbs

@@ -0,0 +1,12 @@
+module Datadog
+  module Tracing
+    module Distributed
+      @propagation: Propagation
+
+      def inject: (TraceDigest digest, ::Hash[untyped, untyped] data) -> (bool | nil)
+
+      def extract: (::Hash[untyped, untyped] data) -> (TraceDigest | nil)
+    end
+  end
+end
+

spec/datadog/tracing/distributed_spec.rb

@@ -0,0 +1,33 @@
+require 'datadog/tracing/distributed'
+
+RSpec.describe Datadog::Tracing::Distributed do
+  context 'integration test' do
+    before { Datadog.configure {} }
+
+    describe '#inject' do
+      subject(:inject) { described_class.inject(digest, data) }
+      let(:trace_id) { Datadog::Tracing::Utils::TraceId.next_id }
+      let(:span_id) { Datadog::Tracing::Utils.next_id }
+      let(:digest) do
+        Datadog::Tracing::TraceDigest.new(trace_id: trace_id, span_id: span_id)
+      end
+      let(:data) { {} }
+
+      it 'injects distributed headers' do
+        inject
+        expect(data).to include('x-datadog-trace-id')
+        expect(data).to include('x-datadog-parent-id')
+      end
+    end
+
+    describe '#extract' do
+      subject(:extract) { described_class.extract(data) }
+
+      let(:data) { { 'x-datadog-trace-id' => '1', 'x-datadog-parent-id' => '2' } }
+
+      it 'extracts distributed headers' do
+        is_expected.to be_a_kind_of(Datadog::Tracing::TraceDigest)
+      end
+    end
+  end
+end