[fix] rack middleware close trace via body proxy

lib/ddtrace/contrib/delayed_job/plugin.rb

@@ -9,7 +9,7 @@ module DelayedJob
       # DelayedJob plugin that instruments invoke_job hook
       class Plugin < Delayed::Plugin
         def self.instrument_invoke(job, &block)
-          return block.call(job) unless tracer && tracer.enabled
+          return yield(job) unless tracer && tracer.enabled
 
           tracer.trace(Ext::SPAN_JOB, service: configuration[:service_name], resource: job_name(job),
                                       on_error: configuration[:error_handler]) do |span|
@@ -29,7 +29,7 @@ def self.instrument_invoke(job, &block)
         end
 
         def self.instrument_enqueue(job, &block)
-          return block.call(job) unless tracer && tracer.enabled
+          return yield(job) unless tracer && tracer.enabled
 
           tracer.trace(Ext::SPAN_ENQUEUE, service: configuration[:client_service_name], resource: job_name(job)) do |span|
             set_sample_rate(span)

lib/ddtrace/contrib/rack/middlewares.rb

@@ -86,7 +86,21 @@ def call(env)
 
           # call the rest of the stack
           status, headers, response = @app.call(env)
-          [status, headers, response]
+
+          # some middleware like Rack::CommonLogger wraps body via Rack::BodyProxy
+          # Rack::BodyProxy is a correct mechanism to finish span after all middleware call
+          # so Rack::CommonLogger will have trace_id & span_id tags
+
+          body = if defined?(::Rack::BodyProxy)
+                   ::Rack::BodyProxy.new(response) do
+                     close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
+                   end
+                 else
+                   close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
+                   response
+                 end
+
+          [status, headers, body]
 
         # rubocop:disable Lint/RescueException
         # Here we really want to catch *any* exception, not only StandardError,
@@ -98,9 +112,13 @@ def call(env)
           # catch exceptions that may be raised in the middleware chain
           # Note: if a middleware catches an Exception without re raising,
           # the Exception cannot be recorded here.
-          request_span.set_error(e) unless request_span.nil?
+          request_span.set_error(e) if request_span
+
+          close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
           raise e
-        ensure
+        end
+
+        def close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
           if request_span
             # Rack is a really low level interface and it doesn't provide any
             # advanced functionality like routers. Because of that, we assume that

lib/ddtrace/contrib/rack/patcher.rb

@@ -14,6 +14,7 @@ def target_version
 
         def patch
           # Patch middleware
+          require 'rack'
           require_relative 'middlewares'
         end
       end