feat: support context propagation through links

Author: gwuah

instrumentation/rack/lib/opentelemetry/instrumentation/rack/instrumentation.rb

@@ -26,6 +26,7 @@ class Instrumentation < OpenTelemetry::Instrumentation::Base
         option :record_frontend_span,     default: false, validate: :boolean
         option :untraced_endpoints,       default: [],    validate: :array
         option :url_quantization,         default: nil,   validate: :callable
+        option :propagate_with_link,      default: nil,   validate: :callable
         option :untraced_requests,        default: nil,   validate: :callable
         option :response_propagators,     default: [],    validate: :array
         # This option is only valid for applications using Rack 2.0 or greater

instrumentation/rack/lib/opentelemetry/instrumentation/rack/middlewares/event_handler.rb

@@ -55,7 +55,15 @@ def on_start(request, _)
             return if untraced_request?(request.env)
 
             parent_context = extract_remote_context(request)
-            span = create_span(parent_context, request)
+            links = nil
+
+            fn = propagate_with_link
+            if fn&.call(request.env)
+              links = prepare_span_links(parent_context)
+              parent_context = OpenTelemetry::Context.empty
+            end
+
+            span = create_span(parent_context, request, links)
             request.env[TOKENS_KEY] = register_current_span(span)
           rescue StandardError => e
             OpenTelemetry.handle_error(exception: e)
@@ -224,6 +232,10 @@ def url_quantization
             config[:url_quantization]
           end
 
+          def propagate_with_link
+            config[:propagate_with_link]
+          end
+
           def response_propagators
             config[:response_propagators]
           end
@@ -253,17 +265,23 @@ def register_current_span(span)
             contexts.map { |context| OpenTelemetry::Context.attach(context) }
           end
 
-          def create_span(parent_context, request)
+          def create_span(parent_context, request, links)
             span = tracer.start_span(
               create_request_span_name(request),
               with_parent: parent_context,
               kind: :server,
-              attributes: request_span_attributes(request.env)
+              attributes: request_span_attributes(request.env),
+              links: links
             )
             request_start_time = OpenTelemetry::Instrumentation::Rack::Util::QueueTime.get_request_start(request.env)
             span.add_event('http.proxy.request.started', timestamp: request_start_time) unless request_start_time.nil?
             span
           end
+
+          def prepare_span_links(ctx)
+            span_context = OpenTelemetry::Trace.current_span(ctx).context
+            span_context.valid? ? [OpenTelemetry::Trace::Link.new(span_context)] : []
+          end
         end
       end
     end

instrumentation/rack/test/opentelemetry/instrumentation/rack/instrumentation_test.rb

@@ -30,6 +30,7 @@
       _(instrumentation.config[:record_frontend_span]).must_equal false
       _(instrumentation.config[:untraced_endpoints]).must_be_empty
       _(instrumentation.config[:url_quantization]).must_be_nil
+      _(instrumentation.config[:propagate_with_link]).must_be_nil
       _(instrumentation.config[:untraced_requests]).must_be_nil
       _(instrumentation.config[:response_propagators]).must_be_empty
       _(instrumentation.config[:use_rack_events]).must_equal true

instrumentation/rack/test/opentelemetry/instrumentation/rack/middlewares/event_handler_test.rb

@@ -24,6 +24,7 @@
       allowed_request_headers: allowed_request_headers,
       allowed_response_headers: allowed_response_headers,
       url_quantization: url_quantization,
+      propagate_with_link: propagate_with_link,
       response_propagators: response_propagators,
       enabled: instrumentation_enabled,
       use_rack_events: true
@@ -51,6 +52,7 @@
   let(:allowed_response_headers) { nil }
   let(:response_propagators) { nil }
   let(:url_quantization) { nil }
+  let(:propagate_with_link) { nil }
   let(:headers) { {} }
   let(:app) do
     Rack::Builder.new.tap do |builder|
@@ -410,4 +412,46 @@ def inject(carrier)
       _(proxy_event).must_be_nil
     end
   end
+
+  describe 'link propagation' do
+    describe 'without link propagation fn' do
+      it 'the root span has no links' do
+        get '/url'
+
+        _(rack_span.name).must_equal 'HTTP GET'
+        _(rack_span.total_recorded_links).must_equal(0)
+      end
+    end
+
+    describe 'with link propagation fn that returns false' do
+      let(:propagate_with_link) do
+        ->(_env) { false }
+      end
+
+      it 'has no links' do
+        get '/url'
+
+        _(rack_span.name).must_equal 'HTTP GET'
+        _(rack_span.total_recorded_links).must_equal(0)
+      end
+    end
+
+    describe 'with link propagation fn that returns true' do
+      let(:propagate_with_link) do
+        ->(env) { env['PATH_INFO'].start_with?('/url') }
+      end
+
+      it 'has links' do
+        trace_id = '618c54694e838292271da0ba122547e9'
+        span_id = 'd408cc622ee29ce0'
+        header 'traceparent', "00-#{trace_id}-#{span_id}-01"
+        get '/url'
+
+        _(rack_span.name).must_equal 'HTTP GET'
+        _(rack_span.total_recorded_links).must_equal(1)
+        _(rack_span.links[0].span_context.trace_id.unpack1('H*')).must_equal(trace_id)
+        _(rack_span.links[0].span_context.span_id.unpack1('H*')).must_equal(span_id)
+      end
+    end
+  end
 end