feat: delay sampling decision to allow sampling on tags and resource

lib/datadog/tracing/distributed/propagation.rb

@@ -44,6 +44,7 @@ def initialize(
         # DEV-2.0: if needed.
         # DEV-2.0: Ideally, we'd have a separate stream to report tracer errors and never
         # DEV-2.0: touch the active span.
+        # DEV-3.0: Sample trace here instead of when generating digest.
         #
         # @param digest [TraceDigest]
         # @param data [Hash]

lib/datadog/tracing/sampling/matcher.rb

@@ -44,6 +44,10 @@ def self.glob_to_regex(glob)
           Regexp.new(glob, Regexp::IGNORECASE)
         end
 
+        def match_all?
+          pattern == MATCH_ALL_PATTERN
+        end
+
         # Returns `true` for any input
         MATCH_ALL = Class.new do
           def match?(_other)
@@ -99,6 +103,10 @@ def match?(trace)
         def tags_match?(trace)
           @tags.all? do |name, matcher|
             tag = trace.get_tag(name)
+            # Floats: Matching floating point values with a non-zero decimal part is not supported.
+            # For floating point values with a non-zero decimal part, any all * pattern always returns true.
+            # Other patterns always return false.
+            return false if tag.is_a?(Float) && !tag.integer? && !matcher.match_all?
 
             # Format metrics as strings, to allow for partial number matching (/4.*/ matching '400', '404', etc.).
             # Because metrics are floats, we use the '%g' format specifier to avoid trailing zeros, which

lib/datadog/tracing/trace_operation.rb

@@ -2,7 +2,7 @@
 
 require_relative '../core/environment/identity'
 require_relative '../core/utils'
-
+require_relative 'tracer'
 require_relative 'event'
 require_relative 'metadata/tagging'
 require_relative 'sampling/ext'
@@ -75,7 +75,9 @@ def initialize(
         metrics: nil,
         trace_state: nil,
         trace_state_unknown_fields: nil,
-        remote_parent: false
+        remote_parent: false,
+        tracer: nil
+
       )
         # Attributes
         @id = id || Tracing::Utils::TraceId.next_id
@@ -98,6 +100,7 @@ def initialize(
         @profiling_enabled = profiling_enabled
         @trace_state = trace_state
         @trace_state_unknown_fields = trace_state_unknown_fields
+        @tracer = tracer
 
         # Generic tags
         set_tags(tags) if tags
@@ -161,6 +164,23 @@ def resource
         @resource || (root_span && root_span.resource)
       end
 
+      # When retrieving tags or metrics we need to include root span tags for sampling purposes
+      def get_tag(key)
+        super || (root_span && root_span.get_tag(key))
+      end
+
+      def get_metric(key)
+        super || (root_span && root_span.get_metric(key))
+      end
+
+      def tags
+        all_tags = {}
+        all_tags.merge(root_span.tags) if root_span
+        all_tags.merge(@tags)
+        all_tags.merge(@metrics)
+        all_tags
+      end
+
       # Returns true if the resource has been explicitly set
       #
       # @return [Boolean]
@@ -284,10 +304,14 @@ def flush!
       # Returns a set of trace headers used for continuing traces.
       # Used for propagation across execution contexts.
       # Data should reflect the active state of the trace.
+      # DEV-3.0: Sampling is a side effect of generating the digest.
+      # We should move the sample call to inject and right before moving to new contexts(threads, forking etc.)
       def to_digest
         # Resolve current span ID
         span_id = @active_span && @active_span.id
         span_id ||= @parent_span_id unless finished?
+        # sample the trace_operation with the tracer
+        @tracer&.sample_trace(self) unless sampling_priority
 
         TraceDigest.new(
           span_id: span_id,

lib/datadog/tracing/tracer.rb

@@ -262,6 +262,17 @@ def continue_trace!(digest, key = nil, &block)
         context.activate!(trace, &block)
       end
 
+      # Sample a span, tagging the trace as appropriate.
+      def sample_trace(trace_op)
+        begin
+          @sampler.sample!(trace_op)
+        rescue StandardError => e
+          SAMPLE_TRACE_LOG_ONLY_ONCE.run do
+            Datadog.logger.warn { "Failed to sample trace: #{e.class.name} #{e} at #{Array(e.backtrace).first}" }
+          end
+        end
+      end
+
       # @!visibility private
       # TODO: make this private
       def trace_completed
@@ -331,12 +342,14 @@ def build_trace(digest = nil)
             trace_state: digest.trace_state,
             trace_state_unknown_fields: digest.trace_state_unknown_fields,
             remote_parent: digest.span_remote,
+            tracer: self
           )
         else
           TraceOperation.new(
             hostname: hostname,
             profiling_enabled: profiling_enabled,
             remote_parent: false,
+            tracer: self
           )
         end
       end
@@ -347,7 +360,6 @@ def bind_trace_events!(trace_op)
         events.span_before_start.subscribe do |event_span_op, event_trace_op|
           event_trace_op.service ||= @default_service
           event_span_op.service ||= @default_service
-          sample_trace(event_trace_op) if event_span_op && event_span_op.parent_id == 0
         end
 
         events.span_finished.subscribe do |event_span, event_trace_op|
@@ -463,17 +475,6 @@ def subscribe_trace_deactivation!(context, trace, original_trace)
         end
       end
 
-      # Sample a span, tagging the trace as appropriate.
-      def sample_trace(trace_op)
-        begin
-          @sampler.sample!(trace_op)
-        rescue StandardError => e
-          SAMPLE_TRACE_LOG_ONLY_ONCE.run do
-            Datadog.logger.warn { "Failed to sample trace: #{e.class.name} #{e} at #{Array(e.backtrace).first}" }
-          end
-        end
-      end
-
       SAMPLE_TRACE_LOG_ONLY_ONCE = Core::Utils::OnlyOnce.new
       private_constant :SAMPLE_TRACE_LOG_ONLY_ONCE
 
@@ -492,6 +493,7 @@ def sample_span(trace_op, span)
 
       # Flush finished spans from the trace buffer, send them to writer.
       def flush_trace(trace_op)
+        sample_trace(trace_op) unless trace_op.sampling_priority
         begin
           trace = @trace_flush.consume!(trace_op)
           write(trace) if trace && !trace.empty?

spec/datadog/tracing/integration_spec.rb

@@ -222,7 +222,11 @@ def agent_receives_span_step3(previous_success)
         @span = trace.spans[0]
       end
 
-      tracer.trace('my.op').finish
+      tracer.trace('my.op', service: 'my.service') do |span|
+        span.set_tag('tag', 'tag_value')
+        span.set_tag('tag2', 'tag_value2')
+        span.resource = 'my.resource'
+      end
 
       try_wait_until { tracer.writer.stats[:traces_flushed] >= 1 }
 
@@ -319,6 +323,87 @@ def agent_receives_span_step3(previous_success)
           it_behaves_like 'sampling decision', '-3'
         end
 
+        context 'with a matching resource name' do
+          include_context 'DD_TRACE_SAMPLING_RULES configuration' do
+            let(:rule) { { resource: 'my.resource', sample_rate: 1.0 } }
+          end
+
+          it_behaves_like 'flushed trace'
+          it_behaves_like 'priority sampled', Datadog::Tracing::Sampling::Ext::Priority::USER_KEEP
+          it_behaves_like 'rule sampling rate metric', 1.0
+          it_behaves_like 'rate limit metric', 1.0
+          it_behaves_like 'sampling decision', '-3'
+        end
+
+        context 'with a matching service name' do
+          include_context 'DD_TRACE_SAMPLING_RULES configuration' do
+            let(:rule) { { service: 'my.service', sample_rate: 1.0 } }
+          end
+
+          it_behaves_like 'flushed trace'
+          it_behaves_like 'priority sampled', Datadog::Tracing::Sampling::Ext::Priority::USER_KEEP
+          it_behaves_like 'rule sampling rate metric', 1.0
+          it_behaves_like 'rate limit metric', 1.0
+          it_behaves_like 'sampling decision', '-3'
+        end
+
+        context 'with matching tags' do
+          include_context 'DD_TRACE_SAMPLING_RULES configuration' do
+            let(:rule) { { tags: { tag: 'tag_value', tag2: 'tag_value2' }, sample_rate: 1.0 } }
+          end
+
+          it_behaves_like 'flushed trace'
+          it_behaves_like 'priority sampled', Datadog::Tracing::Sampling::Ext::Priority::USER_KEEP
+          it_behaves_like 'rule sampling rate metric', 1.0
+          it_behaves_like 'rate limit metric', 1.0
+          it_behaves_like 'sampling decision', '-3'
+        end
+
+        context 'with matching tags and matching service and matching resource' do
+          include_context 'DD_TRACE_SAMPLING_RULES configuration' do
+            let(:rule) do
+              { resource: 'my.resource', service: 'my.service', tags: { tag: 'tag_value', tag2: 'tag_value2' },
+                sample_rate: 1.0 }
+            end
+          end
+
+          it_behaves_like 'flushed trace'
+          it_behaves_like 'priority sampled', Datadog::Tracing::Sampling::Ext::Priority::USER_KEEP
+          it_behaves_like 'rule sampling rate metric', 1.0
+          it_behaves_like 'rate limit metric', 1.0
+          it_behaves_like 'sampling decision', '-3'
+        end
+
+        context 'with not matching tags and matching service and matching resource' do
+          include_context 'DD_TRACE_SAMPLING_RULES configuration' do
+            let(:rule) do
+              { resource: 'my.resource', service: 'my.service', tags: { tag: 'wrong_tag_value' },
+                sample_rate: 1.0 }
+            end
+          end
+
+          it_behaves_like 'flushed trace'
+          it_behaves_like 'priority sampled', Datadog::Tracing::Sampling::Ext::Priority::AUTO_KEEP
+          it_behaves_like 'rule sampling rate metric', nil # Rule is not applied
+          it_behaves_like 'rate limit metric', nil # Rate limiter is never reached, thus has no value to provide
+          it_behaves_like 'sampling decision', '-0'
+        end
+
+        context 'drop with matching tags and matching service and matching resource' do
+          include_context 'DD_TRACE_SAMPLING_RULES configuration' do
+            let(:rule) do
+              { resource: 'my.resource', service: 'my.service', tags: { tag: 'tag_value' },
+                sample_rate: 0 }
+            end
+          end
+
+          it_behaves_like 'flushed trace'
+          it_behaves_like 'priority sampled', Datadog::Tracing::Sampling::Ext::Priority::USER_REJECT
+          it_behaves_like 'rule sampling rate metric', 0.0
+          it_behaves_like 'rate limit metric', nil # Rate limiter is never reached, thus has no value to provide
+          it_behaves_like 'sampling decision', nil
+        end
+
         context 'with low sample rate' do
           let(:rule) { Datadog::Tracing::Sampling::SimpleRule.new(sample_rate: Float::MIN) }
 

spec/datadog/tracing/trace_operation_spec.rb

@@ -136,6 +136,22 @@
     end
 
     context 'given' do
+      context ':trace_operation_samples' do
+        let(:tracer) { instance_double(Datadog::Tracing::Tracer) }
+        let(:trace_op) { described_class.new(tracer: tracer) }
+
+        describe '#to_digest' do
+          before do
+            allow(tracer).to receive(:sample_trace)
+          end
+
+          it 'calls tracer.sample_trace' do
+            expect(tracer).to receive(:sample_trace).with(trace_op)
+            trace_op.to_digest
+          end
+        end
+      end
+
       context ':agent_sample_rate' do
         subject(:options) { { agent_sample_rate: agent_sample_rate } }
         let(:agent_sample_rate) { 0.5 }