[sinatra] tag span with request id if it is available in response headers

docs/GettingStarted.md

@@ -980,6 +980,7 @@ Where `options` is an optional `Hash` that accepts the following parameters:
 | ``resource_script_names`` | Prepend resource names with script name | ``false`` |
 | ``distributed_tracing`` | Enables [distributed tracing](#distributed-tracing) so that this service trace is connected with a trace of another service if tracing headers are received | `false` |
 | ``tracer`` | A ``Datadog::Tracer`` instance used to instrument the application. Usually you don't need to set that. | ``Datadog.tracer`` |
+| ``headers`` | Hash of HTTP request or response headers to add as tags to the `sinatra.request`. Accepts `request` and `response` keys with Array values e.g. `['Last-Modified']`. Adds `http.request.headers.*` and `http.response.headers.*` tags respectively. | ``{ response: ['Content-Type', 'X-Request-ID'] }`` |
 
 ### Sucker Punch
 

lib/ddtrace/contrib/sinatra/request_span.rb

@@ -0,0 +1,39 @@
+module Datadog
+  module Contrib
+    module Sinatra
+      # Middleware used for automatically tagging configured headers and handle request span
+      module RequestSpan
+        SINATRA_REQUEST_SPAN = 'datadog.sinatra_request_span'.freeze
+        SINATRA_REQUEST_TRACE_NAME = 'sinatra.request'.freeze
+
+        module_function
+
+        def span!(env)
+          env[SINATRA_REQUEST_SPAN] ||= build_span(env)
+        end
+
+        def fetch_span(env)
+          env[SINATRA_REQUEST_SPAN]
+        end
+
+        def build_span(env)
+          tracer = configuration[:tracer]
+          distributed_tracing = configuration[:distributed_tracing]
+
+          if distributed_tracing && tracer.provider.context.trace_id.nil?
+            context = HTTPPropagator.extract(env)
+            tracer.provider.context = context if context.trace_id
+          end
+
+          tracer.trace(SINATRA_REQUEST_TRACE_NAME,
+                       service: configuration[:service_name],
+                       span_type: Datadog::Ext::HTTP::TYPE)
+        end
+
+        def configuration
+          Datadog.configuration[:sinatra]
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/sinatra/tracer.rb

@@ -1,16 +1,18 @@
-
 require 'sinatra/base'
 
 require 'ddtrace/ext/app_types'
 require 'ddtrace/ext/errors'
 require 'ddtrace/ext/http'
 require 'ddtrace/propagation/http_propagator'
 
+require 'ddtrace/contrib/sinatra/tracer_middleware'
+require 'ddtrace/contrib/sinatra/request_span'
+
 sinatra_vs = Gem::Version.new(Sinatra::VERSION)
 sinatra_min_vs = Gem::Version.new('1.4.0')
 if sinatra_vs < sinatra_min_vs
   raise "sinatra version #{sinatra_vs} is not supported yet " \
-        + "(supporting versions >=#{sinatra_min_vs})"
+          + "(supporting versions >=#{sinatra_min_vs})"
 end
 
 Datadog::Tracer.log.info("activating instrumentation for sinatra #{sinatra_vs}")
@@ -21,6 +23,10 @@ module Sinatra
       # Datadog::Contrib::Sinatra::Tracer is a Sinatra extension which traces
       # requests.
       module Tracer
+        DEFAULT_HEADERS = {
+          response: %w[Content-Type X-Request-ID]
+        }.freeze
+
         include Base
         register_as :sinatra
 
@@ -32,6 +38,7 @@ module Tracer
         option :tracer, default: Datadog.tracer
         option :resource_script_names, default: false
         option :distributed_tracing, default: false
+        option :headers, default: DEFAULT_HEADERS
 
         def route(verb, action, *)
           # Keep track of the route name when the app is instantiated for an
@@ -49,8 +56,6 @@ def route(verb, action, *)
           super
         end
 
-        # rubocop:disable Metrics/AbcSize
-        # rubocop:disable Metrics/MethodLength
         def self.registered(app)
           ::Sinatra::Base.module_eval do
             def render(engine, data, *)
@@ -71,52 +76,30 @@ def render(engine, data, *)
             end
           end
 
+          app.use TracerMiddleware
+
           app.before do
             return unless Datadog.configuration[:sinatra][:tracer].enabled
 
-            if instance_variable_defined? :@datadog_request_span
-              if @datadog_request_span
-                Datadog::Tracer.log.error('request span active in :before hook')
-                @datadog_request_span.finish()
-                @datadog_request_span = nil
-              end
-            end
-
-            tracer = Datadog.configuration[:sinatra][:tracer]
-            distributed_tracing = Datadog.configuration[:sinatra][:distributed_tracing]
-
-            if distributed_tracing && tracer.provider.context.trace_id.nil?
-              context = HTTPPropagator.extract(request.env)
-              tracer.provider.context = context if context.trace_id
-            end
-
-            span = tracer.trace('sinatra.request',
-                                service: Datadog.configuration[:sinatra][:service_name],
-                                span_type: Datadog::Ext::HTTP::TYPE)
+            span = RequestSpan.span!(request.env)
             span.set_tag(Datadog::Ext::HTTP::URL, request.path)
             span.set_tag(Datadog::Ext::HTTP::METHOD, request.request_method)
-
-            @datadog_request_span = span
           end
 
           app.after do
             return unless Datadog.configuration[:sinatra][:tracer].enabled
 
-            span = @datadog_request_span
-            begin
-              unless span
-                Datadog::Tracer.log.error('missing request span in :after hook')
-                return
-              end
+            span = RequestSpan.fetch_span(request.env)
 
-              span.resource = "#{request.request_method} #{@datadog_route}"
-              span.set_tag('sinatra.route.path', @datadog_route)
-              span.set_tag(Datadog::Ext::HTTP::STATUS_CODE, response.status)
-              span.set_error(env['sinatra.error']) if response.server_error?
-              span.finish()
-            ensure
-              @datadog_request_span = nil
+            unless span
+              Datadog::Tracer.log.error('missing request span in :after hook')
+              return
             end
+
+            span.resource = "#{request.request_method} #{@datadog_route}"
+            span.set_tag('sinatra.route.path', @datadog_route)
+            span.set_tag(Datadog::Ext::HTTP::STATUS_CODE, response.status)
+            span.set_error(env['sinatra.error']) if response.server_error?
           end
         end
       end

lib/ddtrace/contrib/sinatra/tracer_middleware.rb

@@ -0,0 +1,74 @@
+require 'ddtrace/contrib/sinatra/request_span'
+
+module Datadog
+  module Contrib
+    module Sinatra
+      # Middleware used for automatically tagging configured headers and handle request span
+      class TracerMiddleware
+        def initialize(app)
+          @app = app
+        end
+
+        def call(env)
+          span = RequestSpan.span!(env)
+
+          # Request headers
+          parse_request_headers(env).each do |name, value|
+            span.set_tag(name, value) if span.get_tag(name).nil?
+          end
+
+          status, headers, response_body = @app.call(env)
+
+          # Response headers
+          parse_response_headers(headers).each do |name, value|
+            span.set_tag(name, value) if span.get_tag(name).nil?
+          end
+
+          [status, headers, response_body]
+        ensure
+          span.finish
+        end
+
+        private
+
+        def parse_request_headers(env)
+          {}.tap do |result|
+            whitelist = configuration[:headers][:request] || []
+            whitelist.each do |header|
+              rack_header = header_to_rack_header(header)
+              if env.key?(rack_header)
+                result[Datadog::Ext::HTTP::RequestHeaders.to_tag(header)] = env[rack_header]
+              end
+            end
+          end
+        end
+
+        def parse_response_headers(headers)
+          {}.tap do |result|
+            whitelist = configuration[:headers][:response] || []
+            whitelist.each do |header|
+              if headers.key?(header)
+                result[Datadog::Ext::HTTP::ResponseHeaders.to_tag(header)] = headers[header]
+              else
+                # Try a case-insensitive lookup
+                uppercased_header = header.to_s.upcase
+                matching_header = headers.keys.find { |h| h.upcase == uppercased_header }
+                if matching_header
+                  result[Datadog::Ext::HTTP::ResponseHeaders.to_tag(header)] = headers[matching_header]
+                end
+              end
+            end
+          end
+        end
+
+        def configuration
+          Datadog.configuration[:sinatra]
+        end
+
+        def header_to_rack_header(name)
+          "HTTP_#{name.to_s.upcase.gsub(/[-\s]/, '_')}"
+        end
+      end
+    end
+  end
+end

test/contrib/sinatra/tracer_test.rb

@@ -4,6 +4,7 @@
 class TracerTest < TracerTestBase
   class TracerTestApp < Sinatra::Application
     get '/request' do
+      headers['X-Request-Id'] = request.env['HTTP_X_REQUEST_ID']
       'hello world'
     end
 
@@ -83,10 +84,10 @@ def test_distributed_request
     # Enable distributed tracing
     Datadog.configuration.use(:sinatra, distributed_tracing: true)
 
-    response = get  '/request', {},
-                    'HTTP_X_DATADOG_TRACE_ID' => '1',
-                    'HTTP_X_DATADOG_PARENT_ID' => '2',
-                    'HTTP_X_DATADOG_SAMPLING_PRIORITY' => Datadog::Ext::Priority::USER_KEEP.to_s
+    response = get '/request', {},
+                   'HTTP_X_DATADOG_TRACE_ID' => '1',
+                   'HTTP_X_DATADOG_PARENT_ID' => '2',
+                   'HTTP_X_DATADOG_SAMPLING_PRIORITY' => Datadog::Ext::Priority::USER_KEEP.to_s
 
     assert_equal(200, response.status)
 
@@ -237,4 +238,61 @@ def test_literal_template
     assert_equal(0, root.status)
     assert_nil(root.parent)
   end
+
+  def test_tagging_default_connection_headers
+    request_id = SecureRandom.uuid
+    get '/request', {}, 'HTTP_X_REQUEST_ID' => request_id
+
+    assert_equal(200, last_response.status)
+
+    spans = @writer.spans
+    assert_equal(1, spans.length)
+
+    span = spans[0]
+    assert_equal('sinatra', span.service)
+    assert_equal('GET /request', span.resource)
+    assert_equal('GET', span.get_tag(Datadog::Ext::HTTP::METHOD))
+    assert_equal('/request', span.get_tag(Datadog::Ext::HTTP::URL))
+    assert_equal(Datadog::Ext::HTTP::TYPE, span.span_type)
+    assert_equal(request_id, span.get_tag('http.response.headers.x_request_id'))
+    assert_equal('text/html;charset=utf-8', span.get_tag('http.response.headers.content_type'))
+
+    assert_equal(0, span.status)
+    assert_nil(span.parent)
+  end
+
+  def test_tagging_configured_connection_headers
+    Datadog.configuration.use(:sinatra,
+                              headers: {
+                                response: ['Content-Type'],
+                                request: ['X-Request-Header']
+                              })
+
+    request_headers = {
+      'HTTP_X_REQUEST_HEADER' => 'header_value',
+      'HTTP_X_HEADER' => "don't tag me"
+    }
+
+    get '/request#foo?a=1', {}, request_headers
+
+    assert_equal(200, last_response.status)
+
+    spans = @writer.spans
+    assert_equal(1, spans.length)
+
+    span = spans[0]
+    assert_equal('sinatra', span.service)
+    assert_equal('GET /request', span.resource)
+    assert_equal('GET', span.get_tag(Datadog::Ext::HTTP::METHOD))
+    assert_equal('/request', span.get_tag(Datadog::Ext::HTTP::URL))
+    assert_equal(Datadog::Ext::HTTP::TYPE, span.span_type)
+    assert_equal('header_value', span.get_tag('http.request.headers.x_request_header'))
+    assert_equal('text/html;charset=utf-8', span.get_tag('http.response.headers.content_type'))
+    assert_nil(span.get_tag('http.request.headers.x_header'))
+
+    assert_equal(0, span.status)
+    assert_nil(span.parent)
+  ensure
+    Datadog.configuration.use(:sinatra, headers: Datadog::Contrib::Sinatra::Tracer::DEFAULT_HEADERS)
+  end
 end