Merge pull request #3180 from newrelic/hybrid-agent-starting-transaction-tests

Enable "Starting transaction tests" in hybrid agent CAT

Author: kaylareopelle

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

@@ -7,14 +7,10 @@ module Agent
     module OpenTelemetry
       module Trace
         class Span < ::OpenTelemetry::Trace::Span
-          attr_reader :context
+          attr_accessor :finishable
 
-          def initialize(segment:, transaction:)
-            @context = ::OpenTelemetry::Trace::SpanContext.new(
-              trace_id: transaction.trace_id,
-              span_id: segment.guid,
-              trace_flags: 1
-            )
+          def finish(end_timestamp: nil)
+            finishable&.finish
           end
         end
       end

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

@@ -12,22 +12,65 @@ def initialize(name = nil, version = nil)
             @version = version || ''
           end
 
+          def start_span(name, with_parent: nil, attributes: nil, links: nil, start_timestamp: nil, kind: nil)
+            parent_span_context = ::OpenTelemetry::Trace.current_span(with_parent).context
+
+            finishable = if can_start_transaction?(parent_span_context)
+              return if internal_span_kind_with_invalid_parent?(kind, parent_span_context)
+
+              nr_obj = NewRelic::Agent::Tracer.start_transaction_or_segment(name: name, category: :otel)
+              add_remote_partent_span_context_to_txn(nr_obj, parent_span_context)
+              nr_obj
+            else
+              NewRelic::Agent::Tracer.start_segment(name: name)
+            end
+
+            span = get_span_from_finishable(finishable)
+            span.finishable = finishable
+            span
+          end
+
           def in_span(name, attributes: nil, links: nil, start_timestamp: nil, kind: nil)
-            case kind
-            when :internal
-              begin
-                return yield unless NewRelic::Agent::Tracer.current_transaction
+            span = start_span(name, attributes: attributes, links: links, start_timestamp: start_timestamp, kind: kind)
+            begin
+              yield
+            rescue => e
+              # TODO: Update for segment errors if finishable is a segment
+              NewRelic::Agent.notice_error(e)
+              raise
+            end
+          ensure
+            span&.finish
+          end
 
-                segment = NewRelic::Agent::Tracer.start_segment(name: name)
+          private
 
-                NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
-              ensure
-                segment&.finish
-              end
+          def get_span_from_finishable(finishable)
+            case finishable
+            when NewRelic::Agent::Transaction
+              finishable.segments.first.instance_variable_get(:@otel_span)
+            when NewRelic::Agent::Transaction::Segment
+              finishable.instance_variable_get(:@otel_span)
             else
-              NewRelic::Agent.logger.debug("Span kind: #{kind} is not supported yet")
+              NewRelic::Agent.logger.warn('Tracer#get_span_from_finishable failed to get span from finishable - finishable is not a transaction or segment')
+              nil
             end
           end
+
+          def can_start_transaction?(parent_span_context)
+            parent_span_context.remote? || !parent_span_context.valid?
+          end
+
+          def internal_span_kind_with_invalid_parent?(kind, parent_span_context)
+            !parent_span_context.valid? && kind == :internal
+          end
+
+          def add_remote_partent_span_context_to_txn(txn, parent_span_context)
+            return unless txn.is_a?(NewRelic::Agent::Transaction) && parent_span_context.remote?
+
+            txn.trace_id = parent_span_context.trace_id
+            txn.parent_span_id = parent_span_context.span_id
+          end
         end
       end
     end

lib/new_relic/agent/opentelemetry/transaction_patch.rb

@@ -13,29 +13,16 @@ def initialize(_category, _options)
           super
         end
 
-        def otel_current_span_key
-          # CURRENT_SPAN_KEY is a private constant
-          ::OpenTelemetry::Trace.const_get(:CURRENT_SPAN_KEY)
-        end
-
         def set_current_segment(new_segment)
           @current_segment_lock.synchronize do
-            # detach the current token, if one is present
             unless opentelemetry_context.empty?
               ::OpenTelemetry::Context.detach(opentelemetry_context[otel_current_span_key])
             end
 
-            # create an otel span for the new segment
-            span = Trace::Span.new(segment: new_segment, transaction: new_segment.transaction)
-
-            # set otel's current span to the newly created otel span
+            span = find_or_create_span(new_segment)
             ctx = ::OpenTelemetry::Context.current.set_value(otel_current_span_key, span)
-
-            # attach the token generated from updating the current span
             token = ::OpenTelemetry::Context.attach(ctx)
 
-            # update our context tracking hash to correlate the context token
-            # with the otel_current_span_key
             opentelemetry_context[otel_current_span_key] = token
           end
 
@@ -51,6 +38,32 @@ def remove_current_segment_by_thread_id(id)
 
           super
         end
+
+        private
+
+        def find_or_create_span(segment)
+          if segment.instance_variable_defined?(:@otel_span)
+            segment.instance_variable_get(:@otel_span)
+          else
+            span = Trace::Span.new(span_context: span_context_from_segment(segment))
+            segment.instance_variable_set(:@otel_span, span)
+            span
+          end
+        end
+
+        def span_context_from_segment(segment)
+          ::OpenTelemetry::Trace::SpanContext.new(
+            trace_id: segment.transaction.trace_id,
+            span_id: segment.guid,
+            trace_flags: ::OpenTelemetry::Trace::TraceFlags::SAMPLED,
+            remote: false
+          )
+        end
+
+        def otel_current_span_key
+          # CURRENT_SPAN_KEY is a private constant
+          ::OpenTelemetry::Trace.const_get(:CURRENT_SPAN_KEY)
+        end
       end
     end
   end

test/multiverse/suites/hybrid_agent/Envfile

@@ -8,4 +8,5 @@ end
 
 gemfile <<~RB
   gem 'opentelemetry-api'
+  gem 'rack'
 RB

test/multiverse/suites/hybrid_agent/commands.rb

@@ -10,11 +10,20 @@ def do_work_in_span(span_name:, span_kind:, &block)
   end
 
   def do_work_in_span_with_remote_parent(span_name:, span_kind:, &block)
-    span = @tracer.start_span(span_name, kind: span_kind)
-    span.context.instance_variable_set(:@remote, true)
+    context = OpenTelemetry::Context.current
+    span_context = OpenTelemetry::Trace::SpanContext.new(
+      trace_id: 'ba8bc8cc6d062849b0efcf3c169afb5a',
+      span_id: '6d3efb1b173fecfa',
+      trace_flags: '01',
+      remote: true
+    )
+
+    span = OpenTelemetry::Trace.non_recording_span(span_context)
+    parent_context = OpenTelemetry::Trace.context_with_span(span, parent_context: context)
+
+    span = @tracer.start_span(span_name, with_parent: parent_context, kind: span_kind)
     yield if block
-  ensure
-    span&.finish
+    span.finish
   end
 
   def do_work_in_span_with_inbound_context(span_name:, span_kind:, trace_id_in_header:,

test/multiverse/suites/hybrid_agent/hybrid_agent_test.rb

@@ -32,9 +32,10 @@ def teardown
   # until the full suite can be run on the CI
   def focus_tests
     %w[
-      does_not_create_segment_without_a_transaction
       creates_opentelemetry_segment_in_a_transaction
       creates_new_relic_span_as_child_of_opentelemetry_span
+      does_not_create_segment_without_a_transaction
+      starting_transaction_tests
     ]
   end
 
@@ -51,7 +52,7 @@ def focus_tests
           parse_operation(o)
         end
 
-        harvest_and_verify_agent_output(test_case['agentOutput'])
+        verify_agent_output(test_case['agentOutput'])
       else
         skip('marked pending by exclusion from #focus_tests')
       end

test/multiverse/suites/hybrid_agent/parsing_helpers.rb

@@ -3,14 +3,6 @@
 # frozen_string_literal: true
 
 module ParsingHelpers
-  def harvest_and_verify_agent_output(agent_output)
-    txns = harvest_transaction_events![1]
-    spans = harvest_span_events![1]
-
-    verify_agent_output(txns, agent_output, 'transactions')
-    verify_agent_output(spans, agent_output, 'spans')
-  end
-
   def parse_operation(operation)
     command = parse_command(operation['command'])
     parameters = parse_parameters(operation['parameters'])
@@ -123,47 +115,48 @@ def matches_assertion(parameters)
     assert_equal expected, actual, "Expected #{object} to equal #{expected}"
   end
 
-  def verify_agent_output(harvest, output, type)
-    puts "Agent Output: #{type}: #{output[type]}" if ENV['ENABLE_OUTPUT']
+  def verify_agent_output(output)
+    verify_transaction_output(output['transactions'])
+    verify_span_output(output['spans'])
+  end
 
-    if output[type].empty?
-      assert_empty harvest, "Agent output expected no #{type}. Found: #{harvest}"
-    else
-      outputs = prepare_keys(output[type])
-      events = harvest.flatten.select { |a| a.key?('type') }
+  def verify_transaction_output(output)
+    txns = harvest_transaction_events![1].flatten.map { |t| t['name'] }.compact
 
-      events.each_with_index do |event, i|
-        output = outputs[i]
-        if output && event
-          msg = "Agent output for #{type.capitalize} wasn't found in the harvest." \
-            "\nHarvest: #{event}\nAgent output: #{output}"
+    assert_equal output.length, txns.length, 'Wrong number of transactions found'
 
-          if output['parent_name']
-            result = events.find { |e| e['guid'] == event['parentId'] }.dig('name')
+    output.each do |txn|
+      assert_includes(txns, txn['name'], "Transaction name #{txn['name']} missing from output")
+    end
+  end
 
-            assert_equal output['parent_name'], result, msg
-            output.delete('parent_name')
-          end
+  def verify_span_output(output)
+    spans = harvest_span_events![1].flatten.reject { |h| h.empty? }
 
-          assert event >= output, msg
-        end
+    output.each do |expected|
+      actual = spans.find { |s| s['name'] == expected['name'] }
+
+      if expected['category']
+        assert_equal expected['category'], actual['category'], 'Unexpected category'
+      end
+
+      if expected['entryPoint']
+        assert_equal expected['entryPoint'], actual['nr.entryPoint'], 'Unexpected entryPoint'
+      end
+
+      expected['attributes']&.each do |expected_key, expected_value|
+        assert_equal expected_value, actual['attributes'][expected_key], 'Unexpected attribute'
+      end
+
+      if expected['parentName']
+        result = spans.find { |s| s['guid'] == actual['parentId'] }.dig('name')
+
+        assert_equal expected['parentName'], result, 'Unexpected parent name'
       end
     end
   end
 
   def snake_sub_downcase(key)
     key.gsub(/(.)([A-Z])/, '\1_\2').downcase
   end
-
-  def prepare_keys(output)
-    output.map do |o|
-      o.transform_keys do |k|
-        if k == 'entryPoint'
-          k = 'nr.entryPoint'
-        else
-          snake_sub_downcase(k)
-        end
-      end
-    end
-  end
 end

test/multiverse/suites/hybrid_agent/span_test.rb

@@ -7,23 +7,34 @@ module Agent
     module OpenTelemetry
       module Trace
         class SpanTest < Minitest::Test
-          Segment = Struct.new(:guid)
+          Segment = Struct.new(:guid, :transaction)
           Transaction = Struct.new(:trace_id)
 
-          def teardown
-            NewRelic::Agent.instance.transaction_event_aggregator.reset!
-            NewRelic::Agent.instance.span_event_aggregator.reset!
+          def test_finish_does_not_fail_if_no_finishable_present
+            span = NewRelic::Agent::OpenTelemetry::Trace::Span.new
+
+            assert_nil span.finishable
+            assert_nil span.finish
           end
 
-          def test_span_has_context
-            segment = Segment.new('123')
-            transaction = Transaction.new('456')
+          def test_finishable_can_finish_transactions
+            txn = NewRelic::Agent::Tracer.start_transaction_or_segment(category: :web, name: 'test')
+            span = NewRelic::Agent::OpenTelemetry::Trace::Span.new
+            span.finishable = txn
+            span.finish
+
+            assert_predicate span.finishable, :finished?
+            assert_predicate txn, :finished?
+          end
 
-            span = Span.new(segment: segment, transaction: transaction)
+          def test_finishable_can_finish_segments
+            segment = NewRelic::Agent::Transaction::Segment.new
+            span = NewRelic::Agent::OpenTelemetry::Trace::Span.new
+            span.finishable = segment
+            span.finish
 
-            assert_equal '123', span.context.span_id
-            assert_equal '456', span.context.trace_id
-            assert_equal 1, span.context.trace_flags
+            assert_predicate span.finishable, :finished?
+            assert_predicate segment, :finished?
           end
         end
       end

test/multiverse/suites/hybrid_agent/tracer_test.rb

@@ -16,14 +16,6 @@ def teardown
             NewRelic::Agent.instance.span_event_aggregator.reset!
           end
 
-          def test_in_span_logs_when_span_kind_unknown
-            NewRelic::Agent.stub(:logger, NewRelic::Agent::MemoryLogger.new) do
-              @tracer.in_span('fruit', kind: :mango) { 'yep' }
-
-              assert_logged(/Span kind: mango is not supported yet/)
-            end
-          end
-
           def test_in_span_creates_segment_when_span_kind_internal
             txn = in_transaction do
               @tracer.in_span('fruit', kind: :internal) { 'seeds' }
@@ -47,6 +39,17 @@ def test_in_span_captures_error_when_span_kind_internal
             assert_transaction_noticed_error txn, 'RuntimeError'
           end
 
+          def test_start_span_assigns_finishable_to_transaction
+            otel_span = @tracer.start_span('otel_api_span')
+            otel_finishable = otel_span.finishable
+
+            assert_instance_of NewRelic::Agent::Transaction, otel_finishable, "OTel span's finishable should be an NR Transaction"
+
+            otel_span.finish
+
+            assert_predicate otel_finishable, :finished?, 'OTel span should finish NR transaction'
+          end
+
           private
 
           def assert_logged(expected)