Sampling-Propagation:Export sampling mechanism in distributed trace tags

lib/datadog/tracing/distributed/headers/datadog.rb

@@ -43,13 +43,14 @@ def extract(env)
               # DEV: `Parser#id` will not return 0
               return unless (trace_id && parent_id) || (origin && trace_id)
 
-              trace_distributed_tags = extract_tags(headers)
+              trace_distributed_tags, sampling_mechanism = extract_tags(headers)
 
               # Return new trace headers
               TraceDigest.new(
                 span_id: parent_id,
                 trace_id: trace_id,
                 trace_origin: origin,
+                trace_sampling_mechanism: sampling_mechanism,
                 trace_sampling_priority: sampling_priority,
                 trace_distributed_tags: trace_distributed_tags,
               )
@@ -64,15 +65,23 @@ def extract(env)
             # DEV: Ideally, we'd have a dedicated error reporting stream for all of ddtrace.
             # DEV: The same comment applies to the {.extract_tags}.
             def inject_tags(digest, env)
-              return if digest.trace_distributed_tags.nil? || digest.trace_distributed_tags.empty?
+              if (digest.trace_distributed_tags.nil? || digest.trace_distributed_tags.empty?) &&
+                  digest.trace_sampling_mechanism.nil?
+                return
+              end
 
               if ::Datadog.configuration.tracing.x_datadog_tags_max_length <= 0
                 active_trace = Tracing.active_trace
                 active_trace.set_tag('_dd.propagation_error', 'disabled') if active_trace
                 return
               end
 
-              encoded_tags = DatadogTagsCodec.encode(digest.trace_distributed_tags)
+              tags = digest.trace_distributed_tags || {}
+              if digest.trace_sampling_mechanism
+                # Digest's tags are a frozen Hash, we have to create a copy here.
+                tags = tags.merge(Metadata::Ext::Distributed::TAG_DECISION_MAKER => "-#{digest.trace_sampling_mechanism}")
+              end
+              encoded_tags = DatadogTagsCodec.encode(tags)
 
               if encoded_tags.size > ::Datadog.configuration.tracing.x_datadog_tags_max_length
                 active_trace = Tracing.active_trace
@@ -122,14 +131,30 @@ def extract_tags(headers)
               tags = DatadogTagsCodec.decode(tags_header)
               # Only extract keys with the expected Datadog prefix
               tags.select! { |key, _| key.start_with?(Tracing::Metadata::Ext::Distributed::TAGS_PREFIX) }
-              tags
+
+              sampling_mechanism = extract_sampling_mechanism(tags[Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER])
+
+              [tags, sampling_mechanism]
             rescue => e
               active_trace = Tracing.active_trace
               active_trace.set_tag('_dd.propagation_error', 'decoding_error') if active_trace
               ::Datadog.logger.warn(
                 "Failed to extract x-datadog-tags: #{e.class.name} #{e.message} at #{Array(e.backtrace).first}"
               )
             end
+
+            # This tag is of the format `part1-sampling_mechanism`.
+            # `part1` is currently ignored.
+            # This method returns the second part, the `sampling_mechanism`
+            # which is always an Integer.
+            # If we can't find a valid `sampling_mechanism`, returns `nil`.
+            # @return [Integer, nil]
+            def extract_sampling_mechanism(decision_maker)
+              return unless decision_maker
+
+              _, sampling_mechanism = decision_maker.split('-')
+              Integer(sampling_mechanism) rescue nil
+            end
           end
         end
       end

spec/datadog/tracing/distributed/headers/datadog_spec.rb

@@ -164,6 +164,44 @@ def env_header(name)
             end
           end
         end
+
+        context 'with trace_sampling_mechanism' do
+          let(:digest) do
+            Datadog::Tracing::TraceDigest.new(trace_distributed_tags: tags, trace_sampling_mechanism: sampling_mechanism)
+          end
+
+          let(:sampling_mechanism) { 1 }
+
+          context 'conflicting with trace_distributed_tags' do
+            let(:tags) { { '_dd.p.dm' => '-0' } }
+
+            it 'overrides tag with explicit trace_sampling_mechanism set' do
+              is_expected.to include('x-datadog-tags' => '_dd.p.dm=-1')
+            end
+          end
+
+          context 'non-conflicting with trace_distributed_tags' do
+            let(:tags) { { key: 'value' } }
+
+            it 'merges both values' do
+              is_expected.to include('x-datadog-tags' => 'key=value,_dd.p.dm=-1')
+            end
+          end
+        end
+      end
+
+      context 'with trace_sampling_mechanism' do
+        let(:digest) { Datadog::Tracing::TraceDigest.new(trace_sampling_mechanism: sampling_mechanism) }
+
+        context 'nil' do
+          let(:sampling_mechanism) { nil }
+          it { is_expected.to_not include('x-datadog-tags') }
+        end
+
+        context 'a valid value' do
+          let(:sampling_mechanism) { 1 }
+          it { is_expected.to include('x-datadog-tags' => '_dd.p.dm=-1') }
+        end
       end
     end
   end
@@ -248,6 +286,18 @@ def env_header(name)
           it { is_expected.to eq('_dd.p.key' => 'value') }
         end
 
+        context '{ _dd.p.dm: "-1" }' do
+          let(:tags) { '_dd.p.dm=-1' }
+          it { is_expected.to eq('_dd.p.dm' => '-1') }
+          it { expect(extract.trace_sampling_mechanism).to eq(1) }
+        end
+
+        context '{ _dd.p.dm: "invalid value" }' do
+          let(:tags) { '_dd.p.dm=invalid value' }
+          it { is_expected.to eq('_dd.p.dm' => 'invalid value') }
+          it { expect(extract.trace_sampling_mechanism).to be_nil }
+        end
+
         context 'within an active trace' do
           before do
             allow(Datadog::Tracing).to receive(:active_trace).and_return(active_trace)