Merge pull request #3368 from newrelic/extend-missing-tracestate-header-sampling-to-otel

Extend sampling decision logic to OTel

Author: kaylareopelle

lib/new_relic/agent/agent.rb

@@ -100,6 +100,12 @@ def init_components
         @monotonic_gc_profiler = VM::MonotonicGCProfiler.new
         @adaptive_sampler = AdaptiveSampler.new(Agent.config[:sampling_target],
           Agent.config[:sampling_target_period_in_seconds])
+        # TODO: update the configs used here to match new sampler-specific options
+        # TODO: Consider refactoring to only create the sampler if 'adaptive' is the sampler's setting
+        @adaptive_sampler_remote_parent_sampled = AdaptiveSampler.new(Agent.config[:sampling_target],
+          Agent.config[:sampling_target_period_in_seconds])
+        @adaptive_sampler_remote_parent_not_sampled = AdaptiveSampler.new(Agent.config[:sampling_target],
+          Agent.config[:sampling_target_period_in_seconds])
         @serverless_handler = ServerlessHandler.new
         @opentelemetry_bridge = OpenTelemetryBridge.new
       end
@@ -178,6 +184,8 @@ module InstanceMethods
         attr_reader :transaction_event_recorder
         attr_reader :attribute_filter
         attr_reader :adaptive_sampler
+        attr_reader :adaptive_sampler_remote_parent_sampled
+        attr_reader :adaptive_sampler_remote_parent_not_sampled
         attr_reader :serverless_handler
 
         def transaction_event_aggregator

lib/new_relic/agent/opentelemetry/trace/tracer.rb

@@ -19,7 +19,9 @@ def start_span(name, with_parent: nil, attributes: nil, links: nil, start_timest
               return if internal_span_kind_with_invalid_parent?(kind, parent_otel_context)
 
               nr_item = NewRelic::Agent::Tracer.start_transaction_or_segment(name: name, category: :otel)
+
               add_remote_context_to_txn(nr_item, parent_otel_context)
+
               nr_item
             else
               NewRelic::Agent::Tracer.start_segment(name: name)
@@ -93,18 +95,24 @@ def set_tracestate(distributed_tracer, otel_context)
           def set_nr_trace_state(distributed_tracer, otel_context)
             distributed_tracer.instance_variable_set(:@trace_state_payload, otel_context.tracestate)
             distributed_tracer.parent_transaction_id = distributed_tracer.trace_state_payload.transaction_id
-            distributed_tracer.determine_sampling_decision(otel_context.tracestate, otel_context.trace_flags)
+            trace_flags = parse_trace_flags(otel_context.trace_flags)
+
+            distributed_tracer.determine_sampling_decision(otel_context.tracestate, trace_flags)
           end
 
           def set_otel_trace_state(distributed_tracer, otel_context)
             nr_entry = otel_context.tracestate.value(Transaction::TraceContext::AccountHelpers.trace_state_entry_key)
-            return unless nr_entry
-
-            nr_payload = NewRelic::Agent::TraceContextPayload.from_s(nr_entry)
-            distributed_tracer.instance_variable_set(:@trace_state_payload, nr_payload)
-            distributed_tracer.parent_transaction_id = distributed_tracer.trace_state_payload.transaction_id
             trace_flags = parse_trace_flags(otel_context.trace_flags)
-            distributed_tracer.determine_sampling_decision(nr_payload, trace_flags)
+
+            if nr_entry
+              nr_payload = NewRelic::Agent::TraceContextPayload.from_s(nr_entry)
+
+              distributed_tracer.instance_variable_set(:@trace_state_payload, nr_payload)
+              distributed_tracer.parent_transaction_id = distributed_tracer.trace_state_payload.transaction_id
+              distributed_tracer.determine_sampling_decision(nr_payload, trace_flags)
+            else
+              distributed_tracer.determine_sampling_decision(NewRelic::Agent::TraceContextPayload::INVALID, trace_flags)
+            end
           end
 
           def parse_trace_flags(trace_flags)

lib/new_relic/agent/transaction/trace_context.rb

@@ -157,32 +157,42 @@ def determine_sampling_decision(payload, trace_flags)
             set_priority_and_sampled(
               NewRelic::Agent.config[:'distributed_tracing.sampler.remote_parent_sampled'],
               NewRelic::Agent.config[:'distributed_tracing.sampler.remote_parent_sampled.trace_id_ratio_based.ratio'],
-              payload
+              payload,
+              trace_flags
             )
           elsif trace_flags == '00'
             set_priority_and_sampled(
               NewRelic::Agent.config[:'distributed_tracing.sampler.remote_parent_not_sampled'],
               NewRelic::Agent.config[:'distributed_tracing.sampler.remote_parent_not_sampled.trace_id_ratio_based.ratio'],
-              payload
+              payload,
+              trace_flags
             )
           else
-            use_nr_tracestate_sampled(payload)
+            use_nr_tracestate_sampled(payload, trace_flags)
           end
         rescue
-          use_nr_tracestate_sampled(payload)
+          use_nr_tracestate_sampled(payload, trace_flags)
         end
 
-        def use_nr_tracestate_sampled(payload)
-          unless payload.sampled.nil?
+        def use_nr_tracestate_sampled(payload, trace_flags)
+          if payload.sampled.nil?
+            if trace_flags == '01'
+              transaction.sampled = NewRelic::Agent.instance.adaptive_sampler_remote_parent_sampled.sampled?
+            elsif trace_flags == '00'
+              transaction.sampled = NewRelic::Agent.instance.adaptive_sampler_remote_parent_not_sampled.sampled?
+            end
+
+            transaction.priority = transaction.default_priority
+          else
             transaction.sampled = payload.sampled
             transaction.priority = payload.priority if payload.priority
           end
         end
 
-        def set_priority_and_sampled(sampler, ratio, payload)
+        def set_priority_and_sampled(sampler, ratio, payload, trace_flags = nil)
           case sampler
           when 'adaptive'
-            use_nr_tracestate_sampled(payload)
+            use_nr_tracestate_sampled(payload, trace_flags)
           when 'always_on'
             transaction.sampled = true
             transaction.priority = 2.0

test/new_relic/agent/distributed_tracing/trace_context_cross_agent_test.rb

@@ -43,6 +43,8 @@ def self.focus_tests
 
             define_method("test_#{test_case['test_name']}") do
               NewRelic::Agent.instance.adaptive_sampler.stubs(:sampled?).returns(test_case['force_adaptive_sampled_true'])
+              NewRelic::Agent.instance.adaptive_sampler_remote_parent_sampled.stubs(:sampled?).returns(test_case['force_adaptive_sampled_true'])
+              NewRelic::Agent.instance.adaptive_sampler_remote_parent_not_sampled.stubs(:sampled?).returns(test_case['force_adaptive_sampled_true'])
 
               config = {
                 :'distributed_tracing.enabled' => true,