feat: lazy sampling

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/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

@@ -331,12 +331,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 +349,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|
@@ -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) }