Improve rack resource names

lib/ddtrace/contrib/rack/middlewares.rb

@@ -13,16 +13,6 @@ module Rack
       # application. If request tags are not set by the app, they will be set using
       # information available at the Rack level.
       class TraceMiddleware
-        include Base
-        register_as :rack
-
-        option :tracer, default: Datadog.tracer
-        option :service_name, default: 'rack', depends_on: [:tracer] do |value|
-          get_option(:tracer).set_service_info(value, 'rack', Ext::AppTypes::WEB)
-          value
-        end
-        option :distributed_tracing, default: false
-
         def initialize(app)
           @app = app
         end
@@ -79,7 +69,7 @@ def call(env)
           # the result for this request; `resource` and `tags` are expected to
           # be set in another level but if they're missing, reasonable defaults
           # are used.
-          request_span.resource = "#{env['REQUEST_METHOD']} #{status}".strip unless request_span.resource
+          request_span.resource ||= resource_name_for(env, status)
           if request_span.get_tag(Datadog::Ext::HTTP::METHOD).nil?
             request_span.set_tag(Datadog::Ext::HTTP::METHOD, env['REQUEST_METHOD'])
           end
@@ -105,6 +95,14 @@ def call(env)
           # memory leaks.
           tracer.provider.context = Datadog::Context.new
         end
+
+        def resource_name_for(env, status)
+          if Datadog.configuration[:rack][:middleware_names]
+            "#{env['RESPONSE_MIDDLEWARE']}##{env['REQUEST_METHOD']}"
+          else
+            "#{env['REQUEST_METHOD']} #{status}".strip
+          end
+        end
       end
     end
   end

lib/ddtrace/contrib/rack/patcher.rb

@@ -0,0 +1,58 @@
+module Datadog
+  module Contrib
+    module Rack
+      # Provides instrumentation for `rack`
+      module Patcher
+        include Base
+        register_as :rack
+        option :tracer, default: Datadog.tracer
+        option :distributed_tracing, default: false
+        option :middleware_names, default: false
+        option :application
+        option :service_name, default: 'rack', depends_on: [:tracer] do |value|
+          get_option(:tracer).set_service_info(value, 'rack', Ext::AppTypes::WEB)
+          value
+        end
+
+        module_function
+
+        def patch
+          return true if patched?
+
+          require_relative 'middlewares'
+          @patched = true
+
+          return unless get_option(:middleware_names)
+
+          top = get_option(:application) || rails_app
+          retain_middleware_name(top)
+        end
+
+        def patched?
+          @patched ||= false
+        end
+
+        def rails_app
+          return unless Datadog.registry[:rails].compatible?
+          ::Rails.application.app
+        end
+
+        def retain_middleware_name(middleware)
+          return unless middleware && middleware.respond_to?(:call)
+
+          middleware.singleton_class.class_eval do
+            alias_method :__call, :call
+
+            def call(env)
+              env['RESPONSE_MIDDLEWARE'] = self.class.to_s
+              __call(env)
+            end
+          end
+
+          following = middleware.instance_variable_get('@app')
+          retain_middleware_name(following)
+        end
+      end
+    end
+  end
+end

lib/ddtrace/monkey.rb

@@ -4,6 +4,7 @@
 # because patchers do not include any 3rd party module nor even our
 # patching code, which is required on demand, when patching.
 require 'ddtrace/contrib/base'
+require 'ddtrace/contrib/rack/patcher'
 require 'ddtrace/contrib/rails/patcher'
 require 'ddtrace/contrib/active_record/patcher'
 require 'ddtrace/contrib/elasticsearch/patcher'

test/contrib/rack/resource_name_test.rb

@@ -0,0 +1,64 @@
+require_relative 'helpers'
+
+class ResourceNameTest < Minitest::Test
+  include Rack::Test::Methods
+  attr_reader :app
+
+  def setup
+    @previous_configuration = Datadog.configuration[:rack].to_h
+    @tracer = get_test_tracer
+    @app = Rack::Builder.new do
+      use Datadog::Contrib::Rack::TraceMiddleware
+      use AuthMiddleware
+      run BottomMiddleware.new
+    end.to_app
+
+    remove_patch!(:rack)
+    Datadog.configuration.use(
+      :rack,
+      middleware_names: true,
+      tracer: @tracer,
+      application: @app
+    )
+  end
+
+  def teardown
+    Datadog.configuration.use(:rack, @previous_configuration)
+  end
+
+  def test_resource_name_full_chain
+    get '/', {}, 'HTTP_AUTH_TOKEN' => '1234'
+
+    spans = @tracer.writer.spans
+    assert(last_response.ok?)
+    assert_equal(1, spans.length)
+    assert_match(/BottomMiddleware#GET/, spans[0].resource)
+  end
+
+  def test_resource_name_short_circuited_request
+    get '/', {}, 'HTTP_AUTH_TOKEN' => 'Wrong'
+
+    spans = @tracer.writer.spans
+    refute(last_response.ok?)
+    assert_equal(1, spans.length)
+    assert_match(/AuthMiddleware#GET/, spans[0].resource)
+  end
+
+  class AuthMiddleware
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      return [401, {}, []] if env['HTTP_AUTH_TOKEN'] != '1234'
+
+      @app.call(env)
+    end
+  end
+
+  class BottomMiddleware
+    def call(_)
+      [200, {}, []]
+    end
+  end
+end

test/helper.rb

@@ -237,3 +237,9 @@ def try_wait_until(options = {})
     attempts -= 1
   end
 end
+
+def remove_patch!(integration)
+  Datadog
+    .registry[integration]
+    .instance_variable_set('@patched', false)
+end

test/monkey_test.rb

@@ -13,6 +13,7 @@
 class MonkeyTest < Minitest::Test
   def test_autopatch_modules
     expected = {
+      rack: false,
       rails: true,
       elasticsearch: true,
       http: true,
@@ -42,7 +43,7 @@ def test_patch_module
     assert_equal(false, Datadog::Contrib::Grape::Patcher.patched?)
     assert_equal(false, Datadog::Contrib::Aws::Patcher.patched?)
     assert_equal(false, Datadog::Contrib::ActiveRecord::Patcher.patched?)
-    assert_equal({ rails: false, elasticsearch: false, http: false, redis: false, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false }, Datadog::Monkey.get_patched_modules())
+    assert_equal({ rails: false, elasticsearch: false, http: false, redis: false, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false, rack: false }, Datadog::Monkey.get_patched_modules())
 
     Datadog::Monkey.patch_module(:redis)
     assert_equal(false, Datadog::Contrib::Elasticsearch::Patcher.patched?)
@@ -52,7 +53,7 @@ def test_patch_module
     assert_equal(false, Datadog::Contrib::Aws::Patcher.patched?)
     assert_equal(false, Datadog::Contrib::ActiveRecord::Patcher.patched?)
     refute(Datadog::Contrib::Faraday::Patcher.patched?)
-    assert_equal({ rails: false, elasticsearch: false, http: false, redis: true, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false }, Datadog::Monkey.get_patched_modules())
+    assert_equal({ rails: false, elasticsearch: false, http: false, redis: true, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false, rack: false }, Datadog::Monkey.get_patched_modules())
 
     # now do it again to check it's idempotent
     Datadog::Monkey.patch_module(:redis)
@@ -63,7 +64,7 @@ def test_patch_module
     assert_equal(false, Datadog::Contrib::Aws::Patcher.patched?)
     assert_equal(false, Datadog::Contrib::ActiveRecord::Patcher.patched?)
     refute(Datadog::Contrib::Faraday::Patcher.patched?)
-    assert_equal({ rails: false, elasticsearch: false, http: false, redis: true, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false }, Datadog::Monkey.get_patched_modules())
+    assert_equal({ rails: false, elasticsearch: false, http: false, redis: true, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false, rack: false }, Datadog::Monkey.get_patched_modules())
 
     Datadog::Monkey.patch(elasticsearch: true, redis: true)
     assert_equal(true, Datadog::Contrib::Elasticsearch::Patcher.patched?)
@@ -72,7 +73,7 @@ def test_patch_module
     assert_equal(false, Datadog::Contrib::Grape::Patcher.patched?)
     assert_equal(false, Datadog::Contrib::Aws::Patcher.patched?)
     assert_equal(false, Datadog::Contrib::ActiveRecord::Patcher.patched?)
-    assert_equal({ rails: false, elasticsearch: true, http: false, redis: true, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false }, Datadog::Monkey.get_patched_modules())
+    assert_equal({ rails: false, elasticsearch: true, http: false, redis: true, grape: false, faraday: false, aws: false, sucker_punch: false, active_record: false, mongo: false, dalli: false, resque: false, rack: false }, Datadog::Monkey.get_patched_modules())
 
     # verify that active_record is not auto patched by default
     Datadog::Monkey.patch_all()
@@ -82,7 +83,7 @@ def test_patch_module
     assert_equal(false, Datadog::Contrib::Grape::Patcher.patched?)
     assert_equal(true, Datadog::Contrib::Aws::Patcher.patched?)
     assert_equal(false, Datadog::Contrib::ActiveRecord::Patcher.patched?)
-    assert_equal({ rails: false, elasticsearch: true, http: true, redis: true, grape: false, faraday: true, aws: true, sucker_punch: true, active_record: false, mongo: false, dalli: true, resque: true }, Datadog::Monkey.get_patched_modules())
+    assert_equal({ rails: false, elasticsearch: true, http: true, redis: true, grape: false, faraday: true, aws: true, sucker_punch: true, active_record: false, mongo: false, dalli: true, resque: true, rack: false }, Datadog::Monkey.get_patched_modules())
 
     Datadog::Monkey.patch_module(:active_record)
     assert_equal(true, Datadog::Contrib::Elasticsearch::Patcher.patched?)
@@ -91,6 +92,6 @@ def test_patch_module
     assert_equal(false, Datadog::Contrib::Grape::Patcher.patched?)
     assert_equal(true, Datadog::Contrib::Aws::Patcher.patched?)
     assert_equal(true, Datadog::Contrib::ActiveRecord::Patcher.patched?)
-    assert_equal({ rails: false, elasticsearch: true, http: true, redis: true, grape: false, faraday: true, aws: true, sucker_punch: true, active_record: true, mongo: false, dalli: true, resque: true }, Datadog::Monkey.get_patched_modules())
+    assert_equal({ rails: false, elasticsearch: true, http: true, redis: true, grape: false, faraday: true, aws: true, sucker_punch: true, active_record: true, mongo: false, dalli: true, resque: true, rack: false }, Datadog::Monkey.get_patched_modules())
   end
 end